"""Walk-forward validation for time series models."""

from dataclasses import dataclass, field
from typing import (
    Optional, Union, List, Dict, Any, Type, Callable, Tuple
)

import numpy as np
import pandas as pd

from ad_inventory_forecast.core.base import BaseForecaster
from ad_inventory_forecast.backtesting.metrics import (
    calculate_mase,
    calculate_smape,
    calculate_mape,
    calculate_rmse,
    calculate_mae,
    calculate_wmape,
    calculate_bias,
)


@dataclass
class FoldResult:
    """
    Result of a single validation fold.

    Attributes
    ----------
    fold : int
        Fold number (0-indexed).
    train_start : pd.Timestamp
        Start date of training period.
    train_end : pd.Timestamp
        End date of training period.
    test_start : pd.Timestamp
        Start date of test period.
    test_end : pd.Timestamp
        End date of test period.
    n_train : int
        Number of training observations.
    n_test : int
        Number of test observations.
    actual : pd.Series
        Actual values in test period.
    predicted : pd.Series
        Predicted values for test period.
    metrics : Dict[str, float]
        Calculated metrics for this fold.
    """

    fold: int
    train_start: pd.Timestamp
    train_end: pd.Timestamp
    test_start: pd.Timestamp
    test_end: pd.Timestamp
    n_train: int
    n_test: int
    actual: pd.Series
    predicted: pd.Series
    metrics: Dict[str, float]


@dataclass
class BacktestResult:
    """
    Aggregated result of walk-forward validation.

    Attributes
    ----------
    fold_results : List[FoldResult]
        Results for each fold.
    metrics : Dict[str, float]
        Aggregated metrics across all folds.
    metrics_std : Dict[str, float]
        Standard deviation of metrics across folds.
    n_folds : int
        Number of folds used.
    horizon : int
        Forecast horizon used.
    model_params : Dict[str, Any]
        Parameters of the model used.
    """

    fold_results: List[FoldResult]
    metrics: Dict[str, float] = field(default_factory=dict)
    metrics_std: Dict[str, float] = field(default_factory=dict)
    n_folds: int = 0
    horizon: int = 0
    model_params: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        """Calculate aggregated metrics after initialization."""
        if self.fold_results and not self.metrics:
            self._aggregate_metrics()

    def _aggregate_metrics(self):
        """Aggregate metrics across folds."""
        self.n_folds = len(self.fold_results)

        if self.n_folds == 0:
            return

        # Collect metrics from each fold
        metric_values = {}
        for fold in self.fold_results:
            for name, value in fold.metrics.items():
                if name not in metric_values:
                    metric_values[name] = []
                if not np.isnan(value) and not np.isinf(value):
                    metric_values[name].append(value)

        # Calculate mean and std
        for name, values in metric_values.items():
            if values:
                self.metrics[name] = np.mean(values)
                self.metrics_std[name] = np.std(values)
            else:
                self.metrics[name] = np.nan
                self.metrics_std[name] = np.nan

    def summary(self) -> str:
        """Return a summary string of the backtest results."""
        lines = [
            f"Walk-Forward Validation Results",
            f"{'=' * 40}",
            f"Folds: {self.n_folds}",
            f"Horizon: {self.horizon} periods",
            f"",
            f"Aggregated Metrics (mean ± std):",
        ]

        for name, value in sorted(self.metrics.items()):
            std = self.metrics_std.get(name, 0)
            lines.append(f"  {name}: {value:.4f} ± {std:.4f}")

        return "\n".join(lines)

    def to_dataframe(self) -> pd.DataFrame:
        """Convert fold results to DataFrame."""
        records = []
        for fold in self.fold_results:
            record = {
                "fold": fold.fold,
                "train_start": fold.train_start,
                "train_end": fold.train_end,
                "test_start": fold.test_start,
                "test_end": fold.test_end,
                "n_train": fold.n_train,
                "n_test": fold.n_test,
            }
            record.update(fold.metrics)
            records.append(record)

        return pd.DataFrame(records)


