"""
Statistical analyses not covered by models.py.

S1  Information-theoretic model comparison
    For every element x treatment series, the conservative null (Lambda = 1, the
    concentration rise is entirely matrix shrinkage) is compared against the
    liberation model of Eq. (9). Small-sample corrected Akaike information
    criterion and Akaike weights quantify the support for liberation without
    relying on a significance threshold - appropriate here because the series
    have only six points and no replicate-level error term for Cr, Cd and Pb.

S2  Bootstrap confidence intervals on the liberation factor
    Non-parametric bootstrap over element x treatment units within each group,
    which requires no assumption about the (partly unavailable) analytical error.

S3  Correlation structure and ordination of the mature products
    Spearman rank correlation (robust at n = 5) plus principal component
    analysis of the standardised 90-day composition.
"""
import json
import numpy as np
import ministats as stats
from data import MEANS, DAYS, TRT, PL_FRAC, R

D = np.array(DAYS, float)
DT = D - D[0]
MATRIX_NOFE = ["Zn", "Cu", "Mn", "Ni"]
TOXIC = ["Cr", "Cd", "Pb"]
ALL = ["Zn", "Cu", "Fe", "Mn", "Ni"] + TOXIC
M = json.load(open("models.json"))
LAM = M["lambda_scenarios"]["S2 30 % DM loss"]
out = {}

# ═══════════════════════════════════════════════════════ S1 model comparison
def aicc(rss, n, k):
    """Small-sample corrected Akaike information criterion for a least-squares
    fit with Gaussian errors, Eq. (11) of the paper.

        AICc = n ln(RSS/n) + 2k + 2k(k+1)/(n-k-1)

    n is the number of points in the series (six sampling occasions) and k the
    number of estimated parameters including the residual variance: k = 1 for
    the conservative null, which estimates only the variance, and k = 3 for the
    liberation model, which estimates Lambda_inf, the rate constant k, and the
    variance. The correction term matters here: with n = 6 the small-sample
    penalty is far from negligible, and plain AIC would favour the richer model
    too readily."""
    if n - k - 1 <= 0:
        return np.inf                                   # correction undefined
    return n * np.log(rss / n) + 2 * k + 2 * k * (k + 1) / (n - k - 1)


def fit_liberation(traj):
    """Least-squares fit of the first-order liberation model, Eq. (10),

        Lambda(t) = Lambda_inf - (Lambda_inf - 1) exp(-k * dt)

    to one element x treatment trajectory. Returns (RSS, Lambda_inf, k).

    The fit is done by exhaustive search on a 500 x 320 grid rather than by a
    gradient method. With only six points the residual surface has a long
    curved valley in which Lambda_inf and k trade off against each other, and an
    iterative optimiser can stop at different places depending on its starting
    value. A grid search has no starting value, so the result is reproducible
    exactly. The Lambda_inf axis runs from 1 (no liberation) to three times the
    largest observed value, and k is searched on a log scale from about
    5e-4 to 0.3 per day, which brackets half-times from roughly two days to
    four years and so cannot clip the optimum at either end."""
    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)
            rss = float(((pred - traj) ** 2).sum())
            if rss < best[0]:
                best = (rss, float(Li), float(k))
    return best


print("=" * 84)
print("S1  MODEL COMPARISON: conservative null (Lambda = 1) vs liberation model")
print("=" * 84)
print(f"{'Element':<8}{'Trt':<5}{'RSS null':>11}{'RSS lib':>11}"
      f"{'AICc null':>11}{'AICc lib':>11}{'dAICc':>9}{'w(lib)':>9}  support")
