"""Cohort-component population accounting and net-migration residual.

Meseguer (2010, eq. 22) builds population by single year of age with the
cohort-component identity, using Jan-1 populations P, calendar-year deaths D,
births B, and net migration M:

    P_{t+1}(0)   = B_t          - D_t(0)   + M_t(0)      (new births -> age 0)
    P_{t+1}(a+1) = P_t(a)       - D_t(a)   + M_t(a)      (survive one year older)

Net migration M is *not* modelled or published; here we derive it as the residual
that closes the identity, given the observed population, deaths and births:

    closed_{t+1}(a) = "survivors if M = 0"
    M_t(a)          = P_{t+1}(a) - closed_{t+1}(a)

The residual equals true net migration plus census/estimation error, and (because
HMD deaths are calendar-year age *squares* rather than cohort Lexis parallelograms)
a small period-vs-cohort timing approximation. Magnitudes and the age profile are
nonetheless the recognisable US migration signal.
"""

import numpy as np
import pandas as pd


def net_migration(pop, deaths, births, sex):
    """Derive net migration by age for one sex via the cohort-component residual.

    Parameters
    ----------
    pop, deaths : dict sex -> DataFrame[Year, Age]  (Jan-1 population; year deaths).
    births : DataFrame indexed by Year with a column per sex ('Female'/'Male'/'Total').
    sex : 'female' | 'male' | 'total'.

    Returns
    -------
    (mig, closed) : DataFrames indexed by the *start* year t (population is dated
        t+1), columns = age reached at t+1 (0..open). `mig` is net migration during
        year t; `closed` is the migration-free survivor projection.
    """
    P = pop[sex]                       # Jan-1 population, Year index is a string
    D = deaths[sex]                    # calendar-year deaths, integer Year index
    B = births[sex.capitalize()]
    ages = P.columns.to_numpy()
    n = len(ages)
    open_i = n - 1                     # last column is the open age group (110+)

    # US territorial break: Jan-1 1959 appears twice. '1959-' = old territory
    # (contiguous US, matches 1958); '1959+' = new territory (incl. AK & HI,
    # matches 1960). Bridge the step so the ~0.8M reclassification is not counted
    # as migration: end the 1958->1959 step at '1959-', start the 1959->1960 step
    # at '1959+'.
    def start_lbl(t):  # Jan-1 of year t (start of the interval)
        return "1959+" if t == 1959 else str(t)

    def end_lbl(t):    # Jan-1 of year t+1 (end of the interval)
        return "1959-" if t + 1 == 1959 else str(t + 1)

    years = [t for t in D.index
             if start_lbl(t) in P.index and end_lbl(t) in P.index and t in B.index]
    mig_rows, closed_rows = {}, {}
    for t in years:
        Pt = P.loc[start_lbl(t)].to_numpy()
        Pt1 = P.loc[end_lbl(t)].to_numpy()
        Dt = D.loc[t].to_numpy()
        # Deaths in HMD are calendar-year age *squares*; the cohort that ages from
        # a-1 (Jan 1 t) to a (Jan 1 t+1) dies in the upper triangle of square a-1
        # and the lower triangle of square a, i.e. ~half of each -> 0.5*(D[a-1]+D[a]).
        # Year-t births lose ~the lower triangle of the age-0 square (~0.5*D[0]).
        closed = np.empty(n)
        closed[0] = B.loc[t] - 0.5 * Dt[0]                             # surviving births -> age 0
        closed[1:open_i] = Pt[0:open_i - 1] - 0.5 * (Dt[0:open_i - 1] + Dt[1:open_i])
        # open group: survivors of the last single age and of the open group itself
        closed[open_i] = (Pt[open_i - 1] + Pt[open_i]) - (0.5 * Dt[open_i - 1] + Dt[open_i])
        closed_rows[t] = closed
        mig_rows[t] = Pt1 - closed

    mig = pd.DataFrame(mig_rows, index=ages).T
    closed = pd.DataFrame(closed_rows, index=ages).T
    mig.columns.name = closed.columns.name = "Age"
    return mig, closed


def project_closed(pop, deaths, births, sex, base_year=1960, end_year=None):
    """Roll population forward from a base year with NO migration (closed population).

    Uses the observed base-year population, then ages it forward year by year with
    observed births and deaths only (same Lexis half-split as `net_migration`). The
    result is what the population *would* be if net migration were zero, so the gap
    between this and the published population equals the cumulative net migration.

    base_year can be any data year (default 1933, the first). Rolling across 1959
    is fine mechanically (the roll uses its own carried population, not the split
    1959 rows); the only effect is that the ~0.8M Alaska/Hawaii population added at
    the 1959 territorial break is not in the migration-free roll, so it shows up in
    the published-minus-derived gap after 1959 (a small, one-off territorial term on
    top of true net migration).
    Returns DataFrame indexed by Year, columns = Age.
    """
    P, Dd = pop[sex], deaths[sex]
    B = births[sex.capitalize()]
    ages = P.columns.to_numpy()
    n = len(ages)
    open_i = n - 1
    if end_year is None:
        end_year = int(max(Dd.index))

    cur = P.loc[str(base_year)].to_numpy().astype(float)
    rows = {base_year: cur.copy()}
    for t in range(base_year, end_year):
        Dt = Dd.loc[t].to_numpy()
        nxt = np.empty(n)
        nxt[0] = B.loc[t] - 0.5 * Dt[0]
        nxt[1:open_i] = cur[0:open_i - 1] - 0.5 * (Dt[0:open_i - 1] + Dt[1:open_i])
        nxt[open_i] = (cur[open_i - 1] + cur[open_i]) - (0.5 * Dt[open_i - 1] + Dt[open_i])
        cur = np.clip(nxt, 0, None)
        rows[t + 1] = cur.copy()

    out = pd.DataFrame(rows, index=ages).T
    out.columns.name = "Age"
    return out