#!/usr/bin/env python3
"""
Reanalysis of public soil shotgun metagenomes (Berkelmann et al. 2020, Environmental Microbiome 15:5;
NCBI SRA PRJNA562410): rainforest vs jungle rubber, rubber and oil palm in Jambi, Sumatra.

Author: Mc Arthur L. Cababan, PhD
Inputs (supplementary files of Berkelmann et al. 2020):
  --s5  40793_2020_353_MOESM17_ESM.txt   KEGG functional counts (Table S5)
  --s4  40793_2020_353_MOESM16_ESM.txt   taxonomic counts (Table S4)
  --s3  40793_2020_353_MOESM15_ESM.xlsx  soil properties (Table S3)
  --s1  40793_2020_353_MOESM1_ESM.csv    read counts (Table S1)
Requires: pandas numpy scipy scikit-learn matplotlib openpyxl
"""
import argparse, os, warnings
import numpy as np, pandas as pd
from scipy import stats
from sklearn.manifold import MDS
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
warnings.filterwarnings("ignore")

SEED = 2026
RNG = np.random.default_rng(SEED)
LU = {"F": "Rainforest", "J": "Jungle rubber", "R": "Rubber", "O": "Oil palm"}
LU_ORDER = ["Rainforest", "Jungle rubber", "Rubber", "Oil palm"]
LS = {"B": "Bukit Duabelas", "H": "Harapan"}
# validated categorical palette (dataviz reference slots 1-4); entity-fixed across all figures
COL = {"Rainforest": "#2a78d6", "Jungle rubber": "#eb6834", "Rubber": "#1baf7a", "Oil palm": "#eda100"}
INK, MUTED, GRID = "#1a1a19", "#5f5e5a", "#e4e3df"
plt.rcParams.update({"font.family": "DejaVu Sans", "font.size": 9, "axes.edgecolor": MUTED,
                     "axes.labelcolor": INK, "xtick.color": MUTED, "ytick.color": MUTED,
                     "axes.spines.top": False, "axes.spines.right": False, "savefig.dpi": 600})

# Marker genes (KEGG orthologs). mcrA = methyl-coenzyme M reductase (K00399).
GENES = [
    ("Nitrogen cycle", "nifH", "K02588"), ("Nitrogen cycle", "amoA/pmoA", "K10944"),
    ("Nitrogen cycle", "narG", "K00370"), ("Nitrogen cycle", "nirK", "K00368"),
    ("Nitrogen cycle", "nirS", "K15864"), ("Nitrogen cycle", "nosZ", "K00376"),
    ("Methane cycle", "mmoX", "K16157"), ("Methane cycle", "mmoY", "K16158"),
    ("Methane cycle", "mxaF", "K14028"), ("Methane cycle", "mcrA", "K00399"),
    ("Motility", "fliG", "K02410"), ("Motility", "fliH", "K02411"), ("Motility", "fliK", "K02414"),
    ("Motility", "fliM", "K02416"), ("Motility", "fliN", "K02417"),
    ("Chemotaxis", "cheA", "K03407"), ("Chemotaxis", "cheW", "K03408"), ("Chemotaxis", "cheX", "K03409"),
    ("Chemotaxis", "cheY", "K03413"), ("Chemotaxis", "cheR", "K00575"),
    ("Secretion", "virB4 (T4SS)", "K03199"), ("Secretion", "virD4 (T4SS)", "K03205"),
    ("Secretion", "vgrG (T6SS)", "K11904"), ("Secretion", "hcp (T6SS)", "K11903"),
]
# KEGG level-3 pathways, exact names as they appear in Table S5
PATHWAYS = ["Nitrogen metabolism", "Methane metabolism", "Carbon fixation pathways in prokaryotes",
            "Sulfur metabolism", "Photosynthesis", "Bacterial chemotaxis", "Flagellar assembly",
            "Quorum sensing", "Bacterial secretion system"]
# Berkelmann et al. (2020) Fig. S7: DESeq2 log2 fold changes vs rainforest (J, R, O) per landscape,
# transcribed from the published figure, used only to validate this reanalysis.
PUBLISHED_L2FC = {
    "Bacterial chemotaxis": {"B": [-0.185, -0.18, -0.131], "H": [-0.448, -0.535, -0.476]},
    "Flagellar assembly": {"B": [-0.31, -0.452, -0.44], "H": [-0.644, -0.907, -0.744]},
    "Quorum sensing": {"B": [-0.029, -0.169, -0.01], "H": [-0.096, -0.044, -0.003]},
    "Carbon fixation pathways in prokaryotes": {"B": [-0.001, 0.015, 0.073], "H": [-0.001, 0.107, 0.115]},
    "Methane metabolism": {"B": [-0.009, -0.007, 0.02], "H": [-0.029, 0.032, 0.06]},
    "Nitrogen metabolism": {"B": [-0.043, -0.141, -0.06], "H": [-0.073, -0.033, 0.005]},
    "Sulfur metabolism": {"B": [-0.044, -0.113, -0.068], "H": [-0.143, -0.194, -0.149]},
    "Photosynthesis": {"B": [-0.014, 0.016, 0.093], "H": [-0.035, 0.064, 0.087]},
}


