"""Demographic reconciliation and hierarchical forecasting."""

from typing import Optional, Union, Dict, List, Literal
from dataclasses import dataclass

import numpy as np
import pandas as pd


@dataclass
class ReconciliationResult:
    """
    Result of demographic reconciliation.

    Attributes
    ----------
    total_forecast : pd.Series
        Total (aggregate) forecast.
    segment_forecasts : pd.DataFrame
        Forecasts for each segment.
    ratios : pd.DataFrame
        Segment ratios used for reconciliation.
    """

    total_forecast: pd.Series
    segment_forecasts: pd.DataFrame
    ratios: pd.DataFrame


class DemographicReconciler:
    """
    Reconcile total forecasts into demographic segments.

    Uses historical segment breakdowns to distribute total forecasts
    across demographic segments (e.g., age groups, genders, regions).

    Supports multiple methods:
    - Top-Down: Distribute total using smoothed segment ratios
    - TSB (Teunter-Syntetos-Babai): Handles intermittent/sparse segments
    - Proportional: Simple historical average ratios

    Parameters
    ----------
    method : {'top_down', 'tsb', 'proportional'}, default 'top_down'
        Reconciliation method:
        - 'top_down': Exponentially smooth segment ratios
        - 'tsb': TSB method for sparse segments
        - 'proportional': Use simple historical averages
    smoothing_alpha : float, default 0.3
        Smoothing parameter for ratio estimation.
    min_ratio : float, default 0.001
        Minimum segment ratio (prevents zero allocations).

    Examples
    --------
    >>> # Historical segment data
    >>> dates = pd.date_range('2024-01-01', periods=90, freq='D')
    >>> segment_data = pd.DataFrame({
    ...     'total': np.random.randint(1000, 2000, 90),
    ...     'segment_a': np.random.randint(400, 800, 90),
    ...     'segment_b': np.random.randint(300, 600, 90),
    ...     'segment_c': np.random.randint(200, 400, 90),
    ... }, index=dates)
    >>>
    >>> reconciler = DemographicReconciler(method='top_down')
    >>> reconciler.fit(segment_data, total_col='total')
    >>>
    >>> # Reconcile a total forecast
    >>> total_forecast = pd.Series([1500, 1600, 1700], index=pd.date_range('2024-04-01', periods=3))
    >>> result = reconciler.reconcile(total_forecast)
    >>> print(result.segment_forecasts)

    Notes
    -----
    For sparse segments (many zeros), use method='tsb' which models
    both the demand size and demand probability separately.
    """

    def __init__(
        self,
        method: Literal["top_down", "tsb", "proportional"] = "top_down",
        smoothing_alpha: float = 0.3,
        min_ratio: float = 0.001,
    ):
        """Initialize the reconciler."""
        self.method = method
        self.smoothing_alpha = smoothing_alpha
        self.min_ratio = min_ratio

        # Fitted attributes
        self.segment_columns_ = None
        self.total_column_ = None
        self.smoothed_ratios_ = None
        self.historical_ratios_ = None
        self._tsb_params = {}

    def fit(
        self,
        data: pd.DataFrame,
        total_col: str = "total",
        segment_cols: Optional[List[str]] = None,
    ) -> "DemographicReconciler":
        """
        Fit the reconciler to historical segment data.

        Parameters
        ----------
        data : pd.DataFrame
            Historical data with total and segment columns.
        total_col : str, default 'total'
            Name of the total column.
        segment_cols : List[str], optional
            Names of segment columns. If None, uses all columns except total.

        Returns
        -------
        self : DemographicReconciler
            Fitted reconciler.
        """
        self.total_column_ = total_col

        if segment_cols is None:
            segment_cols = [col for col in data.columns if col != total_col]

        self.segment_columns_ = segment_cols

        # Calculate historical ratios
        self.historical_ratios_ = self._calculate_ratios(data)

        # Apply method-specific fitting
        if self.method == "top_down":
            self.smoothed_ratios_ = self._smooth_ratios(self.historical_ratios_)
        elif self.method == "tsb":
            self._fit_tsb(data)
            self.smoothed_ratios_ = self._get_tsb_ratios()
        elif self.method == "proportional":
            # Use simple averages
            self.smoothed_ratios_ = pd.Series(
                self.historical_ratios_.mean(),
                index=self.segment_columns_,
            )
        else:
            raise ValueError(f"Unknown method: {self.method}")

        return self

    def reconcile(
        self,
        total_forecast: pd.Series,
        return_ratios: bool = False,
    ) -> ReconciliationResult:
        """
        Reconcile a total forecast into segment forecasts.

        Parameters
        ----------
        total_forecast : pd.Series
            Total forecast to distribute.
        return_ratios : bool, default False
            Whether to return the ratios used.

        Returns
        -------
        result : ReconciliationResult
            Reconciliation result with segment forecasts.
        """
        if self.smoothed_ratios_ is None:
            raise ValueError("Reconciler must be fitted before reconciling")

        # Get ratios for forecast period
        if isinstance(self.smoothed_ratios_, pd.Series):
            # Static ratios
            ratios = pd.DataFrame(
                [self.smoothed_ratios_.values] * len(total_forecast),
                index=total_forecast.index,
                columns=self.segment_columns_,
            )
        else:
            # Time-varying ratios (use last known)
            last_ratios = self.smoothed_ratios_.iloc[-1]
            ratios = pd.DataFrame(
                [last_ratios.values] * len(total_forecast),
                index=total_forecast.index,
                columns=self.segment_columns_,
            )

        # Normalize ratios to sum to 1
        ratio_sums = ratios.sum(axis=1)
        ratios = ratios.div(ratio_sums, axis=0)

        # Apply minimum ratio
        ratios = ratios.clip(lower=self.min_ratio)
        ratios = ratios.div(ratios.sum(axis=1), axis=0)

        # Distribute total forecast
        segment_forecasts = ratios.mul(total_forecast, axis=0)

        return ReconciliationResult(
            total_forecast=total_forecast,
            segment_forecasts=segment_forecasts,
            ratios=ratios,
        )

    def _calculate_ratios(self, data: pd.DataFrame) -> pd.DataFrame:
        """Calculate historical segment ratios."""
        ratios = pd.DataFrame(index=data.index, columns=self.segment_columns_)

        totals = data[self.total_column_]

        for col in self.segment_columns_:
            # Avoid division by zero
            with np.errstate(divide="ignore", invalid="ignore"):
                ratios[col] = data[col] / totals
                ratios[col] = ratios[col].replace([np.inf, -np.inf], np.nan)
                ratios[col] = ratios[col].fillna(0)

        return ratios

    def _smooth_ratios(self, ratios: pd.DataFrame) -> pd.DataFrame:
        """Apply exponential smoothing to ratios."""
        smoothed = pd.DataFrame(index=ratios.index, columns=ratios.columns)

        for col in ratios.columns:
            smoothed[col] = ratios[col].ewm(alpha=self.smoothing_alpha).mean()

        return smoothed

    def _fit_tsb(self, data: pd.DataFrame) -> None:
        """Fit TSB (Teunter-Syntetos-Babai) model for sparse segments."""
        for col in self.segment_columns_:
            segment_values = data[col].values
            total_values = data[self.total_column_].values

            # Calculate demand probability (fraction of non-zero periods)
            non_zero = segment_values > 0
            demand_prob = non_zero.mean()

            # Smooth demand probability
            prob_smooth = self._exponential_smooth(
                non_zero.astype(float), self.smoothing_alpha
            )

            # Calculate average demand size (when non-zero)
            if non_zero.sum() > 0:
                demand_size = segment_values[non_zero].mean()
                # Smooth demand ratio
                ratio_when_nonzero = np.where(
                    total_values > 0,
                    segment_values / total_values,
                    0,
                )
                ratio_when_nonzero = np.where(non_zero, ratio_when_nonzero, np.nan)
                ratio_smooth = self._exponential_smooth_nonzero(
                    ratio_when_nonzero, self.smoothing_alpha
                )
            else:
                demand_size = 0
                ratio_smooth = 0

            self._tsb_params[col] = {
                "demand_prob": prob_smooth,
                "demand_size": demand_size,
                "ratio_smooth": ratio_smooth,
            }

    def _get_tsb_ratios(self) -> pd.Series:
        """Get TSB-based ratios."""
        ratios = {}

        for col in self.segment_columns_:
            params = self._tsb_params[col]
            # TSB forecast = probability × size ratio
            ratios[col] = params["demand_prob"] * params["ratio_smooth"]

        ratios_series = pd.Series(ratios)

        # Normalize
        total = ratios_series.sum()
        if total > 0:
            ratios_series = ratios_series / total

        return ratios_series

    @staticmethod
    def _exponential_smooth(values: np.ndarray, alpha: float) -> float:
        """Apply exponential smoothing and return last value."""
        result = values[0]
        for v in values[1:]:
            result = alpha * v + (1 - alpha) * result
        return result

    @staticmethod
    def _exponential_smooth_nonzero(values: np.ndarray, alpha: float) -> float:
        """Apply exponential smoothing, ignoring NaN values."""
        result = None
        for v in values:
            if not np.isnan(v):
                if result is None:
                    result = v
                else:
                    result = alpha * v + (1 - alpha) * result
        return result if result is not None else 0

    def get_segment_ratios(self) -> pd.Series:
        """
        Get current segment ratios.

        Returns
        -------
        ratios : pd.Series
            Segment ratios (sum to 1).
        """
        if isinstance(self.smoothed_ratios_, pd.Series):
            ratios = self.smoothed_ratios_.copy()
        else:
            ratios = self.smoothed_ratios_.iloc[-1].copy()

        # Normalize
        return ratios / ratios.sum()

    def top_down_forecast(
        self,
        total_forecast: pd.Series,
        historical_breakdowns: pd.DataFrame,
        segment_col: str = "segment_id",
        value_col: str = "impressions",
    ) -> pd.DataFrame:
        """
        SDK-compliant method to distribute total forecast to segments.

        This is a convenience method that combines fit() and reconcile()
        for the common use case of distributing a total forecast.

        Parameters
        ----------
        total_forecast : pd.Series
            Total forecast to distribute across segments.
        historical_breakdowns : pd.DataFrame
            Historical data with segment breakdown. Must have columns
            for segment identifier and values.
        segment_col : str, default 'segment_id'
            Column name for segment identifier.
        value_col : str, default 'impressions'
            Column name for values (impressions).

        Returns
        -------
        segment_forecasts : pd.DataFrame
            DataFrame with forecasts for each segment. Has columns:
            - 'date': Forecast date
            - 'segment_id': Segment identifier
            - 'predicted_impressions': Forecasted impressions

        Examples
        --------
        >>> reconciler = DemographicReconciler(method='tsb')
        >>> segment_forecasts = reconciler.top_down_forecast(
        ...     total_forecast=total_preds,
        ...     historical_breakdowns=segments_history
        ... )
        >>> print(segment_forecasts[segment_forecasts['segment_id'] == 'M_18-34'].head())
        """
        # Pivot historical data to wide format if needed
        if segment_col in historical_breakdowns.columns:
            # Data is in long format - pivot to wide
            if historical_breakdowns.index.name != "date":
                if "date" in historical_breakdowns.columns:
                    historical_breakdowns = historical_breakdowns.set_index("date")

            pivot_data = historical_breakdowns.pivot_table(
                index=historical_breakdowns.index,
                columns=segment_col,
                values=value_col,
                aggfunc="sum",
            ).fillna(0)

            # Calculate total from segments
            pivot_data["total"] = pivot_data.sum(axis=1)

            # Fit the reconciler
            self.fit(pivot_data, total_col="total")
        else:
            # Assume data is already in wide format
            # Find total column
            if "total" in historical_breakdowns.columns:
                self.fit(historical_breakdowns, total_col="total")
            else:
                raise ValueError(
                    f"Cannot find segment column '{segment_col}' or 'total' column. "
                    "Provide data with segment_id column or pre-pivoted data with total."
                )

        # Reconcile
        result = self.reconcile(total_forecast)

        # Convert to long format for SDK compatibility
        segment_forecasts = result.segment_forecasts.reset_index()
        segment_forecasts = segment_forecasts.melt(
            id_vars=["index"],
            var_name="segment_id",
            value_name="predicted_impressions",
        )
        segment_forecasts = segment_forecasts.rename(columns={"index": "date"})

        return segment_forecasts