rows = {}
for m in ALL:
    rows[m] = {}
    for i in range(3):
        # Observed liberation trajectory, Eq. (7): the concentration ratio
        # relative to the first sampling, with the matrix-shrinkage term
        # exp(-lambda*dt) divided out. What remains is the part of the rise that
        # dry-matter loss cannot account for. Evaluated under the mid-range
        # 30% dry-matter-loss scenario, which is the conservative choice: it
        # removes more of the rise than the tracer-based estimate does.
        c = np.array(MEANS[m][i], float)
        traj = (c / c[0]) * np.exp(-LAM * DT)
        # Conservative null: the element is fully recovered at every occasion,
        # so Lambda(t) == 1 throughout and the residual sum of squares is simply
        # the squared departure of the observed trajectory from unity.
        rss0 = float(((traj - 1.0) ** 2).sum())
        rss1, Li, k = fit_liberation(traj)
        a0, a1 = aicc(rss0, 6, 1), aicc(rss1, 6, 3)
        d = a1 - a0
        # Akaike weight of the liberation model: its share of the total
        # evidence across the two candidates, w = exp(-dAICc/2)/(1+exp(-dAICc/2)),
        # written in logistic form for numerical stability at large |dAICc|.
        w = 1 / (1 + np.exp(0.5 * d))
        sup = ("decisive" if d < -10 else "strong" if d < -4 else
               "weak" if d < 2 else "none (null favoured)")
        rows[m][TRT[i]] = dict(rss0=rss0, rss1=rss1, aicc0=a0, aicc1=a1,
                               dAICc=d, w=w, support=sup)
        print(f"{m:<8}{'T'+str(i+1):<5}{rss0:>11.4f}{rss1:>11.4f}"
              f"{a0:>11.2f}{a1:>11.2f}{d:>9.2f}{w:>9.3f}  {sup}")
    print()
out["model_comparison"] = rows
tox_d = [rows[m][TRT[i]]["dAICc"] for m in TOXIC for i in range(3)]
mat_d = [rows[m][TRT[i]]["dAICc"] for m in MATRIX_NOFE for i in range(3)]
print(f"  mean dAICc, Cr/Cd/Pb          : {np.mean(tox_d):+.1f} "
      f"(liberation favoured in {sum(1 for x in tox_d if x < -4)}/9 series)")
print(f"  mean dAICc, conservative set  : {np.mean(mat_d):+.1f} "
      f"(liberation favoured in {sum(1 for x in mat_d if x < -4)}/12 series)")
out["dAICc_toxic_mean"] = float(np.mean(tox_d))
out["dAICc_matrix_mean"] = float(np.mean(mat_d))
out["n_toxic_strong"] = int(sum(1 for x in tox_d if x < -4))
out["n_matrix_strong"] = int(sum(1 for x in mat_d if x < -4))

# ═══════════════════════════════════════════════ S2 bootstrap CI on Lambda
print("\n" + "=" * 84)
print("S2  BOOTSTRAP 95 % CONFIDENCE INTERVALS ON THE LIBERATION FACTOR")
print("=" * 84)
# Resampling unit is the element x treatment combination, not the replicate:
# replicate-level values do not survive for Cr, Cd and Pb, and the quantity of
# interest is the group mean across elements. Because each element contributes
# one value to each treatment the units are not fully independent, so these
# intervals describe the observed set of elements rather than a wider
# population; the paper states this explicitly in Sect. 4.7.
rng = np.random.default_rng(20240115)   # fixed seed: intervals reproduce exactly
NB = 20000                              # resamples
print(f"{'Scenario':<20}{'group':<20}{'mean':>8}{'95 % CI':>20}")
boot = {}
for sc in M["lambda_scenarios"]:
    boot[sc] = {}
    for nm, grp in (("conservative set", MATRIX_NOFE), ("Cr, Cd, Pb", TOXIC)):
        v = np.array([M["LIB"][sc][m][i] for m in grp for i in range(3)])
        idx = rng.integers(0, len(v), size=(NB, len(v)))
        bs = v[idx].mean(axis=1)
        lo, hi = np.percentile(bs, [2.5, 97.5])
        boot[sc][nm] = dict(mean=float(v.mean()), lo=float(lo), hi=float(hi))
        print(f"{sc.split(None,1)[1]:<20}{nm:<20}{v.mean():>8.2f}"
              f"{f'[{lo:.2f}, {hi:.2f}]':>20}")
    print()
