from pathlib import Path
from typing import Tuple

import matplotlib.pyplot as plt
import mlflow
import mlflow.sklearn
import numpy as np
import pandas as pd
import seaborn as sns
import yaml

from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
    f1_score,
    precision_score,
    recall_score,
)
from typing_extensions import Annotated

from src.classifiers import CatBoostClassifier, get_model_strategy
from src.fine_tuning import ModelTuner
from src.logger import setup_logger
from src.mlflow_utils import create_mlflow_experiment, model_signature
from src.utils import save_yaml_config, yaml_config_loader

logger = setup_logger(__name__)
# mlflow.set_tracking_uri("../mlruns")


class ModelTrainer:
    def __init__(self, config_path: str = "configs/config.yaml"):
        self.config = self.load_config(config_path)
        self.data_paths = self.config["data_paths"]
        self.mlflow_config = self.config["mlflow"]
        self.processed_data_dir = Path(self.data_paths["processed_data_dir"])
        self.model_output_path = Path("models")
        self.training_config: dict = self.config["model_training"]
        self.model = self.training_config["model"]
        self.n_trials = self.training_config["n_trials"]

    @staticmethod
    def load_config(config_path: str) -> dict:
        return yaml_config_loader(config_path)

    def load_data(
        self,
    ) -> Tuple[
        Annotated[pd.DataFrame, "X_train"],
        Annotated[pd.Series, "y_train"],
        Annotated[pd.DataFrame, "X_test"],
        Annotated[pd.Series, "y_test"],
    ]:
        try:
            X_train = pd.read_csv(self.processed_data_dir / "SMOTE/X_train_processed.csv")
            y_train = pd.read_csv(self.processed_data_dir / "SMOTE/y_train_processed.csv")
            X_test = pd.read_csv(self.processed_data_dir / "SMOTE/X_test_processed.csv")
            y_test = pd.read_csv(self.processed_data_dir / "SMOTE/y_test_processed.csv")

            # Convert y_train and y_test to 1-dimensional arrays
            y_train = y_train.values.ravel()
            y_test = y_test.values.ravel()

            logger.info(f"{X_train.shape=}, {y_train.shape=}, {X_test.shape=}, {y_test.shape=}")

            logger.info("Data loaded successfully from processed data paths")
            return X_train, y_train, X_test, y_test
        except Exception as e:
            logger.error(f"Error loading data: {e}")
            raise

    def train_and_evaluate_model(self) -> None:
        X_train, y_train, X_test, y_test = self.load_data()

        # Start MLflow run
        # mlflow.set_tracking_uri(self.mlflow_config['tracking_uri'])
        # mlflow.set_experiment(self.mlflow_config['experiment_name'])
        experiment_id = create_mlflow_experiment(
            experiment_name=self.mlflow_config["experiment_name"],
            artifact_location=self.mlflow_config["artifact_location"],
            tags=self.mlflow_config["tags"],
        )

        with mlflow.start_run(run_name=self.training_config["run_name"], experiment_id=experiment_id):
            # Train model
            try:
                model_params = self.training_config.get("model_params", None)
                if model_params:
                    model = get_model_strategy(self.model, **model_params)
                    logger.info("Parameters loaded from config file")
                else:
                    _, best_params, _ = self.find_best_hyperparameters()
                    model = get_model_strategy(self.model, **best_params)
                    logger.info("Parameters tuning using Optuna")
            except Exception as e:
                logger.error(f"Model parameters are not properly configured. {e}")
                raise

            model.fit(X_train, y_train)

            # Make predictions
            y_pred = model.predict(X_test)

            # Log model
            # signature = mlflow.models.infer_signature(X_test, y_pred)
            signature = model_signature(model, X_test, y_pred)
            if isinstance(model, CatBoostClassifier):
                mlflow.catboost.log_model(model, "model", signature=signature)
            else:
                mlflow.sklearn.log_model(model, "model", signature=signature)

            # Calculate metrics
            accuracy = accuracy_score(y_test, y_pred)
            f1 = f1_score(y_test, y_pred, average="weighted")
            precision = precision_score(y_test, y_pred, average="weighted")
            recall = recall_score(y_test, y_pred, average="weighted")
            report = classification_report(y_test, y_pred)

            logger.info(f"Accuracy: {accuracy}")
            logger.info(f"F1 Score: {f1}")
            logger.info(f"Precision: {precision}")
            logger.info(f"Recall: {recall}")
            logger.info(f"Classification Report: \n{report}")

            # Log parameters, metrics, and model
            mlflow.log_param("X_train_shape", X_train.shape[0])
            mlflow.log_param("X_test_shape", X_test.shape[0])

            mlflow.log_metric("accuracy", accuracy)
            mlflow.log_metric("f1_score", f1)
            mlflow.log_metric("precision", precision)
            mlflow.log_metric("recall", recall)
            mlflow.log_text(report, "classification_report.txt")

            confusion_matrix_plot = self.plot_confusion_matrix(y_test, y_pred)
            mlflow.log_figure(confusion_matrix_plot, "confusion_matrix.png")

            mlflow.log_param("model_type", self.model)
            mlflow.log_params(model.get_params())

            # log config file
            mlflow.log_artifact("configs/config.yaml")
            save_yaml_config(self.training_config, "configs/generated_config.yaml")
            mlflow.log_artifact("configs/generated_config.yaml")
            mlflow.log_artifact(f"configs/model_training/{self.model}.yaml")

        try:
            # Save model locally
            self.model_output_path.mkdir(exist_ok=True, parents=True)
            model_output_file = self.model_output_path / f"{self.model}"
            mlflow.sklearn.save_model(model, model_output_file)
            logger.info(f"Model saved to {model_output_file}")
        except mlflow.exceptions.MlflowException:
            logger.warning("Model not saved. Model already exists.")

    def plot_confusion_matrix(
        self, y_true: pd.DataFrame | np.ndarray, y_pred: pd.DataFrame | np.ndarray, figsize: tuple[int, int] = (10, 10)
    ) -> plt.Figure:
        cm = confusion_matrix(y_true, y_pred, labels=np.unique(y_true))
        cm_sum = np.sum(cm, axis=1, keepdims=True)
        cm_perc = cm / cm_sum.astype(float) * 100
        annot = np.empty_like(cm).astype(str)
        nrows, ncols = cm.shape
        for i in range(nrows):
            for j in range(ncols):
                c = cm[i, j]
                p = cm_perc[i, j]
                if i == j:
                    s = cm_sum[i]
                    annot[i, j] = f"{p:.3f}%\n{c}/{s[0]}"
                elif c == 0:
                    annot[i, j] = ""
                else:
                    annot[i, j] = f"{p:.3f}%\n{c}"

        cm_df = pd.DataFrame(cm, index=np.unique(y_true), columns=np.unique(y_true))
        cm_df.index.name = "Actual"
        cm_df.columns.name = "Predicted"

        fig, ax = plt.subplots(figsize=figsize)
        sns.heatmap(
            cm_df,
            cmap="YlGnBu",
            annot=annot,
            fmt="",
            square=True,
            cbar=True,
            linewidths=0.5,
            ax=ax,
        )
        return fig

    def find_best_hyperparameters(self) -> tuple:
        X_train, y_train, X_test, y_test = self.load_data()

        tuner = ModelTuner(
            model=self.model,
            X_train=X_train,
            y_train=y_train,
            X_test=X_test,
            y_test=y_test,
            n_trials=self.n_trials,
        )

        best_trial, best_params, best_values = tuner.fine_tune()
        self.save_best_hyperparameters(best_params)
        return best_trial, best_params, best_values

    def save_best_hyperparameters(self, best_params: dict) -> None:
        self.training_config["model_params"] = best_params
        with open(f"configs/model_training/{self.model}.yaml", "w+") as file:
            yaml.dump(self.training_config, file)
        logger.info("Best hyperparameters saved to best_hyperparameters.yaml")


if __name__ == "__main__":
    trainer = ModelTrainer()
    trainer.train_and_evaluate_model()

    # x_train, y_train, X_test, y_test = trainer.load_data()
    # x_train.info()
