from __future__ import annotations
from scipy.stats import randint, uniform, loguniform
from sklearn.base import BaseEstimator, RegressorMixin, clone
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import Ridge, ElasticNet
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
from sklearn.model_selection import GridSearchCV, RandomizedSearchCV, KFold
import numpy as np


class PhysicalTargetScoringWrapper(BaseEstimator, RegressorMixin):
    """Wrap an estimator so CV scoring can use physical-unit MAE.

    For precipitation, the model is still trained to predict the transformed
    target pr_factor. During GridSearchCV/RandomizedSearchCV, a physical scale
    column pr_o_cli is appended to X only for scoring. The wrapped estimator
    strips that column before fit/predict, while the scorer computes
    MAE(pr_bc, pr_o) = MAE(max(pr_factor_hat*pr_o_cli, 0), pr_o_factor*pr_o_cli).

    When predicting operational/evaluation grids, X does not contain the scale
    column; the wrapper detects this and passes X directly to the base estimator.
    """
    def __init__(self, estimator=None, has_physical_scale: bool = False, clip_nonnegative: bool = True):
        self.estimator = estimator
        self.has_physical_scale = has_physical_scale
        self.clip_nonnegative = clip_nonnegative

    def _split_X(self, X):
        X = np.asarray(X)
        if self.has_physical_scale:
            base_n = getattr(self, "base_n_features_in_", None)
            if base_n is None:
                # During fit, X includes the appended scale column.
                return X[:, :-1], X[:, -1]
            if X.shape[1] == int(base_n) + 1:
                return X[:, :-1], X[:, -1]
        return X, None

    def fit(self, X, y, **fit_params):
        X0, _ = self._split_X(X)
        self.base_n_features_in_ = X0.shape[1]
        self.estimator_ = clone(self.estimator)
        self.estimator_.fit(X0, y, **fit_params)
        return self

    def predict(self, X):
        X0, _ = self._split_X(X)
        return self.estimator_.predict(X0)


def physical_neg_mae_scorer(estimator, X, y):
    """Negative MAE in physical units for CV model selection."""
    pred = estimator.predict(X)
    y = np.asarray(y, dtype=float)
    if getattr(estimator, "has_physical_scale", False):
        scale = np.asarray(X, dtype=float)[:, -1]
        pred_phys = pred.astype(float) * scale
        y_phys = y * scale
        if getattr(estimator, "clip_nonnegative", True):
            pred_phys = np.where(np.isfinite(pred_phys), np.maximum(pred_phys, 0.0), np.nan)
        err = np.abs(pred_phys - y_phys)
    else:
        err = np.abs(pred.astype(float) - y)
    return -float(np.nanmean(err))


def estimator(method: str, random_seed: int = 42):
    m = method.lower()
    if m == 'ridge':
        return Pipeline([('scaler', StandardScaler()), ('model', Ridge())])
    if m == 'elasticnet':
        return Pipeline([('scaler', StandardScaler()), ('model', ElasticNet(max_iter=20000, random_state=random_seed))])
    if m == 'rf':
        return RandomForestRegressor(n_jobs=1, random_state=random_seed)
    if m == 'gbm':
        return GradientBoostingRegressor(random_state=random_seed)
    if m == 'xgboost':
        try:
            from xgboost import XGBRegressor
        except Exception as exc:
            raise ImportError('xgboost is not installed. Install xgboost or remove xgboost from config.') from exc
        return XGBRegressor(objective='reg:squarederror', n_jobs=1, random_state=random_seed, tree_method='hist', verbosity=0)
    raise ValueError(f'Unknown AI model: {method}')


