import pathlib
import click
import os
import glob
import psutil
import itertools
import numpy as np
import pandas as pd

from datetime import datetime
from sklearn.model_selection import KFold
from sklearn.linear_model._coordinate_descent import _alpha_grid

from pilot import DEFAULT_DF_SETTINGS
from benchmark_util import *


def print_with_timestamp(message):
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{timestamp}] {message}")


OUTPUTFOLDER = pathlib.Path(__file__).parent / "Output"
DATAFOLDER = pathlib.Path(__file__).parent / "Data"

df_setting_alpha01 = dict(
    zip(
        DEFAULT_DF_SETTINGS.keys(),
        1 + 0.01 * (np.array(list(DEFAULT_DF_SETTINGS.values())) - 1),
    )
)

df_setting_alpha5 = dict(
    zip(
        DEFAULT_DF_SETTINGS.keys(),
        1 + 0.5 * (np.array(list(DEFAULT_DF_SETTINGS.values())) - 1),
    )
)

df_setting_alpha01_no_blin = df_setting_alpha01.copy()
df_setting_alpha01_no_blin["blin"] = -1

df_setting_alpha5_no_blin = df_setting_alpha5.copy()
df_setting_alpha5_no_blin["blin"] = -1

df_setting_no_blin = DEFAULT_DF_SETTINGS.copy()
df_setting_no_blin["blin"] = -1


