1. marrts

"""
Metric layer. Each function turns the 8M-row panel into a mart of a few
thousand rows. The report layer never touches the panel — it reads marts.

That separation is what makes month N+1 cheap: append one partition, rebuild
the panel, rebuild the marts, and every chart re-renders from the same numbers.
"""
from __future__ import annotations

import numpy as np
import pandas as pd

from . import config as C

TOTAL = "__all__"


def _write(df: pd.DataFrame, name: str) -> pd.DataFrame:
    C.MART_DIR.mkdir(parents=True, exist_ok=True)
    df.to_parquet(C.MART_DIR / f"{name}.parquet", index=False, compression="zstd")
    return df


def read(name: str) -> pd.DataFrame:
    return pd.read_parquet(C.MART_DIR / f"{name}.parquet")


def _long_by_dim(panel: pd.DataFrame, agg_fn, dims=None) -> pd.DataFrame:
    """Run one aggregation for the total and once per dimension, stacked long."""
    dims = list(dims or C.DIMENSIONS)
    out = []
    tot = agg_fn(panel.groupby("period", observed=True))
    tot.insert(0, "dim_value", TOTAL)
    tot.insert(0, "dim_name", TOTAL)
    out.append(tot.reset_index())
    for d in dims:
        if d not in panel.columns:
            continue
        part = agg_fn(panel.groupby(["period", d], observed=True))
        part = part.reset_index().rename(columns={d: "dim_value"})
        part["dim_value"] = part["dim_value"].astype("string")
        part.insert(0, "dim_name", d)
        out.append(part)
    return pd.concat(out, ignore_index=True)


# ==========================================================================
# 1. Portfolio stock
# ==========================================================================
def portfolio_monthly(panel: pd.DataFrame) -> pd.DataFrame:
    def agg(g):
        return pd.DataFrame({
            "n_loans": g["anketa"].size(),
            "n_clients": g["contragent"].nunique(),
            "exposure": g["exposure"].sum(),
            "npl_exposure": g["npl_exposure"].sum(),
            "watch_exposure": g["watch_exposure"].sum(),
            "npl_gross_exposure": g["npl_gross_exposure"].sum(),
            "exposure_gross_of_insurance": g["exposure_gross_of_insurance"].sum(),
            "reserve": g["reserve_summa_uzs"].sum(),
            "insurance_paid": g["straxovka_summa_uzs"].sum(),
            "n_npl": g["is_npl"].sum(),
            "n_watch": g["is_watch"].sum(),
            "n_in_court": g["in_court"].sum(),
            "n_court_pipeline": g["court_pipeline"].sum(),
            "n_insurance_saved": g["insurance_saved_npl"].sum(),
            "repaid": g["pogashen_tekushiy_summa_uzs"].sum(),
            "interest_accrued": g["interest_accrued"].sum(),
            "written_off": g["spisat_95413_summa_uzs"].sum(),
        })

    m = _long_by_dim(panel, agg)
    m["npl_ratio"] = m["npl_exposure"] / m["exposure"].replace(0, np.nan)
    m["watch_ratio"] = m["watch_exposure"] / m["exposure"].replace(0, np.nan)
    m["npl_ratio_gross_of_insurance"] = (
        m["npl_gross_exposure"] / m["exposure_gross_of_insurance"].replace(0, np.nan))
    m["insurance_npl_relief_pp"] = (
        m["npl_ratio_gross_of_insurance"] - m["npl_ratio"]) * 100
    m["coverage_ratio"] = m["reserve"] / m["npl_exposure"].replace(0, np.nan)
    m["loans_per_client"] = m["n_loans"] / m["n_clients"].replace(0, np.nan)
    m["avg_ticket"] = m["exposure"] / m["n_loans"].replace(0, np.nan)
    return _write(m, "portfolio_monthly")


