from typing import List

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sb
from sklearn.decomposition import PCA
from sklearn.ensemble import RandomForestRegressor
from sklearn.linear_model import LinearRegression, Ridge
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVR
from statsmodels.nonparametric.smoothers_lowess import lowess

from .utils import plot_general

# from xgboost import XGBRegressor


class SimpleML:
    use_lowess = False

    def __init__(s, X: np.ndarray, Y: np.ndarray):
        s.X, s.Y = X, Y
        s.F = s.X.shape[-1]

    def sanity_check(s):
        # 1. scatter plot
        plots = {}
        for i in range(s.F):
            x = s.X[:, i]
            plots[f"x{i}_hist"] = x
            plots[f"x{i}_vs_y"] = {"x": x, "y": s.Y}
            if s.use_lowess:
                r = lowess(s.Y, x).T
                plots[f"x{i}_vs_y_lowess"] = {"x": r[0], "y": r[1]}
        C = 3 if s.use_lowess else 2
        plots["y_hist"] = s.Y
        plot_general(plots, "ml_1_scatter", C=C)

        # 2. correlation heatmap
        df = pd.DataFrame(s.X, columns=[f"x{i}" for i in range(s.F)])
        df["Y"] = s.Y
        plt.figure(figsize=(10, 10))
        sb.heatmap(df.corr(), annot=True)
        plt.savefig("ml_2_corr_heatmap")
        plt.close()

        # 3. PCA
        pca = PCA(n_components=2)
        x = pca.fit_transform(s.X).T
        plt.scatter(x[0], x[1], c=s.Y, s=2)
        plt.colorbar(label="Y")
        plt.xlabel("pca_x0")
        plt.ylabel("pca_x1")
        plt.title("PCA colored by Y")
        plt.savefig("ml_3_pca")
        plt.close()

        # 4. linear regression
        m = LinearRegression()
        m.fit(s.X, s.Y)
        print(f"LinearRegression score: {m.score(s.X, s.Y)}")

    def fit(s, models=None):
        X1, X2, Y1, Y2 = train_test_split(s.X, s.Y, test_size=0.4)
        scaler = StandardScaler()
        X1 = scaler.fit_transform(X1)
        X2 = scaler.transform(X2)
        if models is None:
            models: List[SVR] = [
                Ridge(),
                RandomForestRegressor(),
                SVR(),
                # XGBRegressor(),
            ]
        for m in models:
            m.fit(X1, Y1)
            name = m.__class__.__name__.replace("Regressor", "")
            print(
                f"{name:20} scores: train {m.score(X1, Y1):.3f}, test {m.score(X2, Y2):.3f}"
            )
