"""
gtm_engine_slice.py  —  the minimal RUNNABLE slice of the GTM world model.

This is NOT the v2.0 spec (that doc is unchanged). This is the first piece of
actual machinery: the smallest subset that can be instantiated, simulated, and
falsified today. It deliberately includes only what is measurable from a monthly
MRR series, and REFUSES where the observable window is too short. No twelve-factor
Phi, no game theory, no category layer — those are the additions every critique
asked for and that would move us away from runnability, not toward it.

What is what (the epistemic contract, enforced in code):
  IDENTITY   (exact)     : the MRR walk. M_t = M_{t-1} + New - Churn - Contraction + Expansion.
                           Reconstructed and asserted every period. Not estimated.
  LATENT     (ex-ante)   : Phi_0 in [0,1]. A PMF prior FIXED before the retention
                           window is observed. It parameterizes the PRIOR over
                           monthly revenue retention. It is never fit to the outcome.
                           It is therefore FALSIFIABLE: if realized retention lands
                           outside the band Phi_0 predicted, Phi_0 was a failed
                           prediction and is flagged.
  ESTIMATED  (posterior) : monthly revenue-retention rate p, with a Beta posterior
                           whose width is tied to the NUMBER OF MONTHS observed.
  REFUSED                : LTV when the posterior on customer lifetime is too wide,
                           i.e. retention has not matured. The "noise floor" as a
                           gate, not a disclaimer. Expansion-inclusive LTV is also
                           deliberately not computed (that is where naive models
                           overstate 2-3x); NRR is reported separately instead.

Run:  python3 gtm_engine_slice.py
Deps: numpy only.
"""

from dataclasses import dataclass
import numpy as np

rng = np.random.default_rng(7)


# ----------------------------------------------------------------------------
# IDENTITY LAYER  — the MRR walk. This is exact. It is the only thing we trust
# without estimation. Everything else is inferred from it.
# ----------------------------------------------------------------------------
@dataclass
class MonthlyFlows:
    new: float           # exogenous acquisition (driven by funnel/spend; not modeled here)
    expansion: float     # upsell on the retained base
    contraction: float   # downgrade on the retained base
    churn: float         # gross revenue lost from the retained base


def mrr_walk(prev_mrr: float, f: MonthlyFlows) -> float:
    """The sole identity. Returns M_t and asserts the books balance."""
    m_t = prev_mrr + f.new + f.expansion - f.contraction - f.churn
    # the identity must hold to the penny; this is not a model, it is accounting
    assert abs((m_t) - (prev_mrr + f.new + f.expansion - f.contraction - f.churn)) < 1e-9
    return m_t


def retained_fraction(prev_mrr: float, f: MonthlyFlows) -> float:
    """Gross revenue retention on the prior base (downside only; excludes new + expansion).
    This is the quantity Phi is a prior about, and the quantity LTV depends on."""
    if prev_mrr <= 0:
        return np.nan
    return (prev_mrr - f.churn - f.contraction) / prev_mrr


# ----------------------------------------------------------------------------
# LATENT LAYER  — Phi_0 as an EX-ANTE prior. Mapping Phi_0 -> Beta prior on the
# monthly retention rate is a modeling CHOICE, stated openly. The point is only
# that Phi_0 is fixed up front and can be contradicted by data.
# ----------------------------------------------------------------------------
def phi0_to_retention_prior(phi0: float, prior_months: float = 4.0):
    """Phi_0 in [0,1] -> Beta(alpha0, beta0) prior on monthly revenue retention p.

    prior mean retention runs 0.90 (phi0=0) .. 0.99 (phi0=1).
    prior_months = strength of the prior in effective months of evidence.
    """
    phi0 = float(np.clip(phi0, 0.0, 1.0))
    prior_mean = 0.90 + 0.09 * phi0
    alpha0 = prior_mean * prior_months
    beta0 = (1.0 - prior_mean) * prior_months
    return alpha0, beta0, prior_mean


