import IMLearn.learners.regressors.linear_regression
from IMLearn.learners.regressors import PolynomialFitting
from IMLearn.utils import split_train_test

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.io as pio

pio.templates.default = "simple_white"

dic = {1: 0, 2: 31, 3: 59, 4: 90, 5: 120, 6: 151, 7: 181, 8: 212, 9: 243, 10: 273, 11: 304, 12: 334}


def load_data(filename: str) -> pd.DataFrame:
    """
    Load city daily temperature dataset and preprocess data.
    Parameters
    ----------
    filename: str
        Path to house prices dataset

    Returns
    -------
    Design matrix and response vector (Temp)
    """
    df = pd.read_csv(filename, parse_dates=['Date']).dropna().drop_duplicates()

    df = df[df['Month'].isin(range(1, 13))]

    feb_cond = (df['Month'] == 2) & (((df['Day'].isin(range(30))) & (df['Year'] % 4 == 0)) | ((df['Year'] % 4 != 0) & (df['Day'].isin(range(29)))))

    month30 = (df['Month'].isin([4, 6, 9, 11]) & df['Day'].isin(range(1, 31)))

    month31 = (df['Month'].isin([1, 3, 5, 7, 8, 10, 12]) & df['Day'].isin(range(1, 32)))
    df = df[month31 | month30 | feb_cond]
    tmp = df.apply(lambda row: row.Day + dic[row.Month], axis=1)
    df['DayOfYear'] = tmp
    # print(df.shape)

    df = df[df['Temp'] > -72]


    df['Year'] = df['Year'].astype(str)
    print(df.shape)
    return df


if __name__ == '__main__':
    np.random.seed(0)
    # Question 1 - Load and preprocessing of city temperature dataset
    filename = '../datasets/City_Temperature.csv'
    df = load_data(filename)
    X, y = df.drop('Temp', 1), df['Temp']
    print(df.size)
    # Question 2 - Exploring data for specific country
    israel_data = df[df['Country'] == 'Israel']
    fig = px.scatter(israel_data, x='DayOfYear', y='Temp', color='Year',
                     title='Daily Change of Temperature \'DayOfYear\' Attribute (different colors for different years)')
    fig.show()

    israel_by_month = israel_data.groupby(['Month']).agg({'Temp': 'std'}).reset_index()
    fig = px.bar(israel_by_month, x='Month', y='Temp',
                 labels={'Month': 'Month',
                         'Temperature': 'the Standard Deviation of Temperature'},
                 title='The Standard deviation of the Daily Temperatures Grouped by Month')
    fig.show()

    # Question 3 - Exploring differences between countries
    country_month = df.groupby(['Country', 'Month']).agg(avg=('Temp', 'mean'), std=('Temp', 'std')).reset_index()
    fig = px.line(country_month, x='Month', y='avg', color='Country', error_y='std',
                  labels={'Month': 'Month', 'Temp_average': 'Temperatures\' Average for Each Month',
                          'Temp_std': 'Temperatures\' Standard Deviation for Each Month'},
                  title='the Temperatures\' Average for Each Month By Country With Standard Deviation Error')
    fig.show()

    # Question 4 - Fitting model for different values of `k`
    temp_X, temp_y = israel_data.drop('Temp', 1), israel_data['Temp']
    Day_Of_Year = israel_data['DayOfYear']
    new_X, new_y = israel_data.drop('DayOfYear', 1), israel_data['DayOfYear']
    train_X, train_y, test_X, test_y = split_train_test(israel_data.DayOfYear.to_frame(), israel_data.Temp)
    train_X = train_X.DayOfYear
    test_X = test_X.DayOfYear
    loss = []
    K = range(1, 11)
    for k in K:
        poly_fit = PolynomialFitting(k=k)
        poly_fit.fit(train_X.to_numpy(), train_y.to_numpy())
        losss = poly_fit.loss(test_X.to_numpy(), test_y.to_numpy())
        # print(losss)
        losss = np.round(losss, decimals=2)
        loss.append(losss)
    print(list(zip(K, loss)))
    fig = px.bar(israel_data, x=list(K), y=loss, title="the test error recorded for each value of k",
                 labels={'y': "loss of the model", 'x': "k - degree value"})
    fig.show()

    # Question 5 - Evaluating fitted model on different countries

    k = 5
    country_loss = {}
    temp_PF = PolynomialFitting(k)
    temp_PF.fit(israel_data.DayOfYear.to_numpy(), israel_data.Temp.to_numpy())

    for country in set(df.Country):
        if country != "Israel":
            country_test = df[df['Country'] == country]
            print(country_test.head())
            temp_error = temp_PF.loss(country_test.DayOfYear.to_numpy(), country_test.Temp.to_numpy())
            country_loss[country] = temp_error

    fig = px.bar(df, x=country_loss.keys(), y=country_loss.values(), title="the model’s error over each of the countries",
                 labels={'y': "loss of the model", 'x': "country"})

    fig.show()
