"""
Mechanistic and risk models.

M1  Mass-conservative partition model
    ------------------------------------------------------------------
    Metal mass in a static pile is conserved. The measured pseudo-total
    concentration can therefore change for only two reasons:

      (i)  the matrix loses dry mass          m(t) = m0 exp(-lambda t)
      (ii) the fraction of the metal inventory that the digestion actually
           recovers changes                   f(t)

    C(t) = M f(t) / m(t)   =>   C(t)/C(t0) = LAMBDA(t) exp(lambda dt)

    where  LAMBDA(t) = f(t)/f(t0)  is the LIBERATION FACTOR: the part of the
    concentration rise that matrix loss cannot explain. Rearranged,

           LAMBDA(t) = [C(t)/C(t0)] exp(-lambda dt)

    This is exact and requires no curve fitting once lambda is specified.
    LAMBDA = 1 means fully conservative behaviour. LAMBDA > 1 means metal that
    the digestion could not reach at t0 has become reachable.

    A lower bound on the initially unrecovered ("occluded") fraction follows:
           theta0 >= 1 - 1/LAMBDA

    lambda is not identifiable from concentration data alone, so it is treated
    as a controlled parameter under three scenarios spanning the entire
    physically plausible range of composting dry-matter loss.

M2  Kinetics of liberation - first-order approach to an asymptote.
M3  Monte-Carlo compliance model.
M4  Long-term soil accumulation model.
M5  Geochemical indices (CF, Igeo, PLI, Hakanson RI).
"""
import numpy as np, json
import ministats as stats
from data import MEANS, DAYS, TRT, PL_FRAC, R

D = np.array(DAYS, float)
DT = D - D[0]
TSPAN = DT[-1]
MATRIX = ["Zn", "Cu", "Fe", "Mn", "Ni"]
TOXIC = ["Cr", "Cd", "Pb"]
ALL = MATRIX + TOXIC
REF = ["Zn", "Mn"]          # conservative reference set (see note below)

BACKGROUND = {"Zn": 70.0, "Cu": 30.0, "Fe": 35000.0, "Mn": 490.0,
              "Ni": 29.0, "Cr": 60.0, "Cd": 0.41, "Pb": 27.0}   # Kabata-Pendias (2011)
TR = {"Zn": 1, "Cr": 2, "Cu": 5, "Ni": 5, "Pb": 5, "Cd": 30}    # Hakanson (1980)

out = {}
EF = lambda m, i: MEANS[m][i][5] / MEANS[m][i][0]

# =====================================================================
# lambda scenarios
# =====================================================================
print("=" * 80)
print("M1  DRY-MATTER LOSS SCENARIOS")
print("=" * 80)
print("  Fe is excluded from the reference set: it is the only element enriched in the")
print("  vegetable-waste endmember and therefore tracks mineral (soil/dust) input, not")
print("  the organic matrix. Zn and Mn are the most temporally stable elements and are")
print("  used as conservative tracers for the lower-bound scenario.\n")

lam_ref = {}
for i, t in enumerate(TRT):
    e = np.mean([np.log(EF(m, i)) for m in REF]) / TSPAN
    lam_ref[t] = float(e)
lam_lo = float(np.mean(list(lam_ref.values())))
lam_30 = -np.log(1 - 0.30) / TSPAN
lam_50 = -np.log(1 - 0.50) / TSPAN
SCEN = {"S1 tracer-based": lam_lo, "S2 30 % DM loss": lam_30, "S3 50 % DM loss": lam_50}
for nm, L in SCEN.items():
    print(f"  {nm:<20} lambda = {L:.5f} /d   ->  {(1-np.exp(-L*TSPAN))*100:5.1f} % "
          f"dry-matter loss over the sampling window")
out["lambda_scenarios"] = SCEN

