"""H3 — Concentration and the long tail of merchant catalogue sizes (S_A, all stores with >= 1 listing).
Gini, top-1 % / top-10 % shares (store bootstrap CIs), rank-size slope, discrete power-law vs log-normal (CSN MLE + Vuong LR).
Within-region Gini (S_R). Preregistered decision: Gini >= 0.70 and top-10 % share >= 0.50.
"""
from __future__ import annotations
import json
import numpy as np, pandas as pd
from scipy import stats, optimize, special
from src.common import ROOT, SEED, get_logger, ledger, write_json, PROC_DIR, EXP_DIR, P

log = get_logger("models.h3")
PROC = PROC_DIR


def gini(x):
    x = np.sort(np.asarray(x, float)); n = len(x)
    return float((2 * np.arange(1, n + 1) - n - 1).dot(x) / (n * x.sum()))


def top_share(x, q):
    x = np.sort(np.asarray(x, float))[::-1]; k = max(1, int(np.ceil(q * len(x)))); return float(x[:k].sum() / x.sum())


def powerlaw_fit(x):
    """Discrete power law with x_min chosen by KS minimisation (Clauset, Shalizi & Newman 2009)."""
    x = np.asarray(x, int); best = None
    for xmin in np.unique(x)[:-10]:
        t = x[x >= xmin]; n = len(t)
        if n < 50: break
        nll = lambda a: -(n * -np.log(special.zeta(a, xmin)) - a * np.log(t).sum()) if a > 1.0001 else 1e18
        r = optimize.minimize_scalar(nll, bounds=(1.01, 6), method="bounded"); a = r.x
        # KS distance
        xs = np.arange(xmin, t.max() + 1); cdf_m = np.cumsum(xs ** -a) / special.zeta(a, xmin)
        emp = np.searchsorted(np.sort(t), xs, side="right") / n
        ks = np.max(np.abs(emp - cdf_m))
        if best is None or ks < best["ks"]:
            best = {"xmin": int(xmin), "alpha": float(a), "ks": float(ks), "n_tail": int(n), "ll_pl": float(-nll(a))}
    return best


def lognormal_vs_powerlaw(x, fit):
    """Log-normal fitted by MAXIMUM LIKELIHOOD on the same truncated tail (x >= x_min), discrete support, compared with the
    discrete power law by a Vuong likelihood-ratio test (referee 1: untruncated moments overstated the power-law advantage)."""
    t = np.asarray(x, int); t = t[t >= fit["xmin"]]; xmin = fit["xmin"]; xs = np.arange(xmin, t.max() * 3 + 1)
    def disc_ll(params):
        mu, ls = params; sig = np.exp(ls)
        pm = stats.lognorm.cdf(xs + 0.5, s=sig, scale=np.exp(mu)) - stats.lognorm.cdf(xs - 0.5, s=sig, scale=np.exp(mu))
        pm = pm / pm.sum(); lp = np.log(np.maximum(pm[t - xmin], 1e-300)); return lp
    r = optimize.minimize(lambda p_: -disc_ll(p_).sum(), x0=[np.log(t).mean(), np.log(np.log(t).std() + 1e-9)], method="Nelder-Mead")
    ll_ln = disc_ll(r.x); mu, sig = r.x[0], float(np.exp(r.x[1]))
    a = fit["alpha"]; ll_pl = -a * np.log(t) - np.log(special.zeta(a, xmin))
    d = ll_pl - ll_ln; R = d.sum(); se = d.std(ddof=1) * np.sqrt(len(d))
    z = float(R / se) if se > 0 else None
    return {"vuong_R": float(R), "vuong_z": z, "vuong_p": float(2 * (1 - stats.norm.cdf(abs(z)))) if z is not None else None,
            "favoured": ("power-law" if R > 0 else "log-normal") if (z is not None and abs(z) > 1.96) else "indistinguishable",
            "lognormal_mu": float(mu), "lognormal_sigma": sig, "ll_lognormal": float(ll_ln.sum()), "ll_powerlaw": float(ll_pl.sum()), "method": "truncated discrete log-normal MLE on the same tail"}


def main():
    sp = pd.read_parquet(PROC / "store_panel.parquet"); rng = np.random.default_rng(SEED)
    x = sp.n_listings.values
    B = 2000; g_b = np.empty(B); t1 = np.empty(B); t10 = np.empty(B)
    for b in range(B):
        d = rng.choice(x, len(x), replace=True); g_b[b] = gini(d); t1[b] = top_share(d, 0.01); t10[b] = top_share(d, 0.10)
    ranks = np.arange(1, len(x) + 1); xs = np.sort(x)[::-1]
    slope = float(np.polyfit(np.log(ranks), np.log(xs), 1)[0])
    fit = powerlaw_fit(x); lv = lognormal_vs_powerlaw(x, fit)
    res = {"stores": int(len(x)), "listings": int(x.sum()), "gini": round(gini(x), 4), "gini_ci": [round(float(np.percentile(g_b, 2.5)), 4), round(float(np.percentile(g_b, 97.5)), 4)],
           "top1pct_share": round(top_share(x, 0.01), 4), "top1_ci": [round(float(np.percentile(t1, 2.5)), 4), round(float(np.percentile(t1, 97.5)), 4)],
           "top10pct_share": round(top_share(x, 0.10), 4), "top10_ci": [round(float(np.percentile(t10, 2.5)), 4), round(float(np.percentile(t10, 97.5)), 4)],
           "median_catalogue": float(np.median(x)), "p90": float(np.percentile(x, 90)), "p99": float(np.percentile(x, 99)), "max": int(x.max()),
           "rank_size_slope_loglog": round(slope, 4), "powerlaw": fit, "lognormal_vs_powerlaw": lv,
           "decision_prereg": bool(gini(x) >= 0.70 and top_share(x, 0.10) >= 0.50)}
    # within-region
    wr = {}
    for r, g in sp[sp.region_final != "Inconnue"].groupby("region_final"):
        v = g.n_listings.values
        if len(v) < 5: wr[r] = {"stores": int(len(v)), "gini": None}; continue
        gb = [gini(rng.choice(v, len(v), replace=True)) for _ in range(500)]
        wr[r] = {"stores": int(len(v)), "listings": int(v.sum()), "gini": round(gini(v), 3), "gini_ci": [round(float(np.percentile(gb, 2.5)), 3), round(float(np.percentile(gb, 97.5)), 3)],
                 "top10pct_share": round(top_share(v, 0.10), 3), "median": float(np.median(v))}
    res["within_region"] = wr
    pd.DataFrame({"rank": ranks, "n_listings": xs}).to_csv(PROC / "rank_size.csv", index=False)
    write_json(res, EXP_DIR / "h3_longtail.json", "h3"); ledger("models.h3_longtail", **{k: v for k, v in res.items() if k != "within_region"})
    log.info(json.dumps({k: v for k, v in res.items() if k != "within_region"}))


if __name__ == "__main__":
    main()