# ----------------------------------------------------------------- statistics helpers
def bh(p):
    p = np.asarray(p, float); out = np.full_like(p, np.nan); ok = ~np.isnan(p)
    v = p[ok]; n = len(v)
    if n:
        o = np.argsort(v); r = v[o] * n / np.arange(1, n + 1)
        r = np.minimum.accumulate(r[::-1])[::-1]; a = np.empty(n); a[o] = np.minimum(r, 1); out[ok] = a
    return out


def design(meta, interaction=False):
    cols = [np.ones(len(meta))]
    for lu in LU_ORDER[1:]:
        cols.append((meta.landuse == lu).values.astype(float))
    h = (meta.landscape == "Harapan").values.astype(float); cols.append(h)
    if interaction:
        for lu in LU_ORDER[1:]:
            cols.append((meta.landuse == lu).values * h)
    return np.column_stack(cols)


def ols(y, X):
    b, *_ = np.linalg.lstsq(X, y, rcond=None)
    res = y - X @ b; df = len(y) - X.shape[1]; s2 = res @ res / df
    se = np.sqrt(np.diag(s2 * np.linalg.inv(X.T @ X)))
    return b, se, df, res @ res


def landuse_model(y, meta):
    """ln(y) ~ land use + landscape. Returns per land use: estimate, 95% CI, p; interaction p; residual normality p."""
    X = design(meta); b, se, df, rss0 = ols(y, X)
    Xi = design(meta, True); _, _, dfi, rss1 = ols(y, Xi)
    F = ((rss0 - rss1) / 3) / (rss1 / dfi); p_int = stats.f.sf(F, 3, dfi)
    tcrit = stats.t.ppf(0.975, df); out = {}
    for i, lu in enumerate(LU_ORDER[1:], start=1):
        t = b[i] / se[i]
        out[lu] = (b[i], b[i] - tcrit * se[i], b[i] + tcrit * se[i], 2 * stats.t.sf(abs(t), df))
    resid = y - X @ b
    return out, p_int, stats.shapiro(resid).pvalue


def strat_perm(y, meta, lu, n_perm=9999):
    """Permutation test (labels shuffled within landscape) of the landscape-averaged mean difference."""
    sel = meta.landuse.isin(["Rainforest", lu]).values
    yy, g, s = y[sel], (meta.landuse[sel] == lu).values, meta.landscape[sel].values
    def stat(gg):
        return np.mean([yy[(s == k) & gg].mean() - yy[(s == k) & ~gg].mean() for k in np.unique(s)])
    obs = stat(g); cnt = 0
    for _ in range(n_perm):
        gp = g.copy()
        for k in np.unique(s):
            idx = np.where(s == k)[0]; gp[idx] = RNG.permutation(g[idx])
        cnt += abs(stat(gp)) >= abs(obs) - 1e-12
    return (cnt + 1) / (n_perm + 1)


def bray(X):
    X = np.asarray(X, float); D = np.zeros((len(X), len(X)))
    for i in range(len(X)):
        D[i] = np.abs(X[i] - X).sum(1) / (X[i] + X).sum(1)
    return D


def permanova(D, g, strata=None, n_perm=9999):
    g = np.asarray(g); n = len(g); D2 = D ** 2; iu = np.triu_indices(n, 1); sst = D2[iu].sum() / n
    def F(gg):
        ssw = 0.0
        for l in np.unique(gg):
            i = np.where(gg == l)[0]; ssw += D2[np.ix_(i, i)][np.triu_indices(len(i), 1)].sum() / len(i)
        a = len(np.unique(gg)); return ((sst - ssw) / (a - 1)) / (ssw / (n - a)), (sst - ssw) / sst
    f0, r2 = F(g); c = 0
    for _ in range(n_perm):
        gp = g.copy()
        if strata is None:
            gp = RNG.permutation(g)
        else:
            for k in np.unique(strata):
                i = np.where(strata == k)[0]; gp[i] = RNG.permutation(g[i])
        c += F(gp)[0] >= f0
    return f0, r2, (c + 1) / (n_perm + 1)


def shannon(c):
    c = c[c > 0]; p = c / c.sum(); return -(p * np.log(p)).sum()


# ----------------------------------------------------------------- loading
def load_soil(path):
    x = pd.read_excel(path, header=1).iloc[:32].copy()
    x.columns = ["sample", "P_mg_kg", "N_mg_kg", "C_g_kg", "C_N", "pH", "Al_g_m2", "Ca_g_m2", "Fe_g_m2",
                 "K_g_m2", "Mg_g_m2", "Mn_g_m2", "Na_g_m2", "ECEC_mmolc_kg", "Base_saturation_pct",
                 "Basal_respiration", "Microbial_biomass_C", "H2O_pct", "Land_use_index"]
    x = x.set_index("sample")
    log = []
    # Repair values whose decimal point was lost in the published spreadsheet (x1000 inflation).
    for col, lim in (("C_g_kg", 1000), ("N_mg_kg", 1e5), ("P_mg_kg", 1000)):
        bad = x[col] >= lim
        for s in x.index[bad]:
            log.append([s, col, x.at[s, col], x.at[s, col] / 1000])
        x.loc[bad, col] = x.loc[bad, col] / 1000
    x["CN_check"] = (x.C_g_kg * 1000 / x.N_mg_kg) / x.C_N
    return x, pd.DataFrame(log, columns=["sample", "variable", "published_value", "corrected_value"])