# ==========================================================================
# 2. Origination
# ==========================================================================
def origination_monthly(panel: pd.DataFrame) -> pd.DataFrame:
    first_seen = panel.groupby("contragent", observed=True)["period"].min()
    first_loan = panel.groupby("anketa", observed=True)["period"].min()

    new = panel[panel["is_disbursed_this_month"]].copy()
    new["amount"] = new["vidacha_tekushiy_summa_uzs"]
    new["is_new_client"] = new["contragent"].map(first_seen).eq(new["period"])
    new["is_first_loan_month"] = new["anketa"].map(first_loan).eq(new["period"])

    def agg(g):
        return pd.DataFrame({
            "n_loans": g["anketa"].nunique(),
            "n_clients": g["contragent"].nunique(),
            "amount": g["amount"].sum(),
            "n_new_clients": g["is_new_client"].sum(),
        })

    m = _long_by_dim(new, agg)
    m["avg_ticket"] = m["amount"] / m["n_loans"].replace(0, np.nan)
    m["loans_per_client"] = m["n_loans"] / m["n_clients"].replace(0, np.nan)
    m["repeat_client_share"] = 1 - m["n_new_clients"] / m["n_clients"].replace(0, np.nan)

    # Left truncation: in the first observed month everyone looks "new".
    first_period = int(panel["period"].min())
    m["truncation_warning"] = m["period"] <= first_period
    return _write(m, "origination_monthly")


# ==========================================================================
# 3. Vintages  (GL30+@3MOB and friends)
# ==========================================================================
def vintages(panel: pd.DataFrame, segments=("__all__", "avto", "department")) -> pd.DataFrame:
    first_period = int(panel["period"].min())
    p = panel[panel["vintage"].notna() & (panel["vintage"] >= first_period)].copy()
    p = p[p["mob"].between(0, 36)]

    orig = (p[p["is_disbursed_this_month"]]
            .groupby("anketa", observed=True)["vidacha_tekushiy_summa_uzs"].max())
    p["orig_amount"] = p["anketa"].map(orig).fillna(p["exposure"])

    rows = []
    for thr, _ in {(t, m) for t, m in C.VINTAGE_MARKS}:
        p[f"ever_{thr}_by_mob"] = p["peak_dpd_to_date"] > thr
        p[f"at_{thr}"] = p["dpd"] > thr

    for seg in segments:
        keys = ["vintage", "mob"] if seg == "__all__" else ["vintage", "mob", seg]
        g = p.groupby(keys, observed=True)
        d = pd.DataFrame({
            "n_loans": g["anketa"].size(),
            "orig_amount": g["orig_amount"].sum(),
        })
        for thr in sorted({t for t, _ in C.VINTAGE_MARKS}):
            d[f"n_ever_{thr}"] = g[f"ever_{thr}_by_mob"].sum()
            d[f"n_at_{thr}"] = g[f"at_{thr}"].sum()
            d[f"amt_ever_{thr}"] = g.apply(
                lambda x, t=thr: x.loc[x[f"ever_{t}_by_mob"], "orig_amount"].sum(),
                include_groups=False)
        d = d.reset_index()
        if seg == "__all__":
            d["dim_name"], d["dim_value"] = TOTAL, TOTAL
        else:
            d = d.rename(columns={seg: "dim_value"})
            d["dim_name"] = seg
            d["dim_value"] = d["dim_value"].astype("string")
        rows.append(d)

    m = pd.concat(rows, ignore_index=True)
    for thr in sorted({t for t, _ in C.VINTAGE_MARKS}):
        m[f"gl{thr}_ever_rate"] = m[f"n_ever_{thr}"] / m["n_loans"].replace(0, np.nan)
        m[f"gl{thr}_at_rate"] = m[f"n_at_{thr}"] / m["n_loans"].replace(0, np.nan)
        m[f"gl{thr}_ever_rate_amt"] = m[f"amt_ever_{thr}"] / m["orig_amount"].replace(0, np.nan)
    m["is_reliable"] = m["n_loans"] >= C.MIN_COHORT_SIZE
    return _write(m, "vintages")


def vintage_marks(vint: pd.DataFrame) -> pd.DataFrame:
    """Flatten the vintage surface into the GLx@yMOB table stakeholders ask for."""
    out = []
    for thr, mob in C.VINTAGE_MARKS:
        s = vint[vint["mob"] == mob].copy()
        s["metric"] = f"GL{thr}+@{mob}MOB"
        s["rate"] = s[f"gl{thr}_ever_rate"]
        s["rate_amt"] = s[f"gl{thr}_ever_rate_amt"]
        out.append(s[["dim_name", "dim_value", "vintage", "mob", "metric",
                      "rate", "rate_amt", "n_loans", "is_reliable"]])
    return _write(pd.concat(out, ignore_index=True), "vintage_marks")


