"""
Functions for comparing baseline ML models using cross-validation.
"""

import matplotlib.pyplot as plt
import pandas as pd
from sklearn.ensemble import (
    AdaBoostClassifier,
    AdaBoostRegressor,
    BaggingClassifier,
    BaggingRegressor,
    GradientBoostingClassifier,
    GradientBoostingRegressor,
    RandomForestClassifier,
    RandomForestRegressor,
    VotingClassifier,
    VotingRegressor,
)
from sklearn.linear_model import LinearRegression, LogisticRegression
from sklearn.model_selection import KFold, cross_val_score
from sklearn.neighbors import KNeighborsClassifier
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor


class CvCompareOutput:
    """
    Output object returned by cv_compare.

    Attributes
    ----------
    names : tuple
        Model names.
    cv_scoring : str
        Scoring metric used.
    model_task : str
        Classification or regression.
    scores : tuple
        Cross-validation scores for each model.
    summary_table : pd.DataFrame
        Mean and std of CV scores, sorted descending.
    raw_table : pd.DataFrame
        Raw CV scores for each model.
    """

    def __init__(self, names, cv_scoring, model_task, scores, summary_table, raw_table):
        self.names = names
        self.cv_scoring = cv_scoring
        self.model_task = model_task
        self.scores = scores
        self.summary_table = summary_table
        self.raw_table = raw_table

    def plot(self):
        """
        Display a boxplot of cross-validation scores for each model.
        """
        plt.boxplot(self.scores, tick_labels=self.names)
        plt.xticks(rotation=45)
        plt.ylabel(self.cv_scoring)
        plt.title(f"{self.model_task} comparison (KFold CV)")
        plt.tight_layout()
        plt.show()


def cv_compare(
    X,
    y,
    model_task="classification",
    models=None,
    n_splits=6,
    add_models=None,
    random_state=1,
    scaler=True,
    cv_scoring=None,
):
    """
    Parameters
    ----------
    X : array-like
        Feature matrix
    y : array-like
        Target vector
    model_task : str, optional
        Either "classification" or "regression" (Default value = "classification")
    models : list, optional
        Custom list of sklearn models (Default value = None - a predefined set of models)
    n_splits : int, optional
        Determines the number of folds in the KFold cross validation (Default value = 6)
    add_models : list, optional
        Custom list of sklearn models, further to predefined set (Default value = None)
    random_state : int, optional
        Seed (Default value = 1)
    scaler : bool, optional
        Whether to apply StandardScaler (Default value = True)
    cv_scoring : str, optional
        Custom performance metric (Default value = None)

    Returns
    -------
    CvCompareOutput
        Object containing scores, tables and a plot method (boxplot comparison of models)
    """
    if model_task not in ("classification", "regression"):
        raise ValueError(f"model_task must be either 'classification' or 'regression', got {model_task}")

    if len(X) != len(y):
        raise ValueError(f"len(X) != len(y) : len(X) = {len(X)}, len(y) = {len(y)}")

    if len(X) < n_splits:
        raise ValueError(f"Fewer sample observations than n_splits : len(X) = {len(X)}, n_splits = {n_splits}")

    if models is not None and not isinstance(models, list):
        raise TypeError(f"Expected list of 'models', got {type(models).__name__}")

    if add_models is not None and not isinstance(add_models, list):
        raise TypeError(f"Expected list of 'add_models', got {type(add_models).__name__}")

    # Define scoring methodology
    if cv_scoring is None:
        if model_task == "regression":
            cv_scoring = "neg_root_mean_squared_error"
        else:
            cv_scoring = "accuracy"

    # Define raw models
    if models is not None:
        raw_models = models
    elif model_task == "classification":
        raw_models = [
            KNeighborsClassifier(),
            LogisticRegression(max_iter=1000, random_state=random_state),
            DecisionTreeClassifier(random_state=random_state),
            RandomForestClassifier(random_state=random_state),
            BaggingClassifier(
                estimator=DecisionTreeClassifier(random_state=random_state),
                random_state=random_state,
            ),
            BaggingClassifier(
                estimator=RandomForestClassifier(n_estimators=50, random_state=random_state),
                random_state=random_state,
            ),
            AdaBoostClassifier(random_state=random_state),
            GradientBoostingClassifier(subsample=0.8, max_features="sqrt", random_state=random_state),
            VotingClassifier(
                [
                    (
                        "lr",
                        LogisticRegression(max_iter=1000, random_state=random_state),
                    ),
                    ("dt", DecisionTreeClassifier(random_state=random_state)),
                    ("rf", RandomForestClassifier(random_state=random_state)),
                ]
            ),
        ]
    else:
        raw_models = [
            LinearRegression(),
            DecisionTreeRegressor(random_state=random_state),
            RandomForestRegressor(random_state=random_state),
            BaggingRegressor(
                estimator=DecisionTreeRegressor(random_state=random_state),
                random_state=random_state,
            ),
            BaggingRegressor(
                estimator=RandomForestRegressor(n_estimators=50, random_state=random_state),
                random_state=random_state,
            ),
            AdaBoostRegressor(random_state=random_state),
            GradientBoostingRegressor(subsample=0.8, max_features=0.2, random_state=random_state),
            VotingRegressor(
                [
                    ("lr", LinearRegression()),
                    ("dt", DecisionTreeRegressor(random_state=random_state)),
                    ("rf", RandomForestRegressor(random_state=random_state)),
                ]
            ),
        ]
    if add_models is not None:
        raw_models += add_models

    # Define list of model name and pipeline tuples
    models = [
        (
            type(m).__name__,
            Pipeline([("scaler", StandardScaler()), ("model", m)] if scaler else [("model", m)]),
        )
        for m in raw_models
    ]

    # Generate results
    results = []
    kf = KFold(n_splits=n_splits, random_state=random_state, shuffle=True)
    for name, pipeline in models:
        cv_results = cross_val_score(pipeline, X, y, cv=kf, scoring=cv_scoring)
        results.append((name, cv_results))
    names, scores = zip(*results, strict=False)

    # Generate results DF
    summary_table = pd.DataFrame(
        {
            "model": names,
            f"mean {cv_scoring}": [s.mean() for s in scores],
            "std": [s.std() for s in scores],
        }
    ).sort_values(f"mean {cv_scoring}", ascending=False)

    raw_table = pd.DataFrame(
        {
            "model": names,
            f"{cv_scoring}": scores,
        }
    )

    # Return vars
    return CvCompareOutput(
        names=names,
        scores=scores,
        cv_scoring=cv_scoring,
        model_task=model_task,
        summary_table=summary_table,
        raw_table=raw_table,
    )
