"""Hyperparameter optimization using Optuna with TimeSeriesSplit.

Automatically finds optimal hyperparameters for each model type.
Results are cached so we don't re-optimize on every retrain.
"""
from __future__ import annotations
import logging
import os
import numpy as np
import pandas as pd
from typing import Optional

logger = logging.getLogger(__name__)


def _cfg(key: str, fallback):
    try:
        from flask import current_app
        return current_app.config.get(key, fallback)
    except RuntimeError:
        return fallback


class HyperOptimizer:
    """Optuna-based hyperparameter optimization for tree models."""

    def __init__(self):
        self._cache: dict[str, dict] = {}

    def optimize(
        self, feature_df: pd.DataFrame, target_col: str = 'target_return',
        model_type: str = 'xgboost', n_trials: int | None = None,
    ) -> dict:
        """Run Optuna optimization for a model type.

        Args:
            feature_df: Feature DataFrame with target column
            target_col: Name of target column
            model_type: 'xgboost', 'lightgbm', 'catboost', or 'gb'
            n_trials: Number of Optuna trials (default from config)

        Returns:
            dict of best hyperparameters
        """
        try:
            import optuna
            optuna.logging.set_verbosity(optuna.logging.WARNING)
        except ImportError:
            logger.warning('Optuna not installed, using default params')
            return self._default_params(model_type)

        if n_trials is None:
            n_trials = _cfg('ML_OPTUNA_TRIALS', 50)

        # Check cache
        cache_key = f'{model_type}_{len(feature_df)}'
        if cache_key in self._cache:
            return self._cache[cache_key]

        X = feature_df.drop(columns=[target_col]).select_dtypes(include=[np.number])
        y = feature_df[target_col].values

        mask = ~(np.isnan(X.values).any(axis=1) | np.isnan(y))
        X, y = X[mask], y[mask]

        if len(X) < 100:
            return self._default_params(model_type)

        def objective(trial):
            from sklearn.model_selection import TimeSeriesSplit
            from sklearn.preprocessing import StandardScaler
            from sklearn.metrics import mean_absolute_error

            params = self._suggest_params(trial, model_type)
            model = self._create_model(model_type, params)
            if model is None:
                return float('inf')

            n_splits = _cfg('ML_CV_SPLITS', 5)
            tscv = TimeSeriesSplit(n_splits=n_splits)
            scores = []

            for train_idx, val_idx in tscv.split(X):
                X_train, X_val = X.iloc[train_idx].values, X.iloc[val_idx].values
                y_train, y_val = y[train_idx], y[val_idx]

                scaler = StandardScaler()
                X_train_s = scaler.fit_transform(X_train)
                X_val_s = scaler.transform(X_val)

                try:
                    model_copy = self._create_model(model_type, params)
                    model_copy.fit(X_train_s, y_train)
                    preds = model_copy.predict(X_val_s)

                    # Combined objective: MAE + directional accuracy
                    mae = mean_absolute_error(y_val, preds)
                    dir_acc = np.mean(np.sign(preds) == np.sign(y_val))
                    # Minimize: MAE - 0.5 * dir_acc (balance both)
                    scores.append(mae - 0.5 * dir_acc)
                except Exception:
                    scores.append(float('inf'))

            return float(np.mean(scores))

        study = optuna.create_study(direction='minimize')
        study.optimize(objective, n_trials=n_trials, show_progress_bar=False)

        best_params = study.best_params
        self._cache[cache_key] = best_params

        logger.info(f'Optuna best params for {model_type}: {best_params}')
        return best_params

    def _suggest_params(self, trial, model_type: str) -> dict:
        """Suggest hyperparameters for an Optuna trial."""
        params = {
            'n_estimators': trial.suggest_int('n_estimators', 100, 800, step=50),
            'max_depth': trial.suggest_int('max_depth', 4, 12),
            'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.15, log=True),
            'subsample': trial.suggest_float('subsample', 0.6, 1.0),
        }

        if model_type in ('xgboost', 'lightgbm'):
            params['colsample_bytree'] = trial.suggest_float('colsample_bytree', 0.5, 1.0)
            params['reg_alpha'] = trial.suggest_float('reg_alpha', 0.0, 2.0)
            params['reg_lambda'] = trial.suggest_float('reg_lambda', 0.5, 3.0)

        if model_type == 'catboost':
            params['l2_leaf_reg'] = trial.suggest_float('l2_leaf_reg', 1.0, 10.0)

        return params

    def _create_model(self, model_type: str, params: dict):
        """Create model instance with given params."""
        n_est = params.get('n_estimators', 300)
        max_d = params.get('max_depth', 8)
        lr = params.get('learning_rate', 0.05)
        sub = params.get('subsample', 0.8)

        if model_type == 'xgboost':
            try:
                from xgboost import XGBRegressor
                return XGBRegressor(
                    n_estimators=n_est, max_depth=max_d,
                    learning_rate=lr, subsample=sub,
                    colsample_bytree=params.get('colsample_bytree', 0.8),
                    reg_alpha=params.get('reg_alpha', 0.1),
                    reg_lambda=params.get('reg_lambda', 1.0),
                    random_state=42, verbosity=0,
                )
            except ImportError:
                pass

        elif model_type == 'lightgbm':
            try:
                from lightgbm import LGBMRegressor
                return LGBMRegressor(
                    n_estimators=n_est, max_depth=max_d,
                    learning_rate=lr, subsample=sub,
                    colsample_bytree=params.get('colsample_bytree', 0.8),
                    reg_alpha=params.get('reg_alpha', 0.1),
                    reg_lambda=params.get('reg_lambda', 1.0),
                    random_state=42, verbose=-1,
                )
            except ImportError:
                pass

        elif model_type == 'catboost':
            try:
                from catboost import CatBoostRegressor
                return CatBoostRegressor(
                    iterations=n_est, depth=min(max_d, 10),
                    learning_rate=lr, random_seed=42, verbose=0, logging_level='Silent',
                    l2_leaf_reg=params.get('l2_leaf_reg', 3.0),
                )
            except ImportError:
                pass

        from sklearn.ensemble import GradientBoostingRegressor
        return GradientBoostingRegressor(
            n_estimators=n_est, max_depth=max_d,
            learning_rate=lr, subsample=sub, random_state=42,
        )

    def _default_params(self, model_type: str) -> dict:
        """Return sensible defaults when Optuna is not available."""
        return {
            'n_estimators': _cfg('ML_TREE_ESTIMATORS', 500),
            'max_depth': _cfg('ML_MAX_DEPTH', 8),
            'learning_rate': 0.05,
            'subsample': 0.8,
        }