# ==========================================================================
# 4. Transition matrix, roll rates, cure rates
# ==========================================================================
def transitions(panel: pd.DataFrame) -> pd.DataFrame:
    p = panel[panel["next_bucket"].notna()].copy()
    p["from_bucket"] = p["bucket"].astype("string")
    p["to_bucket"] = p["next_bucket"].astype("string")

    g = p.groupby(["period", "from_bucket", "to_bucket"], observed=True)
    m = pd.DataFrame({"n_loans": g["anketa"].size(),
                      "exposure": g["exposure"].sum()}).reset_index()
    denom = m.groupby(["period", "from_bucket"], observed=True)[["n_loans", "exposure"]].transform("sum")
    m["rate_n"] = m["n_loans"] / denom["n_loans"].replace(0, np.nan)
    m["rate_amt"] = m["exposure"] / denom["exposure"].replace(0, np.nan)
    return _write(m, "transitions")


def cure_rates(panel: pd.DataFrame) -> pd.DataFrame:
    """
    Cure = a loan that was in bucket B this month and is current next month.
    Split out separately: 'partial cure' (improved but not current) and
    'exit' (left the book — repaid, written off, or sold, which is a very
    different story from cured and must not be mixed in).
    """
    order = {b: i for i, b in enumerate(C.BUCKET_ORDER)}
    p = panel[panel["next_bucket"].notna() & (panel["bucket"].astype("string") != "CUR")].copy()
    p["from_bucket"] = p["bucket"].astype("string")
    p["to_bucket"] = p["next_bucket"].astype("string")
    fo = p["from_bucket"].map(order)
    to = p["to_bucket"].map(order)

    p["outcome"] = np.select(
        [p["to_bucket"].eq("EXIT"), p["to_bucket"].eq("CUR"), to < fo, to == fo],
        ["exit", "cured", "improved", "flat"], default="worsened")

    g = p.groupby(["period", "from_bucket", "outcome"], observed=True)
    m = pd.DataFrame({"n_loans": g["anketa"].size(),
                      "exposure": g["exposure"].sum()}).reset_index()
    denom = m.groupby(["period", "from_bucket"], observed=True)[["n_loans", "exposure"]].transform("sum")
    m["rate_n"] = m["n_loans"] / denom["n_loans"].replace(0, np.nan)
    m["rate_amt"] = m["exposure"] / denom["exposure"].replace(0, np.nan)
    return _write(m, "cure_rates")


# ==========================================================================
# 5. Observed probability of default
# ==========================================================================
def pd_observed(panel: pd.DataFrame, horizon: int | None = None) -> pd.DataFrame:
    """
    Not a model score — the realised default frequency.

    Population: performing loans in month t. Event: default at any point in
    t+1 .. t+H. Observations without a full H-month window and without an
    early default are censored out, otherwise recent months look artificially
    safe purely because the future has not happened yet.
    """
    H = horizon or C.PD_HORIZON_MONTHS
    p = panel[~panel["is_npl"]].copy()
    last_mi = float(panel["mi"].max())

    defaulted = p["months_to_npl"].le(H) & p["months_to_npl"].gt(0)
    censored = (~defaulted) & (p["mi"] + H > last_mi)
    p = p[~censored].copy()
    p["defaulted_in_horizon"] = defaulted[p.index]

    def agg(g):
        return pd.DataFrame({
            "n_obs": g["anketa"].size(),
            "n_default": g["defaulted_in_horizon"].sum(),
            "exposure": g["exposure"].sum(),
            "exposure_default": g.apply(
                lambda x: x.loc[x["defaulted_in_horizon"], "exposure"].sum(),
                include_groups=False),
        })

    m = _long_by_dim(p, agg)
    by_bucket = agg(p.groupby(["period", "bucket"], observed=True)).reset_index()
    by_bucket = by_bucket.rename(columns={"bucket": "dim_value"})
    by_bucket["dim_value"] = by_bucket["dim_value"].astype("string")
    by_bucket["dim_name"] = "bucket"
    m = pd.concat([m, by_bucket], ignore_index=True)

    m["pd"] = m["n_default"] / m["n_obs"].replace(0, np.nan)
    m["pd_amt"] = m["exposure_default"] / m["exposure"].replace(0, np.nan)
    m["horizon_months"] = H
    m["is_reliable"] = m["n_obs"] >= C.MIN_COHORT_SIZE
    return _write(m, "pd_observed")


