from __future__ import annotations
import numpy as np
import pandas as pd
from sklearn import datasets
from IMLearn.metrics import mean_square_error
from IMLearn.utils import split_train_test
from IMLearn.model_selection import cross_validate
from IMLearn.learners.regressors import PolynomialFitting, LinearRegression, RidgeRegression
from sklearn.linear_model import Lasso

from utils import *
import plotnine as gg


def select_polynomial_degree(n_samples: int = 100, noise: float = 5):
    """
    Simulate data from a polynomial model and use cross-validation to select the best fitting degree

    Parameters
    ----------
    n_samples: int, default=100
        Number of samples to generate

    noise: float, default = 5
        Noise level to simulate in responses
    """

    # Question 1 - Generate dataset for model f(x)=(x+3)(x+2)(x+1)(x-1)(x-2) + eps for eps Gaussian noise
    # and split into training- and testing portions

    def f(x):
        return (x + 3) * (x + 2) * (x + 1) * (x - 1) * (x - 2)

    X = np.linspace(-1.2, 2, n_samples)
    y = f(X) + np.random.normal(0, noise, n_samples)
    train_X, train_y, test_X, test_y = split_train_test(pd.DataFrame(X), pd.Series(y), train_proportion=(2 / 3))

    df_train = pd.DataFrame({"x": train_X.squeeze(), "y": train_y, "type": "Train"})
    df_test = pd.DataFrame({"x": test_X.squeeze(), "y": test_y, "type": "test"})
    x_stat = np.linspace(-1.4, 2, 100)
    df_stat = pd.DataFrame({"x": x_stat, "y": f(x_stat), "type": "Model"})
    df = pd.concat([df_test, df_train])
    title = f"f(x) = (x+3)(x+2)(x+1)(x-1)(x-2) + Gaussian noise ~ N(0,{noise})"

    p = gg.ggplot() + \
        gg.geom_point(df, gg.aes("x", "y", color="type")) + \
        gg.geom_line(df_stat, gg.aes("x", "y")) + \
        gg.theme_bw() + \
        gg.ggtitle(title)
    # print(p)
    gg.ggsave(filename=f'../../IML/ex5/plots/{title}.png', plot=p, verbose=False)

    # Question 2 - Perform CV for polynomial fitting with degrees 0,1,...,10
    train_err = []
    validation_err = []
    for k in range(11):
        pf = PolynomialFitting(k)
        train_score, validation_score = cross_validate(pf, train_X.to_numpy(), train_y.to_numpy(), mean_square_error)
        train_err.append(train_score)
        validation_err.append(validation_score)

    df1 = pd.DataFrame({"k": range(11), "avg error": train_err, "type": "train error"})
    df2 = pd.DataFrame({"k": range(11), "avg error": validation_err, "type": "validation error"})
    df = pd.concat([df1, df2])
    title = f" Cross Validation for Polynomial Fitting Over Different Degrees k"
    p = gg.ggplot(df, gg.aes("k", "avg error", color="type")) + \
        gg.geom_point() + \
        gg.theme_bw() + gg.scale_x_continuous(breaks=range(11)) + \
        gg.labs(y="Average training and validation errors",
                title=f"{title} \nWith Noise: {noise}, Num of samples: {n_samples}")
    gg.ggsave(filename=f'../../IML/ex5/plots/{title} {noise} {n_samples}.png', plot=p, verbose=False)

    # Question 3 - Using best value of k, fit a k-degree polynomial model and report test error
    best_k = np.argmin(np.array(validation_err))
    pf = PolynomialFitting(int(best_k))
    pf.fit(train_X.to_numpy(), train_y.to_numpy())
    y_pred = pf.predict(test_X.to_numpy())
    print("best k =", best_k)
    print("Test = ", round(mean_square_error(test_y.to_numpy(), y_pred), 2))
    print("Validation = ", round(validation_err[best_k], 2))


