1. enrich

"""
Derived fields. Every one of them is computed exactly once, here, so that
"NPL" means the same thing in the vintage table and in the branch ranking.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

from . import config as C
from . import io_layer as io


def as_of(report_date) -> pd.Timestamp:
    """report_date is the *stamp*; the data describes the last day before it."""
    return pd.Timestamp(report_date) - pd.Timedelta(days=1)


def period_of(report_date) -> int:
    d = as_of(report_date)
    return d.year * 100 + d.month


def add_period_cols(df: pd.DataFrame) -> pd.DataFrame:
    a = df["report_date"].map(as_of)
    df["as_of"] = a
    df["period"] = (a.dt.year * 100 + a.dt.month).astype("int32")
    v = df["vidacha_sana"]
    df["vintage"] = (v.dt.year * 100 + v.dt.month).astype("Int32")
    df["mob"] = ((a.dt.year - v.dt.year) * 12 + (a.dt.month - v.dt.month)).astype("Int32")
    return df


def bucket_of(dpd: pd.Series) -> pd.Series:
    edges = [b[0] for b in C.DPD_BUCKETS[1:]]
    labels = [b[2] for b in C.DPD_BUCKETS]
    idx = np.searchsorted(np.array(edges), dpd.to_numpy(), side="left")
    return pd.Categorical([labels[i] for i in idx],
                          categories=C.BUCKET_ORDER, ordered=True)


def enrich(df: pd.DataFrame, fx: pd.DataFrame, fx_mode: str = "nominal") -> pd.DataFrame:
    df = add_period_cols(df)

    rate = io.fx_factor(df, fx, fx_mode)
    df["fx_rate"] = rate
    for c in ("brutto_95_summa", "brutto_summa", "reserve_summa",
              "vidacha_tekushiy_summa", "pogashen_tekushiy_summa",
              "straxovka_summa", "straxovka_brutto_95_summa",
              "spisat_95413_summa", "all_time_given_amount"):
        df[c + "_uzs"] = df[c] * rate

    df["exposure"] = df[C.EXPOSURE_COL + "_uzs"]
    df["interest_accrued"] = (
        df["protsent_16309_summa"] + df["protsent_16377_summa"] + df["protsent_16379_summa"]
    ) * rate

    df["dpd"] = df["maks_dni"].clip(lower=0)
    df["bucket"] = bucket_of(df["dpd"])

    df["is_npl"] = df["kategoriya_npl_90"].isin(C.NPL_CATEGORIES)
    # Counterfactual: what the book would look like if the insurer had not paid.
    df["is_npl_gross_of_insurance"] = df["kategoriya_straxovka"].isin(C.NPL_CATEGORIES)
    df["insurance_saved_npl"] = df["is_npl_gross_of_insurance"] & ~df["is_npl"]
    df["exposure_gross_of_insurance"] = df["straxovka_brutto_95_summa_uzs"].where(
        df["straxovka_brutto_95_summa_uzs"] > 0, df["exposure"])

    df["is_watch"] = (~df["is_npl"]) & df["dpd"].between(C.WATCH_DPD_MIN, C.WATCH_DPD_MAX)
    df["is_disbursed_this_month"] = df["vidacha_tekushiy_summa_uzs"] > 0
    df["in_court"] = df["sud"].isin(["official"])
    df["court_pipeline"] = df["sud"].isin(["neofficial"])
    df["written_off"] = df["spisat_95413_summa"] > 0

    df["npl_exposure"] = np.where(df["is_npl"], df["exposure"], 0.0)
    df["watch_exposure"] = np.where(df["is_watch"], df["exposure"], 0.0)
    df["npl_gross_exposure"] = np.where(
        df["is_npl_gross_of_insurance"], df["exposure_gross_of_insurance"], 0.0)
    return df


PANEL_COLS = [
    "anketa", "contragent", "period", "vintage", "mob", "dpd", "bucket",
    "is_npl", "is_watch", "exposure", "filial", "department", "passport",
    "avto", "valuta", "sud", "is_disbursed_this_month",
    "vidacha_tekushiy_summa_uzs", "pogashen_tekushiy_summa_uzs",
    "interest_accrued", "all_time_given_amount_uzs", "written_off",
    "is_npl_gross_of_insurance", "insurance_saved_npl", "in_court",
    "court_pipeline", "reserve_summa_uzs", "straxovka_summa_uzs",
    "exposure_gross_of_insurance", "npl_exposure", "watch_exposure",
    "npl_gross_exposure", "spisat_95413_summa_uzs",
]


def month_index(period: pd.Series | np.ndarray) -> np.ndarray:
    """YYYYMM -> a dense month counter, so month arithmetic is just +/- 1."""
    p = np.asarray(period, dtype="float64")
    return (p // 100) * 12 + (p % 100)


_SOURCE_COLS_FOR_PANEL = [
    "anketa", "contragent", "filial", "department", "passport", "avto",
    "valuta", "sud", "vidacha_sana", "maks_dni", "kategoriya_npl_90",
    "kategoriya_straxovka", "brutto_95_summa", "reserve_summa", "brutto_summa",
    "vidacha_tekushiy_summa", "pogashen_tekushiy_summa", "spisat_95413_summa",
    "spisat_91501_summa", "straxovka_summa", "straxovka_brutto_95_summa",
    "protsent_16309_summa", "protsent_16377_summa", "protsent_16379_summa",
    "all_time_given_amount",
]


def build_panel(fx_mode: str = "nominal", dates=None) -> pd.DataFrame:
    """
    The loan-month panel: one slim row per (anketa, period).

    Everything time-dependent — vintages, roll rates, cure rates, observed PD —
    is derived from this single object rather than re-reading the lake, which is
    what keeps the numbers reconcilable with each other.
    """
    fx = io.load_fx()
    frames = []
    for rd, snap in io.iter_snapshots(columns=_SOURCE_COLS_FOR_PANEL, dates=dates):
        frames.append(enrich(snap, fx, fx_mode)[PANEL_COLS])
    if not frames:
        raise RuntimeError("lake is empty — run ingest() first")

    panel = pd.concat(frames, ignore_index=True)
    panel = panel.sort_values(["anketa", "period"], kind="stable").reset_index(drop=True)

    g = panel.groupby("anketa", sort=False)
    panel["mi"] = month_index(panel["period"])
    panel["ever_npl_to_date"] = g["is_npl"].cummax()
    panel["peak_dpd_to_date"] = g["dpd"].cummax()
    panel["obs_index"] = g.cumcount()
    panel["is_first_obs"] = panel["obs_index"] == 0

    prev_mi = g["mi"].shift(1)
    next_mi = g["mi"].shift(-1)
    contiguous_back = prev_mi.eq(panel["mi"] - 1)
    contiguous_fwd = next_mi.eq(panel["mi"] + 1)

    panel["prev_bucket"] = g["bucket"].shift(1).where(contiguous_back)
    panel["prev_dpd"] = g["dpd"].shift(1).where(contiguous_back)
    panel["prev_exposure"] = g["exposure"].shift(1).where(contiguous_back)
    panel["prev_is_npl"] = g["is_npl"].shift(1).where(contiguous_back)
    panel["entered_this_month"] = ~contiguous_back

    # Forward view. A loan with no observation next month has left the book:
    # that is a real outcome, not missing data, so it gets its own state.
    panel["next_bucket"] = (g["bucket"].shift(-1).astype("object")
                            .where(contiguous_fwd, "EXIT"))
    panel["next_is_npl"] = g["is_npl"].shift(-1).where(contiguous_fwd)
    panel["exited_after"] = ~contiguous_fwd
    panel["is_last_obs"] = g.cumcount(ascending=False) == 0

    # First month at or after this one in which the loan is in default.
    # Reverse cumulative minimum inside each loan — one pass, no Python loop.
    big = 10**9
    tmp = panel["mi"].where(panel["is_npl"], big)
    panel["next_npl_mi"] = (tmp.iloc[::-1]
                            .groupby(panel["anketa"].iloc[::-1], sort=False)
                            .cummin().iloc[::-1])
    panel["months_to_npl"] = (panel["next_npl_mi"] - panel["mi"]).where(
        panel["next_npl_mi"] < big)

    C.PANEL_PATH.parent.mkdir(parents=True, exist_ok=True)
    panel.to_parquet(C.PANEL_PATH, index=False, compression="zstd")
    return panel


def load_panel() -> pd.DataFrame:
    return pd.read_parquet(C.PANEL_PATH)