out["bootstrap"] = boot
sc0 = list(M["lambda_scenarios"])[0]
print(f"  Under the tracer scenario the two intervals are disjoint: "
      f"conservative [{boot[sc0]['conservative set']['lo']:.2f}, "
      f"{boot[sc0]['conservative set']['hi']:.2f}] vs toxic "
      f"[{boot[sc0]['Cr, Cd, Pb']['lo']:.2f}, {boot[sc0]['Cr, Cd, Pb']['hi']:.2f}]")

# ═══════════════════════════════════ S3 correlation structure and ordination
print("\n" + "=" * 84)
print("S3  SPEARMAN RANK CORRELATION OF THE MATURE PRODUCTS (n = 5 treatments)")
print("=" * 84)


def rank(a):
    """Ranks with ties averaged, as Spearman's coefficient requires. Written out
    rather than taken from a library so that the archive has no SciPy
    dependency and the tie rule is visible."""
    a = np.asarray(a, float)
    order = a.argsort()
    r = np.empty(len(a), float)
    i = 0
    s = a[order]
    while i < len(s):
        j = i
        while j + 1 < len(s) and s[j + 1] == s[i]:
            j += 1
        r[order[i:j + 1]] = (i + j) / 2 + 1
        i = j + 1
    return r


# Mature-product composition: the day-90 mean of each element in each of the
# five blends, giving a 5 (treatment) x 8 (element) matrix. Spearman rather than
# Pearson because n = 5 and rank correlation does not assume linearity.
X = np.array([[MEANS[m][i][5] for m in ALL] for i in range(5)], float)
Rk = np.array([rank(X[:, j]) for j in range(len(ALL))]).T
C = np.corrcoef(Rk, rowvar=False)
print("      " + "".join(f"{m:>7}" for m in ALL))
for a, m in enumerate(ALL):
    print(f"{m:>5} " + "".join(f"{C[a, b]:>7.2f}" for b in range(len(ALL))))
pairs = [(ALL[a], ALL[b], C[a, b]) for a in range(len(ALL))
         for b in range(a + 1, len(ALL))]
neg = [p for p in pairs if p[2] < 0]
print(f"\n  {len(pairs) - len(neg)}/{len(pairs)} element pairs positively correlated; "
      f"all {len(neg)} negative pairs involve Fe: "
      f"{all(p[0] == 'Fe' or p[1] == 'Fe' for p in neg)}")
out["spearman"] = C.tolist()

print("\n" + "=" * 84)
print("S3b PRINCIPAL COMPONENT ANALYSIS (correlation matrix, standardised)")
print("=" * 84)
# Standardise each element to zero mean and unit variance before the
# eigendecomposition. Without this the ordination would be dominated by Zn and
# Mn purely because they are present at the highest concentrations; the question
# here is the pattern of covariation, not the absolute amounts. Equivalent to a
# PCA on the correlation matrix.
Z = (X - X.mean(axis=0)) / X.std(axis=0, ddof=1)
Cm = np.corrcoef(Z, rowvar=False)
w, V = np.linalg.eigh(Cm)
o = w.argsort()[::-1]
w, V = w[o], V[:, o]
expl = w / w.sum() * 100
print(f"  variance explained: PC1 {expl[0]:.1f} %, PC2 {expl[1]:.1f} %, "
      f"cumulative {expl[0]+expl[1]:.1f} %")
print(f"\n  {'Element':<9}{'PC1':>8}{'PC2':>8}")
for a, m in enumerate(ALL):
    print(f"  {m:<9}{V[a, 0]:>8.3f}{V[a, 1]:>8.3f}")
scores = Z @ V[:, :2]
print(f"\n  {'Treatment':<12}{'PC1':>9}{'PC2':>9}")
for i in range(5):
    print(f"  {TRT[i]:<12}{scores[i, 0]:>9.2f}{scores[i, 1]:>9.2f}")
r_pc1, p_pc1 = stats.pearsonr(PL_FRAC, scores[:, 0])
print(f"\n  PC1 score vs poultry-litter fraction: r = {r_pc1:+.4f}, P = {p_pc1:.4f}")
out["pca"] = dict(explained=expl[:4].tolist(), loadings=V[:, :2].tolist(),
                  scores=scores.tolist(), r_pc1=float(r_pc1), p_pc1=float(p_pc1))

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