def load_functions(path):
    s = pd.read_csv(path, sep="\t", skiprows=1)
    samples = [c for c in s.columns if c not in ("KO", "pathway")]
    s["L3"] = s.pathway.str.split(";").str[2]
    s = s[s.KO != "Unclassified"]                                 # reads without a KEGG assignment
    ko = s.drop_duplicates("KO").set_index("KO")[samples]        # each KO once (rows repeat per pathway)
    pw = {p: s[s.L3 == p].drop_duplicates("KO")[samples].sum() for p in PATHWAYS}
    return ko, pd.DataFrame(pw).T, samples


def load_taxa(path):
    t = pd.read_csv(path, sep="\t", skiprows=1)
    samples = [c for c in t.columns if c not in ("#OTU ID", "taxonomy")]
    parts = t.taxonomy.str.split("; ", expand=True)
    t["domain"] = parts[0]
    t["order"] = parts[0] + "; " + parts[1].fillna("NA") + "; " + parts[3].fillna("NA")
    return t, samples


# ----------------------------------------------------------------- main analysis
def main():
    ap = argparse.ArgumentParser()
    for k in ("s5", "s4", "s3", "s1", "out"):
        ap.add_argument("--" + k, required=k != "s4")
    a = ap.parse_args(); out = a.out; os.makedirs(out, exist_ok=True)

    ko, pw, samples = load_functions(a.s5)
    raw = pd.read_csv(a.s5, sep="\t", skiprows=1)
    unclassified = raw[raw.KO == "Unclassified"].iloc[0][samples].astype(float)
    meta = pd.DataFrame({"landscape": [LS[s[0]] for s in samples], "landuse": [LU[s[1]] for s in samples]},
                        index=samples)
    soil, fixes = load_soil(a.s3); soil = soil.loc[samples]
    fixes.to_csv(f"{out}/S1_soil_value_corrections.csv", index=False)
    soil.round(4).to_csv(f"{out}/S2_soil_properties_corrected.csv")
    reads = pd.read_csv(a.s1, skiprows=1).set_index("Sample")["Read count (after QF)"].loc[samples]
    total = ko.sum()                                            # KO-annotated reads per sample
    # DESeq2-style median-of-ratios size factors (KOs detected in every sample), scaled so the
    # effective library size equals the geometric mean of annotated reads. Robust to compositional shifts.
    full = ko[(ko > 0).all(1)]
    lg = np.log(full); sf = np.exp((lg.sub(lg.mean(1), axis=0)).median())
    lib = sf * np.exp(np.log(total).mean())
    pd.DataFrame({"KO_annotated_reads": total, "size_factor": sf, "effective_library": lib}).to_csv(
        f"{out}/S7_normalization_size_factors.csv")

    # ---------- Table 1: dataset summary
    t1 = pd.DataFrame({"landscape": meta.landscape, "landuse": meta.landuse, "reads_after_QC": reads,
                       "KO_annotated_reads": total, "unassigned_reads": unclassified,
                       "pct_KO_assigned": total / (total + unclassified) * 100,
                       "KOs_detected": (ko > 0).sum(), "soil_pH": soil.pH, "soil_C_g_kg": soil.C_g_kg,
                       "C_N": soil.C_N, "P_mg_kg": soil.P_mg_kg, "Base_saturation_pct": soil.Base_saturation_pct})
    t1.to_csv(f"{out}/S3_sample_summary.csv")
    def ms(x): return f"{x.mean():.2f} ± {x.std(ddof=1):.2f}"
    rows = []
    for lu in LU_ORDER:
        g = t1[t1.landuse == lu]
        pct = f"{g.pct_KO_assigned.mean():.1f} ± {g.pct_KO_assigned.std():.1f}"
        rows.append([lu, len(g), f"{g.reads_after_QC.mean() / 1e6:.1f} ± {g.reads_after_QC.std() / 1e6:.1f}",
                     f"{g.KO_annotated_reads.mean() / 1e6:.2f} ± {g.KO_annotated_reads.std() / 1e6:.2f}", pct,
                     f"{g.KOs_detected.mean():.0f} ± {g.KOs_detected.std():.0f}", ms(g.soil_pH),
                     f"{g.soil_C_g_kg.mean():.1f} ± {g.soil_C_g_kg.std():.1f}", ms(g.C_N),
                     f"{g.P_mg_kg.mean():.1f} ± {g.P_mg_kg.std():.1f}",
                     f"{g.Base_saturation_pct.mean():.1f} ± {g.Base_saturation_pct.std():.1f}"])
    t1s = pd.DataFrame(rows, columns=["Land use", "n", "Reads after QC (M)", "KO-annotated reads (M)", "Reads assigned to a KO (%)",
                                      "KOs detected", "Soil pH", "Soil C (g/kg)", "C:N", "P (mg/kg)",
                                      "Base saturation (%)"])
    t1s.to_csv(f"{out}/Table1_dataset_and_soil_summary.csv", index=False)
    soil_tests = []
    for v in ["pH", "C_g_kg", "C_N", "P_mg_kg", "Base_saturation_pct", "Ca_g_m2", "Fe_g_m2"]:
        est, pint, _ = landuse_model(np.log(soil[v].values), meta)
        for lu in LU_ORDER[1:]:
            soil_tests.append([v, lu, *est[lu]])
    st = pd.DataFrame(soil_tests, columns=["variable", "landuse", "LRR", "CI_low", "CI_high", "p"])
    st["p_BH"] = bh(st.p); st["pct_change"] = (np.exp(st.LRR) - 1) * 100
    st.round(4).to_csv(f"{out}/S4_soil_landuse_effects.csv", index=False)

    # ---------- functional diversity (Shannon, rarefied 50x)
    depth = int(total.min()); H = {}
    pool_ko = ko.values.astype(np.int64)
    for j, smp in enumerate(samples):
        v = pool_ko[:, j]; vals = []
        for _ in range(50):
            vals.append(shannon(RNG.multivariate_hypergeometric(v, depth).astype(float)))
        H[smp] = np.mean(vals)
    meta["H_KO"] = pd.Series(H)
    est, pint, pnorm = landuse_model(meta.H_KO.values, meta)
    kw = stats.kruskal(*[meta.H_KO[meta.landuse == lu] for lu in LU_ORDER])
    div = pd.DataFrame([[lu, meta.H_KO[meta.landuse == lu].mean(), meta.H_KO[meta.landuse == lu].std(),
                         *(est[lu] if lu != "Rainforest" else [np.nan] * 4)] for lu in LU_ORDER],
                       columns=["landuse", "mean_H", "sd_H", "diff_vs_forest", "CI_low", "CI_high", "p"])
    div["p_BH"] = bh(div.p)
    div.round(5).to_csv(f"{out}/Table3_functional_diversity.csv", index=False)
    with open(f"{out}/stats_functional_diversity.txt", "w") as fh:
        fh.write(f"Rarefaction depth {depth} KO-annotated reads, 50 iterations\n"
                 f"Kruskal-Wallis H = {kw.statistic:.3f}, df = 3, p = {kw.pvalue:.4g}\n"
                 f"Land use x landscape interaction p = {pint:.4g}; Shapiro residuals p = {pnorm:.3g}\n")

    # ---------- composition: NMDS + PERMANOVA + envfit
    rel = (ko / total).T; rel = rel.loc[:, (ko > 0).sum(1).values >= 4]
    D = bray(rel.values)
    best = None
    for seed in range(10):
        m = MDS(n_components=2, metric=False, dissimilarity="precomputed", random_state=seed, n_init=1,
                max_iter=3000, eps=1e-9, normalized_stress=True)
        p = m.fit_transform(D)
        if best is None or m.stress_ < best[1]:
            best = (p, m.stress_)
    pts, stress = best
    lu_arr, ls_arr = meta.landuse.values, meta.landscape.values
    perm = {"Land use (within landscape)": permanova(D, lu_arr, ls_arr),
            "Landscape": permanova(D, ls_arr)}
    pair = []
    for lu in LU_ORDER[1:]:
        sel = np.isin(lu_arr, ["Rainforest", lu])
        f, r2, p = permanova(D[np.ix_(sel, sel)], lu_arr[sel], ls_arr[sel])
        pair.append(["Rainforest vs " + lu, f, r2, p])
    pair = pd.DataFrame(pair, columns=["contrast", "pseudo_F", "R2", "p"]); pair["p_BH"] = bh(pair.p)
    pair.round(4).to_csv(f"{out}/Table4b_pairwise_PERMANOVA.csv", index=False)
    pd.DataFrame([[k, *v] for k, v in perm.items()], columns=["term", "pseudo_F", "R2", "p"]).round(4) \
        .to_csv(f"{out}/Table4a_PERMANOVA.csv", index=False)
    env = []
    envvars = ["pH", "C_N", "C_g_kg", "N_mg_kg", "P_mg_kg", "Base_saturation_pct", "ECEC_mmolc_kg", "Ca_g_m2",
               "Fe_g_m2", "Al_g_m2", "K_g_m2", "Mg_g_m2", "Mn_g_m2", "Basal_respiration", "Microbial_biomass_C",
               "H2O_pct", "Land_use_index"]
    for v in envvars:
        z = soil[v].values.astype(float)
        def fit(zz):
            Xz = np.c_[np.ones(32), zz]; B = np.linalg.lstsq(Xz, pts, rcond=None)[0]
            r = pts - Xz @ B; return 1 - (r ** 2).sum() / ((pts - pts.mean(0)) ** 2).sum(), B[1]
        r2, vec = fit(z); c = sum(fit(RNG.permutation(z))[0] >= r2 for _ in range(9999))
        env.append([v, r2, (c + 1) / 10000, *(vec / np.linalg.norm(vec))])
    env = pd.DataFrame(env, columns=["variable", "r2", "p", "NMDS1", "NMDS2"]); env["p_BH"] = bh(env.p)
    env = env.sort_values("r2", ascending=False); env.round(4).to_csv(f"{out}/Table5_envfit_soil.csv", index=False)

    # ---------- marker genes and pathways
    rows = []
    def analyse(kind, group, name, counts):
        y = np.log((counts + 0.5) / lib * 1e6)
        est, pint, pnorm = landuse_model(y.values, meta)
        cpm = counts / lib * 1e6
        for lu in LU_ORDER[1:]:
            pp = strat_perm(y.values, meta, lu, 4999)
            rows.append([kind, group, name, lu, cpm[meta.landuse == "Rainforest"].mean(),
                         cpm[meta.landuse == lu].mean(), *est[lu], pp, pint, pnorm])
    excluded = []
    for grp, g, k in GENES:
        c = ko.loc[k] if k in ko.index else pd.Series(0, index=samples)
        if (c > 0).sum() < 24 or c.sum() < 100:      # detected in >= 75% of samples and >= 100 reads in total
            excluded.append([g, k, int(c.sum()), int((c > 0).sum())]); continue
        analyse("gene", grp, f"{g} ({k})", c)
    pd.DataFrame(excluded, columns=["gene", "KO", "total_reads_all_samples", "samples_detected"]).to_csv(
        f"{out}/S8_genes_excluded_low_detection.csv", index=False)
    for p in PATHWAYS:
        analyse("pathway", "KEGG level 3", p, pw.loc[p])
    res = pd.DataFrame(rows, columns=["kind", "group", "feature", "landuse", "mean_normCPM_rainforest",
                                      "mean_normCPM_landuse", "LRR", "CI_low", "CI_high", "p_model", "p_permutation",
                                      "p_interaction_landscape", "p_shapiro_resid"])
    for kind in ("gene", "pathway"):
        i = res.kind == kind
        res.loc[i, "p_BH"] = bh(res.loc[i, "p_model"]); res.loc[i, "p_perm_BH"] = bh(res.loc[i, "p_permutation"])
    res["pct_change"] = (np.exp(res.LRR) - 1) * 100
    res.round(5).to_csv(f"{out}/Table2_marker_genes_and_pathways.csv", index=False)

    # validation against published DESeq2 log2FC (per landscape)
    val = []
    for p, d in PUBLISHED_L2FC.items():
        for lsk, lsn in LS.items():
            for lu, pub in zip(LU_ORDER[1:], d[lsk]):
                c = pw.loc[p]; f = meta.landscape == lsn
                ours = np.log2((c[f & (meta.landuse == lu)] / lib[f & (meta.landuse == lu)]).mean() /
                               (c[f & (meta.landuse == "Rainforest")] / lib[f & (meta.landuse == "Rainforest")]).mean())
                val.append([p, lsn, lu, pub, ours])
    val = pd.DataFrame(val, columns=["pathway", "landscape", "landuse", "published_log2FC", "this_study_log2FC"])
    r_val = stats.pearsonr(val.published_log2FC, val.this_study_log2FC)
    sign = (np.sign(val.published_log2FC) == np.sign(val.this_study_log2FC))[val.published_log2FC.abs() > 0.05].mean()
    val.round(4).to_csv(f"{out}/S5_validation_vs_published.csv", index=False)

    # ---------- soil pH as moderator
    mod = []
    for _, r in res.drop_duplicates("feature").iterrows():
        f = r.feature
        counts = pw.loc[f] if r.kind == "pathway" else ko.loc[f.split("(")[-1].strip(")")] \
            if f.split("(")[-1].strip(")") in ko.index else pd.Series(0, index=samples)
        y = np.log((counts + 0.5) / lib * 1e6)
        rho, p = stats.spearmanr(soil.pH, y)
        X = np.c_[np.ones(32), soil.pH.values, (meta.landscape == "Harapan").values]
        b, se, df, _ = ols(y.values, X); pt = 2 * stats.t.sf(abs(b[1] / se[1]), df)
        mod.append([r.kind, f, rho, p, b[1], pt])
    mod = pd.DataFrame(mod, columns=["kind", "feature", "spearman_rho_pH", "p_spearman", "slope_per_pH_unit",
                                     "p_slope_adj_landscape"])
    mod["p_BH"] = bh(mod.p_spearman); mod.round(5).to_csv(f"{out}/Table6_soil_pH_associations.csv", index=False)

    # ---------- taxonomy (orders)
    tax_res = None
    if a.s4:
        t, ts = load_taxa(a.s4); t = t[t.domain != "Unclassified_tax"]
        tot_t = t[ts].sum()
        orders = t.groupby("order")[ts].sum()
        bact = orders[orders.index.str.startswith("Bacteria")]
        top = (bact / tot_t).mean(1).sort_values(ascending=False)
        top = top[~top.index.str.endswith("; NA")].index[:12]
        trows = []
        for o in top:
            y = np.log(bact.loc[o] / tot_t * 100); est, pint, _ = landuse_model(y.values, meta)
            for lu in LU_ORDER[1:]:
                trows.append([o.split("; ")[1], o.split("; ")[2], lu,
                              (bact.loc[o] / tot_t * 100)[meta.landuse == "Rainforest"].mean(),
                              (bact.loc[o] / tot_t * 100)[meta.landuse == lu].mean(), *est[lu], pint])
        tax_res = pd.DataFrame(trows, columns=["phylum", "order", "landuse", "pct_reads_rainforest",
                                               "pct_reads_landuse", "LRR", "CI_low", "CI_high", "p", "p_interaction"])
        tax_res["p_BH"] = bh(tax_res.p); tax_res["pct_change"] = (np.exp(tax_res.LRR) - 1) * 100
        tax_res.round(5).to_csv(f"{out}/Table7_bacterial_orders.csv", index=False)
        burk = [i for i in bact.index if i.endswith("; Burkholderiales")][0]
        b_rel = bact.loc[burk] / tot_t * 100
        link = []
        for nm, ser in [("Flagellar assembly", pw.loc["Flagellar assembly"]), ("Bacterial chemotaxis", pw.loc["Bacterial chemotaxis"]),
                        ("Bacterial secretion system", pw.loc["Bacterial secretion system"]),
                        ("nirS (K15864)", ko.loc["K15864"]), ("nifH (K02588)", ko.loc["K02588"])]:
            rho, p = stats.spearmanr(b_rel, ser / lib * 1e6); link.append(["Burkholderiales", nm, rho, p])
        link = pd.DataFrame(link, columns=["order", "function", "spearman_rho", "p"]); link["p_BH"] = bh(link.p)
        link.round(5).to_csv(f"{out}/S6_Burkholderiales_function_links.csv", index=False)
        meta["Burkholderiales_pct"] = b_rel

    meta.round(5).to_csv(f"{out}/S0_sample_metadata_and_diversity.csv")
    with open(f"{out}/stats_summary.txt", "w") as fh:
        fh.write(f"Samples: {len(samples)}; KOs in table: {len(ko)}\n")
        for k, (f, r2, p) in perm.items():
            fh.write(f"PERMANOVA {k}: pseudo-F = {f:.3f}, R2 = {r2:.3f}, p = {p:.4f}\n")
        fh.write(f"NMDS stress (Kruskal stress-1) = {stress:.3f}\n")
        fh.write(f"Validation vs published Fig. S7: Pearson r = {r_val[0]:.3f} (p = {r_val[1]:.2g}, n = {len(val)}); "
                 f"sign agreement for |log2FC| > 0.05: {sign * 100:.0f}%\n")
        fh.write(f"Soil corrections applied: {len(fixes)} values; C:N cross-check range "
                 f"{soil.CN_check.min():.3f}-{soil.CN_check.max():.3f}\n")

    figures(out, meta, pts, stress, perm, env, res, soil, ko, pw, lib, val, r_val, tax_res)
    print(open(f"{out}/stats_summary.txt").read())


