"""Year-over-Year copy forecaster for spike-seasonal daily series."""

import warnings
from typing import Optional, Tuple, Union

import numpy as np
import pandas as pd
from scipy.stats import norm

from ad_inventory_forecast.core.base import BaseForecaster


class YoYCopyForecaster(BaseForecaster):
    """
    Year-over-Year copy forecaster for spike-seasonal daily series.

    Forecasts by copying last year's daily actuals scaled by a YoY growth
    factor. Preserves spike shapes implicitly via the template year, making
    it well-suited for high-volatility series with recurring seasonal spikes
    (e.g., sports events, holidays).

    Parameters
    ----------
    yoy_method : str, default "median"
        Aggregation method for computing the YoY growth factor from
        (iso-week, day-of-week) matched pairs. One of "median" or "mean".
    growth_clip : tuple, default (0.5, 3.0)
        (min, max) bounds for clamping the computed YoY ratio. Guards
        against extreme values from data anomalies.
    min_obs : int, default 365
        Minimum number of daily observations required to fit. Raises
        ValueError if len(y) < min_obs.
    """

    def __init__(
        self,
        yoy_method: str = "median",
        growth_clip: Tuple[float, float] = (0.5, 3.0),
        min_obs: int = 365,
    ):
        super().__init__()
        self.yoy_method = yoy_method
        self.growth_clip = growth_clip
        self.min_obs = min_obs

    def fit(
        self,
        y: Union[pd.Series, np.ndarray],
        X=None,
    ) -> "YoYCopyForecaster":
        """
        Fit the forecaster to training data.

        Parameters
        ----------
        y : pd.Series
            Daily time series with DatetimeIndex.
        X : ignored

        Returns
        -------
        self
        """
        y = self._validate_y(y)

        # Require daily frequency
        freq = self._infer_frequency(y)
        if freq != "D":
            raise ValueError(
                f"YoYCopyForecaster requires daily data, got frequency '{freq}'. "
                "Resample to daily before fitting."
            )

        if len(y) < self.min_obs:
            raise ValueError(
                f"YoYCopyForecaster requires >= {self.min_obs} daily observations, "
                f"got {len(y)}."
            )

        self.freq_ = freq
        self.n_obs_ = len(y)
        self.training_series_ = y

        # --- Compute YoY growth factor ---
        # Match each date in the last 365 days to its equivalent 52 iso-weeks prior
        recent = y.iloc[-365:]
        ratios = []

        for date, value in recent.items():
            if value <= 0:
                continue
            iso = date.isocalendar()
            target_week = iso.week
            target_dow = iso.weekday  # 1=Mon … 7=Sun

            # Find matching dates in training data that are ~52 weeks earlier
            prior_window_start = date - pd.Timedelta(weeks=54)
            prior_window_end = date - pd.Timedelta(weeks=50)
            prior_candidates = y.loc[prior_window_start:prior_window_end]

            for prior_date, prior_value in prior_candidates.items():
                if prior_value <= 0:
                    continue
                prior_iso = prior_date.isocalendar()
                if (
                    prior_iso.week == target_week
                    and prior_iso.weekday == target_dow
                ):
                    ratios.append(value / prior_value)
                    break  # one match per recent date is sufficient

        if not ratios:
            warnings.warn(
                "YoYCopyForecaster: no prior-year matches found — "
                "only 1 year of data available. Defaulting yoy_growth_factor_ to 1.0.",
                UserWarning,
                stacklevel=2,
            )
            self.yoy_growth_factor_ = 1.0
        else:
            if self.yoy_method == "mean":
                raw_factor = float(np.mean(ratios))
            else:
                raw_factor = float(np.median(ratios))

            self.yoy_growth_factor_ = float(
                np.clip(raw_factor, self.growth_clip[0], self.growth_clip[1])
            )

        # --- Build template keyed by (iso_week, day_of_week) ---
        # Use last 366 days so we cover a full leap year
        template_series = y.iloc[-366:]
        template = {}
        for date, value in template_series.items():
            iso = date.isocalendar()
            key = (iso.week, iso.weekday)
            # Prefer more recent observations when there are duplicates
            template[key] = (date, float(value))

        self.template_ = template

        # --- Residual CV via leave-one-year-out on template year ---
        residuals = []
        prior_year = y.iloc[: -366] if len(y) > 366 else None

        if prior_year is not None and len(prior_year) >= 1:
            for date, actual in template_series.items():
                if actual <= 0:
                    continue
                iso = date.isocalendar()
                key = (iso.week, iso.weekday)

                # Find prior year match
                pw_start = date - pd.Timedelta(weeks=54)
                pw_end = date - pd.Timedelta(weeks=50)
                candidates = prior_year.loc[
                    pw_start:pw_end
                ] if pw_start <= prior_year.index[-1] else pd.Series(dtype=float)

                prior_val = None
                for pd_date, pv in candidates.items():
                    if pv <= 0:
                        continue
                    p_iso = pd_date.isocalendar()
                    if p_iso.week == key[0] and p_iso.weekday == key[1]:
                        prior_val = float(pv)
                        break

                if prior_val is not None:
                    predicted = prior_val * self.yoy_growth_factor_
                    residuals.append(actual - predicted)

        if len(residuals) >= 2:
            template_mean = float(template_series.mean())
            self.residual_cv_ = (
                float(np.std(residuals)) / template_mean
                if template_mean > 0
                else 0.1
            )
        else:
            # Default to a modest CV when we can't compute residuals
            self.residual_cv_ = 0.10

        self.is_fitted_ = True
        return self

    def predict(
        self,
        horizon: int,
        return_conf_int: bool = False,
        alpha: float = 0.05,
        return_df: bool = True,
    ) -> Union[pd.Series, pd.DataFrame, Tuple[pd.Series, pd.DataFrame]]:
        """
        Generate daily forecasts by copying scaled last-year actuals.

        Parameters
        ----------
        horizon : int
            Number of daily periods to forecast.
        return_conf_int : bool, default False
            Legacy parameter; ignored when return_df=True.
        alpha : float, default 0.05
            Significance level for confidence intervals (0.05 → 95% CI).
        return_df : bool, default True
            Return DataFrame with standard 3-column format.

        Returns
        -------
        pd.DataFrame or pd.Series
        """
        self._check_is_fitted()

        last_date = self.training_series_.index[-1]
        forecast_index = pd.date_range(
            start=last_date + pd.Timedelta(days=1),
            periods=horizon,
            freq="D",
        )

        z = float(norm.ppf(1 - alpha / 2))
        forecasts = []
        lowers = []
        uppers = []

        for d in forecast_index:
            iso = d.isocalendar()
            key = (iso.week, iso.weekday)

            if key in self.template_:
                _, template_val = self.template_[key]
            else:
                # Fall back to nearest calendar date within ±3 days
                template_val = self._nearest_template_value(d)

            point = template_val * self.yoy_growth_factor_
            lower = point * max(0.0, 1.0 - z * self.residual_cv_)
            upper = point * (1.0 + z * self.residual_cv_)

            forecasts.append(point)
            lowers.append(lower)
            uppers.append(upper)

        forecast_series = pd.Series(forecasts, index=forecast_index, name="predicted_impressions")

        if not return_df:
            if return_conf_int:
                ci = pd.DataFrame(
                    {"lower": lowers, "upper": uppers}, index=forecast_index
                )
                return forecast_series, ci
            return forecast_series

        df = pd.DataFrame(
            {
                "predicted_impressions": forecasts,
                "lower_bound": lowers,
                "upper_bound": uppers,
            },
            index=forecast_index,
        )
        return df

    def _nearest_template_value(self, d: pd.Timestamp) -> float:
        """
        Find the nearest template value within ±3 days when exact iso-key
        lookup fails (e.g., iso-week 53 not in template).
        """
        for delta in range(1, 4):
            for sign in (1, -1):
                candidate = d + pd.Timedelta(days=delta * sign)
                c_iso = candidate.isocalendar()
                key = (c_iso.week, c_iso.weekday)
                if key in self.template_:
                    _, val = self.template_[key]
                    return val

        # Last resort: return the overall mean of the template
        vals = [v for _, v in self.template_.values()]
        return float(np.mean(vals)) if vals else 0.0