# ----------------------------------------------------------------------------
# ESTIMATION LAYER — Beta posterior on monthly retention. Each observed MONTH
# contributes a fixed effective weight, so posterior width shrinks with sqrt(T):
# few months -> wide -> refusal; many months -> tight -> forecast allowed.
# (A production version would use a hierarchical / state-space model on logit p;
# this keeps it conjugate, transparent, and honest about where uncertainty comes
# from — observed DURATION, not dollar magnitude within a month.)
# ----------------------------------------------------------------------------
def update_retention_posterior(alpha0, beta0, retained_fracs, weight_per_month=1.0):
    fracs = np.asarray([x for x in retained_fracs if np.isfinite(x)], dtype=float)
    alpha = alpha0 + weight_per_month * fracs.sum()
    beta = beta0 + weight_per_month * (len(fracs) - fracs.sum())
    return alpha, beta


def posterior_samples(alpha, beta, n=40000):
    return rng.beta(alpha, beta, size=n)


def summary(samples, lo=5, hi=95):
    return {
        "mean": float(np.mean(samples)),
        "ci_lo": float(np.percentile(samples, lo)),
        "ci_hi": float(np.percentile(samples, hi)),
    }


# ----------------------------------------------------------------------------
# REFUSAL GATE — the heart of the discipline. LTV = ARPA * margin / (1 - p),
# geometric lifetime. If the posterior on LTV is too wide relative to its
# median, retention has not matured: REFUSE rather than guess.
# ----------------------------------------------------------------------------
@dataclass
class LTVResult:
    refused: bool
    reason: str = ""
    median: float = float("nan")
    ci_lo: float = float("nan")
    ci_hi: float = float("nan")


def ltv_with_refusal(p_samples, arpa, gross_margin,
                     max_rel_ci_width=1.0, min_months=6, n_months_observed=0):
    loss_rate = np.clip(1.0 - p_samples, 1e-4, 1.0)        # guard divide-by-zero
    lifetime = 1.0 / loss_rate
    ltv = arpa * gross_margin * lifetime
    med = float(np.median(ltv))
    lo, hi = np.percentile(ltv, [5, 95])
    rel_width = (hi - lo) / med if med > 0 else np.inf

    if n_months_observed < min_months:
        return LTVResult(True,
                         f"REFUSED: only {n_months_observed} months observed "
                         f"(< {min_months} required). Retention not matured.")
    if rel_width > max_rel_ci_width:
        return LTVResult(True,
                         f"REFUSED: LTV 90% CI spans {rel_width:.1f}x the median "
                         f"(> {max_rel_ci_width:.1f}x cap). Estimate too uncertain to act on.")
    return LTVResult(False, "", med, float(lo), float(hi))


# ----------------------------------------------------------------------------
# FALSIFIABILITY — does realized retention land where Phi_0 predicted? If the
# posterior mean falls outside the Phi_0-implied prior band, Phi_0 was a failed
# ex-ante prediction. This is what converts Phi from "explains anything" into
# "can be wrong".
# ----------------------------------------------------------------------------
def phi0_calibration_check(phi0, posterior_p_samples, prior_months=4.0):
    a0, b0, prior_mean = phi0_to_retention_prior(phi0, prior_months)
    prior_band = (float(rng.beta(a0, b0, 40000).std()))
    lo, hi = prior_mean - 2 * prior_band, prior_mean + 2 * prior_band
    post_mean = float(np.mean(posterior_p_samples))
    ok = lo <= post_mean <= hi
    return ok, prior_mean, (lo, hi), post_mean


# ----------------------------------------------------------------------------
# DEMO — synthetic monthly series. Shows: the identity balancing, the refusal
# gate triggering on a short window and releasing on a long one, NRR reported
# separately, and Phi_0 caught when it is miscalibrated.
# ----------------------------------------------------------------------------
def synth_series(n_months, true_churn, true_contr, true_exp, base0=100_000.0,
                 new_per_month=8_000.0, noise=0.15):
    """Generate a monthly MRR series and its retained fractions."""
    mrr = base0
    fracs, walk = [], [mrr]
    for _ in range(n_months):
        churn = max(0.0, rng.normal(true_churn, true_churn * noise)) * mrr
        contr = max(0.0, rng.normal(true_contr, true_contr * noise)) * mrr
        exp = max(0.0, rng.normal(true_exp, true_exp * noise)) * mrr
        f = MonthlyFlows(new=new_per_month, expansion=exp, contraction=contr, churn=churn)
        fr = retained_fraction(mrr, f)
        mrr = mrr_walk(mrr, f)          # identity advances + self-checks
        fracs.append(fr); walk.append(mrr)
    return walk, fracs


