import sys
import optuna
import numpy as np
from sklearn.model_selection import cross_val_score, KFold
from sklearn.utils.multiclass import type_of_target

# Regression
from xgboost import XGBRegressor
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor, AdaBoostRegressor
from sklearn.linear_model import LinearRegression, Ridge, Lasso, ElasticNet
from sklearn.svm import SVR
from sklearn.neighbors import KNeighborsRegressor

# Classification
from xgboost import XGBClassifier
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier, AdaBoostClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.naive_bayes import GaussianNB
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier

from ..exception.exception import ProjectXecption
import warnings, logging

import warnings
warnings.filterwarnings("ignore")
import warnings
import logging
import os

warnings.filterwarnings("ignore", category=UserWarning)
logging.getLogger("xgboost").setLevel(logging.CRITICAL)
os.environ["XGB_WARN_EARLY_STOPPING"] = "false"



class ModelTrainer:
    def __init__(self, n_trials=None, cv_splits=None, fast_mode=None, use_gpu=None):
        self.n_trials = n_trials if n_trials!= None else 50
        self.cv_splits = cv_splits if cv_splits != None else 5
        self.cv = KFold(n_splits=self.cv_splits, shuffle=True, random_state=42)
        self.best_model = None
        self.best_score = -np.inf
        self.task = None
        self.fast_mode = fast_mode if fast_mode != None else False
        self.use_gpu = use_gpu if use_gpu != None else False

        # ✅ Full model lists
        self.regression_models = [
            "XGBRegressor", "RandomForest", "GradientBoosting",
            "Ridge", "Lasso", "ElasticNet", "SVR", "KNN", "LinearRegression"
        ]
        self.classification_models = [
            "XGBoost", "RandomForest", "GradientBoosting", "AdaBoost",
            "Logistic", "SVC", "NaiveBayes", "KNN", "DecisionTree"
        ]

    def _detect_task(self, y):
        target_type = type_of_target(y)
        self.task = "regression" if target_type in ["continuous", "continuous-multioutput"] else "classification"
        print(f"[AUTO] Detected Task: {self.task}")

    def _gpu_params(self):
        return {"tree_method": "gpu_hist"} if self.use_gpu else {"tree_method": "hist"}

    def _build_best_model(self, model_name: str, best_params: dict):
        """Rebuild the best model from Optuna best params."""
        if self.task == "regression":
            if model_name == "XGBRegressor":
                return XGBRegressor(**best_params, **self._gpu_params())
            elif model_name == "RandomForest":
                return RandomForestRegressor(**best_params)
            elif model_name == "GradientBoosting":
                return GradientBoostingRegressor(**best_params)
            elif model_name == "Ridge":
                return Ridge(**best_params)
            elif model_name == "Lasso":
                return Lasso(**best_params)
            elif model_name == "ElasticNet":
                return ElasticNet(**best_params)
            elif model_name == "SVR":
                return SVR(**best_params)
            elif model_name == "KNN":
                return KNeighborsRegressor(**best_params)
            else:
                return LinearRegression(**best_params)
        else:
            if model_name == "XGBoost":
                return XGBClassifier(**best_params, **self._gpu_params(), eval_metric="mlogloss")
            elif model_name == "RandomForest":
                return RandomForestClassifier(**best_params)
            elif model_name == "GradientBoosting":
                return GradientBoostingClassifier(**best_params)
            elif model_name == "AdaBoost":
                return AdaBoostClassifier(**best_params, algorithm="SAMME")
            elif model_name == "Logistic":
                return LogisticRegression(**best_params)
            elif model_name == "SVC":
                return SVC(**best_params)
            elif model_name == "NaiveBayes":
                return GaussianNB()
            elif model_name == "KNN":
                return KNeighborsClassifier(**best_params)
            elif model_name == "DecisionTree":
                return DecisionTreeClassifier(**best_params)

    def _get_model(self, trial):
        models_list = self.regression_models if self.task == "regression" else self.classification_models
        if self.fast_mode:
            models_list = ["Logistic", "Ridge", "RandomForest", "XGBRegressor", "XGBoost"]

        model_name = trial.suggest_categorical("model", models_list)

        # ✅ Same as before but added GPU & early stopping support
        if self.task == "regression":
            if model_name == "XGBRegressor":
                return XGBRegressor(
                    n_estimators=trial.suggest_int("n_estimators", 50, 300),
                    learning_rate=trial.suggest_float("learning_rate", 0.01, 0.3),
                    max_depth=trial.suggest_int("max_depth", 3, 12),
                    subsample=trial.suggest_float("subsample", 0.5, 1.0),
                    colsample_bytree=trial.suggest_float("colsample_bytree", 0.5, 1.0),
                    eval_metric="rmse",
                    random_state=42
                )
            elif model_name == "RandomForest":
                return RandomForestRegressor(
                    n_estimators=trial.suggest_int("n_estimators", 50, 200),
                    max_depth=trial.suggest_int("max_depth", 3, 15),
                    min_samples_split=trial.suggest_int("min_samples_split", 2, 10),
                    random_state=42
                )
            elif model_name == "Ridge":
                return Ridge(alpha=trial.suggest_float("alpha", 0.01, 10.0))
            else:
                return LinearRegression()
        else:
            if model_name == "XGBoost":
                return XGBClassifier(
                    n_estimators=trial.suggest_int("n_estimators", 50, 300),
                    learning_rate=trial.suggest_float("learning_rate", 0.01, 0.3),
                    max_depth=trial.suggest_int("max_depth", 3, 12),
                    subsample=trial.suggest_float("subsample", 0.5, 1.0),
                    colsample_bytree=trial.suggest_float("colsample_bytree", 0.5, 1.0),
                    eval_metric="mlogloss",
                    random_state=42
                )
            elif model_name == "RandomForest":
                return RandomForestClassifier(
                    n_estimators=trial.suggest_int("n_estimators", 50, 200),
                    max_depth=trial.suggest_int("max_depth", 3, 15),
                    min_samples_split=trial.suggest_int("min_samples_split", 2, 10),
                    random_state=42
                )
            elif model_name == "Logistic":
                return LogisticRegression(
                    C=trial.suggest_float("C", 0.01, 10.0),
                    max_iter=300
                )
            else:
                return GaussianNB()

    def _objective(self, trial, X, y):
        try:
            model = self._get_model(trial)
            scoring = "r2" if self.task == "regression" else "accuracy"
            # ✅ Use subset of data for faster tuning if fast_mode=True
            if self.fast_mode and X.shape[0] > 20000:
                subset = np.random.choice(len(y), 20000, replace=False)
                X, y = X.iloc[subset], y.iloc[subset]
            score = cross_val_score(model, X, y, cv=self.cv, scoring=scoring, n_jobs=-1).mean()
            trial.report(score, step=0)
            if trial.should_prune():
                raise optuna.TrialPruned()
            return score
        except Exception as e:
            raise ProjectXecption(e, sys)

    def fit(self, X_train, y_train):
        try:
            self._detect_task(y_train)
            study = optuna.create_study(
                direction="maximize",
                pruner=optuna.pruners.MedianPruner(n_warmup_steps=5)
            )
            study.optimize(
                lambda trial: self._objective(trial, X_train, y_train),
                n_trials=self.n_trials,
                timeout=600,
                n_jobs=-1,
                show_progress_bar=True
            )
            print("✅ Best Trial:", study.best_trial.params)
            print("✅ Best Score:", study.best_value)
            best_params = study.best_trial.params
            best_model_name = best_params.pop("model")
            best_model = self._build_best_model(best_model_name, best_params)
            best_model.fit(X_train, y_train)
            self.best_model = best_model
            self.best_score = study.best_value
            return best_model
        except Exception as e:
            raise ProjectXecption(e, sys)