"""Fertility: 7-group aggregation and the Lee-Tuljapurkar benchmark.

Meseguer (2010) models fertility for 7 age groups within a Lee-Carter age
decomposition, but forecasts the total fertility rate (TFR) as a **constrained
ARMA(1,1)** whose unconditional mean is fixed at an ultimate value F* (he uses
F*=2, the SSA Alternative-2 assumption) -- the Lee (1993) / Lee-Tuljapurkar (1994)
approach, since fertility is mean-reverting rather than trending.
"""

import numpy as np
import pandas as pd
from scipy.optimize import brentq
from statsmodels.tsa.arima.model import ARIMA

from hmd_io import asfr_by_age
from mortality import lee_carter

FERT_LABELS = ["<20", "20-24", "25-29", "30-34", "35-39", "40-44", "45+"]
_EDGES = [25, 30, 35, 40, 45]           # upper edges of the middle groups


def _gid(age):
    if age < 20:
        return 0
    for i, e in enumerate(_EDGES, 1):
        if age < e:
            return i
    return 6


def group_asfr(path):
    """Aggregate single-year period ASFR into the 7 fertility groups.

    Group rate = sum of the annual age-specific rates in the group, so
    sum over groups = TFR. Returns DataFrame[Year, group label].
    """
    asfr = asfr_by_age(path)
    gid = np.array([_gid(a) for a in asfr.columns.to_numpy()])
    return pd.DataFrame({FERT_LABELS[k]: asfr.values[:, gid == k].sum(1) for k in range(7)},
                        index=asfr.index)


def tfr_from_groups(G):
    """TFR = sum of the 7 group rates."""
    return G.sum(axis=1)


def tfr_to_k(a, b, target):
    """Invert TFR(k)=sum_g exp(a_g+b_g k)=target for the Lee-Carter index k."""
    f = lambda k: np.sum(np.exp(a + b * k)) - target
    return brentq(f, -80, 80)


def constrained_arma_forecast(tfr, Fstar, H, alpha=0.10):
    """Constrained ARMA(1,1) for the TFR with fixed ultimate mean F* (Meseguer eq. 8).

    Fit a zero-mean ARMA(1,1) to (TFR - F*), so forecasts revert to F*.
    Returns (mean, lo, hi, params).
    """
    x = tfr.values - Fstar
    m = ARIMA(x, order=(1, 0, 1), trend="n").fit()
    fc = m.get_forecast(H)
    ci = fc.conf_int(alpha=alpha)
    return (fc.predicted_mean + Fstar, ci[:, 0] + Fstar, ci[:, 1] + Fstar, m.params)


def constrained_arma_draws(tfr, Fstar, H, ndraw, seed=0):
    """Sample TFR paths from the constrained ARMA(1,1) (mean fixed at F*).

    Simulates x_t = phi x_{t-1} + e_t + theta e_{t-1} forward from the last fitted
    state, then adds F*. Returns (ndraw, H) TFR path draws — the fertility analogue
    of the BVAR posterior-predictive draws, for full uncertainty propagation.
    """
    x = tfr.values - Fstar
    res = ARIMA(x, order=(1, 0, 1), trend="n").fit()
    phi = res.arparams[0] if len(res.arparams) else 0.0
    theta = res.maparams[0] if len(res.maparams) else 0.0
    sigma = float(np.sqrt(res.params[-1]))
    x_last, e_last = x[-1], res.resid[-1]
    rng = np.random.default_rng(seed)
    out = np.empty((ndraw, H))
    for d in range(ndraw):
        xp, ep = x_last, e_last
        for h in range(H):
            e = rng.normal(0, sigma)
            xp = phi * xp + e + theta * ep
            out[d, h] = xp + Fstar
            ep = e
    return out


def rw_drift_forecast(tfr, H):
    """Unconstrained random-walk-with-drift TFR forecast (for contrast)."""
    d = (tfr.values[-1] - tfr.values[0]) / (len(tfr) - 1)
    return tfr.values[-1] + np.arange(1, H + 1) * d


def reconstruct_group_rates(base_rates, b, tfr_path):
    """Distribute a forecast TFR path across the 7 groups using the LC age pattern.

    Jump-off adjusted (Lee-Miller style): anchored to the observed last-year rates
    `base_rates` rather than the LC-fitted ones, so the forecast is continuous with
    the data. For each year we scale
        f_g(t) = base_g * exp(b_g * delta_t),
    solving delta_t so that sum_g f_g(t) equals the forecast TFR. At the jump-off
    (TFR = sum base_g) delta = 0, reproducing the observed rates exactly.
    Returns an (H x 7) matrix.
    """
    base_rates, b = np.asarray(base_rates), np.asarray(b)
    out = np.empty((len(tfr_path), len(base_rates)))
    for h, T in enumerate(tfr_path):
        f = lambda d: np.sum(base_rates * np.exp(b * d)) - T
        delta = brentq(f, -80, 80)
        out[h] = base_rates * np.exp(b * delta)
    return out