def select_regularization_parameter(n_samples: int = 50, n_evaluations: int = 500):
    """
    Using sklearn's diabetes dataset use cross-validation to select the best fitting regularization parameter
    values for Ridge and Lasso regressions

    Parameters
    ----------
    n_samples: int, default=50
        Number of samples to generate

    n_evaluations: int, default = 500
        Number of regularization parameter values to evaluate for each of the algorithms
    """
    # Question 6 - Load diabetes dataset and split into training and testing portions
    X, y = datasets.load_diabetes(return_X_y=True, as_frame=True)
    train_X, train_y, test_X, test_y = X.iloc[:50, :], y[:50], X.iloc[50:, ], y[50:]

    # Question 7 - Perform CV for different values of the regularization parameter for Ridge and Lasso regressions
    for name, learner, ran in [("Ridge", RidgeRegression, np.linspace(0.001, 0.05, 500)),
                               ("Lasso", Lasso, np.linspace(0.001, 0.5, 500))]:
        train_err = []
        validation_err = []
        for lam in ran:
            rg = learner(lam)
            train_score, validation_score = cross_validate(rg, train_X.to_numpy(), train_y.to_numpy(),
                                                           mean_square_error)
            train_err.append(train_score)
            validation_err.append(validation_score)
        df1 = pd.DataFrame({"lambda": ran, "avg error": train_err, "type": "train error"})
        df2 = pd.DataFrame({"lambda": ran, "avg error": validation_err, "type": "validation error"})
        df = pd.concat([df1, df2])
        title = f"{name} Regularization Cross Validate Over Different Lambda"
        p = gg.ggplot(df, gg.aes("lambda", "avg error", color="type")) + \
            gg.geom_line() + \
            gg.theme_bw() + gg.labs(y="Average training and validation errors", title=title)
        gg.ggsave(filename=f'../../IML/ex5/plots/{title}.png', plot=p, verbose=False)

        # Question 8 - Compare best Ridge model, best Lasso model and Least Squares model
        best_lam = np.argmin(np.array(validation_err))
        rg = learner(ran[best_lam])
        rg.fit(train_X.to_numpy(), train_y.to_numpy())
        y_pred = rg.predict(test_X.to_numpy())
        print(f"best lambda {name} = {round(ran[best_lam], 3)}")
        print(f"Test MSE {name} = {round(mean_square_error(test_y.to_numpy(), y_pred), 2)}")
    lr = LinearRegression()
    lr.fit(train_X.to_numpy(), train_y.to_numpy())
    print("Linear Regression Loss = ", lr.loss(test_X.to_numpy(), test_y.to_numpy()))


if __name__ == '__main__':
    np.random.seed(0)
    select_polynomial_degree()
    select_polynomial_degree(noise=0)
    select_polynomial_degree(n_samples=1500, noise=10)
    select_regularization_parameter()