# ==========================================================================
# 6. NPL bridge  (the waterfall the board actually wants)
# ==========================================================================
def npl_bridge(panel: pd.DataFrame) -> pd.DataFrame:
    """
    closing = opening + inflow - cured - exited + residual

    The residual is amortisation, partial repayment and FX revaluation on loans
    that stayed in default. It is computed as a balancing figure on purpose:
    the bridge must tie to the stock exactly, or stakeholders stop trusting it.
    """
    rows = []
    periods = sorted(panel["period"].unique())
    npl_stock = (panel.groupby("period", observed=True)["npl_exposure"].sum())

    p = panel[panel["prev_is_npl"].notna()].copy()
    inflow = p[p["is_npl"] & ~p["prev_is_npl"].astype(bool)].groupby("period", observed=True)["exposure"].sum()
    cured = p[~p["is_npl"] & p["prev_is_npl"].astype(bool)].groupby("period", observed=True)["prev_exposure"].sum()

    exited = (panel[panel["is_npl"] & panel["exited_after"]]
              .assign(next_period=lambda d: d["period"].map(_next_period))
              .groupby("next_period", observed=True)["exposure"].sum())

    for i, per in enumerate(periods):
        if i == 0:
            continue
        prev = periods[i - 1]
        opening = float(npl_stock.get(prev, 0.0))
        closing = float(npl_stock.get(per, 0.0))
        inf = float(inflow.get(per, 0.0))
        cur = float(cured.get(per, 0.0))
        exi = float(exited.get(per, 0.0))
        rows.append({
            "period": per, "opening": opening, "inflow": inf,
            "cured": -cur, "exited": -exi,
            "residual": closing - (opening + inf - cur - exi),
            "closing": closing,
        })
    return _write(pd.DataFrame(rows), "npl_bridge")


def _next_period(p: int) -> int:
    y, m = divmod(int(p), 100)
    return (y + 1) * 100 + 1 if m == 12 else p + 1


# ==========================================================================
# 7. Client burden, concentration, realised term, cost of credit
# ==========================================================================
def client_burden(panel: pd.DataFrame) -> pd.DataFrame:
    cnt = (panel.groupby(["period", "contragent"], observed=True)
           .agg(n_loans=("anketa", "nunique"), exposure=("exposure", "sum"))
           .reset_index())
    cnt["burden"] = pd.cut(cnt["n_loans"], [0, 1, 2, 3, 4, 10**9],
                           labels=["1", "2", "3", "4", "5+"])
    g = cnt.groupby(["period", "burden"], observed=True)
    m = pd.DataFrame({"n_clients": g["contragent"].size(),
                      "n_loans": g["n_loans"].sum(),
                      "exposure": g["exposure"].sum()}).reset_index()
    denom = m.groupby("period", observed=True)[["n_clients", "exposure"]].transform("sum")
    m["client_share"] = m["n_clients"] / denom["n_clients"]
    m["exposure_share"] = m["exposure"] / denom["exposure"]
    return _write(m, "client_burden")


def concentration(panel: pd.DataFrame, top_n: int = 20) -> pd.DataFrame:
    rows = []
    for per, g in panel.groupby("period", observed=True):
        by_client = g.groupby("contragent", observed=True)["exposure"].sum().sort_values(ascending=False)
        tot = by_client.sum()
        shares = (g.groupby("filial", observed=True)["exposure"].sum() / tot).fillna(0)
        rows.append({
            "period": per,
            f"top{top_n}_share": by_client.head(top_n).sum() / tot if tot else np.nan,
            "top1_share": by_client.head(1).sum() / tot if tot else np.nan,
            "hhi_filial": float((shares ** 2).sum() * 10_000),
            "n_clients": int(by_client.size),
        })
    return _write(pd.DataFrame(rows), "concentration")