class WalkForwardValidator:
    """
    Walk-forward (rolling window) validation for time series models.

    Implements proper time series cross-validation that respects temporal
    ordering. Training data always precedes test data.

    Parameters
    ----------
    n_folds : int, default 5
        Number of validation folds.
    horizon : int, default 30
        Forecast horizon (number of periods to predict).
    min_train_size : int, default 180
        Minimum number of training observations required.
    step_size : int, default 7
        Number of periods to advance between folds.
        Smaller values give more folds but with overlapping test periods.
    expanding : bool, default True
        If True, use expanding window (training grows with each fold).
        If False, use sliding window (fixed training size).
    gap : int, default 0
        Gap between training and test sets (useful for models that
        require future data for training).
    primary_metric : str, default 'smape'
        Primary metric for model comparison.

    Attributes
    ----------
    fold_boundaries_ : List[Tuple]
        Boundaries of each fold after fitting.

    Examples
    --------
    >>> from ad_inventory_forecast.models import InventoryForecaster
    >>> validator = WalkForwardValidator(n_folds=5, horizon=7)
    >>> result = validator.validate(InventoryForecaster, y)
    >>> print(result.summary())

    Notes
    -----
    Walk-forward validation is the gold standard for time series model
    evaluation because it:
    1. Respects temporal ordering (no data leakage)
    2. Tests the model's ability to forecast into the future
    3. Provides multiple forecast origins for robust evaluation
    """

    def __init__(
        self,
        n_folds: int = 5,
        horizon: int = 30,
        min_train_size: int = 180,
        step_size: int = 7,
        expanding: bool = True,
        gap: int = 0,
        primary_metric: str = "smape",
    ):
        """Initialize the validator."""
        self.n_folds = n_folds
        self.horizon = horizon
        self.min_train_size = min_train_size
        self.step_size = step_size
        self.expanding = expanding
        self.gap = gap
        self.primary_metric = primary_metric

        self.fold_boundaries_ = None

    def validate(
        self,
        model: Union[Type[BaseForecaster], BaseForecaster],
        y: pd.Series,
        model_params: Optional[Dict[str, Any]] = None,
        verbose: bool = False,
    ) -> BacktestResult:
        """
        Run walk-forward validation on a model.

        Parameters
        ----------
        model : class or BaseForecaster instance
            Model class to instantiate, or a fitted model instance.
            If class, will be instantiated with model_params.
        y : pd.Series
            Time series with DatetimeIndex.
        model_params : dict, optional
            Parameters to pass to model constructor.
        verbose : bool, default False
            Whether to print progress.

        Returns
        -------
        result : BacktestResult
            Validation results with metrics for each fold.
        """
        if not isinstance(y.index, pd.DatetimeIndex):
            raise ValueError("Series must have DatetimeIndex")

        # Generate fold boundaries
        self.fold_boundaries_ = self._generate_folds(y)

        if len(self.fold_boundaries_) < 1:
            raise ValueError(
                f"Insufficient data for validation. "
                f"Need at least {self.min_train_size + self.horizon} observations."
            )

        fold_results = []
        model_params = model_params or {}

        for fold_idx, (train_end_idx, test_start_idx, test_end_idx) in enumerate(
            self.fold_boundaries_
        ):
            if verbose:
                print(f"Fold {fold_idx + 1}/{len(self.fold_boundaries_)}")

            # Determine training window
            if self.expanding:
                train_start_idx = 0
            else:
                train_start_idx = max(0, train_end_idx - self.min_train_size)

            # Extract train and test data
            y_train = y.iloc[train_start_idx:train_end_idx]
            y_test = y.iloc[test_start_idx:test_end_idx]

            # Skip if test is empty
            if len(y_test) == 0:
                continue

            # Instantiate and fit model
            if isinstance(model, type):
                model_instance = model(**model_params)
            else:
                # Clone the model
                model_instance = model.__class__(**model.get_params())

            try:
                model_instance.fit(y_train)
                predictions = model_instance.predict(len(y_test))
            except Exception as e:
                if verbose:
                    print(f"  Fold {fold_idx + 1} failed: {e}")
                continue

            # Align predictions with actuals
            if isinstance(predictions, pd.Series):
                pred_values = predictions.values
            else:
                pred_values = predictions

            actual_values = y_test.values
            n_compare = min(len(actual_values), len(pred_values))
            actual_values = actual_values[:n_compare]
            pred_values = pred_values[:n_compare]

            # Calculate metrics
            metrics = self._calculate_fold_metrics(
                actual_values, pred_values, y_train.values
            )

            fold_result = FoldResult(
                fold=fold_idx,
                train_start=y_train.index[0],
                train_end=y_train.index[-1],
                test_start=y_test.index[0],
                test_end=y_test.index[min(n_compare - 1, len(y_test) - 1)],
                n_train=len(y_train),
                n_test=n_compare,
                actual=pd.Series(actual_values, index=y_test.index[:n_compare]),
                predicted=pd.Series(pred_values, index=y_test.index[:n_compare]),
                metrics=metrics,
            )
            fold_results.append(fold_result)

        return BacktestResult(
            fold_results=fold_results,
            horizon=self.horizon,
            model_params=model_params,
        )

    def _generate_folds(self, y: pd.Series) -> List[Tuple[int, int, int]]:
        """
        Generate fold boundaries.

        Returns list of (train_end_idx, test_start_idx, test_end_idx) tuples.
        """
        n = len(y)
        folds = []

        # Calculate how much space we need
        # Last fold: train ends at n - horizon, test is [n - horizon, n)
        # First fold: train ends at min_train_size, test is [min_train_size, min_train_size + horizon)

        # Work backwards from the end
        test_end = n
        fold_count = 0

        while fold_count < self.n_folds:
            test_start = test_end - self.horizon
            train_end = test_start - self.gap

            # Check if we have enough training data
            if train_end < self.min_train_size:
                break

            folds.append((train_end, test_start, test_end))
            fold_count += 1

            # Move window back
            test_end -= self.step_size

        # Reverse to get chronological order
        folds.reverse()

        return folds

    def _calculate_fold_metrics(
        self,
        actual: np.ndarray,
        predicted: np.ndarray,
        training_series: np.ndarray,
    ) -> Dict[str, float]:
        """Calculate all metrics for a fold."""
        metrics = {}

        # Basic metrics
        metrics["mae"] = calculate_mae(actual, predicted)
        metrics["rmse"] = calculate_rmse(actual, predicted)
        metrics["smape"] = calculate_smape(actual, predicted)
        metrics["wmape"] = calculate_wmape(actual, predicted)
        metrics["bias"] = calculate_bias(actual, predicted)

        # MAPE (may fail with zeros)
        try:
            metrics["mape"] = calculate_mape(actual, predicted)
        except ValueError:
            metrics["mape"] = np.nan

        # MASE (requires training series)
        try:
            # Use period=7 for daily data as default
            metrics["mase"] = calculate_mase(
                actual, predicted, training_series, period=7
            )
        except Exception:
            metrics["mase"] = np.nan

        return metrics

    def plot_folds(self, y: pd.Series):
        """
        Visualize the fold structure.

        Parameters
        ----------
        y : pd.Series
            Original time series.

        Returns
        -------
        fig : matplotlib.figure.Figure
            Figure showing fold structure.
        """
        try:
            import matplotlib.pyplot as plt
            import matplotlib.patches as mpatches
        except ImportError:
            raise ImportError("matplotlib is required for plotting")

        if self.fold_boundaries_ is None:
            self.fold_boundaries_ = self._generate_folds(y)

        fig, ax = plt.subplots(figsize=(14, 4 + len(self.fold_boundaries_) * 0.5))

        # Plot the time series at the top
        ax.plot(range(len(y)), y.values, color="gray", alpha=0.5, linewidth=0.5)

        # Plot fold structure
        colors_train = plt.cm.Blues(np.linspace(0.3, 0.7, len(self.fold_boundaries_)))
        colors_test = plt.cm.Oranges(np.linspace(0.4, 0.8, len(self.fold_boundaries_)))

        y_positions = np.linspace(
            y.min() - (y.max() - y.min()) * 0.1,
            y.min() - (y.max() - y.min()) * 0.5,
            len(self.fold_boundaries_),
        )

        for i, (train_end, test_start, test_end) in enumerate(self.fold_boundaries_):
            train_start = 0 if self.expanding else max(0, train_end - self.min_train_size)

            # Training bar
            ax.barh(
                y_positions[i],
                train_end - train_start,
                left=train_start,
                height=(y.max() - y.min()) * 0.03,
                color=colors_train[i],
                edgecolor="black",
                linewidth=0.5,
            )

            # Test bar
            ax.barh(
                y_positions[i],
                test_end - test_start,
                left=test_start,
                height=(y.max() - y.min()) * 0.03,
                color=colors_test[i],
                edgecolor="black",
                linewidth=0.5,
            )

            ax.text(
                -len(y) * 0.02,
                y_positions[i],
                f"Fold {i + 1}",
                ha="right",
                va="center",
                fontsize=9,
            )

        # Legend
        train_patch = mpatches.Patch(color=colors_train[0], label="Training")
        test_patch = mpatches.Patch(color=colors_test[0], label="Test")
        ax.legend(handles=[train_patch, test_patch], loc="upper right")

        ax.set_xlabel("Observation Index")
        ax.set_ylabel("Value")
        ax.set_title(
            f"Walk-Forward Validation Structure\n"
            f"({len(self.fold_boundaries_)} folds, "
            f"horizon={self.horizon}, "
            f"{'expanding' if self.expanding else 'sliding'} window)"
        )

        plt.tight_layout()
        return fig