# =====================================================================
# M1  liberation factor LAMBDA (exact)
# =====================================================================
print("\n" + "=" * 80)
print("M1  LIBERATION FACTOR  LAMBDA = EF x exp(-lambda dt)   [1.00 = conservative]")
print("=" * 80)
LIB = {}
for nm, L in SCEN.items():
    corr = np.exp(-L * TSPAN)
    print(f"\n  --- {nm} (correction factor {corr:.3f}) ---")
    print(f"  {'Element':<9}" + "".join(f"{t.split()[0]:>9}" for t in TRT) +
          f"{'T1-T3 mean':>13}{'theta0 min':>12}")
    LIB[nm] = {}
    for m in ALL:
        v = [EF(m, i) * corr for i in range(5)]
        LIB[nm][m] = v
        m3 = float(np.mean(v[:3]))
        th = max(0.0, 1 - 1 / m3)
        print(f"  {m:<9}" + "".join(f"{x:>9.2f}" for x in v) +
              f"{m3:>13.2f}{th*100:>11.1f}%")
    a = [LIB[nm][m][i] for m in MATRIX if m != "Fe" for i in range(3)]
    b = [LIB[nm][m][i] for m in TOXIC for i in range(3)]
    U, p = stats.mannwhitneyu(a, b)
    print(f"  {'':<9}matrix (excl. Fe) {np.mean(a):.2f} vs toxic {np.mean(b):.2f}   "
          f"Mann-Whitney U = {U:.0f}, P = {p:.2g}")
    out.setdefault("lambda_test", {})[nm] = dict(matrix=float(np.mean(a)),
                                                 toxic=float(np.mean(b)),
                                                 U=float(U), p=float(p))
out["LIB"] = LIB

# =====================================================================
# M2  kinetics of liberation
# =====================================================================
def fit_kinetics(traj, dt):
    """traj = LAMBDA(t) series. Fit LAMBDA(t) = Li - (Li-1) exp(-k dt)."""
    best = (np.inf, np.nan, np.nan)
    for Li in np.linspace(1.0, max(1.05, traj.max() * 3.0), 500):
        for k in np.logspace(-3.3, -0.5, 320):
            pred = Li - (Li - 1) * np.exp(-k * dt)
            sse = float(((pred - traj) ** 2).sum())
            if sse < best[0]:
                best = (sse, float(Li), float(k))
    sse, Li, k = best
    pred = Li - (Li - 1) * np.exp(-k * dt)
    denom = float(((traj - traj.mean()) ** 2).sum())
    r2 = 1 - sse / denom if denom > 0 else np.nan
    return Li, k, r2, pred

print("\n" + "=" * 80)
print("M2  KINETICS OF LIBERATION  (scenario S2, 30 % dry-matter loss)")
print("=" * 80)
L2 = SCEN["S2 30 % DM loss"]
print(f"  {'Element':<9}{'Trt':<6}{'LAMBDA_inf':>12}{'k (1/d)':>10}{'t-half (d)':>12}{'R2':>8}")
KIN = {}
for m in TOXIC + ["Ni", "Cu"]:
    KIN[m] = {}
    for i in range(3):
        c = np.array(MEANS[m][i], float)
        traj = (c / c[0]) * np.exp(-L2 * DT)
        Li, k, r2, pred = fit_kinetics(traj, DT)
        KIN[m][TRT[i]] = dict(Linf=Li, k=k, r2=r2,
                              traj=traj.tolist(), pred=pred.tolist())
        print(f"  {m:<9}{TRT[i].split()[0]:<6}{Li:>12.2f}{k:>10.4f}"
              f"{np.log(2)/k:>12.1f}{r2:>8.3f}")
    print()
out["kinetics"] = {m: {t: {kk: vv for kk, vv in KIN[m][t].items()}
                       for t in KIN[m]} for m in KIN}

kt = [KIN[m][TRT[i]]["k"] for m in TOXIC for i in range(3)]
print(f"  liberation rate constant across Cr, Cd, Pb (T1-T3): "
      f"k = {np.mean(kt):.4f} +/- {np.std(kt, ddof=1):.4f} /d, "
      f"half-time {np.log(2)/np.mean(kt):.0f} d")
