from typing import Tuple

import numpy as np
import optuna
import pandas as pd

from sklearn.metrics import f1_score, recall_score
from typing_extensions import Annotated

from src.classifiers import get_model_strategy
from src.logger import setup_logger

logger = setup_logger(__name__)


class ModelTuner:
    def __init__(
        self,
        model: str,
        X_train: pd.DataFrame,
        y_train: np.ndarray,
        X_test: pd.DataFrame,
        y_test: np.ndarray,
        n_trials: int = 100,
    ) -> None:
        self.model = model
        self.X_train = X_train
        self.X_test = X_test
        self.y_train = y_train
        self.y_test = y_test
        self.n_trials = n_trials

    def objective(self, trial: optuna.Trial) -> Tuple[float, float]:
        model_strategy = get_model_strategy(self.model)
        params = model_strategy.get_tuning_params(trial)
        model = model_strategy.set_params(**params)

        logger.info(f"Fine tuning model with params: {params}")

        try:
            model.fit(self.X_train, self.y_train)
            y_pred = model.predict(self.X_test)
            logger.info(f"Fine tuning model with params: {params}")
        except Exception as e:
            logger.error(f"Error in Fine tuning process > objective function, {e}")
            raise

        recall_value = recall_score(self.y_test, y_pred, average="weighted")
        f1_score_value = f1_score(self.y_test, y_pred, average="weighted")

        return recall_value, f1_score_value

    def fine_tune(
        self,
    ) -> Tuple[
        Annotated[optuna.trial.FrozenTrial, "best_trial"],
        Annotated[dict, "best_params"],
        Annotated[dict, "best_values"],
    ]:
        logger.info(">>>>>> Finding best parameters...")
        study = optuna.create_study(directions=["maximize", "maximize"])
        study.optimize(self.objective, n_trials=self.n_trials, n_jobs=-1, show_progress_bar=True)

        best_trial = study.best_trials[0]
        best_params = best_trial.params
        best_values = best_trial.values

        logger.info(
            f"Fine tuning done...,\nBest trial: {best_trial}\nBest values: {best_values}\nBest params: {best_params}"
        )

        return best_trial, best_params, best_values