class CrostonForecaster:
    """
    Croston's method for intermittent demand.

    Suitable for segments with many zero values (sparse demand).

    Parameters
    ----------
    alpha : float, default 0.3
        Smoothing parameter.

    Examples
    --------
    >>> values = pd.Series([0, 0, 5, 0, 0, 0, 3, 0, 0, 4, 0, 0])
    >>> forecaster = CrostonForecaster()
    >>> forecaster.fit(values)
    >>> forecast = forecaster.predict(horizon=5)
    """

    def __init__(self, alpha: float = 0.3):
        """Initialize Croston forecaster."""
        self.alpha = alpha

        self.level_ = None
        self.interval_ = None

    def fit(self, y: pd.Series) -> "CrostonForecaster":
        """
        Fit Croston's method.

        Parameters
        ----------
        y : pd.Series
            Time series (can contain zeros).

        Returns
        -------
        self : CrostonForecaster
            Fitted forecaster.
        """
        values = y.values

        # Find non-zero values and their intervals
        non_zero_idx = np.where(values > 0)[0]

        if len(non_zero_idx) == 0:
            self.level_ = 0
            self.interval_ = 1
            return self

        # Initialize with first non-zero
        self.level_ = values[non_zero_idx[0]]
        self.interval_ = non_zero_idx[0] + 1 if non_zero_idx[0] > 0 else 1

        # Update with subsequent non-zeros
        for i in range(1, len(non_zero_idx)):
            current_idx = non_zero_idx[i]
            prev_idx = non_zero_idx[i - 1]

            # Demand size
            demand = values[current_idx]
            self.level_ = self.alpha * demand + (1 - self.alpha) * self.level_

            # Inter-demand interval
            interval = current_idx - prev_idx
            self.interval_ = self.alpha * interval + (1 - self.alpha) * self.interval_

        return self

    def predict(self, horizon: int) -> pd.Series:
        """
        Predict future demand rate.

        Parameters
        ----------
        horizon : int
            Number of periods.

        Returns
        -------
        forecast : pd.Series
            Constant forecast (demand rate).
        """
        # Croston forecast = level / interval
        if self.interval_ > 0:
            demand_rate = self.level_ / self.interval_
        else:
            demand_rate = 0

        return pd.Series([demand_rate] * horizon, name="forecast")