out["k_toxic_mean"] = float(np.mean(kt))

# =====================================================================
# M3  Monte-Carlo compliance
# =====================================================================
print("\n" + "=" * 80)
print("M3  MONTE-CARLO COMPLIANCE MODEL (Cd, EU 2019/1009 limit 1.5 mg/kg DM)")
print("=" * 80)
rng = np.random.default_rng(20240115)
NS = 200_000
y = np.array([MEANS["Cd"][i][5] for i in range(5)], float)
x = np.array(PL_FRAC, float)
sl, ic, r, p, se_sl = stats.linregress(x, y)
resid = y - (ic + sl * x)
s_res = float(np.sqrt((resid ** 2).sum() / 3))
se_ic = s_res * np.sqrt(1 / 5 + x.mean() ** 2 / float(((x - x.mean()) ** 2).sum()))
CV_AN = 0.208
# The intercept and slope of an ordinary least-squares fit are correlated,
# here strongly so: corr(b0, b1) = -xbar / sqrt(Sxx/n + xbar^2) = -0.82.
# Drawing them independently would therefore overstate the uncertainty of the
# fitted mean. The regression component is instead drawn directly from the
# sampling distribution of the fitted mean at each inclusion rate, whose
# standard error s*sqrt(1/n + (f - xbar)^2 / Sxx) already carries that
# covariance.
xbar = float(x.mean())
Sxx = float(((x - xbar) ** 2).sum())
CORR_B0B1 = -xbar / np.sqrt(Sxx / len(x) + xbar ** 2)
grid = np.linspace(0, 1, 201)
pex = np.empty_like(grid)
for j, f in enumerate(grid):
    se_fit = s_res * np.sqrt(1 / len(x) + (f - xbar) ** 2 / Sxx)
    mu = (ic + sl * f) + rng.standard_normal(NS) * se_fit
    pex[j] = float((mu * (1 + rng.standard_normal(NS) * CV_AN) > 1.5).mean())
print(f"  Cd = {sl:.3f} f_PL {ic:+.3f}   (r = {r:.4f}, P = {p:.4f})")
print(f"  residual SD = {s_res:.4f}; analytical CV = {CV_AN*100:.1f} %\n")
print(f"  {'target P(exceed)':>18}{'max poultry litter':>22}")
for tgt in (0.05, 0.10, 0.25, 0.50):
    k_ = int(np.argmin(np.abs(pex - tgt)))
    print(f"  {tgt:>17.0%}{grid[k_]*100:>21.1f} %")
print(f"\n  {'PL %':>6}{'P(exceed)':>12}")
for f in (0, 10, 20, 25, 30, 40, 50, 60, 75, 100):
    print(f"  {f:>6}{pex[int(f*2)]:>12.3f}")
out["mc"] = dict(grid=grid.tolist(), p_exceed=pex.tolist(), slope=sl,
                 intercept=ic, se_slope=se_sl, se_int=float(se_ic),
                 corr_b0b1=float(CORR_B0B1), s_res=s_res, cv=CV_AN)

# =====================================================================
# M4  soil accumulation
# =====================================================================
print("\n" + "=" * 80)
print("M4  SOIL ACCUMULATION UNDER REPEATED APPLICATION")
print("=" * 80)
SOIL_MASS = 0.20 * 1300 * 1e4 / 1000
CD_LIMIT, CD_BG = 1.0, 0.41
print(f"  plough layer 0-20 cm, bulk density 1.3 Mg/m3 -> {SOIL_MASS:,.0f} Mg soil/ha")
print(f"  target soil Cd limit {CD_LIMIT} mg/kg (EU 86/278/EEC, pH 6-7); "
      f"background {CD_BG} mg/kg\n")
