"""Abridged (Meseguer) mortality: age grouping, abridged life table, Lee-Carter.

Meseguer (2010) models all-cause mortality for each sex over 22 age groups:
    age 0, 1-4, 5-9, 10-14, ..., 95-99, 100+.
This module aggregates the single-year HMD death rates into those groups, builds
an abridged period life table (for life expectancy), and fits the Lee-Carter
benchmark model.
"""

import numpy as np
import pandas as pd

from lifetable import a0_andreev_kingkade

# Meseguer's 22 groups as (left age, width n); the open group uses width np.inf.
GROUP_LO = np.array([0, 1] + list(range(5, 100, 5)) + [100])
GROUP_N = np.array([1, 4] + [5] * len(range(5, 100, 5)) + [np.inf])
GROUP_LABEL = (["0", "1-4"] +
               [f"{a}-{a+4}" for a in range(5, 100, 5)] + ["100+"])


def aggregate_to_groups(mx, deaths, pop, sex):
    """Aggregate single-year m_x to the 22 groups: n_mx = sum(deaths) / sum(exposure).

    Exposure is recovered exactly from HMD's own rates as E_x = D_x / m_x (falling
    back to Jan-1 population where m_x = 0). Returns DataFrame[Year, group label].
    """
    Dx, Mx = deaths[sex], mx[sex]
    # Jan-1 population aligned to the mortality years (drop the 1959-/1959+ split).
    Px = pop[sex].copy()
    Px = Px[[str(y).isdigit() for y in Px.index]]
    Px.index = Px.index.astype(int)
    Px = Px.reindex(Mx.index)

    Ex = (Dx / Mx).where(Mx > 0, Px)          # person-years exposure
    ages = Mx.columns.to_numpy()

    rows_m = {}
    for lo, n, lab in zip(GROUP_LO, GROUP_N, GROUP_LABEL):
        hi = ages.max() + 1 if np.isinf(n) else lo + n
        cols = ages[(ages >= lo) & (ages < hi)]
        rows_m[lab] = Dx[cols].sum(axis=1) / Ex[cols].sum(axis=1)
    out = pd.DataFrame(rows_m)
    out.columns.name = "group"
    return out


def _nax(nmx, sex):
    """Average years lived in each interval by those dying in it (separation factors)."""
    a = np.empty(len(nmx))
    a[0] = a0_andreev_kingkade(nmx[0], sex)   # age 0
    a[1] = 1.5                                # 1-4 (n=4), standard approx
    a[2:-1] = 2.5                             # 5-year groups
    a[-1] = 1.0 / nmx[-1] if nmx[-1] > 0 else 2.5   # open group
    return a


def life_table_abridged(nmx, sex, radix=100000.0):
    """Abridged period life table from grouped rates n_mx. Returns dict incl. 'ex'."""
    nmx = np.asarray(nmx, float)
    n = GROUP_N.copy()
    ax = _nax(nmx, sex)

    qx = np.empty(len(nmx))
    qx[:-1] = (n[:-1] * nmx[:-1]) / (1 + (n[:-1] - ax[:-1]) * nmx[:-1])
    qx[-1] = 1.0                              # open group: everyone dies
    k = len(nmx)
    lx = np.empty(k); lx[0] = radix
    for i in range(1, k):
        lx[i] = lx[i - 1] * (1 - qx[i - 1])
    dx = np.empty(k); dx[:-1] = lx[:-1] - lx[1:]; dx[-1] = lx[-1]

    Lx = np.empty(k)
    Lx[:-1] = n[:-1] * lx[1:] + ax[:-1] * dx[:-1]
    Lx[-1] = lx[-1] / nmx[-1] if nmx[-1] > 0 else lx[-1] * ax[-1]
    Tx = np.cumsum(Lx[::-1])[::-1]
    ex = Tx / lx
    return {"nmx": nmx, "qx": qx, "lx": lx, "ex": ex, "e0": ex[0]}


