import plotly.io

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

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

pio.templates.default = "simple_white"


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]
    """
    # reading the csv
    df = pd.read_csv(filename).dropna().drop_duplicates()
    # preprocessing and removing bad data

    df = df.drop("id", axis=1).drop("date", axis=1).drop("long", axis=1).drop("lat", axis=1)
    # encoding the zipcode so instead of 70 unique zipcodes in one column we will have 70 columns
    df = pd.get_dummies(data=df, prefix="zc", columns=['zipcode'], drop_first=True)

    # removing rows that have "bad" values
    df.query("price > 0", inplace=True)
    df.query("sqft_lot15 > 0", inplace=True)

    """
    example of adding information that has been deducted from the data
    adding a column to the data frame that determines if a house needs renovation based on the years since its 
    renovation or if the house was never renovated, based on the years since its creation.
    """
    df["needs_renovation"] = np.where(df["yr_renovated"] > 0, 2022 - df["yr_renovated"], 2022 - df["yr_built"])
    # lets pre-determine that a house needs renovation every 20 years
    df["needs_renovation"] = np.where(df["needs_renovation"] >= 20, 1, 0)
    # now that we used these columns, we dont need them anymore
    df = df.drop("yr_renovated", axis=1)
    df = df.drop("yr_built", axis=1)

    # removing suspicious lines e.g. a house with 33 rooms
    df = df[df["bedrooms"] < 20]
    df = df[df["sqft_lot"] < 1250000]
    df = df[df["sqft_lot15"] < 500000]

    df.insert(0, "intercept", 1, True)
    return df.drop("price", axis=1), df["price"]


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
    """
    # dropping the zipcode columns since they will most likely have very low correlation because they were split
    X = X[X.columns.drop(list(X.filter(regex='zc')))]
    # intercept has 0 correlation
    X = X.drop("intercept", axis=1)
    i = 0
    y_std = np.std(y)
    # iterating over the columns and calculating the pearson correlation coeficient
    for col in X:
        if i == X.shape[1]:
            break
        x = np.array(X[col].values)
        corr = np.cov(x, y)[0][1] / (np.std(x) * y_std)
        # ive added a trendline because it helps better understand the relationship between a feature and the response
        fig = px.scatter(x=x, y=y, trendline='ols', trendline_color_override='black')
        fig.update_xaxes(title="Feature values")
        fig.update_yaxes(title="Response")
        fig.update_layout(title=f"Feature: {col}\n, Correlation: {round(corr, 3)}")
        # saving the image as col.png where col is the name of the feature
        fig.write_image(output_path + col + ".png")
        i += 1


if __name__ == '__main__':
    np.random.seed(0)
    # Question 1 - Load and preprocessing of housing prices dataset
    X, y = load_data("../datasets/house_prices.csv")
    # Question 2 - Feature evaluation with respect to response
    feature_evaluation(X, y)

    # Question 3 - Split samples into training- and testing sets.

    train_x, train_y, test_x, test_y = split_train_test(X, y, 0.75)

    # 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)

    # using pd.concat to join the training data so that when sampling %p of the data the X and y would be in sync
    train_df = pd.concat([train_x, train_y], axis=1)
    lin = LinearRegression(False)
    avg_loss = np.ndarray(shape=(91,))
    avg_loss_std = np.ndarray(shape=(91,))
    # iterating over the range [10,100]
    for i in range(10, 101):
        loss = np.ndarray(shape=(10,))
        p = i / 100
        # training and testing 10 times, taking the loss value each time
        for j in range(10):
            curr_train_df = train_df.sample(frac=p)
            curr_y = curr_train_df["price"]
            curr_x = curr_train_df.drop("price", axis=1)
            lin.fit(curr_x, curr_y)
            loss[j] = lin.loss(np.array(test_x), np.array(test_y))
        # calculating the mean and standard deviation of the loss
        avg_loss[i - 10] = loss.mean()
        avg_loss_std[i - 10] = loss.std()
    # plotting the graph (which is 3 graphs together, the upper error range, the lower error range, and the MSE
    ps = [i / 100 for i in range(10, 101)]
    y_upper = avg_loss + 2 * avg_loss_std
    y_lower = avg_loss - 2 * avg_loss_std
    fig = go.Figure([go.Scatter(x=ps, y=avg_loss, line=dict(color='rgb(0,100,80)'), mode='lines'),
                     go.Scatter(
                         x=ps,
                         y=y_upper,
                         mode='lines',
                         fill='tonexty',
                         fillcolor='rgba(0,100,80,0.2)',
                         line=dict(color='rgba(255,255,255,0)'),
                         hoverinfo="skip",
                         showlegend=False),
                     go.Scatter(
                         x=ps,
                         y=y_lower,
                         mode='lines',
                         fill='tonexty',
                         fillcolor='rgba(0,100,80,0.2)',
                         line=dict(color='rgba(255,255,255,0)'),
                         hoverinfo="skip",
                         showlegend=False)
                     ])
    fig.update_xaxes(title="% of training set used")
    fig.update_yaxes(title="MSE with 2σ error ribbon")
    fig.update_layout(title="MSE distribution with a 2σ error ribbon as a function of % of training data used:")
    fig.show()