def realised_term(panel: pd.DataFrame) -> pd.DataFrame:
    """
    There is no maturity date in the tape, so 'how long do loans actually live'
    is answered from behaviour: the MOB at which a loan leaves the book, split
    by whether it left clean or in default.
    """
    last = panel[panel["is_last_obs"] & panel["exited_after"]].copy()
    last = last[last["mob"].notna()]
    last["closure_type"] = np.where(
        last["is_npl"] | last["written_off"], "closed_bad", "closed_clean")
    g = last.groupby(["period", "closure_type"], observed=True)
    m = pd.DataFrame({
        "n_loans": g["anketa"].size(),
        "avg_life_months": g["mob"].mean(),
        "median_life_months": g["mob"].median(),
        "exposure_at_exit": g["exposure"].sum(),
    }).reset_index()
    # Ignore the final snapshot: everything looks like it "exited" there.
    m = m[m["period"] < panel["period"].max()]
    return _write(m, "realised_term")


def cost_of_credit(panel: pd.DataFrame) -> pd.DataFrame:
    """
    Proxy for total cost of credit: interest accrued over the life of a vintage
    against the amount disbursed. Fees, insurance premia and commissions are not
    in the tape, so this is a floor, not the full TCC.
    """
    orig = (panel[panel["is_disbursed_this_month"]]
            .groupby(["anketa", "vintage"], observed=True)["vidacha_tekushiy_summa_uzs"]
            .max().reset_index())
    interest = panel.groupby("anketa", observed=True)["interest_accrued"].sum()
    orig["interest_total"] = orig["anketa"].map(interest).fillna(0)
    life = panel.groupby("anketa", observed=True)["mob"].max()
    orig["life_months"] = orig["anketa"].map(life)

    g = orig.groupby("vintage", observed=True)
    m = pd.DataFrame({
        "n_loans": g["anketa"].size(),
        "disbursed": g["vidacha_tekushiy_summa_uzs"].sum(),
        "interest": g["interest_total"].sum(),
        "avg_life_months": g["life_months"].mean(),
    }).reset_index()
    m["tcc_proxy"] = m["interest"] / m["disbursed"].replace(0, np.nan)
    m["tcc_annualised"] = m["tcc_proxy"] / (m["avg_life_months"] / 12).replace(0, np.nan)
    return _write(m, "cost_of_credit")


def insurance_effect(panel: pd.DataFrame) -> pd.DataFrame:
    g = panel.groupby(["period", "department"], observed=True)
    m = pd.DataFrame({
        "exposure": g["exposure"].sum(),
        "npl_exposure": g["npl_exposure"].sum(),
        "npl_gross_exposure": g["npl_gross_exposure"].sum(),
        "exposure_gross": g["exposure_gross_of_insurance"].sum(),
        "insurance_paid": g["straxovka_summa_uzs"].sum(),
        "n_saved": g["insurance_saved_npl"].sum(),
    }).reset_index()
    m["npl_ratio"] = m["npl_exposure"] / m["exposure"].replace(0, np.nan)
    m["npl_ratio_gross"] = m["npl_gross_exposure"] / m["exposure_gross"].replace(0, np.nan)
    m["relief_pp"] = (m["npl_ratio_gross"] - m["npl_ratio"]) * 100
    return _write(m, "insurance_effect")


def build_all(panel: pd.DataFrame) -> dict[str, pd.DataFrame]:
    out = {
        "portfolio_monthly": portfolio_monthly(panel),
        "origination_monthly": origination_monthly(panel),
        "transitions": transitions(panel),
        "cure_rates": cure_rates(panel),
        "pd_observed": pd_observed(panel),
        "npl_bridge": npl_bridge(panel),
        "client_burden": client_burden(panel),
        "concentration": concentration(panel),
        "realised_term": realised_term(panel),
        "cost_of_credit": cost_of_credit(panel),
        "insurance_effect": insurance_effect(panel),
    }
    v = vintages(panel)
    out["vintages"] = v
    out["vintage_marks"] = vintage_marks(v)
    return out