def cross_validate(
    model: Union[Type[BaseForecaster], BaseForecaster],
    y: pd.Series,
    n_folds: int = 5,
    horizon: int = 30,
    min_train_size: int = 180,
    model_params: Optional[Dict[str, Any]] = None,
    return_predictions: bool = False,
) -> Union[Dict[str, float], Tuple[Dict[str, float], pd.DataFrame]]:
    """
    Convenience function for walk-forward cross-validation.

    Parameters
    ----------
    model : class or BaseForecaster instance
        Model to validate.
    y : pd.Series
        Time series with DatetimeIndex.
    n_folds : int, default 5
        Number of validation folds.
    horizon : int, default 30
        Forecast horizon.
    min_train_size : int, default 180
        Minimum training observations.
    model_params : dict, optional
        Model constructor parameters.
    return_predictions : bool, default False
        Whether to return all predictions.

    Returns
    -------
    metrics : dict
        Aggregated metrics.
    predictions : pd.DataFrame, optional
        All predictions (only if return_predictions=True).

    Examples
    --------
    >>> from ad_inventory_forecast.models import InventoryForecaster
    >>> metrics = cross_validate(InventoryForecaster, y, n_folds=5)
    >>> print(f"SMAPE: {metrics['smape']:.2f}%")
    """
    validator = WalkForwardValidator(
        n_folds=n_folds,
        horizon=horizon,
        min_train_size=min_train_size,
    )

    result = validator.validate(model, y, model_params=model_params)

    if return_predictions:
        # Combine all predictions
        all_preds = []
        for fold in result.fold_results:
            df = pd.DataFrame({
                "fold": fold.fold,
                "actual": fold.actual,
                "predicted": fold.predicted,
            })
            all_preds.append(df)

        predictions_df = pd.concat(all_preds, ignore_index=False)
        return result.metrics, predictions_df

    return result.metrics


