import math

from IMLearn.utils import split_train_test
from IMLearn.learners.regressors import LinearRegression

from typing import NoReturn
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import plotly.express as px
import plotly.io as pio
import datetime as dt

pio.templates.default = "simple_white"

ZIP_1 = {98001, 98002, 98003, 98004, 98005, 98006, 98007, 98008, 98009, 98010, 98011}
ZIP_2 = {98020, 98021, 98022, 98023, 98024, 98025, 98026, 98027, 98028, 98029, 98030, 98031}
ZIP_3 = {98032, 98033, 98034, 98035, 98036, 98037, 98038, 98039, 98040, 98041}
ZIP_4 = {98100, 98101, 98102, 98103, 98104, 98105, 98106, 98107, 98108, 98109, 98110}
ZIP_5 = {98111, 98112, 98113, 98114, 98115, 98116, 98117, 98118, 98119, 98120}


def create_new_features(data):
    data["common_bedrooms"] = (1.65 <= data["bedrooms"]) & (6.60 >= data["bedrooms"])
    data["common_bathrooms"] = (0.80 <= data["bathrooms"]) & (4.80 >= data["bathrooms"])
    data["common_sqft_living"] = (290 <= data["sqft_living"]) & (5590 >= data["sqft_living"])
    data["common_sqft_lot"] = (520 <= data["sqft_lot"]) & (83061.95 >= data["sqft_lot"])
    data["common_floors"] = (1 <= data["floors"]) & (2.13 >= data["floors"])
    data["common_sqft_above"] = (290 <= data["sqft_above"]) & (4394 >= data["sqft_above"])
    data["common_sqft_basement"] = (0 <= data["sqft_basement"]) & (1446 >= data["sqft_basement"])
    data["common_long"] = (-122.46 <= data["long"]) & (-121.74 >= data["long"])
    data["common_sqft_living15"] = (689.55 <= data["sqft_living15"]) & (4176.15 >= data["sqft_living15"])
    data["common_sqft_lot15"] = (651 <= data["sqft_lot15"]) & (87705.90 >= data["sqft_lot15"])
    data["common_yr_renovated"] = (0 <= data["yr_renovated"]) & (175 >= data["yr_renovated"])


def handle_zipcode(data):
    data['zip1'] = (data["zipcode"].isin(ZIP_1))
    data['zip2'] = (data["zipcode"].isin(ZIP_2))
    data['zip3'] = (data["zipcode"].isin(ZIP_3))
    data['zip4'] = (data["zipcode"].isin(ZIP_4))
    data['zip5'] = (data["zipcode"].isin(ZIP_5))
    other_func = lambda x: not (x in ZIP_1 or x in ZIP_2 or x in ZIP_3 or x in ZIP_4 or x in ZIP_5)
    data['other'] = data["zipcode"].apply(other_func)


def handle_date(data):
    data['date'] = data['date'].str.replace('T000000', '')
    data['date'] = data['date'].fillna(0)
    data['date'] = data['date'].astype(int)
    # data['date'] = data['date'].map(dt.datetime.toordinal)


def load_data(filename: str):
    """
    Load house prices dataset and preprocess data.
    Parameters
    ----------
    filename: str
        Path to house prices dataset

    Returns
    -------
    Design matrix and response vector (prices) - either as a single
    DataFrame or a Tuple[DataFrame, Series]
    """
    data = pd.read_csv(filename).drop_duplicates()
    data.fillna(data.mean(), inplace=True)
    y = data.pop("price")

    handle_zipcode(data)
    handle_date(data)
    create_new_features(data)

    data = data.drop(columns=["bedrooms", "bathrooms", "sqft_living", "sqft_lot", "floors",
                              "sqft_above", "sqft_basement", "long", "zipcode", "sqft_living15", "sqft_lot15",
                              "yr_renovated"])


    return data, y


def feature_evaluation(X: pd.DataFrame, y: pd.Series, output_path: str = ".") -> NoReturn:
    """
    Create scatter plot between each feature and the response.
        - Plot title specifies feature name
        - Plot title specifies Pearson Correlation between feature and response
        - Plot saved under given folder with file name including feature name
    Parameters
    ----------
    X : DataFrame of shape (n_samples, n_features)
        Design matrix of regression problem

    y : array-like of shape (n_samples, )
        Response vector to evaluate against

    output_path: str (default ".")
        Path to folder in which plots are saved
    """

    for column in X.columns:
        sigma_x = math.sqrt(X[column].var())
        sigma_y = math.sqrt(y.var())
        cov = np.cov(X[column], y)
        pearson_correlation = cov[0][1] / (sigma_x * sigma_y)
        fig = go.Figure([go.Scatter(x=X[column], y=y, name="pearson correlation",
                                    mode="markers", line=dict(width=1))],
                        layout=go.Layout(title= "Pearson Correlation between "+ column +" and response :"+str(pearson_correlation),
                                         xaxis={"title": column},
                                         yaxis={"title": "Response"},
                                         height=400))
        fig.write_image(output_path + "/" + column + "_correlation.png",format='png')


if __name__ == '__main__':
    np.random.seed(0)
    # Question 1 - Load and preprocessing of housing prices dataset
    df, labels = load_data("../datasets/house_prices.csv")
    # print("df_shape: "+str(np.shape(df))+" lables_shape:"+str(np.shape(labels)))

    # Question 2 - Feature evaluation with respect to response
    feature_evaluation(df, labels, "../ex2_plots")
    # print("menash")
    # Question 3 - Split samples into training- and testing sets.
    X_train, y_train, X_test, y_test = split_train_test(df, labels)
    # print("df_shape: "+str(np.shape(X_train))+" lables_shape:"+str(np.shape(y_train)))
    # Question 4 - Fit model over increasing percentages of the overall training data
    # For every percentage p in 10%, 11%, ..., 100%, repeat the following 10 times:
    #   1) Sample p% of the overall training data
    #   2) Fit linear model (including intercept) over sampled set
    #   3) Test fitted model over test set
    #   4) Store average and variance of loss over test set
    # Then plot average loss as function of training size with error ribbon of size (mean-2*std, mean+2*std)
    all_avg = []
    std_up = []
    std_down = []
    for p in range(10,101):
        results = []
        for i in range(10):
            samples = X_train.sample(frac = p/100.0)
            y_sample =y_train[samples.index]
            estimator = LinearRegression().fit(samples, y_sample)
            loss = estimator.loss(X_test,y_test)
            results.append(loss)
        var_loss = np.std(results)
        avg_loss = np.mean(results)
        all_avg.append(avg_loss)
        std_up.append(avg_loss+2*var_loss)
        std_down.append(avg_loss-2*var_loss)

    x=list(range(10,101))
    fig = go.Figure([
        go.Scatter(
            x=x,
            y=all_avg,
            line=dict(color='rgb(0,100,80)'),
            mode='lines'
        ),
        go.Scatter(
            x=x+x[::-1], # x, then x reversed
            y=std_up+std_down[::-1], # upper, then lower reversed
            fill='toself',
            fillcolor='rgba(0,100,80,0.2)',
            line=dict(color='rgba(255,255,255,0)'),
            hoverinfo="skip",
            showlegend=False
        )
    ])
    print("menash")
    fig.write_image("../ex2_plots/per.png",format='png')