def search_space(method: str):
    m = method.lower()
    if m == 'ridge':
        return {'model__alpha': [0.01, 0.1, 1.0, 10.0, 100.0]}
    if m == 'elasticnet':
        return {'model__alpha': [0.0005, 0.001, 0.005, 0.01, 0.05, 0.1], 'model__l1_ratio': [0.1, 0.25, 0.5, 0.75, 0.9]}
    if m == 'rf':
        return {'n_estimators': randint(120, 500), 'max_depth': [6, 10, 14, None], 'min_samples_leaf': randint(2, 12), 'max_features': ['sqrt', 0.6, 0.8, 1.0]}
    if m == 'gbm':
        return {'n_estimators': randint(100, 500), 'learning_rate': loguniform(0.01, 0.12), 'max_depth': randint(2, 5), 'subsample': uniform(0.65, 0.35), 'min_samples_leaf': randint(2, 15)}
    if m == 'xgboost':
        return {'n_estimators': randint(150, 600), 'learning_rate': loguniform(0.01, 0.12), 'max_depth': randint(2, 5), 'subsample': uniform(0.65, 0.35), 'colsample_bytree': uniform(0.65, 0.35), 'reg_lambda': loguniform(0.1, 20.0), 'min_child_weight': randint(1, 8)}
    raise ValueError(method)


def _prefix_space(space: dict, prefix: str) -> dict:
    return {f'{prefix}{k}': v for k, v in space.items()}


def fit_with_optional_search(X, y, method: str, cfg: dict, physical_scale=None, clip_nonnegative: bool = True):
    """Fit an AI model, optionally tuning hyperparameters by CV.

    physical_scale is used for precipitation: y remains pr_factor, but CV scoring
    is performed after inverse transformation to physical precipitation units.
    For temperature anomaly, physical MAE equals anomaly MAE because the same
    climatological offset is added to prediction and observation.
    """
    seed = int(cfg['training'].get('random_seed', 42))
    base = estimator(method, seed)
    use_physical = physical_scale is not None
    if use_physical:
        scale = np.asarray(physical_scale, dtype=float)
        ok = np.isfinite(scale) & np.isfinite(y) & np.all(np.isfinite(X), axis=1)
        Xfit = np.column_stack([X[ok], scale[ok].astype(np.float32)])
        yfit = np.asarray(y)[ok]
        base_for_search = PhysicalTargetScoringWrapper(base, has_physical_scale=True, clip_nonnegative=clip_nonnegative)
        scoring = physical_neg_mae_scorer
        space = _prefix_space(search_space(method), 'estimator__')
        scoring_name = 'physical_mae_after_inverse_transform'
    else:
        Xfit = X
        yfit = y
        base_for_search = base
        scoring = 'neg_mean_absolute_error'
        space = search_space(method)
        scoring_name = 'mae_on_training_target_units'

    if len(yfit) == 0:
        raise ValueError('No finite samples available for model fitting.')

    if not cfg['training'].get('tune_hyperparameters', True):
        return base_for_search.fit(Xfit, yfit), {'search': 'disabled', 'scoring': scoring_name}

    cvn = min(int(cfg['training'].get('cv_folds', 3)), max(2, len(yfit) // 10))
    if cvn < 2:
        return base_for_search.fit(Xfit, yfit), {'search': 'skipped_small_sample', 'scoring': scoring_name}

    cv = KFold(n_splits=cvn, shuffle=True, random_state=seed)
    if method.lower() in {'ridge', 'elasticnet'}:
        search = GridSearchCV(base_for_search, space, scoring=scoring, cv=cv, n_jobs=int(cfg['training'].get('n_jobs_cv', 1)), refit=True)
    else:
        search = RandomizedSearchCV(base_for_search, space, n_iter=int(cfg['training'].get('random_search_iter', 8)), scoring=scoring, cv=cv, n_jobs=int(cfg['training'].get('n_jobs_cv', 1)), random_state=seed, refit=True)
    search.fit(Xfit, yfit)
    # Convert prefixed params back to the natural estimator names in metadata.
    bp = dict(search.best_params_)
    if use_physical:
        bp = {k.replace('estimator__', '', 1): v for k, v in bp.items()}
    return search.best_estimator_, {
        'search': search.__class__.__name__,
        'scoring': scoring_name,
        'best_score_neg_mae': float(search.best_score_),
        'best_mae_physical_or_target': float(-search.best_score_),
        'best_params': bp,
    }