class TSBForecaster:
    """
    Teunter-Syntetos-Babai (TSB) method for intermittent demand.

    Improves on Croston by updating probability every period (not just
    on demand occurrence).

    Parameters
    ----------
    alpha : float, default 0.3
        Smoothing parameter for demand size.
    beta : float, default 0.3
        Smoothing parameter for demand probability.

    Examples
    --------
    >>> values = pd.Series([0, 0, 5, 0, 0, 0, 3, 0, 0, 4, 0, 0])
    >>> forecaster = TSBForecaster()
    >>> forecaster.fit(values)
    >>> forecast = forecaster.predict(horizon=5)
    """

    def __init__(self, alpha: float = 0.3, beta: float = 0.3):
        """Initialize TSB forecaster."""
        self.alpha = alpha
        self.beta = beta

        self.level_ = None
        self.prob_ = None

    def fit(self, y: pd.Series) -> "TSBForecaster":
        """
        Fit TSB method.

        Parameters
        ----------
        y : pd.Series
            Time series (can contain zeros).

        Returns
        -------
        self : TSBForecaster
            Fitted forecaster.
        """
        values = y.values

        # Initialize
        non_zero = values > 0
        if non_zero.sum() == 0:
            self.level_ = 0
            self.prob_ = 0
            return self

        self.level_ = values[non_zero].mean()
        self.prob_ = non_zero.mean()

        # Update for each period
        for i, val in enumerate(values):
            # Update probability every period
            indicator = 1.0 if val > 0 else 0.0
            self.prob_ = self.beta * indicator + (1 - self.beta) * self.prob_

            # Update level only when demand occurs
            if val > 0:
                self.level_ = self.alpha * val + (1 - self.alpha) * self.level_

        return self

    def predict(self, horizon: int) -> pd.Series:
        """
        Predict future demand rate.

        Parameters
        ----------
        horizon : int
            Number of periods.

        Returns
        -------
        forecast : pd.Series
            Constant forecast (expected demand).
        """
        # TSB forecast = probability × level
        expected_demand = self.prob_ * self.level_

        return pd.Series([expected_demand] * horizon, name="forecast")
