"""Automatic model selection and tuning."""

from dataclasses import dataclass, field
from typing import Optional, Union, List, Dict, Any, Type
import warnings

import numpy as np
import pandas as pd

from ad_inventory_forecast.core.base import BaseForecaster
from ad_inventory_forecast.backtesting.validator import WalkForwardValidator
from ad_inventory_forecast.backtesting.metrics import calculate_smape, calculate_mase


@dataclass
class TuningResult:
    """
    Result of model tuning/selection.

    Attributes
    ----------
    best_model : BaseForecaster
        Best fitted model.
    best_model_name : str
        Name of the best model.
    best_params : Dict[str, Any]
        Parameters of the best model.
    best_score : float
        Best score achieved.
    all_results : List[Dict]
        Results for all models tried.
    metric : str
        Metric used for selection.
    """

    best_model: BaseForecaster
    best_model_name: str
    best_params: Dict[str, Any]
    best_score: float
    all_results: List[Dict] = field(default_factory=list)
    metric: str = "smape"

    def summary(self) -> str:
        """Return summary of tuning results."""
        lines = [
            "AutoTuner Results",
            "=" * 40,
            f"Best Model: {self.best_model_name}",
            f"Best {self.metric.upper()}: {self.best_score:.4f}",
            "",
            "All Models Ranked:",
        ]

        sorted_results = sorted(self.all_results, key=lambda x: x["score"])
        for i, result in enumerate(sorted_results, 1):
            lines.append(
                f"  {i}. {result['name']}: {self.metric}={result['score']:.4f}"
            )

        return "\n".join(lines)

    def to_dataframe(self) -> pd.DataFrame:
        """Convert results to DataFrame."""
        return pd.DataFrame(self.all_results)