def ex_series(group_mx, sex, age=0):
    """Life expectancy at exact `age` (a group boundary) for each year.

    `age` must be one of the group start ages (0, 1, 5, 10, ..., 95, 100); e.g.
    age=65 gives e65 (period life expectancy at exact age 65).
    """
    idx = int(np.where(GROUP_LO == age)[0][0])
    return pd.Series({y: life_table_abridged(group_mx.loc[y].to_numpy(), sex)["ex"][idx]
                      for y in group_mx.index}, name=f"e{age}")


def e0_series(group_mx, sex):
    """Life expectancy at birth for each year (convenience wrapper for ex_series)."""
    return ex_series(group_mx, sex, 0).rename("e0")


def lee_carter(log_m):
    """Fit Lee-Carter log m_{t,x} = a_x + b_x k_t by SVD (constraints sum b=1, sum k=0).

    Parameters
    ----------
    log_m : DataFrame[Year, age group] of log central death rates (the sample).

    Returns
    -------
    dict with Series a (by group), b (by group), k (by year).
    """
    groups = log_m.columns
    years = log_m.index
    a = log_m.mean(axis=0)                     # a_x = mean over time
    Z = (log_m - a).to_numpy()                 # centred, T x A
    U, S, Vt = np.linalg.svd(Z, full_matrices=False)
    b = Vt[0]
    k = U[:, 0] * S[0]
    c = b.sum()
    b = b / c
    k = k * c
    kmean = k.mean()                           # enforce sum(k)=0, fold into a
    a = a + pd.Series(b * kmean, index=groups)
    k = k - kmean
    return {"a": a, "b": pd.Series(b, index=groups), "k": pd.Series(k, index=years)}


def forecast_lee_carter(fit, horizon):
    """Random-walk-with-drift forecast of k and the implied log-mortality.

    Returns (k_forecast Series indexed 1..horizon-offset, drift, sigma) plus a
    function to turn a k value into a group-rate vector. The caller supplies the
    forecast years.
    """
    k = fit["k"].to_numpy()
    T = len(k)
    drift = (k[-1] - k[0]) / (T - 1)
    resid = np.diff(k) - drift
    sigma = resid.std(ddof=1)
    h = np.arange(1, horizon + 1)
    kf = k[-1] + h * drift
    se = sigma * np.sqrt(h)                     # RW-with-drift forecast s.e.
    return {"k": kf, "se": se, "drift": drift, "sigma": sigma, "h": h}


def rates_from_k(fit, k_values):
    """Group central death rates exp(a + b*k) for one or many k values."""
    a, b = fit["a"].to_numpy(), fit["b"].to_numpy()
    k_values = np.atleast_1d(k_values)
    return np.exp(a[None, :] + np.outer(k_values, b))   # (len(k), A)


def project_rates(fit, k_values, logm_jumpoff=None):
    """Forecast group rates from projected k.

    With `logm_jumpoff` (the *observed* log-rates in the last fitted year) this
    applies the Lee-Miller (2001) jump-off adjustment: the forecast is anchored to
    the actual last-year rates rather than the LC-fitted ones, so h=0 reproduces
    the observed jump-off year exactly (no discontinuity) and the LC trend rides on
    top. Without it, rates are the plain exp(a + b*k).
    """
    b = fit["b"].to_numpy()
    kv = np.atleast_1d(k_values)
    if logm_jumpoff is None:
        return np.exp(fit["a"].to_numpy()[None, :] + np.outer(kv, b))
    klast = fit["k"].iloc[-1]
    return np.exp(np.asarray(logm_jumpoff)[None, :] + np.outer(kv - klast, b))


# ---------------------------------------------------------------------------
# Li-Lee coherent ("common factor") Lee-Carter
# ---------------------------------------------------------------------------
#   log m_i(x,t) = a_i(x) + B(x) K(t) + b_i(x) k_i(t)
# B(x)K(t) is a factor shared by both sexes (forecast as random walk with drift,
# so the common trend continues); b_i(x)k_i(t) is each sex's deviation, whose
# index k_i(t) is modelled as a stationary AR(1) so the sexes cannot drift apart
# in the long run -- the male/female gap stays bounded (coherence).