print(f"  {'Treatment':<12}{'Cd (mg/kg)':>12}" +
      "".join(f"{f'{a} Mg/ha/yr':>16}" for a in (5, 10, 20)))
acc = {}
for i, t in enumerate(TRT):
    c = MEANS["Cd"][i][5]
    row, cells = {}, ""
    for a in (5, 10, 20):
        d = a * c / SOIL_MASS
        yrs = (CD_LIMIT - CD_BG) / d if d > 0 else np.inf
        row[str(a)] = float(yrs)
        cells += f"{(f'{yrs:,.0f}' if yrs < 1e5 else '>100 000'):>16}"
    acc[t] = row
    print(f"  {t:<12}{c:>12.3f}{cells}")
print("\n  (years of continuous annual application before the plough layer reaches the limit)")
out["accumulation"] = acc

# =====================================================================
# M5  geochemical indices
# =====================================================================
print("\n" + "=" * 80)
print("M5  GEOCHEMICAL INDICES OF THE MATURE PRODUCT")
print("=" * 80)
idx = {}
print(f"  {'':<8}" + "".join(f"{m:>8}" for m in ALL) + f"{'PLI':>8}{'RI':>9}{'class':>15}")
for i, t in enumerate(TRT):
    cf = {m: MEANS[m][i][5] / BACKGROUND[m] for m in ALL}
    igeo = {m: float(np.log2(cf[m] / 1.5)) for m in ALL}
    PL_SET = list(TR)          # PLI over the six regulated pollutants (Fe/Mn excluded)
    pli = float(np.prod([cf[m] for m in PL_SET]) ** (1 / len(PL_SET)))
    ri = float(sum(TR[m] * cf[m] for m in TR))
    cls = ("low" if ri < 150 else "moderate" if ri < 300 else
           "considerable" if ri < 600 else "very high")
    idx[t] = dict(CF=cf, Igeo=igeo, PLI=pli, RI=ri, cls=cls)
    print(f"  {t.split()[0]:<8}" + "".join(f"{cf[m]:>8.2f}" for m in ALL) +
          f"{pli:>8.2f}{ri:>9.1f}{cls:>15}")
print("\n  values are contamination factors CF = C_compost / C_background;"
        " PLI and RI over Zn, Cr, Cu, Ni, Pb, Cd")
print(f"  Cd contribution to RI: " +
      ", ".join(f"{t.split()[0]} {TR['Cd']*idx[t]['CF']['Cd']/idx[t]['RI']*100:.0f} %"
                for t in TRT))
out["indices"] = {t: dict(PLI=idx[t]["PLI"], RI=idx[t]["RI"], cls=idx[t]["cls"],
                          CF=idx[t]["CF"], Igeo=idx[t]["Igeo"]) for t in TRT}

# =====================================================================
# soil-ingress falsification
# =====================================================================
print("\n" + "=" * 80)
print("SOIL-INGRESS FALSIFICATION")
print("=" * 80)
fFe = (MEANS["Fe"][0][5] - MEANS["Fe"][0][0]) / (BACKGROUND["Fe"] - MEANS["Fe"][0][0])
print(f"  mineral fraction implied by the Fe rise: {fFe*100:.3f} % w/w")
si = {}
for m in TOXIC:
    d = MEANS[m][0][5] - MEANS[m][0][0]
    contrib = fFe * (BACKGROUND[m] - MEANS[m][0][0])
    si[m] = dict(obs=float(d), soil=float(contrib), pct=float(contrib / d * 100))
    print(f"  {m}: observed {d:+7.2f} mg/kg | mineral input contributes {contrib:+8.4f} "
          f"({contrib/d*100:6.2f} %)")
print(f"\n  Background Cd ({BACKGROUND['Cd']} mg/kg) is below the initial compost Cd "
      f"({MEANS['Cd'][0][0]} mg/kg):\n  mineral ingress would dilute Cd, not raise it.")
out["soil_ingress"] = dict(fraction=float(fFe), detail=si)

