"""Cohort-component population projection (single year of age, annual steps).

Given a jump-off population by single year of age and sex, forecast single-year
death rates (survival), single-year fertility rates (births), and a net-migration
schedule, iterate the cohort-component identity (Meseguer 2010, eq. 22) forward:

    P_{t+1}(0)   = surviving births              + M(0)
    P_{t+1}(a+1) = P_t(a) * exp(-m_x(a))         + M(a+1)

Births are split by sex with a sex ratio at birth (SRB). The open age group is the
last index (survivors of the last single age plus the open group's own survivors).
"""

import numpy as np


def project(P0, mx, asfr, mig, srb=1.05):
    """Iterate the cohort-component projection.

    Parameters
    ----------
    P0   : dict sex -> (A,) jump-off Jan-1 population by single age (A = open age + 1).
    mx   : dict sex -> (T, A) forecast single-year central death rates for T years.
    asfr : (T, A) forecast single-year fertility rates (female), age-indexed 0..A-1.
    mig  : dict sex -> (A,) net migration by age, added each year (held constant).
    srb  : sex ratio at birth (males per female), default 1.05.

    Returns
    -------
    dict sex -> (T+1, A) population trajectory (row 0 = jump-off).
    """
    sexes = list(P0)
    A = len(P0[sexes[0]])
    T = mx[sexes[0]].shape[0]
    out = {s: np.zeros((T + 1, A)) for s in sexes}
    for s in sexes:
        out[s][0] = P0[s]

    for t in range(T):
        px = {s: np.exp(-mx[s][t]) for s in sexes}                 # annual survival by age
        Pfem = out["female"][t]
        B = float(np.sum(asfr[t] * Pfem))                          # total births in year t
        Bsex = {"female": B / (1 + srb), "male": B * srb / (1 + srb)}
        for s in sexes:
            nxt = np.empty(A)
            nxt[0] = Bsex[s] * np.exp(-0.5 * mx[s][t, 0]) + mig[s][0]     # surviving births (~half-yr exposure)
            nxt[1:A - 1] = out[s][t][0:A - 2] * px[s][0:A - 2] + mig[s][1:A - 1]   # age a-1 -> a
            nxt[A - 1] = (out[s][t][A - 2] * px[s][A - 2]
                          + out[s][t][A - 1] * px[s][A - 1] + mig[s][A - 1])       # open group
            out[s][t + 1] = np.clip(nxt, 0, None)
    return out


def group_of_age(ages, edges):
    """Map each single age to its abridged-group index given group left-edges."""
    edges = np.asarray(edges)
    return np.clip(np.searchsorted(edges, ages, side="right") - 1, 0, len(edges) - 1)


def disaggregate(base_single, group_factor, gid):
    """Scale a jump-off single-year schedule by each age's group forecast factor.

    base_single : (A,) observed single-year rates at the jump-off year.
    group_factor: (T, G) multiplicative factor per group per projection year.
    gid         : (A,) group index of each single age.
    Returns (T, A) forecast single-year rates.
    """
    return base_single[None, :] * group_factor[:, gid]