#!/usr/bin/env python3
import argparse
import lzma
import os
import pickle
import sys
import urllib.request
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import RidgeCV, LinearRegression
from sklearn.metrics import mean_squared_error

import numpy as np

class Dataset:
    """Dataset půjčovny kol

    Dataset obsahuje tyto informace:
    - Hodina - hodina v rámci dne
    - Teplota - ve stupních Celsia
    - Vlhkost vzduchu - %
    - Rychlost větru - m/s
    - Viditelnost - v jednotkách 10m
    - Rosný bod - ve stupních Celsia
    - Sluneční záření - MJ/m^2
    - Množství srážek - mm
    - Výška sněhové pokrývky - cm
    - Roční období - (0: zima, 1: jaro, 2: léto, 3: podzim)
    - Prázdniny - (0: ne, 1: ano)
    - Fungoval systém - (0: ne, 1: ano)
    - Den   - 1-31
    - Měsíc - 1-12
    - Rok   - 2017 nebo 2018

    Cílová proměnná je počet půjčených kol v dané hodině.
    """
    def __init__(self,
                 name="dataset36-2-S.npz",
                 url="https://ksp.mff.cuni.cz/h/ulohy/36/36-2-S/competition-datasets/"):
        if not os.path.exists(name):
            print("Downloading dataset {}...".format(name), file=sys.stderr)
            urllib.request.urlretrieve(url + name, filename=name)

        # načtení serialovaného datasetu
        dataset = np.load(name, allow_pickle=True)

        self.train_data = dataset['train_data']
        self.test_data = dataset['test_data']
        self.train_target = dataset['train_target']
        # pozor: obsahuje vektor -1
        self.test_target = dataset['test_target']


parser = argparse.ArgumentParser()

parser.add_argument("--seed", default=42, type=int, help="Random seed")


class DayOfWeekTransformer(BaseEstimator, TransformerMixin):
    '''
    Transformer, který z dat vytvoří nový sloupec s číslem dne v týdnu.
    
    Pokud columns=None, tak se předpokládá, že vstupní data jsou ve formátu:
    první sloupec odpovídá dnu v měsíci, druhý měsíci a třetí roku.

    Jinak bere pořadí sloupců zadané v columns.
    '''
    def __init__(self, columns=None, copy=True):
        self.columns = columns
        self.copy = copy

    def fit(self, X, y=None):
        return self

    def transform(self, X, copy=None):
        copy = copy if copy is not None else self.copy
        X = X.copy() if copy else X

        if self.columns is None:
            day = X[:, 0]
            month = X[:, 1]
            year = X[:, 2]
        else:
            day = X[:, self.columns[0]]
            month = X[:, self.columns[1]]
            year = X[:, self.columns[2]]
        
        dates = np.array([year, month, day]).T
        import pandas as pd
        dates = np.apply_along_axis(lambda x: pd.to_datetime("-".join(x.astype(str))), 1, dates)
        import pandas as pd
        df = pd.DataFrame(dates)
        
        # Monday=0, Sunday=6
        day_of_week = df[0].dt.dayofweek
        day_of_week = np.array(day_of_week)
        
        X = np.concatenate((X, day_of_week[:, None]), axis=1)
        return X

class PartOfDayTransformer(BaseEstimator, TransformerMixin):
    '''
    Transformer, který z dat vytvoří nový sloupec s číslem části dne.
    
    Pokud columns=None, tak se předpokládá, že vstupní data jsou ve formátu:
    první sloupec odpovídá hodině v dnu.

    Jinak bere pořadí sloupců zadané v columns.
    '''
    def __init__(self, columns=None, copy=True):
        self.columns = columns
        self.copy = copy

    def fit(self, X, y=None):
        return self

    def transform(self, X, copy=None):
        copy = copy if copy is not None else self.copy
        X = X.copy() if copy else X

        if self.columns is None:
            hour = X[:, 0]
        else:
            hour = X[:, self.columns[0]]
        
        part_of_day = np.zeros_like(hour)
        part_of_day[(hour >= 0) & (hour < 6)] = 0
        part_of_day[(hour >= 6) & (hour < 12)] = 1
        part_of_day[(hour >= 12) & (hour < 18)] = 2
        part_of_day[(hour >= 18) & (hour < 24)] = 3

        X = np.concatenate((X, part_of_day[:, None]), axis=1)
        return X