# ----------------------------------------------------------------- figures (Biology and Fertility of Soils specs)
# Lettering: Arial-metric sans (Liberation Sans) 8-9 pt at final size; widths 84 mm (single) or 174 mm (full);
# no titles inside figures; land use identified by marker shape AND colour (accessibility).
MM = 1 / 25.4
W1, W2 = 84 * MM, 174 * MM
SHAPE = {"Rainforest": "o", "Jungle rubber": "s", "Rubber": "D", "Oil palm": "^"}
plt.rcParams.update({"font.family": "Liberation Sans", "font.size": 8, "axes.labelsize": 8, "xtick.labelsize": 8,
                     "ytick.labelsize": 8, "legend.fontsize": 8, "axes.linewidth": 0.6, "xtick.major.width": 0.6,
                     "ytick.major.width": 0.6, "savefig.dpi": 600, "mathtext.fontset": "custom",
                     "mathtext.it": "Liberation Sans:italic", "mathtext.rm": "Liberation Sans"})
PCT_TICKS = [-75, -50, -25, -10, 0, 10, 25, 50, 100, 200, 300]
NICE = {"pH": "pH", "Base_saturation_pct": "Base saturation", "Land_use_index": "Land-use index", "C_N": "C:N",
        "Fe_g_m2": "Fe", "Ca_g_m2": "Ca", "P_mg_kg": "P", "Basal_respiration": "Basal respiration"}