class AutoTuner:
    """
    Automatic model selection using walk-forward validation.

    Compares multiple forecasting models and selects the best one
    based on backtesting performance.

    Parameters
    ----------
    candidates : List[BaseForecaster or Type[BaseForecaster]], optional
        List of model instances or classes to compare.
        If None, uses default set of models.
    metric : str, default 'smape'
        Primary metric for model selection.
        Options: 'smape', 'mase', 'mape', 'rmse', 'mae'.
    n_folds : int, default 3
        Number of validation folds.
    horizon : int, default 30
        Forecast horizon for evaluation.
    min_train_size : int, default 180
        Minimum training observations.
    refit_on_full : bool, default True
        Whether to refit the best model on full data.
    verbose : bool, default False
        Whether to print progress.

    Attributes
    ----------
    result_ : TuningResult
        Results after calling select_best().
    best_model_ : BaseForecaster
        Best model (alias for result_.best_model).

    Examples
    --------
    >>> from ad_inventory_forecast.models import InventoryForecaster
    >>> from ad_inventory_forecast.models.ets import ETSForecaster
    >>> tuner = AutoTuner(
    ...     candidates=[
    ...         InventoryForecaster(),
    ...         ETSForecaster(seasonal='add', seasonal_periods=7),
    ...     ],
    ...     metric='smape',
    ... )
    >>> result = tuner.select_best(y)
    >>> print(result.summary())
    >>> forecast = result.best_model.predict(horizon=30)

    Notes
    -----
    Models that fail during fitting are skipped with a warning.
    If all models fail, raises a ValueError.
    """

    def __init__(
        self,
        candidates: Optional[List[Union[BaseForecaster, Type[BaseForecaster]]]] = None,
        metric: str = "smape",
        n_folds: int = 3,
        horizon: int = 30,
        min_train_size: int = 180,
        refit_on_full: bool = True,
        verbose: bool = False,
    ):
        """Initialize the AutoTuner."""
        self.candidates = candidates
        self.metric = metric
        self.n_folds = n_folds
        self.horizon = horizon
        self.min_train_size = min_train_size
        self.refit_on_full = refit_on_full
        self.verbose = verbose

        # Results
        self.result_ = None
        self.best_model_ = None

    def select_best(
        self,
        y: pd.Series,
        X: Optional[pd.DataFrame] = None,
    ) -> TuningResult:
        """
        Select the best model using cross-validation.

        Parameters
        ----------
        y : pd.Series
            Time series with DatetimeIndex.
        X : pd.DataFrame, optional
            Exogenous features.

        Returns
        -------
        result : TuningResult
            Tuning results with best model.
        """
        if self.candidates is None:
            candidates = self._get_default_candidates(y)
        else:
            candidates = self.candidates

        # Create validator
        validator = WalkForwardValidator(
            n_folds=self.n_folds,
            horizon=self.horizon,
            min_train_size=self.min_train_size,
        )

        all_results = []
        best_score = np.inf
        best_model = None
        best_name = None
        best_params = {}

        for candidate in candidates:
            # Get model instance and name
            if isinstance(candidate, type):
                model_class = candidate
                model_name = candidate.__name__
                model_params = {}
            else:
                model_class = candidate.__class__
                model_name = f"{candidate.__class__.__name__}"
                model_params = candidate.get_params()

            if self.verbose:
                print(f"Evaluating {model_name}...")

            try:
                # Run validation
                backtest_result = validator.validate(
                    model_class,
                    y,
                    model_params=model_params,
                    verbose=False,
                )

                # Get primary metric
                score = backtest_result.metrics.get(self.metric, np.inf)

                result_dict = {
                    "name": model_name,
                    "score": score,
                    "params": model_params,
                    **backtest_result.metrics,
                }
                all_results.append(result_dict)

                if self.verbose:
                    print(f"  {self.metric}: {score:.4f}")

                if score < best_score:
                    best_score = score
                    best_name = model_name
                    best_params = model_params

                    # Store best model class for refitting
                    if isinstance(candidate, type):
                        best_model = candidate()
                    else:
                        best_model = candidate.__class__(**model_params)

            except Exception as e:
                if self.verbose:
                    print(f"  Failed: {e}")
                warnings.warn(f"Model {model_name} failed: {e}")
                all_results.append({
                    "name": model_name,
                    "score": np.inf,
                    "params": model_params,
                    "error": str(e),
                })

        if best_model is None:
            raise ValueError("All candidate models failed during evaluation")

        # Refit on full data
        if self.refit_on_full:
            if self.verbose:
                print(f"\nRefitting {best_name} on full data...")
            best_model.fit(y, X)

        self.result_ = TuningResult(
            best_model=best_model,
            best_model_name=best_name,
            best_params=best_params,
            best_score=best_score,
            all_results=all_results,
            metric=self.metric,
        )
        self.best_model_ = best_model

        return self.result_

    def _get_default_candidates(self, y: pd.Series) -> List[BaseForecaster]:
        """Get default candidate models based on data characteristics."""
        candidates = []

        # Determine seasonal period
        n = len(y)
        if n >= 365:
            seasonal_periods = [7, 365]
        elif n >= 14:
            seasonal_periods = [7]
        else:
            seasonal_periods = []

        # Always try InventoryForecaster
        try:
            from ad_inventory_forecast.models.decomposition import InventoryForecaster

            candidates.append(
                InventoryForecaster(
                    model="auto",
                    trend="auto",
                    seasonal_periods=seasonal_periods or [7],
                )
            )
        except ImportError:
            pass

        # Try ETS if enough data
        if n >= 14:
            try:
                from ad_inventory_forecast.models.ets import ETSForecaster

                candidates.append(
                    ETSForecaster(
                        trend="add",
                        damped_trend=True,
                        seasonal="add" if seasonal_periods else None,
                        seasonal_periods=seasonal_periods[0] if seasonal_periods else None,
                    )
                )
            except ImportError:
                pass

        # Try SARIMA if enough data
        if n >= 30:
            try:
                from ad_inventory_forecast.models.sarima import SARIMAForecaster

                candidates.append(
                    SARIMAForecaster(
                        auto=True,
                        seasonal_periods=seasonal_periods[0] if seasonal_periods else 1,
                    )
                )
            except ImportError:
                pass

        # Try ML model if enough data
        if n >= 90:
            try:
                from ad_inventory_forecast.models.ml_ensemble import MLEnsembleForecaster

                candidates.append(
                    MLEnsembleForecaster(
                        model_type="xgboost",
                        lag_features=[1, 7, 14],
                        rolling_windows=[7, 14],
                    )
                )
            except ImportError:
                pass

        if not candidates:
            raise ValueError("No default models available")

        return candidates