class BeaufortTransformer(BaseEstimator, TransformerMixin):
    '''
    Transformer, který z rychlosti větru vytvoří stupně Beaufortovy stupnice.
    
    Pokud columns=None, tak se předpokládá, že vstupní data jsou ve formátu:
    první sloupec odpovídá rychlosti větru.

    Jinak bere pořadí sloupců zadané v columns.

    replace=True znamená, jestli má být sloupec den v týdnu nahrazen.
    '''
    def __init__(self, columns=None, copy=True):
        self.columns = columns
        self.copy = copy

    def fit(self, X, y=None):
        return self

    def transform(self, X, copy=None):
        copy = copy if copy is not None else self.copy
        X = X.copy() if copy else X

        if self.columns is None:
            wind_speed = X[:, 0]
        else:
            wind_speed = X[:, self.columns[0]]
        
        beaufort = np.zeros_like(wind_speed)
        # Beaufortova stupnice (z wikipedie)
        beaufort[(wind_speed >= 0) & (wind_speed < 0.3)] = 0
        beaufort[(wind_speed >= 0.3) & (wind_speed < 1.6)] = 1
        beaufort[(wind_speed >= 1.6) & (wind_speed < 3.4)] = 2
        beaufort[(wind_speed >= 3.4) & (wind_speed < 5.5)] = 3
        beaufort[(wind_speed >= 5.5) & (wind_speed < 8)] = 4
        beaufort[(wind_speed >= 8) & (wind_speed < 10.8)] = 5
        beaufort[(wind_speed >= 10.8) & (wind_speed < 13.9)] = 6
        beaufort[(wind_speed >= 13.9) & (wind_speed < 17.2)] = 7
        beaufort[(wind_speed >= 17.2) & (wind_speed < 20.8)] = 8
        beaufort[(wind_speed >= 20.8) & (wind_speed < 24.5)] = 9
        beaufort[(wind_speed >= 24.5) & (wind_speed < 28.5)] = 10
        beaufort[(wind_speed >= 28.5) & (wind_speed < 32.7)] = 11
        beaufort[(wind_speed >= 32.7)] = 12

        X = np.concatenate((X, beaufort[:, None]), axis=1)
        return X

class FogCategoriesTransformer(BaseEstimator, TransformerMixin):
    '''
    Transformer, který z dat vytvoří nové sloupce s kategoriemi mlhy.
    
    Pokud columns=None, tak se předpokládá, že vstupní data jsou ve formátu:
    první sloupec odpovídá viditelnosti.

    Jinak bere pořadí sloupců zadané v columns.
    '''
    def __init__(self, columns=None, copy=True):
        self.columns = columns
        self.copy = copy

    def fit(self, X, y=None):
        return self

    def transform(self, X, copy=None):
        copy = copy if copy is not None else self.copy
        X = X.copy() if copy else X

        if self.columns is None:
            visibility = X[:, 0]
        else:
            visibility = X[:, self.columns[0]]
        
        fog = np.zeros_like(visibility)
        # velmi silná mlha, silná mlha, mírná mlha, slabá mlha, bez mlhy
        fog[(visibility >= 0) & (visibility < 5)] = 1
        fog[(visibility >= 5) & (visibility < 20)] = 2
        fog[(visibility >= 20) & (visibility < 50)] = 3
        fog[(visibility >= 50) & (visibility < 100)] = 4
        fog[(visibility >= 100)] = 5

        X = np.concatenate((X, fog[:, None]), axis=1)
        return X


def main(args):
    # nastavení seedu
    np.random.seed(args.seed)

    # načtení datasetu
    dataset = Dataset()

    # zahození řádků, kde nefungoval systém
    dataset.train_target = dataset.train_target[dataset.train_data[:, 11] != 0]
    dataset.train_data = dataset.train_data[dataset.train_data[:, 11] != 0]

    # TODO: Natrénujte model
    pipeline = Pipeline([
        ("day_of_week", DayOfWeekTransformer(columns=[12, 13, 14])),
        ("part_of_day", PartOfDayTransformer(columns=[0])),
        ("beaufort", BeaufortTransformer(columns=[3])),
        ("fog", FogCategoriesTransformer(columns=[4])),
        ("column_transformer", ColumnTransformer([
            ("one_hot", OneHotEncoder(), [0, 9, 10, 12, 13, 14, 15, 16, 17, 18]),
            # ("robust_scaler", RobustScaler(), list(range(1, 9))),
            ("standard_scaler", StandardScaler(), list(range(1,9))),
            # zahazujeme příznak, jestli fungoval systém
            ('drop', 'drop', [11])
        ], remainder="passthrough")),
        ("polynomial_features", PolynomialFeatures(degree=3, include_bias=False)),
        # ('reg', LinearRegression())
        ("reg", RidgeCV(alphas=[1e-3, 5e-2, 1e-2, 5e-1, 0.1, 1.0, 10.0, 100, 1000, 10000])),
    ])

    pipeline.fit(dataset.train_data, dataset.train_target)
    
    def predict_with_correction(data):
        pred = pipeline.predict(data)
        # explicitní nastavení non func days na 0
        pred[data[:,11] == 0] = 0
        pred[pred < 0] = 0
        return pred

    train_pred = predict_with_correction(dataset.train_data)
    test_pred = predict_with_correction(dataset.test_data)
    print(mean_squared_error(dataset.train_target, train_pred, squared=False))
    
    # Pokud si budete chtít uložit model (nebo celou pipeline),
    # se vám může hodit toto:
    # with lzma.open("competition.model", "wb") as model_file:
    #     pickle.dump(model, model_file)

    # Poté načtení modelu uděláte takto:
    # with lzma.open("competition.model", "rb") as model_file:
    #     model = pickle.load(model_file)

    pred = test_pred

    with open("36-2-S-prediction.txt", "w") as prediction_file:
        for p in pred:
            print(p, file=prediction_file)


if __name__ == "__main__":
    args = parser.parse_args()
    main(args)