udyamlinc-ml/training/train_forecast.py

"""
Quantile heads for the cash-flow forecast.

The problem statement is "AI-Driven **Cash Flow Prediction** & Risk Flagging".
The risk head is a model; the forecast, until now, was
``DeterministicForecastStrategy`` — a median of monthly ledger position with a
capped drift and a band widening as sqrt(h). Honest, explainable, and not AI.
This trains the replacement.

Shape
-----
One row per (enterprise, as-of month end, horizon 1..6):

    features known at T   ->   net ledger position in month T + h

Three models, at alpha = 0.10 / 0.50 / 0.90. **Horizon is a feature, not a
separate model per step**: six models per quantile would be eighteen artifacts
that cannot share the signal they all depend on, and the h=6 model would train on
a sixth of the rows.

Two rules carried over from the deterministic strategy
-----------------------------------------------------
1. Quantile regression does not guarantee ordering — P10 can be predicted above
   P50 on an individual row. The three predictions are sorted per row before
   they leave here. A crossed band is worse than no band.
2. **Missing inputs widen the band; they never lower its centre.** A pessimistic
   P50 built on ignorance manufactures breaches, which is how an early-warning
   system loses an officer's trust in one season.

What the metrics mean
---------------------
Pinball loss, per quantile, against the deterministic strategy reproduced on the
same held-out months. Lower is better. This is the comparison that decides
whether the model earns the switch — and, as with the stress head, it is measured
against a generated ledger, so it measures the generator too.
"""

from __future__ import annotations

import argparse
import json
import pickle
from datetime import datetime
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
import psycopg

DSN = "postgresql://postgres:root@localhost:5432/UdyamLinc_DEMO"

QUANTILES = {"p10": 0.10, "p50": 0.50, "p90": 0.90}
HORIZONS = 6

# Features known at T. The point-in-time discipline is the same as the stress
# head's: `valid_for_date <= T`. Forecast-derived features are excluded for the
# same reason they are there — dscr_* reads the current forecast at every date.
FEATURE_CODES = [
    "feed_cost_change",
    "output_price_change",
    "margin_squeeze",
    "price_z_score",
    "rainfall_anomaly",
    "heat_anomaly",
    "cash_flow_change",
    "emi_delay",
    "inventory_days",
    "feed_staleness_days",
    "days_since_last_entry",
    "emi_to_inflow",
]

# --------------------------------------------------------------------------- #

# Monthly net position, and the trailing statistics a forecaster would have had.
# Everything is computed from months strictly before the as-of month, so a row
# never sees the month it is filed under, let alone the one it predicts.
SPINE_SQL = """
with monthly as (
    select enterprise_id,
           date_trunc('month', entry_date)::date as mo,
           sum(case when kind = 'IN' then amount when kind = 'OUT' then -amount else 0 end) as net,
           sum(case when kind = 'IN' then amount else 0 end) as inflow
    from ledger_entry
    group by 1, 2
),
stats as (
    select m.enterprise_id,
           m.mo as as_of_mo,
           avg(h.net)    as trail_mean,
           stddev_pop(h.net) as trail_sd,
           min(h.net)    as trail_min,
           avg(h.inflow) as trail_inflow,
           count(*)      as trail_months
    from monthly m
    join monthly h
      on h.enterprise_id = m.enterprise_id
     and h.mo <  m.mo
     and h.mo >= m.mo - interval '6 month'
    group by 1, 2
),
obligation as (
    select enterprise_id, as_of_date, sum(instalment_amount) as emi
    from loan_record group by 1, 2
)
select s.enterprise_id,
       s.as_of_mo,
       s.trail_mean, s.trail_sd, s.trail_min, s.trail_inflow, s.trail_months,
       e.sector_id, e.district_id,
       o.emi,
       t.mo   as target_mo,
       t.net  as target_net,
       (extract(year from t.mo) * 12 + extract(month from t.mo))
         - (extract(year from s.as_of_mo) * 12 + extract(month from s.as_of_mo)) as horizon
from stats s
join enterprise e on e.enterprise_id = s.enterprise_id
left join obligation o
       on o.enterprise_id = s.enterprise_id
      and o.as_of_date = (s.as_of_mo + interval '1 month - 1 day')::date
join monthly t
  on t.enterprise_id = s.enterprise_id
 and t.mo >  s.as_of_mo
 and t.mo <= s.as_of_mo + interval '6 month'
where s.trail_months >= 3
"""

FEATURES_SQL = """
select enterprise_id, valid_for_date, feature_code, value
from feature_value
where feature_code = any(%(codes)s)
"""


def build(dsn: str) -> pd.DataFrame:
    with psycopg.connect(dsn) as conn:
        spine = pd.read_sql(SPINE_SQL, conn)
        feats = pd.read_sql(FEATURES_SQL, conn, params={"codes": FEATURE_CODES})

    if spine.empty:
        raise SystemExit("No rows. Has the ledger backfill run?")

    # Features are filed at month end; the spine is keyed on month start.
    feats["as_of_mo"] = pd.to_datetime(feats["valid_for_date"]).values.astype("datetime64[M]")
    wide = feats.pivot_table(
        index=["enterprise_id", "as_of_mo"], columns="feature_code", values="value", aggfunc="first"
    ).reset_index()

    spine["as_of_mo"] = pd.to_datetime(spine["as_of_mo"]).values.astype("datetime64[M]")
    df = spine.merge(wide, on=["enterprise_id", "as_of_mo"], how="left")

    for code in FEATURE_CODES:
        if code not in df.columns:
            df[code] = np.nan
        df[f"{code}__missing"] = df[code].isna().astype("int8")

    numeric = FEATURE_CODES + [
        "trail_mean", "trail_sd", "trail_min", "trail_inflow", "trail_months",
        "emi", "horizon", "target_net", "sector_id", "district_id",
    ]
    for col in numeric:
        df[col] = pd.to_numeric(df[col], errors="coerce")

    df = df.dropna(subset=["target_net", "horizon"])
    return df.sort_values(["as_of_mo", "enterprise_id", "horizon"]).reset_index(drop=True)


