"""
MonthWise Forecaster
====================
A practical, transparent demand forecasting method that treats each calendar
month as its own independent forecasting problem.

Instead of fitting a single model to the entire time series, MonthWise looks at
how many same-calendar-month historical data points exist and applies the
appropriate rule:

    Rule 1 (1 data point)  : Growth projection from single value
    Rule 2 (2 data points) : Trend extrapolation with ±50% cap
    Rule 3 (3+ data points): Linear regression with 1.5× cap

Requires minimum 12 months of monthly history.
"""

import numpy as np
import pandas as pd
from dateutil.relativedelta import relativedelta
from scipy import stats


class MonthWiseForecaster:
    """
    MonthWise Forecaster — calendar-month-aware demand forecasting.

    Parameters
    ----------
    forecast_horizon : int, default=18
        Number of months to forecast into the future.
    rule1_growth_step : float, default=0.10
        Growth increment per future occurrence when only one historical
        same-month value exists. For example, 0.10 means the 1st future
        occurrence is +10%, the 2nd is +20%, the 3rd is +30%, etc.

    Examples
    --------
    >>> from monthwise_forecast import MonthWiseForecaster
    >>> import pandas as pd
    >>>
    >>> # Create monthly series
    >>> dates = pd.date_range("2024-01-01", periods=18, freq="MS")
    >>> values = [10, 12, 8, 15, 20, 18, 11, 14, 9, 16, 22, 19, 12, 15, 10, 17, 24, 21]
    >>> series = pd.Series(values, index=dates)
    >>>
    >>> # Forecast
    >>> f = MonthWiseForecaster(forecast_horizon=18)
    >>> result = f.forecast(series)
    >>> print(result)
    >>>
    >>> # Chunked forecast (6-month buckets)
    >>> chunks = f.forecast_chunked(series, chunk_size=6)
    >>> print(chunks)
    >>>
    >>> # Batch forecast for multiple groups
    >>> df = pd.DataFrame({
    ...     "group": ["A"]*18 + ["B"]*18,
    ...     "month_start": list(dates)*2,
    ...     "units": values + [v*2 for v in values]
    ... })
    >>> batch_result = f.forecast_batch(
    ...     df, group_cols=["group"], date_col="month_start", value_col="units"
    ... )
    >>> print(batch_result)
    """

    # Fixed rules — not configurable
    _MIN_HISTORY = 12
    _RULE2_TREND_CAP = 0.50
    _RULE3_CAP_MULTIPLIER = 1.5

    def __init__(self, forecast_horizon=18, rule1_growth_step=0.10):
        if not isinstance(forecast_horizon, int) or forecast_horizon < 1:
            raise ValueError("forecast_horizon must be a positive integer.")
        if not isinstance(rule1_growth_step, (int, float)) or rule1_growth_step <= 0:
            raise ValueError("rule1_growth_step must be a positive number.")

        self.forecast_horizon = forecast_horizon
        self.rule1_growth_step = float(rule1_growth_step)

    def __repr__(self):
        return (
            f"MonthWiseForecaster("
            f"forecast_horizon={self.forecast_horizon}, "
            f"rule1_growth_step={self.rule1_growth_step})"
        )

    # ------------------------------------------------------------------ #
    #  Internal helpers                                                    #
    # ------------------------------------------------------------------ #

    @staticmethod
    def _to_month_start(ts):
        return pd.Timestamp(ts).to_period("M").to_timestamp()

    @staticmethod
    def _safe_round_int(val):
        return max(0, int(round(val)))

    def _forecast_single_month(self, target_year, historical_same_month_values,
                                max_hist_value, nth_future):
        """
        Forecast one future month.

        Parameters
        ----------
        target_year : int
            The year being forecasted.
        historical_same_month_values : dict
            {year: value} for this calendar month.
        max_hist_value : float
            Max of all historical same-month values (for Rule 3 cap).
        nth_future : int
            Which future occurrence this is (1st, 2nd, 3rd...).

        Returns
        -------
        float (unrounded)
        """
        n = len(historical_same_month_values)

        if n == 0:
            return 0.0

        # ---- Rule 1: single historical value ----
        if n == 1:
            base_val = list(historical_same_month_values.values())[0]
            return base_val * (1 + self.rule1_growth_step * nth_future)

        # ---- Rule 2: two historical values ----
        if n == 2:
            years_sorted = sorted(historical_same_month_values.keys())
            val_old = historical_same_month_values[years_sorted[0]]
            val_new = historical_same_month_values[years_sorted[1]]

            if val_old == 0:
                # Fallback to Rule 1
                return val_new * (1 + self.rule1_growth_step * nth_future)

            pct_change = (val_new - val_old) / val_old
            pct_change = np.clip(pct_change, -self._RULE2_TREND_CAP, self._RULE2_TREND_CAP)

            projected = val_new * ((1 + pct_change) ** nth_future)
            return max(0.0, projected)

        # ---- Rule 3: three or more historical values ----
        years = np.array(sorted(historical_same_month_values.keys()), dtype=float)
        values = np.array(
            [historical_same_month_values[y] for y in sorted(historical_same_month_values.keys())],
            dtype=float,
        )

        slope, intercept, _, _, _ = stats.linregress(years, values)
        projected = intercept + slope * target_year

        projected = max(0.0, projected)
        if max_hist_value > 0:
            projected = min(projected, self._RULE3_CAP_MULTIPLIER * max_hist_value)

        return projected

    # ------------------------------------------------------------------ #
    #  Public API                                                          #
    # ------------------------------------------------------------------ #

    def forecast(self, series):
        """
        Forecast a single monthly time series.

        Parameters
        ----------
        series : pd.Series
            Monthly data indexed by month-start Timestamps. Values should be
            numeric (units, revenue, etc.). Index must be datetime-like.

        Returns
        -------
        pd.DataFrame
            Columns: month_start, forecast

        Raises
        ------
        ValueError
            If series has fewer than 12 months of data.
        TypeError
            If series is not a pandas Series with a datetime index.
        """
        # Validation
        if not isinstance(series, pd.Series):
            raise TypeError("Input must be a pandas Series.")
        if not isinstance(series.index, pd.DatetimeIndex):
            series = series.copy()
            series.index = pd.to_datetime(series.index)
        if len(series) < self._MIN_HISTORY:
            raise ValueError(
                f"Need at least {self._MIN_HISTORY} months of history. "
                f"Got {len(series)}."
            )

        series = series.sort_index().astype(float)
        last_hist_month = series.index.max()

        # Build lookup: {calendar_month: {year: value}}
        month_year_lookup = {}
        for ts, val in series.items():
            m = ts.month
            y = ts.year
            if m not in month_year_lookup:
                month_year_lookup[m] = {}
            month_year_lookup[m][y] = float(val)

        # Track how many times each calendar month has been forecasted
        future_month_count = {}

        results = []
        for i in range(1, self.forecast_horizon + 1):
            fm = self._to_month_start(last_hist_month + relativedelta(months=i))
            cal_month = fm.month
            cal_year = fm.year

            hist_values = month_year_lookup.get(cal_month, {})

            # Track nth future occurrence of this calendar month
            future_month_count[cal_month] = future_month_count.get(cal_month, 0) + 1
            nth_future = future_month_count[cal_month]

            max_hist = max(hist_values.values()) if hist_values else 0.0

            fc = self._forecast_single_month(
                target_year=cal_year,
                historical_same_month_values=hist_values,
                max_hist_value=max_hist,
                nth_future=nth_future,
            )
            results.append((fm, fc))

        # Round at the very end
        df = pd.DataFrame(results, columns=["month_start", "forecast"])
        df["forecast"] = df["forecast"].apply(self._safe_round_int)
        return df

    def forecast_chunked(self, series, chunk_size=6):
        """
        Forecast and aggregate into chunks (e.g., 6-month buckets).

        Parameters
        ----------
        series : pd.Series
            Monthly data indexed by month-start Timestamps.
        chunk_size : int, default=6
            Number of months per chunk.

        Returns
        -------
        pd.DataFrame
            One row with columns like '1-6', '7-12', '13-18', etc.
        """
        fc = self.forecast(series)
        values = fc["forecast"].values

        chunks = {}
        for start in range(0, len(values), chunk_size):
            end = min(start + chunk_size, len(values))
            label = f"{start + 1}-{end}"
            chunks[label] = int(values[start:end].sum())

        return pd.DataFrame([chunks])

    def forecast_batch(self, df, group_cols, date_col="month_start",
                        value_col="units"):
        """
        Forecast for multiple groups in a DataFrame.

        Parameters
        ----------
        df : pd.DataFrame
            Long-format DataFrame with group columns, a date column, and a
            value column.
        group_cols : list of str
            Column name(s) that define each group (e.g., ["material"] or
            ["material", "planning_channel"]).
        date_col : str, default="month_start"
            Name of the date column (must be datetime or convertible).
        value_col : str, default="units"
            Name of the value column to forecast.

        Returns
        -------
        pd.DataFrame
            Long-format DataFrame with columns: *group_cols, month_start,
            forecast. Groups with fewer than 12 months are silently skipped.
        """
        if isinstance(group_cols, str):
            group_cols = [group_cols]

        df = df.copy()
        df[date_col] = pd.to_datetime(df[date_col])

        all_results = []
        groups = df.groupby(group_cols, sort=False)
        skipped = 0

        for keys, g in groups:
            if not isinstance(keys, tuple):
                keys = (keys,)

            g = g.sort_values(date_col)
            series = g.set_index(date_col)[value_col]

            if len(series) < self._MIN_HISTORY:
                skipped += 1
                continue

            fc = self.forecast(series)

            for col, val in zip(group_cols, keys):
                fc[col] = val

            all_results.append(fc)

        if skipped > 0:
            print(f"Note: {skipped} group(s) skipped (fewer than {self._MIN_HISTORY} months).")

        if not all_results:
            return pd.DataFrame(columns=group_cols + ["month_start", "forecast"])

        result = pd.concat(all_results, ignore_index=True)
        return result[group_cols + ["month_start", "forecast"]]

    def forecast_batch_chunked(self, df, group_cols, date_col="month_start",
                                value_col="units", chunk_size=6):
        """
        Forecast for multiple groups and return chunked aggregates.

        Parameters
        ----------
        df : pd.DataFrame
            Long-format DataFrame.
        group_cols : list of str
            Column name(s) that define each group.
        date_col : str, default="month_start"
            Date column name.
        value_col : str, default="units"
            Value column name.
        chunk_size : int, default=6
            Months per chunk.

        Returns
        -------
        pd.DataFrame
            One row per group with columns: *group_cols, '1-6', '7-12', etc.
        """
        if isinstance(group_cols, str):
            group_cols = [group_cols]

        df = df.copy()
        df[date_col] = pd.to_datetime(df[date_col])

        all_rows = []
        groups = df.groupby(group_cols, sort=False)
        skipped = 0

        for keys, g in groups:
            if not isinstance(keys, tuple):
                keys = (keys,)

            g = g.sort_values(date_col)
            series = g.set_index(date_col)[value_col]

            if len(series) < self._MIN_HISTORY:
                skipped += 1
                continue

            chunks = self.forecast_chunked(series, chunk_size=chunk_size)
            for col, val in zip(group_cols, keys):
                chunks[col] = val

            all_rows.append(chunks)

        if skipped > 0:
            print(f"Note: {skipped} group(s) skipped (fewer than {self._MIN_HISTORY} months).")

        if not all_rows:
            return pd.DataFrame()

        result = pd.concat(all_rows, ignore_index=True)
        # Reorder: group cols first
        chunk_cols = [c for c in result.columns if c not in group_cols]
        return result[group_cols + chunk_cols]