def pct_axis(ax, lo, hi):
    ticks = PCT_TICKS if hi - lo < 1.2 else [-75, -50, -25, 0, 100, 300]
    t = [v for v in ticks if np.log(1 + v / 100) >= lo - 0.05 and np.log(1 + v / 100) <= hi + 0.05]
    ax.set_xticks([np.log(1 + v / 100) for v in t], [f"{v:+d}%" if v else "0" for v in t])


def lu_handles(lus=LU_ORDER, ms=5.5):
    return [plt.Line2D([], [], marker=SHAPE[l], ls="", mfc=COL[l], mec=INK, mew=0.5, ms=ms, label=l) for l in lus]


def panel(ax, letter):
    ax.text(-0.02, 1.02, letter, transform=ax.transAxes, fontsize=9, fontweight="bold", ha="right", va="bottom")


def forest(ax, d, keycol, keys, labels):
    for j, lu in enumerate(LU_ORDER[1:]):
        s = d[d.landuse == lu].set_index(keycol).reindex(keys)
        y = np.arange(len(keys)) + (j - 1) * 0.26
        ax.hlines(y, s.CI_low, s.CI_high, color=COL[lu], lw=1.1)
        sg = (s.p_BH < 0.05).values
        ax.scatter(s.LRR[~sg], y[~sg], marker=SHAPE[lu], s=20, facecolors="white", edgecolors=COL[lu], linewidths=1.1, zorder=3)
        ax.scatter(s.LRR[sg], y[sg], marker=SHAPE[lu], s=22, c=COL[lu], edgecolors=INK, linewidths=0.5, zorder=3)
    ax.axvline(0, color=MUTED, lw=0.7, ls="--")
    ax.set_yticks(range(len(keys)), labels); ax.invert_yaxis()
    ax.grid(axis="x", color=GRID, lw=0.5); ax.set_axisbelow(True)
    h = lu_handles(LU_ORDER[1:])
    h += [plt.Line2D([], [], marker="o", ls="", mfc=INK, mec=INK, ms=5, label="BH-adjusted p < 0.05"),
          plt.Line2D([], [], marker="o", ls="", mfc="white", mec=INK, ms=5, label="Not significant")]
    return h