def run_case(name, phi0, true_churn, true_contr, true_exp, n_months,
             arpa=1200.0, margin=0.80):
    print(f"\n{'='*72}\n{name}\n{'='*72}")
    a0, b0, prior_mean = phi0_to_retention_prior(phi0)
    walk, fracs = synth_series(n_months, true_churn, true_contr, true_exp)
    a, b = update_retention_posterior(a0, b0, fracs)
    ps = posterior_samples(a, b)
    s = summary(ps)

    nrr = 1.0 + true_exp - true_contr - true_churn   # reported, not folded into LTV

    print(f"  Phi_0 = {phi0:.2f}  -> ex-ante prior mean monthly retention = {prior_mean:.3f}")
    print(f"  Identity check: MRR ${walk[0]:,.0f} -> ${walk[-1]:,.0f} over {n_months} mo "
          f"(books balanced every month).")
    print(f"  Posterior monthly retention p: mean {s['mean']:.3f}  "
          f"90% CI [{s['ci_lo']:.3f}, {s['ci_hi']:.3f}]   (from {n_months} months)")
    print(f"  NRR (reported separately, NOT in LTV): {nrr:.3f}")

    res = ltv_with_refusal(ps, arpa, margin, n_months_observed=n_months)
    if res.refused:
        print(f"  LTV: {res.reason}")
    else:
        print(f"  LTV: ${res.median:,.0f}  90% CI [${res.ci_lo:,.0f}, ${res.ci_hi:,.0f}]")

    ok, pm, band, post_mean = phi0_calibration_check(phi0, ps)
    verdict = "PASS — Phi_0 was a useful ex-ante signal" if ok else \
              "FAIL — Phi_0 contradicted by data; flag and do not reuse"
    print(f"  Phi_0 falsifiability: predicted ~{pm:.3f} (band {band[0]:.3f}-{band[1]:.3f}), "
          f"realized {post_mean:.3f}  ->  {verdict}")


if __name__ == "__main__":
    # 1) Calibrated, but only 3 months of data -> LTV refused (window too short).
    run_case("CASE 1  calibrated venture, 3 months  (refusal expected)",
             phi0=0.80, true_churn=0.020, true_contr=0.005, true_exp=0.030, n_months=3)

    # 2) Same venture, 18 months -> retention matured -> LTV allowed.
    run_case("CASE 2  same venture, 18 months  (forecast allowed)",
             phi0=0.80, true_churn=0.020, true_contr=0.005, true_exp=0.030, n_months=18)

    # 3) Phi_0 claimed strong fit, reality is bad retention -> Phi_0 caught as a
    #    failed prediction. Phi is falsifiable here, not a universal excuse.
    run_case("CASE 3  Phi_0 overclaims fit, real retention poor  (Phi_0 falsified)",
             phi0=0.85, true_churn=0.060, true_contr=0.010, true_exp=0.010, n_months=18)

    print(f"\n{'='*72}")
    print("What this slice proves it can do (that the spec could not):")
    print("  - advance an exact identity and self-check it,")
    print("  - carry uncertainty and report intervals, not point fictions,")
    print("  - REFUSE a forecast when the observable window is too short,")
    print("  - make Phi falsifiable by fixing it ex ante and testing it.")
    print("What it still cannot do (honest scope): regime switching (Psi),")
    print("  the funnel/bottleneck workflow, brand stock, and any learning rule.")
    print("Those are the next slices — each measurable, each runnable, none added")
    print("until this one holds.")
    print('='*72)
