#!/usr/bin/env python3
"""Two questions the essay must answer before it can say anything.

1. POWER. A firm-level diagnostic failed P49 twice. Does the POOLED panel test
   -- the one actually run -- have power to detect a real margin effect? A null
   result from a test with no power is not evidence of absence.
2. ROBUSTNESS. Does the headline survive dropping the outlier, changing the
   window, and switching estimator?
"""
import json, os, math, random, statistics as st

OUT = os.path.dirname(os.path.abspath(__file__))
d = json.load(open(os.path.join(OUT, "panel_final.json")))
WIN, BAL, P = d["window"], d["balanced"], d["panel"]
PRE, POST = WIN[:4], WIN[-4:]


def deltas(firms):
    out = {}
    for t in firms:
        pre = st.mean(P[t]["quarters"][q]["gross_margin"] for q in PRE)
        post = st.mean(P[t]["quarters"][q]["gross_margin"] for q in POST)
        out[t] = post - pre
    return out


def wilcoxon_p(ds):
    nz = [x for x in ds if x != 0]
    order = sorted(range(len(nz)), key=lambda i: abs(nz[i]))
    rk = [0]*len(nz); i = 0
    while i < len(order):
        j = i
        while j+1 < len(order) and abs(nz[order[j+1]]) == abs(nz[order[i]]): j += 1
        for k in range(i, j+1): rk[order[k]] = (i+j)/2+1
        i = j+1
    Wp = sum(rk[i] for i in range(len(nz)) if nz[i] > 0)
    N = len(nz); mu = N*(N+1)/4; sd = math.sqrt(N*(N+1)*(2*N+1)/24)
    z = (Wp-mu)/sd
    return z, 2*(1-0.5*(1+math.erf(abs(z)/math.sqrt(2))))


# ---------------------------------------------------------------- 1. POWER
print("=== POWER OF THE POOLED PANEL TEST ===")
print("53 firms, paired pre/post 4-quarter means. Firm margin sd across the panel")
print("is ~13pp; within-firm quarter noise ~1pp. How big an average compression")
print("would this design catch?\n")
random.seed(20260810)
REPS = 4000
obs = list(deltas(BAL).values())
resid_sd = st.stdev(obs)
print(f"observed sd of firm-level change = {resid_sd*100:.2f} pp (used as the noise scale)")
print(f"{'true effect':>13}{'detected at p<.05':>20}")
for eff in (0.0, -0.005, -0.01, -0.02, -0.03, -0.05):
    hit = 0
    for _ in range(REPS):
        sim = [random.gauss(eff, resid_sd) for _ in range(53)]
        z, p = wilcoxon_p(sim)
        if p < 0.05 and st.median(sim) < 0:
            hit += 1
    print(f"{eff*100:>11.1f}pp{hit/REPS:>19.1%}")
print("\n-> the design detects a true mean compression of ~2pp or more. It could NOT")
print("   rule out a compression smaller than about 1pp.")

# ------------------------------------------------------------ 2. ROBUSTNESS
print("\n=== ROBUSTNESS OF THE HEADLINE ===")
full = deltas(BAL)
z, p = wilcoxon_p(list(full.values()))
print(f"all 53          median {st.median(full.values())*100:+.2f}pp  "
      f"mean {st.mean(full.values())*100:+.2f}pp  z={z:+.2f} p={p:.4f}  "
      f"up {sum(1 for v in full.values() if v>0)}/53")

no_ai = {k: v for k, v in full.items() if k != "AI"}
z, p = wilcoxon_p(list(no_ai.values()))
print(f"drop C3.ai      median {st.median(no_ai.values())*100:+.2f}pp  "
      f"mean {st.mean(no_ai.values())*100:+.2f}pp  z={z:+.2f} p={p:.4f}  "
      f"up {sum(1 for v in no_ai.values() if v>0)}/52")

# alternative windows
for pre_n, post_n, lbl in [(2, 2, "2-qtr blocks"), (6, 6, "6-qtr blocks")]:
    pre_w, post_w = WIN[:pre_n], WIN[-post_n:]
    dd = {}
    for t in BAL:
        a = st.mean(P[t]["quarters"][q]["gross_margin"] for q in pre_w)
        b = st.mean(P[t]["quarters"][q]["gross_margin"] for q in post_w)
        dd[t] = b-a
    z, p = wilcoxon_p(list(dd.values()))
    print(f"{lbl:16}median {st.median(dd.values())*100:+.2f}pp  "
          f"mean {st.mean(dd.values())*100:+.2f}pp  z={z:+.2f} p={p:.4f}")

# post-2023 only: has it turned?
recent = {}
for t in BAL:
    a = st.mean(P[t]["quarters"][q]["gross_margin"] for q in WIN[8:12])   # 2024
    b = st.mean(P[t]["quarters"][q]["gross_margin"] for q in POST)
    recent[t] = b-a
z, p = wilcoxon_p(list(recent.values()))
print(f"\n2024 -> latest  median {st.median(recent.values())*100:+.2f}pp  "
      f"mean {st.mean(recent.values())*100:+.2f}pp  z={z:+.2f} p={p:.4f}  "
      f"up {sum(1 for v in recent.values() if v>0)}/53")

# aggregate line, last 8 quarters — is the trend rolling over?
print("\nAggregate (revenue-weighted) last 8 quarters:")
for q in WIN[-8:]:
    R = sum(P[t]["quarters"][q]["revenue"] for t in BAL)
    G = sum(P[t]["quarters"][q]["gross_profit"] for t in BAL)
    print(f"  {q}  {G/R*100:6.2f}%")

# how much of total panel revenue do the compressors represent?
comp = [t for t, v in full.items() if v < -0.01]
Rall = sum(P[t]["quarters"][WIN[-1]]["revenue"] for t in BAL)
Rc = sum(P[t]["quarters"][WIN[-1]]["revenue"] for t in comp)
print(f"\nFirms down >1pp: {len(comp)} of 53 = {Rc/Rall:.1%} of panel revenue")
print(f"  {sorted(comp)}")