def _svd_factor(Z):
    """Rank-1 LC factor of a centred T x A matrix: returns (b over A, k over T)."""
    U, S, Vt = np.linalg.svd(Z, full_matrices=False)
    b, k = Vt[0], U[:, 0] * S[0]
    c = b.sum()
    b, k = b / c, k * c
    return b, k - k.mean(), k.mean()


def lee_carter_common(logm_by_sex):
    """Fit the Li-Lee common-factor model over both sexes.

    logm_by_sex : dict sex -> DataFrame[Year, group] (the sample).
    Returns dict with the shared factor (a per sex, B, K) and per-sex (b_i, k_i).
    """
    groups = next(iter(logm_by_sex.values())).columns
    years = next(iter(logm_by_sex.values())).index
    a = {s: lm.mean(axis=0) for s, lm in logm_by_sex.items()}
    Z = {s: (lm - a[s]).to_numpy() for s, lm in logm_by_sex.items()}

    Zbar = sum(Z.values()) / len(Z)                 # common centred surface
    B, K, Kmean = _svd_factor(Zbar)
    a = {s: a[s] + pd.Series(B * Kmean, index=groups) for s in a}   # fold K mean into a_i

    sex = {}
    for s in logm_by_sex:
        R = Z[s] - np.outer(K, B)                   # residual after common factor
        b_i, k_i, kmean_i = _svd_factor(R)
        a[s] = a[s] + pd.Series(b_i * kmean_i, index=groups)
        sex[s] = {"b": pd.Series(b_i, index=groups), "k": pd.Series(k_i, index=years)}

    return {"a": a, "B": pd.Series(B, index=groups), "K": pd.Series(K, index=years),
            "sex": sex, "groups": groups, "years": years}


def fit_ar1(x):
    """Fit x_t = mu + phi (x_{t-1} - mu) + e. Returns (mu, phi, sigma)."""
    x = np.asarray(x, float)
    x0, x1 = x[:-1], x[1:]
    phi = np.cov(x0, x1, ddof=0)[0, 1] / np.var(x0)
    mu = x.mean()
    resid = (x1 - mu) - phi * (x0 - mu)
    return mu, phi, resid.std(ddof=1)


def forecast_common(fit, horizon):
    """Forecast the Li-Lee model: K random-walk-with-drift, each k_i AR(1).

    Returns dict: 'K' (array), and per sex a dict with 'k' (array). Coherence comes
    from k_i reverting to its mean while only the shared K carries a trend.
    """
    K = fit["K"].to_numpy()
    T = len(K)
    drift = (K[-1] - K[0]) / (T - 1)
    h = np.arange(1, horizon + 1)
    Kf = K[-1] + h * drift

    out = {"K": Kf, "h": h, "sex": {}}
    for s, f in fit["sex"].items():
        mu, phi, sig = fit_ar1(f["k"].to_numpy())
        last = f["k"].to_numpy()[-1]
        kf = mu + phi ** h * (last - mu)            # mean-reverting -> coherent
        out["sex"][s] = {"k": kf, "mu": mu, "phi": phi}
    return out


def rates_from_common(fit, fc, sex, step, logm_jumpoff=None):
    """Group rates for one sex at forecast step index `step` (0-based).

    `logm_jumpoff` (observed last-year log-rates) applies the same Lee-Miller
    jump-off adjustment as `project_rates`, for continuity with the observed data.
    """
    B = fit["B"].to_numpy()
    b_i = fit["sex"][sex]["b"].to_numpy()
    Kstep, kstep = fc["K"][step], fc["sex"][sex]["k"][step]
    if logm_jumpoff is None:
        return np.exp(fit["a"][sex].to_numpy() + B * Kstep + b_i * kstep)
    Klast = fit["K"].iloc[-1]
    klast = fit["sex"][sex]["k"].iloc[-1]
    return np.exp(np.asarray(logm_jumpoff) + B * (Kstep - Klast) + b_i * (kstep - klast))