@click.command()
@click.option("--experiment_name", "-e", required=True, help="Name of the experiment")
@click.option(
    "--models",
    "-m",
    multiple=True,
    help="Models to run (cart, pilot, rf, cpf, xgb, ridge, lasso). If not specified, runs all models.",
)
def run_benchmark(experiment_name, models):
    experiment_folder = OUTPUTFOLDER / experiment_name
    experiment_folder.mkdir(exist_ok=True)
    experiment_file = experiment_folder / "results.csv"
    print_with_timestamp(f"Results will be stored in {experiment_file}")

    # If no models specified, run all models
    if not models:
        models = ["cart", "pilot", "rf", "cpf", "xgb", "ridge", "lasso"]
    else:
        # Convert to lowercase for case-insensitive comparison
        models = [m.lower() for m in models]

    print_with_timestamp(f"Running benchmark for models: {', '.join(models)}")

    np.random.seed(42)
    cv = KFold(n_splits=5, shuffle=True, random_state=42)

    if experiment_file.exists():
        print_with_timestamp(f"Loading existing results from {experiment_file}")
        results = pd.read_csv(experiment_file)
        processed_repo_ids = results["id"].unique().astype(str)
        results = results.to_dict("records")
    else:
        results = []
        processed_repo_ids = []

    repo_ids_to_process = [
        pathlib.Path(f).stem
        for f in glob.glob(os.path.join(DATAFOLDER, "*"))
        if pathlib.Path(f).stem not in processed_repo_ids
    ]
    print_with_timestamp(f"Datasets to process: {len(repo_ids_to_process)}")

    for repo_id in repo_ids_to_process:
        print_with_timestamp(repo_id)
        kind, repo_id = repo_id.split("_")
        dataset = load_data(repo_id=repo_id, kind=kind)
        if dataset.n_samples > 2e5:
            print_with_timestamp(f"Skipping large dataset {repo_id}")
            continue
        alphagrid = _alpha_grid(
            dataset.X_oh_encoded,
            dataset.y,
            l1_ratio=1,
            fit_intercept=True,
            eps=1e-3,
            n_alphas=100,
            copy_X=False,
        )
        for i, (train, test) in enumerate(cv.split(dataset.X, dataset.y), start=1):
            print_with_timestamp(f"\tFold {i} / 5")
            print("\tRAM Used (GB):", psutil.virtual_memory()[3] / 1000000000)
            train_dataset = dataset.subset(train)
            test_dataset = dataset.subset(test)

            transformers = fit_transformers(train_dataset)

            for col, transformer in transformers.items():
                train_dataset.apply_transformer(col, transformer)
                test_dataset.apply_transformer(col, transformer)

            # CART
            if "cart" in models:
                for md in [6, 20, None]:
                    model_name = f"CART - max_depth = {md}"
                    print_with_timestamp(f"\t\t{model_name}")
                    r = fit_cart(
                        train_dataset=train_dataset,
                        test_dataset=test_dataset,
                        max_depth=md,
                    )
                    results.append(
                        dict(
                            **dataset.summary(),
                            fold=i,
                            model=model_name,
                            **r.asdict(),
                            max_depth=md,
                        )
                    )

            # PILOT
            if "pilot" in models:
                for (df_name, alpha, df_setting), md in itertools.product(
                    [
                        ("df alpha = 0.01", 0.01, df_setting_alpha01),
                        ("default df", 1, DEFAULT_DF_SETTINGS),
                        ("df alpha = 0.5", 0.5, df_setting_alpha5),
                    ],
                    [6, 20, None],
                ):
                    model_name = f"CPILOT - {df_name} - max_depth = {md}"
                    print_with_timestamp(f"\t\t{model_name}")
                    r = fit_cpilot(
                        train_dataset=train_dataset,
                        test_dataset=test_dataset,
                        max_depth=md,
                        df_settings=df_setting,
                    )
                    results.append(
                        dict(
                            **dataset.summary(),
                            fold=i,
                            model=model_name,
                            **r.asdict(),
                            max_depth=md,
                            df_setting=df_setting,
                            alpha=alpha,
                        )
                    )

            # RF
            if "rf" in models:
                for md, mf, nt in itertools.product([6, 20, None], [0.7, 1.0], [100]):
                    model_name = f"RF - max_depth = {md} - max_features = {mf} - n_estimators = {nt}"
                    print_with_timestamp(f"\t\t{model_name}")
                    r = fit_random_forest(
                        train_dataset=train_dataset,
                        test_dataset=test_dataset,
                        n_estimators=nt,
                        max_depth=md,
                        max_features=mf,
                    )
                    results.append(
                        dict(
                            **dataset.summary(),
                            fold=i,
                            model=model_name,
                            **r.asdict(),
                            max_depth=md,
                            max_features=mf,
                            n_estimators=nt,
                        )
                    )

            # CPF
            if "cpf" in models:
                for j, (
                    (df_name, alpha, df_setting),
                    max_depth,
                    max_features,
                    ntrees,
                ) in enumerate(
                    itertools.product(
                        [
                            # ("default df", 1, DEFAULT_DF_SETTINGS),
                            # ("df alpha = 0.01", 0.01, df_setting_alpha01),
                            (
                                "df alpha = 0.01, no blin",
                                0.01,
                                df_setting_alpha01_no_blin,
                            ),
                            ("df no blin", 1, df_setting_no_blin),
                            # ("df alpha = 0.5", 0.5, df_setting_alpha5),
                            ("df alpha = 0.5, no blin", 0.5, df_setting_alpha5_no_blin),
                        ],
                        [6, 20],
                        [0.7, 1.0],
                        [100],
                    )
                ):
                    model_name = f"CPF - {df_name} - max_depth = {max_depth} - max_node_features = {max_features} - n_estimators = {ntrees}"
                    print_with_timestamp(f"\t\t{model_name}")
                    r = fit_cpilot_forest(
                        train_dataset=train_dataset,
                        test_dataset=test_dataset,
                        n_estimators=ntrees,
                        min_sample_leaf=1,
                        min_sample_alpha=2,
                        min_sample_fit=2,
                        max_depth=max_depth,
                        n_features_node=max_features,
                        df_settings=df_setting,
                        max_pivot=10000,
                    )
                    results.append(
                        dict(
                            **dataset.summary(),
                            fold=i,
                            model=model_name,
                            **r.asdict(),
                            df_setting=df_setting,
                            excl_blin="no_blin" in df_name,
                            alpha=alpha,
                            max_depth=max_depth,
                            max_features=max_features,
                            n_estimators=ntrees,
                        )
                    )

            # XGB
            if "xgb" in models:
                for md, mf, nt in itertools.product([6, 20], [0.7, 1.0], [100]):
                    model_name = f"XGB - max_depth = {md} - max_features = {mf} - n_estimators = {nt}"
                    print_with_timestamp(f"\t\t{model_name}")
                    r = fit_xgboost(
                        train_dataset=train_dataset,
                        test_dataset=test_dataset,
                        max_depth=md,
                        max_node_features=mf,
                        n_estimators=nt,
                    )
                    results.append(
                        dict(
                            **dataset.summary(),
                            fold=i,
                            model=model_name,
                            **r.asdict(),
                            max_depth=md,
                            max_features=mf,
                            n_estimators=nt,
                        )
                    )
            # linear models
            if "ridge" in models or "lasso" in models:
                for alpha in alphagrid:
                    if "ridge" in models:
                        model_name = f"Ridge - alpha = {alpha}"
                        print_with_timestamp(f"\t\t{model_name}")
                        r = fit_ridge(
                            train_dataset=train_dataset,
                            test_dataset=test_dataset,
                            alpha=alpha,
                        )
                        results.append(
                            dict(
                                **dataset.summary(),
                                fold=i,
                                model=model_name,
                                **r.asdict(),
                                alpha=alpha,
                            )
                        )
                    if "lasso" in models:
                        model_name = f"Lasso - alpha = {alpha}"
                        print_with_timestamp(f"\t\t{model_name}")
                        r = fit_lasso(
                            train_dataset=train_dataset,
                            test_dataset=test_dataset,
                            alpha=alpha,
                        )
                        results.append(
                            dict(
                                **dataset.summary(),
                                fold=i,
                                model=model_name,
                                **r.asdict(),
                                alpha=alpha,
                            )
                        )

        pd.DataFrame(results).to_csv(experiment_file, index=False)


if __name__ == "__main__":
    run_benchmark()