def feature_columns(df: pd.DataFrame) -> list[str]:
    cols = FEATURE_CODES + [f"{c}__missing" for c in FEATURE_CODES]
    cols += ["trail_mean", "trail_sd", "trail_min", "trail_inflow", "trail_months",
             "emi", "horizon", "sector_id", "district_id"]
    return [c for c in cols if c in df.columns]


def pinball(y: np.ndarray, pred: np.ndarray, alpha: float) -> float:
    """The loss quantile regression is actually optimising. Lower is better."""
    delta = y - pred
    return float(np.mean(np.maximum(alpha * delta, (alpha - 1) * delta)))


def deterministic_baseline(frame: pd.DataFrame, alpha: float) -> np.ndarray:
    """
    The strategy currently in production, reproduced so the comparison is fair.

    Centre is the trailing mean; the band widens with sqrt(horizon) at the same
    18% base spread the Java implementation uses. If the model cannot beat this,
    the switch is not earned and the config value stays on DETERMINISTIC.
    """
    centre = frame["trail_mean"].to_numpy()
    spread = 0.18 * np.sqrt(frame["horizon"].to_numpy())
    if alpha < 0.5:
        return centre * (1 - spread)
    if alpha > 0.5:
        return centre * (1 + spread)
    return centre


def expanding_folds(months: np.ndarray, n_folds: int = 3, min_train: int = 8):
    uniq = np.array(sorted(pd.unique(months)))
    step = max(1, (len(uniq) - min_train) // max(1, n_folds))
    for k in range(n_folds):
        cut = min_train + k * step
        if cut >= len(uniq):
            break
        test = uniq[cut : cut + step] if k < n_folds - 1 else uniq[cut:]
        yield uniq[:cut], test


def run(out_dir: Path, dsn: str = DSN) -> dict:
    df = build(dsn)
    cols = feature_columns(df)
    months = df["as_of_mo"].to_numpy()

    results: dict = {
        "rows": int(len(df)),
        "enterprises": int(df["enterprise_id"].nunique()),
        "as_of_months": int(df["as_of_mo"].nunique()),
        "horizons": int(df["horizon"].max()),
        "validation": "expanding-window time-series split over as-of months",
        "feature_order": cols,
        "quantiles": {},
    }
    print(f"rows {len(df):,} · enterprises {results['enterprises']} · months {results['as_of_months']}")

    for name, alpha in QUANTILES.items():
        model_losses, base_losses = [], []
        for train_months, test_months in expanding_folds(months):
            tr = df[df["as_of_mo"].isin(train_months)]
            te = df[df["as_of_mo"].isin(test_months)]
            if te.empty or tr.empty:
                continue
            model = lgb.LGBMRegressor(
                objective="quantile", alpha=alpha,
                n_estimators=400, learning_rate=0.05, num_leaves=15,
                min_child_samples=30, subsample=0.85, subsample_freq=1,
                colsample_bytree=0.85, verbose=-1, random_state=42,
            )
            model.fit(tr[cols], tr["target_net"])
            pred = model.predict(te[cols])
            y = te["target_net"].to_numpy()
            model_losses.append(pinball(y, pred, alpha))
            base_losses.append(pinball(y, deterministic_baseline(te, alpha), alpha))

        results["quantiles"][name] = {
            "alpha": alpha,
            "pinball_model": round(float(np.mean(model_losses)), 2),
            "pinball_deterministic": round(float(np.mean(base_losses)), 2),
        }
        q = results["quantiles"][name]
        better = q["pinball_model"] < q["pinball_deterministic"]
        print(f"{name}  pinball model {q['pinball_model']:>10,.0f}  "
              f"deterministic {q['pinball_deterministic']:>10,.0f}  "
              f"{'model wins' if better else 'DETERMINISTIC WINS'}")

    # Refit each quantile on everything for serving.
    final = {}
    for name, alpha in QUANTILES.items():
        model = lgb.LGBMRegressor(
            objective="quantile", alpha=alpha,
            n_estimators=400, learning_rate=0.05, num_leaves=15,
            min_child_samples=30, subsample=0.85, subsample_freq=1,
            colsample_bytree=0.85, verbose=-1, random_state=42,
        )
        model.fit(df[cols], df["target_net"])
        final[name] = model

    wins = sum(1 for q in results["quantiles"].values()
               if q["pinball_model"] < q["pinball_deterministic"])
    results["verdict"] = (
        f"{wins}/3 quantiles beat the deterministic strategy on held-out months"
    )
    print(f"\n{results['verdict']}")

    out_dir.mkdir(parents=True, exist_ok=True)
    with open(out_dir / "forecast_model.pkl", "wb") as fh:
        pickle.dump({"models": final, "features": cols, "horizons": HORIZONS,
                     "version": "lightgbm-forecast-v1"}, fh)
    (out_dir / "forecast_metrics.json").write_text(
        json.dumps(results | {"trainedAt": datetime.now().isoformat()}, indent=2), encoding="utf-8"
    )
    print(f"written: {out_dir / 'forecast_model.pkl'}")
    return results


def main() -> None:
    ap = argparse.ArgumentParser(description="Train the UdyamLinc quantile forecast heads")
    ap.add_argument("--out", default="models")
    ap.add_argument("--dsn", default=DSN)
    args = ap.parse_args()
    run(Path(args.out), args.dsn)


if __name__ == "__main__":
    main()