from abc import ABC, abstractmethod
from typing import Any

import pandas as pd

from catboost import CatBoostClassifier
from optuna import Trial
from sklearn.ensemble import GradientBoostingClassifier, RandomForestClassifier
from sklearn.neural_network import MLPClassifier
from sklearn.svm import SVC


class ClassifierStrategy(ABC):
    def __init__(self) -> None:
        self.model = None

    @abstractmethod
    def get_tuning_params(self, trial: Trial) -> dict:
        pass

    def fit(self, X: pd.DataFrame, y: pd.Series) -> None:
        if self.model:
            self.model.fit(X, y)
        else:
            raise ValueError("Model is not initialized")
        # return self

    def predict(self, X: pd.DataFrame) -> Any:
        if self.model:
            return self.model.predict(X)
        else:
            raise ValueError("Model is not initialized")

    def get_params(self) -> dict:
        if self.model:
            return self.model.get_params()
        else:
            raise ValueError("Model is not initialized")

    def set_params(self, **params: dict) -> Any:
        if self.model:
            self.model.set_params(**params)
            return self
        else:
            raise ValueError("Model is not initialized")


class RandomForestModel(ClassifierStrategy):
    def __init__(self, **params: dict) -> None:
        self.model = RandomForestClassifier(**params)

    def get_tuning_params(self, trial: Trial) -> dict:
        return {
            "n_estimators": trial.suggest_int("n_estimators", 80, 120),
            "max_depth": trial.suggest_int("max_depth", 5, 20),
            "min_samples_leaf": trial.suggest_int("min_samples_leaf", 5, 15),
            "bootstrap": True,
            "n_jobs": -1,
        }


class GradientBoostingModel(ClassifierStrategy):
    def __init__(self, **params: dict) -> None:
        self.model = GradientBoostingClassifier(**params)

    def get_tuning_params(self, trial: Trial) -> dict:
        return {
            "n_estimators": trial.suggest_int("n_estimators", 90, 140),
            "learning_rate": trial.suggest_float("learning_rate", 0.0001, 0.2, log=True),
            "subsample": trial.suggest_float("subsample", 0.01, 1, log=True),
            "min_samples_leaf": trial.suggest_int("min_samples_leaf", 5, 15),
            "max_depth": trial.suggest_int("max_depth", 5, 20),
        }


class SVMModel(ClassifierStrategy):
    def __init__(self, **params: dict) -> None:
        self.model = SVC(**params)

    def get_tuning_params(self, trial: Trial) -> dict:
        return {
            "C": trial.suggest_float("C", 0.01, 100, log=True),
            "kernel": trial.suggest_categorical("kernel", ["rbf", "linear", "sigmoid", "poly"]),
            "gamma": trial.suggest_categorical("gamma", ["auto", "scale"]),
            "probability": True,
        }


class MLPModel(ClassifierStrategy):
    def __init__(self, **params: dict) -> None:
        self.model = MLPClassifier(**params)

    def get_tuning_params(self, trial: Trial) -> dict:
        return {
            "hidden_layer_sizes": trial.suggest_categorical("hidden_layer_sizes", [(64,), (128,), (32, 64), (64, 128)]),
            "activation": trial.suggest_categorical("activation", ["relu", "logistic", "tanh"]),
            "alpha": trial.suggest_float("alpha", 1e-5, 0.1, log=True),
        }


class CatBoostModel(ClassifierStrategy):
    def __init__(self, **params: dict) -> None:
        self.model = CatBoostClassifier(**params)

    def get_tuning_params(self, trial: Trial) -> dict:
        return {
            "learning_rate": trial.suggest_float("learning_rate", 0.0001, 0.2, log=True),
            "depth": trial.suggest_int("depth", 1, 16),
            "l2_leaf_reg": trial.suggest_int("l2_leaf_reg", 1, 10),
            "bagging_temperature": trial.suggest_float("bagging_temperature", 0, 1),
            "verbose": 0,
        }


def get_model_strategy(model_name: str, **params: dict) -> Any:
    if model_name == "RandomForest":
        return RandomForestModel(**params)
    elif model_name == "GradientBoosting":
        return GradientBoostingModel(**params)
    elif model_name == "SVM":
        return SVMModel(**params)
    elif model_name == "MLP":
        return MLPModel(**params)
    elif model_name == "CatBoost":
        return CatBoostModel(**params)
    else:
        raise ValueError(f"Model {model_name} is not supported")