# =====================================================================
# M8  Dose-response of the mature product on litter fraction, with the
#     Holm-Bonferroni adjustment for eight simultaneous tests
# =====================================================================
# Section 3.1. Each element's day-90 concentration is regressed on the poultry
# litter mass fraction. Because one regression is fitted per element, eight
# probabilities are generated from the same dataset and the chance of at least
# one spurious result at the nominal 5% level is far above 5%. Holm's step-down
# procedure controls the family-wise error rate: order the raw probabilities
# ascending, multiply the smallest by 8, the next by 7, and so on, then enforce
# monotonicity so that an adjusted value can never fall below the one before it.
# Holm is uniformly more powerful than plain Bonferroni and, unlike a
# false-discovery-rate method, controls the probability of any false positive,
# which is the appropriate criterion when the claim is that essentially every
# element tracks the blend.
print("\n" + "=" * 80)
print("M8  DOSE-RESPONSE AND HOLM-BONFERRONI ADJUSTMENT")
print("=" * 80)
raw = []
for m in ALL:
    y = [MEANS[m][i][5] for i in range(5)]
    sl_, ic_, r_, p_, se_ = stats.linregress(PL_FRAC, y)
    raw.append(dict(element=m, slope=float(sl_), intercept=float(ic_),
                    r=float(r_), p_raw=float(p_)))
order = sorted(raw, key=lambda d: d["p_raw"])
running = 0.0
for i, d in enumerate(order):
    running = min(1.0, max(running, d["p_raw"] * (len(order) - i)))
    d["p_holm"] = float(running)
print(f"  {'Element':<9}{'slope':>10}{'r':>9}{'P raw':>11}{'P Holm':>11}  verdict")
for d in order:
    print(f"  {d['element']:<9}{d['slope']:>10.3f}{d['r']:>9.3f}"
          f"{d['p_raw']:>11.4f}{d['p_holm']:>11.4f}  "
          f"{'retained' if d['p_holm'] < 0.05 else 'not significant'}")
kept = [d['element'] for d in order if d['p_holm'] < 0.05]
print(f"\n  {len(kept)}/8 elements retained after adjustment; "
      f"dropped: {[d['element'] for d in order if d['p_holm'] >= 0.05]}")
out["dose_response"] = {d["element"]: d for d in order}

# =====================================================================
# M9  Mass-loss ceiling test
# =====================================================================
# Section 3.3, Eq. (9). If an element is conservative and matrix shrinkage is
# the only process acting, its enrichment cannot exceed 1/(1 - L), where L is
# the fractional dry-matter loss. Inverting that gives the loss each element
# would require to explain its own enrichment by shrinkage alone,
# L_required = 1 - 1/EF. Values beyond about 0.5 are not attainable in a
# composting pile, so any element demanding more than that cannot be
# conservative. Evaluated on the mean over T1-T3, the litter-rich blends in
# which every concentration sits well above the detection limit.
print("\n" + "=" * 80)
print("M9  MASS-LOSS CEILING TEST")
print("=" * 80)
for L in (0.30, 0.50):
    print(f"  a pile losing {L*100:.0f}% of its dry matter can enrich an element "
          f"at most {1/(1-L):.2f}-fold")
print(f"\n  {'Element':<9}{'EF (T1-T3)':>12}{'L required':>13}  attainable?")
ceil_out = {}
for m in ALL:
    ef = float(np.mean([EF(m, i) for i in range(3)]))
    Lreq = 1 - 1 / ef
    ceil_out[m] = dict(EF_T1_T3=ef, L_required=float(Lreq))
    print(f"  {m:<9}{ef:>12.2f}{Lreq*100:>12.1f}%  "
          f"{'yes' if Lreq <= 0.50 else 'NO - exceeds any attainable loss'}")
out["mass_loss_ceiling"] = ceil_out

json.dump(out, open("models.json", "w"), indent=1, default=float)
print("\n[saved models.json]")