def time_series_split(
    y: pd.Series,
    n_splits: int = 5,
    test_size: int = 30,
    gap: int = 0,
) -> List[Tuple[np.ndarray, np.ndarray]]:
    """
    Generate train/test splits for time series cross-validation.

    Similar to sklearn's TimeSeriesSplit but returns index arrays.

    Parameters
    ----------
    y : pd.Series
        Time series.
    n_splits : int, default 5
        Number of splits.
    test_size : int, default 30
        Size of each test set.
    gap : int, default 0
        Gap between train and test.

    Yields
    ------
    train_idx : np.ndarray
        Training indices.
    test_idx : np.ndarray
        Test indices.

    Examples
    --------
    >>> for train_idx, test_idx in time_series_split(y, n_splits=3):
    ...     y_train, y_test = y.iloc[train_idx], y.iloc[test_idx]
    """
    n = len(y)
    indices = np.arange(n)

    # Calculate split points
    test_starts = []
    for i in range(n_splits):
        test_end = n - i * test_size
        test_start = test_end - test_size
        if test_start - gap > 0:
            test_starts.append(test_start)

    test_starts.reverse()

    for test_start in test_starts:
        train_end = test_start - gap
        test_end = test_start + test_size

        train_idx = indices[:train_end]
        test_idx = indices[test_start:test_end]

        yield train_idx, test_idx