def quick_select(
    y: pd.Series,
    horizon: int = 30,
    verbose: bool = True,
) -> BaseForecaster:
    """
    Quick model selection with sensible defaults.

    Parameters
    ----------
    y : pd.Series
        Time series with DatetimeIndex.
    horizon : int, default 30
        Forecast horizon.
    verbose : bool, default True
        Whether to print progress.

    Returns
    -------
    model : BaseForecaster
        Best model fitted on full data.

    Examples
    --------
    >>> model = quick_select(y, horizon=30)
    >>> forecast = model.predict(horizon=30)
    """
    tuner = AutoTuner(
        metric="smape",
        n_folds=3,
        horizon=horizon,
        verbose=verbose,
    )

    result = tuner.select_best(y)

    if verbose:
        print(f"\nSelected: {result.best_model_name}")
        print(f"SMAPE: {result.best_score:.2f}%")

    return result.best_model


class GridSearchCV:
    """
    Grid search with cross-validation for a single model.

    Parameters
    ----------
    model_class : Type[BaseForecaster]
        Model class to tune.
    param_grid : Dict[str, List[Any]]
        Parameter grid to search.
    metric : str, default 'smape'
        Metric to optimize.
    n_folds : int, default 3
        Number of CV folds.
    horizon : int, default 30
        Forecast horizon.
    verbose : bool, default False
        Whether to print progress.

    Examples
    --------
    >>> from ad_inventory_forecast.models.ets import ETSForecaster
    >>> param_grid = {
    ...     'trend': ['add', None],
    ...     'seasonal': ['add', 'mul'],
    ...     'damped_trend': [True, False],
    ... }
    >>> gs = GridSearchCV(ETSForecaster, param_grid, n_folds=3)
    >>> result = gs.fit(y)
    >>> print(result.best_params_)
    """

    def __init__(
        self,
        model_class: Type[BaseForecaster],
        param_grid: Dict[str, List[Any]],
        metric: str = "smape",
        n_folds: int = 3,
        horizon: int = 30,
        verbose: bool = False,
    ):
        """Initialize grid search."""
        self.model_class = model_class
        self.param_grid = param_grid
        self.metric = metric
        self.n_folds = n_folds
        self.horizon = horizon
        self.verbose = verbose

        # Results
        self.best_params_ = None
        self.best_score_ = None
        self.best_model_ = None
        self.cv_results_ = None

    def fit(self, y: pd.Series) -> "GridSearchCV":
        """
        Run grid search.

        Parameters
        ----------
        y : pd.Series
            Time series with DatetimeIndex.

        Returns
        -------
        self : GridSearchCV
            Fitted grid search.
        """
        # Generate all parameter combinations
        from itertools import product

        param_names = list(self.param_grid.keys())
        param_values = list(self.param_grid.values())
        all_combinations = list(product(*param_values))

        validator = WalkForwardValidator(
            n_folds=self.n_folds,
            horizon=self.horizon,
        )

        results = []
        best_score = np.inf
        best_params = None

        for i, values in enumerate(all_combinations):
            params = dict(zip(param_names, values))

            if self.verbose:
                print(f"Trying {i + 1}/{len(all_combinations)}: {params}")

            try:
                backtest_result = validator.validate(
                    self.model_class,
                    y,
                    model_params=params,
                )

                score = backtest_result.metrics.get(self.metric, np.inf)

                results.append({
                    "params": params,
                    "score": score,
                    **backtest_result.metrics,
                })

                if score < best_score:
                    best_score = score
                    best_params = params

            except Exception as e:
                results.append({
                    "params": params,
                    "score": np.inf,
                    "error": str(e),
                })

        self.cv_results_ = pd.DataFrame(results)
        self.best_params_ = best_params
        self.best_score_ = best_score

        # Fit best model on full data
        if best_params:
            self.best_model_ = self.model_class(**best_params)
            self.best_model_.fit(y)

        return self