def figures(out, meta, pts, stress, perm, env, res, soil, ko, pw, total, val, r_val, tax_res):
    # Fig. 1 NMDS (full width)
    fig, ax = plt.subplots(figsize=(W2, 105 * MM), layout="constrained")
    for lu in LU_ORDER:
        for ls in ("Bukit Duabelas", "Harapan"):
            s = ((meta.landuse == lu) & (meta.landscape == ls)).values
            if ls == "Bukit Duabelas":
                ax.scatter(pts[s, 0], pts[s, 1], marker=SHAPE[lu], s=34, c=COL[lu], edgecolors=INK, linewidths=0.5, zorder=3)
            else:
                ax.scatter(pts[s, 0], pts[s, 1], marker=SHAPE[lu], s=34, facecolors="white", edgecolors=COL[lu], linewidths=1.3, zorder=3)
    sig = env[env.p_BH < 0.05]; sc = 1.1 * np.abs(pts).max(); placed = []
    for _, r in sig.assign(ang=np.degrees(np.arctan2(sig.NMDS2, sig.NMDS1))).sort_values("ang").iterrows():
        d = np.array([r.NMDS1, r.NMDS2]) * sc * np.sqrt(r.r2)
        ax.annotate("", d, (0, 0), arrowprops=dict(arrowstyle="-|>", color=MUTED, lw=0.9), zorder=2)
        tx, ty = d * 1.06
        while any(abs(ty - py) < 0.09 and abs(tx - px) < 0.6 for px, py in placed):
            ty += 0.09
        placed.append((tx, ty))
        ax.text(tx, ty, NICE.get(r.variable, r.variable), color=INK, va="center", ha="right" if d[0] < 0 else "left",
                zorder=4, bbox=dict(boxstyle="round,pad=0.12", fc="white", ec="none", alpha=0.9))
    ax.set_xlabel("NMDS1"); ax.set_ylabel("NMDS2"); ax.grid(color=GRID, lw=0.5); ax.set_axisbelow(True)
    h = lu_handles()
    h += [plt.Line2D([], [], marker="o", ls="", mfc=INK, mec=INK, ms=5, label="Bukit Duabelas (filled)"),
          plt.Line2D([], [], marker="o", ls="", mfc="white", mec=INK, ms=5, label="Harapan (open)")]
    ax.legend(handles=h, frameon=False, loc="upper left", bbox_to_anchor=(1.01, 1))
    fig.savefig(f"{out}/Fig1_NMDS_functional_composition.png"); plt.close(fig)

    # Fig. 2 functional diversity (single column)
    fig, ax = plt.subplots(figsize=(W1, 70 * MM), layout="constrained")
    for i, lu in enumerate(LU_ORDER):
        v = meta.H_KO[meta.landuse == lu]
        ax.boxplot(v, positions=[i], widths=0.5, showfliers=False, medianprops=dict(color=INK, lw=1),
                   boxprops=dict(color=MUTED, lw=0.7), whiskerprops=dict(color=MUTED, lw=0.7), capprops=dict(color=MUTED, lw=0.7))
        for ls in ("Bukit Duabelas", "Harapan"):
            vv = meta.H_KO[(meta.landuse == lu) & (meta.landscape == ls)]
            x = i + RNG.uniform(-0.13, 0.13, len(vv))
            if ls == "Bukit Duabelas":
                ax.scatter(x, vv, marker=SHAPE[lu], s=18, c=COL[lu], edgecolors=INK, linewidths=0.4, zorder=3)
            else:
                ax.scatter(x, vv, marker=SHAPE[lu], s=18, facecolors="white", edgecolors=COL[lu], linewidths=1.0, zorder=3)
    ax.set_xticks(range(4), [l.replace(" ", "\n") for l in LU_ORDER]); ax.set_ylabel("Shannon index (KEGG orthologs)")
    ax.grid(axis="y", color=GRID, lw=0.5); ax.set_axisbelow(True)
    fig.savefig(f"{out}/Fig2_functional_diversity.png"); plt.close(fig)

    # Fig. 3 pathways, Fig. 4 marker genes (full width)
    for kind, fname, hmm in (("pathway", "Fig3_KEGG_pathways.png", 92), ("gene", "Fig4_marker_genes.png", 225)):
        d = res[res.kind == kind]; feats = list(dict.fromkeys(d.feature))
        fig, ax = plt.subplots(figsize=(W2, hmm * MM), layout="constrained")
        def glab(f):
            g = f.split(" (K")[0]
            if kind != "gene":
                return g
            name, _, extra = g.partition(" (")
            name = name.replace("/", "}/\\mathit{")
            return f"$\\mathit{{{name}}}$" + (f" ({extra}" if extra else "")
        labels = [glab(f) for f in feats]
        h = forest(ax, d, "feature", feats, labels)
        pct_axis(ax, d.CI_low.min(), d.CI_high.max())
        ax.set_xlabel("Change relative to rainforest (log-ratio scale)")
        ax.legend(handles=h, frameon=False, loc="upper left", bbox_to_anchor=(1.01, 1))
        fig.savefig(f"{out}/{fname}"); plt.close(fig)

    # Fig. 5 soil pH vs functions (2 x 2, full width)
    feats = [("pathway", "Flagellar assembly", "Flagellar assembly"), ("pathway", "Bacterial chemotaxis", "Bacterial chemotaxis"),
             ("gene", "nirS (K15864)", "nirS"), ("gene", "nifH (K02588)", "nifH")]
    fig, axs = plt.subplots(2, 2, figsize=(W2, 120 * MM), layout="constrained")
    for ax, (kind, f, lab), letter in zip(axs.flat, feats, "abcd"):
        c = pw.loc[f] if kind == "pathway" else ko.loc[f.split("(")[-1].strip(")")]
        y = c / total * 1e6
        for lu in LU_ORDER:
            s = meta.landuse == lu
            ax.scatter(soil.pH[s], y[s], marker=SHAPE[lu], s=18, c=COL[lu], edgecolors=INK, linewidths=0.4, zorder=3)
        rho, p = stats.spearmanr(soil.pH, y)
        ax.text(0.98, 0.96, f"ρ = {rho:.2f}, p = {p:.3f}" if p >= 0.001 else f"ρ = {rho:.2f}, p < 0.001",
                transform=ax.transAxes, ha="right", va="top")
        lo_, hi_ = y.min(), y.max(); ax.set_ylim(lo_ - 0.06 * (hi_ - lo_), hi_ + 0.22 * (hi_ - lo_))
        ax.set_xlabel("Soil pH")
        ax.set_ylabel(f"{lab} (normalized CPM)", fontstyle="italic" if kind == "gene" else "normal")
        ax.grid(color=GRID, lw=0.5); ax.set_axisbelow(True); panel(ax, letter)
    fig.legend(handles=lu_handles(), frameon=False, loc="outside lower center", ncol=4)
    fig.savefig(f"{out}/Fig5_soil_pH_Sumatra.png"); plt.close(fig)

    # Fig. 6 bacterial orders (full width)
    if tax_res is not None:
        orders = list(dict.fromkeys(tax_res.order))
        fig, ax = plt.subplots(figsize=(W2, 120 * MM), layout="constrained")
        lab = [f"{o} ({tax_res[tax_res.order == o].phylum.iloc[0]})" for o in orders]
        h = forest(ax, tax_res, "order", orders, lab)
        pct_axis(ax, tax_res.CI_low.min(), tax_res.CI_high.max())
        ax.set_xlabel("Change in relative abundance relative to rainforest (log-ratio scale)")
        ax.legend(handles=h, frameon=False, loc="upper left", bbox_to_anchor=(1.01, 1))
        fig.savefig(f"{out}/Fig6_bacterial_orders.png"); plt.close(fig)

    # Online Resource figure S1: validation (single column)
    fig, ax = plt.subplots(figsize=(W1, 84 * MM), layout="constrained")
    for lu in LU_ORDER[1:]:
        for ls in ("Bukit Duabelas", "Harapan"):
            ss = (val.landuse == lu) & (val.landscape == ls)
            if ls == "Bukit Duabelas":
                ax.scatter(val.published_log2FC[ss], val.this_study_log2FC[ss], marker=SHAPE[lu], s=20, c=COL[lu],
                           edgecolors=INK, linewidths=0.4, zorder=3)
            else:
                ax.scatter(val.published_log2FC[ss], val.this_study_log2FC[ss], marker=SHAPE[lu], s=20,
                           facecolors="white", edgecolors=COL[lu], linewidths=1.0, zorder=3)
    lim = [min(val.min(numeric_only=True)) - 0.1, max(val.max(numeric_only=True)) + 0.1]
    ax.plot(lim, lim, color=MUTED, lw=0.7, ls="--"); ax.set_xlim(lim); ax.set_ylim(lim)
    ax.set_xlabel("Published log2 fold change"); ax.set_ylabel("Log2 fold change, this study")
    ax.grid(color=GRID, lw=0.5); ax.set_axisbelow(True)
    ax.legend(handles=lu_handles(LU_ORDER[1:], 5), frameon=False, loc="lower right")
    fig.savefig(f"{out}/FigS1_validation_vs_published.png"); plt.close(fig)


if __name__ == "__main__":
    main()