[((1030, 1061), 'numpy.linspace', 'np.linspace', (['(-1.2)', '(2)', 'n_samples'], {}), '(-1.2, 2, n_samples)\n', (1041, 1061), True, 'import numpy as np\n'), ((1408, 1433), 'numpy.linspace', 'np.linspace', (['(-1.4)', '(2)', '(100)'], {}), '(-1.4, 2, 100)\n', (1419, 1433), True, 'import numpy as np\n'), ((1518, 1548), 'pandas.concat', 'pd.concat', (['[df_test, df_train]'], {}), '([df_test, df_train])\n', (1527, 1548), True, 'import pandas as pd\n'), ((1839, 1916), 'plotnine.ggsave', 'gg.ggsave', ([], {'filename': 'f"""../../IML/ex5/plots/{title}.png"""', 'plot': 'p', 'verbose': '(False)'}), "(filename=f'../../IML/ex5/plots/{title}.png', plot=p, verbose=False)\n", (1848, 1916), True, 'import plotnine as gg\n'), ((2496, 2517), 'pandas.concat', 'pd.concat', (['[df1, df2]'], {}), '([df1, df2])\n', (2505, 2517), True, 'import pandas as pd\n'), ((2910, 3011), 'plotnine.ggsave', 'gg.ggsave', ([], {'filename': 'f"""../../IML/ex5/plots/{title} {noise} {n_samples}.png"""', 'plot': 'p', 'verbose': '(False)'}), "(filename=f'../../IML/ex5/plots/{title} {noise} {n_samples}.png',\n    plot=p, verbose=False)\n", (2919, 3011), True, 'import plotnine as gg\n'), ((4036, 4090), 'sklearn.datasets.load_diabetes', 'datasets.load_diabetes', ([], {'return_X_y': '(True)', 'as_frame': '(True)'}), '(return_X_y=True, as_frame=True)\n', (4058, 4090), False, 'from sklearn import datasets\n'), ((5883, 5901), 'IMLearn.learners.regressors.LinearRegression', 'LinearRegression', ([], {}), '()\n', (5899, 5901), False, 'from IMLearn.learners.regressors import PolynomialFitting, LinearRegression, RidgeRegression\n'), ((6072, 6089), 'numpy.random.seed', 'np.random.seed', (['(0)'], {}), '(0)\n', (6086, 6089), True, 'import numpy as np\n'), ((1077, 1114), 'numpy.random.normal', 'np.random.normal', (['(0)', 'noise', 'n_samples'], {}), '(0, noise, n_samples)\n', (1093, 1114), True, 'import numpy as np\n'), ((1171, 1186), 'pandas.DataFrame', 'pd.DataFrame', (['X'], {}), '(X)\n', (1183, 1186), True, 'import pandas as pd\n'), ((1188, 1200), 'pandas.Series', 'pd.Series', (['y'], {}), '(y)\n', (1197, 1200), True, 'import pandas as pd\n'), ((1802, 1819), 'plotnine.ggtitle', 'gg.ggtitle', (['title'], {}), '(title)\n', (1812, 1819), True, 'import plotnine as gg\n'), ((2075, 2095), 'IMLearn.learners.regressors.PolynomialFitting', 'PolynomialFitting', (['k'], {}), '(k)\n', (2092, 2095), False, 'from IMLearn.learners.regressors import PolynomialFitting, LinearRegression, RidgeRegression\n'), ((2769, 2897), 'plotnine.labs', 'gg.labs', ([], {'y': '"""Average training and validation errors"""', 'title': 'f"""{title} \nWith Noise: {noise}, Num of samples: {n_samples}"""'}), '(y=\'Average training and validation errors\', title=\n    f"""{title} \nWith Noise: {noise}, Num of samples: {n_samples}""")\n', (2776, 2897), True, 'import plotnine as gg\n'), ((3128, 3152), 'numpy.array', 'np.array', (['validation_err'], {}), '(validation_err)\n', (3136, 3152), True, 'import numpy as np\n'), ((4856, 4932), 'pandas.DataFrame', 'pd.DataFrame', (["{'lambda': ran, 'avg error': train_err, 'type': 'train error'}"], {}), "({'lambda': ran, 'avg error': train_err, 'type': 'train error'})\n", (4868, 4932), True, 'import pandas as pd\n'), ((4947, 5037), 'pandas.DataFrame', 'pd.DataFrame', (["{'lambda': ran, 'avg error': validation_err, 'type': 'validation error'}"], {}), "({'lambda': ran, 'avg error': validation_err, 'type':\n    'validation error'})\n", (4959, 5037), True, 'import pandas as pd\n'), ((5047, 5068), 'pandas.concat', 'pd.concat', (['[df1, df2]'], {}), '([df1, df2])\n', (5056, 5068), True, 'import pandas as pd\n'), ((5354, 5431), 'plotnine.ggsave', 'gg.ggsave', ([], {'filename': 'f"""../../IML/ex5/plots/{title}.png"""', 'plot': 'p', 'verbose': '(False)'}), "(filename=f'../../IML/ex5/plots/{title}.png', plot=p, verbose=False)\n", (5363, 5431), True, 'import plotnine as gg\n'), ((1776, 1789), 'plotnine.theme_bw', 'gg.theme_bw', ([], {}), '()\n', (1787, 1789), True, 'import plotnine as gg\n'), ((4350, 4379), 'numpy.linspace', 'np.linspace', (['(0.001)', '(0.05)', '(500)'], {}), '(0.001, 0.05, 500)\n', (4361, 4379), True, 'import numpy as np\n'), ((4430, 4458), 'numpy.linspace', 'np.linspace', (['(0.001)', '(0.5)', '(500)'], {}), '(0.001, 0.5, 500)\n', (4441, 4458), True, 'import numpy as np\n'), ((5281, 5345), 'plotnine.labs', 'gg.labs', ([], {'y': '"""Average training and validation errors"""', 'title': 'title'}), "(y='Average training and validation errors', title=title)\n", (5288, 5345), True, 'import plotnine as gg\n'), ((5552, 5576), 'numpy.array', 'np.array', (['validation_err'], {}), '(validation_err)\n', (5560, 5576), True, 'import numpy as np\n'), ((2701, 2714), 'plotnine.theme_bw', 'gg.theme_bw', ([], {}), '()\n', (2712, 2714), True, 'import plotnine as gg\n'), ((5265, 5278), 'plotnine.theme_bw', 'gg.theme_bw', ([], {}), '()\n', (5276, 5278), True, 'import plotnine as gg\n'), ((1638, 1649), 'plotnine.ggplot', 'gg.ggplot', ([], {}), '()\n', (1647, 1649), True, 'import plotnine as gg\n'), ((1746, 1762), 'plotnine.aes', 'gg.aes', (['"""x"""', '"""y"""'], {}), "('x', 'y')\n", (1752, 1762), True, 'import plotnine as gg\n'), ((2673, 2688), 'plotnine.geom_point', 'gg.geom_point', ([], {}), '()\n', (2686, 2688), True, 'import plotnine as gg\n'), ((5234, 5248), 'plotnine.geom_line', 'gg.geom_line', ([], {}), '()\n', (5246, 5248), True, 'import plotnine as gg\n'), ((1680, 1710), 'plotnine.aes', 'gg.aes', (['"""x"""', '"""y"""'], {'color': '"""type"""'}), "('x', 'y', color='type')\n", (1686, 1710), True, 'import plotnine as gg\n'), ((2621, 2659), 'plotnine.aes', 'gg.aes', (['"""k"""', '"""avg error"""'], {'color': '"""type"""'}), "('k', 'avg error', color='type')\n", (2627, 2659), True, 'import plotnine as gg\n'), ((5173, 5216), 'plotnine.aes', 'gg.aes', (['"""lambda"""', '"""avg error"""'], {'color': '"""type"""'}), "('lambda', 'avg error', color='type')\n", (5179, 5216), True, 'import plotnine as gg\n')]