import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.model_selection import GroupShuffleSplit

from src.training.decision_tree import preprocessor
from src.training.decision_tree.model import Model
from src.training.decision_tree.visualizer import Visualizer
from src.training.simulated_annealing.simulated_annealing import SimulatedAnnealing

# Path setup for saving figures
SAVE_DIR = "./doc/hyperparametersOptimisation/simulatedAnnealing/"
SAVE_PATH = os.path.join(SAVE_DIR, "pairwise_exploration.png")
os.makedirs(SAVE_DIR, exist_ok=True)

# Shared history state
evaluations_history = []

def create_split(x, y, game_ids):

    splitter = GroupShuffleSplit(
        n_splits=1,
        test_size=0.10,
        random_state=42
    )

    train_indices, test_indices = next(splitter.split(x, y, groups=game_ids))

    x_train = x[train_indices]
    x_test = x[test_indices]

    y_train = y[train_indices]
    y_test = y[test_indices]

    train_games = set(game_ids[train_indices])
    test_games = set(game_ids[test_indices])

    print("Train positions:", len(x_train))
    print("Test positions:", len(x_test))
    print("Train games:", len(train_games))
    print("Test games:", len(test_games))
    print("Overlapping games:", len(train_games & test_games))

    return x_train, x_test, y_train, y_test

# Graph for hyperparameter visualisation
def update_live_pairwise_plot(history, save_path):
    """Refreshes the live Seaborn pairplot and exports the updated figure to disk."""
    df = pd.DataFrame(
        [
            {
                "max_depth": e["params"]["max_depth"],
                "min_samples_leaf": e["params"]["min_samples_leaf"],
                # "n_estimators": e["params"]["n_estimators"],
                "max_features": e["params"]["max_features"],
                "test_mse": e["test_mse"],
            }
            for e in history
        ]
    )

    # Close previous figures to avoid accumulating memory and orphaned windows
    plt.close("all")

    # Render pairwise exploration matrix
    grid = sns.pairplot(
        df,
        vars=[
            "max_depth",
            "min_samples_leaf",
            # "n_estimators",
            "max_features"],
        hue="test_mse",
        palette="crest",
        diag_kind="kde",
    )
    grid.figure.suptitle(f"Live SA Search (Iteration {len(history)})", y=1.02)

    # Save to disk
    grid.savefig(save_path, bbox_inches="tight")

    # Force Matplotlib canvas redraw and flush event loop
    grid.figure.canvas.draw()
    grid.figure.canvas.flush_events()
    plt.pause(0.05)

# Callback function for training the model with given hyperparameters
def train_model(parameters):
    # Enable interactive figure mode on first run

    model = Model(
        tree=HistGradientBoostingRegressor(**parameters),
        feature_names=preprocessor.get_feature_names(),
        train_test_split=split,
    )

    model.train()
    train_mae, train_mse, test_mae, test_mse = model.evaluate()
    model.release_training_data()

    # Track current iteration state
    evaluations_history.append(
        {
            "params": parameters,
            "train_mse": train_mse,
            "test_mse": test_mse,
        }
    )

    # Trigger live update & file saving
    update_live_pairwise_plot(evaluations_history, SAVE_PATH)

    return test_mse


if __name__ == "__main__":
    model_path = "./src/training/decision_tree/cache/model.joblib"

    if os.path.exists(model_path):
        model = Model.load(model_path)
    else:
        x, y, game_ids = preprocessor.run("./dataset/output.csv")
        split = create_split(x, y, game_ids)
        del x, y, game_ids
        
        plt.ion()

        base_model_parameters = {
            "max_depth": {"min": 1, "max": 100},
            "min_samples_leaf": {"min": 1, "max": 30},
            # "n_estimators": {"min": 10, "max": 100},
            "max_features": {"min": 0.01, "max": 1.0},
            "random_state": 42,
            # "n_jobs": 15,
        }

        optimization_algorithm = SimulatedAnnealing(
            parameter_space=base_model_parameters,
            initial_temperature=1,
            cooling_rate=0.90,
            max_iterations=100,
        )
        optimization_algorithm.set_training_callback(train_model)

        try:
            best_solution, best_score = optimization_algorithm.optimize()
        except KeyboardInterrupt:
            print("\n[Optimization Interrupted] Early exit triggered by user.")
            
            
        print("\n[Optimization Complete] Training model with best hyperparameters :")
        model = Model(
            tree=HistGradientBoostingRegressor(**best_solution),
            feature_names=preprocessor.get_feature_names(),
            train_test_split=split,
        )
    
        model.train()
        train_mae, train_mse, test_mae, test_mse = model.evaluate()
        model.release_training_data()

        plt.ioff()
        plt.show()

    visualizer = Visualizer(model=model)
    visualizer.show_feature_importance(15)