#!/usr/bin/env python3
"""
Cross-study comparison: oil palm vs forest soil functional potential in Sumatra (Berkelmann et al. 2020)
and Sabah, Borneo (Tripathi et al. 2016). Each site's effect is estimated within its own normalization;
only effect sizes (log response ratios) are combined.

Author: Mc Arthur L. Cababan, PhD
Inputs:
  --s5     Berkelmann et al. 2020 Table S5 (KEGG counts)
  --seed   Tripathi et al. 2016 Dryad "SEED functional profile at subsystem level 3.xlsx"
  --soil   Tripathi et al. 2016 Dryad "Soil chemical variables.xlsx"
"""
import argparse, itertools, os, warnings
import numpy as np, pandas as pd
from scipy import stats
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
warnings.filterwarnings("ignore")
RNG = np.random.default_rng(2026)

INK, MUTED, GRID = "#1a1a19", "#5f5e5a", "#e4e3df"
SITE_COL = {"Sumatra – Bukit Duabelas": "#2a78d6", "Sumatra – Harapan": "#1baf7a", "Sabah, Borneo": "#eb6834"}
MM = 1 / 25.4; W2 = 174 * MM
SITE_SHAPE = {"Sumatra – Bukit Duabelas": "o", "Sumatra – Harapan": "s", "Sabah, Borneo": "^"}
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, "axes.edgecolor": MUTED, "axes.labelcolor": INK,
                     "xtick.color": MUTED, "ytick.color": MUTED, "axes.spines.top": False,
                     "axes.spines.right": False, "savefig.dpi": 600})

# Matched functional categories: KEGG orthologs (Sumatra) and SEED subsystems (Sabah)
KEGG = {
    "Flagellar motility": ("pathway", "Flagellar assembly"),
    "Bacterial chemotaxis": ("pathway", "Bacterial chemotaxis"),
    "Type VI secretion": ("t6ss", None),
    "Denitrification": ("kos", ["K00368", "K15864", "K04561", "K02305", "K00376"]),
    "Nitrogen fixation": ("kos", ["K02588", "K02586", "K02591"]),
}
SEED = {
    "Flagellar motility": ["Flagellum", "Flagellar motility", "Flagellum in Campylobacter"],
    "Bacterial chemotaxis": ["Bacterial Chemotaxis"],
    "Type VI secretion": ["Type VI secretion systems"],
    "Denitrification": ["Denitrification", "Dissimilatory nitrite reductase"],
    "Nitrogen fixation": ["Nitrogen fixation"],
}


def bh(p):
    p = np.asarray(p, float); n = len(p); o = np.argsort(p); r = p[o] * n / np.arange(1, n + 1)
    r = np.minimum.accumulate(r[::-1])[::-1]; a = np.empty(n); a[o] = np.minimum(r, 1); return a


def mor_library(counts):
    """Median-of-ratios size factors from features present in every sample; effective library sizes."""
    full = counts[(counts > 0).all(1)]
    lg = np.log(full); sf = np.exp(lg.sub(lg.mean(1), axis=0).median())
    return sf * np.exp(np.log(counts.sum()).mean())


def effect(y_t, y_c):
    """Difference in mean ln(abundance) = log response ratio, its variance, and an exact permutation p."""
    y_t, y_c = np.asarray(y_t), np.asarray(y_c)
    d = y_t.mean() - y_c.mean()
    v = y_t.var(ddof=1) / len(y_t) + y_c.var(ddof=1) / len(y_c)
    allv = np.r_[y_t, y_c]; n_t = len(y_t); obs = abs(d); c = 0; tot = 0
    for idx in itertools.combinations(range(len(allv)), n_t):
        m = np.zeros(len(allv), bool); m[list(idx)] = True
        c += abs(allv[m].mean() - allv[~m].mean()) >= obs - 1e-12; tot += 1
    return d, v, c / tot


def pool(y, v):
    """Random-effects pooling: DerSimonian-Laird tau2 with Hartung-Knapp-Sidik-Jonkman CI (t, k-1 df)."""
    y, v = np.asarray(y), np.asarray(v); k = len(y); w = 1 / v
    fe = (w * y).sum() / w.sum(); Q = (w * (y - fe) ** 2).sum()
    C = w.sum() - (w ** 2).sum() / w.sum(); tau2 = max(0, (Q - (k - 1)) / C)
    ws = 1 / (v + tau2); mu = (ws * y).sum() / ws.sum()
    se_hk = np.sqrt((ws * (y - mu) ** 2).sum() / ((k - 1) * ws.sum()))
    se_dl = np.sqrt(1 / ws.sum()); se = max(se_hk, se_dl)        # never narrower than the DL interval
    tq = stats.t.ppf(0.975, k - 1); p = 2 * stats.t.sf(abs(mu / se), k - 1)
    I2 = max(0, (Q - (k - 1)) / Q) * 100 if Q > 0 else 0
    return mu, mu - tq * se, mu + tq * se, p, tau2, I2, Q, stats.chi2.sf(Q, k - 1)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--s5", required=True); ap.add_argument("--seed", required=True)
    ap.add_argument("--soil", required=True); ap.add_argument("--s3", required=True, help="Berkelmann Table S3 soil xlsx")
    ap.add_argument("--out", default="results_crossstudy")
    a = ap.parse_args(); os.makedirs(a.out, exist_ok=True)

    # ---------- Sumatra
    s = pd.read_csv(a.s5, sep="\t", skiprows=1); s = s[s.KO != "Unclassified"]
    samples = [c for c in s.columns if c not in ("KO", "pathway")]
    s["L3"] = s.pathway.str.split(";").str[2]; s["desc"] = s.pathway.str.split(";").str[3]
    ko = s.drop_duplicates("KO").set_index("KO")[samples]; lib_su = mor_library(ko)
    def kegg_counts(cat):
        kind, arg = KEGG[cat]
        if kind == "pathway":
            return s[s.L3 == arg].drop_duplicates("KO")[samples].sum()
        if kind == "kos":
            return ko.loc[arg].sum()
        t6 = s[s.desc.str.lower().str.startswith("type vi secretion system", na=False)].drop_duplicates("KO")
        return t6[samples].sum()
    kegg_n_t6 = s[s.desc.str.lower().str.startswith("type vi secretion system", na=False)].KO.nunique()

    # ---------- Sabah
    seed = pd.read_excel(a.seed, header=None)
    names = seed.iloc[0, 1:].astype(str).values; ids = seed.iloc[1:, 0].astype(str).values
    mat = pd.DataFrame(seed.iloc[1:, 1:].values.astype(float).T, columns=ids)   # roles x samples
    soil = pd.read_excel(a.soil, index_col=0); soil.index = soil.index.astype(str); soil = soil.loc[ids]
    lib_sa = mor_library(mat)
    def seed_counts(cat):
        return mat[np.isin(names, SEED[cat])].sum()

    rows, cats = [], list(KEGG)
    for cat in cats:
        cs = kegg_counts(cat); ys = np.log((cs + 0.5) / lib_su * 1e6)
        for land, lab in (("B", "Sumatra – Bukit Duabelas"), ("H", "Sumatra – Harapan")):
            t = [x for x in samples if x[0] == land and x[1] == "O"]; c = [x for x in samples if x[0] == land and x[1] == "F"]
            d, v, p = effect(ys[t], ys[c]); rows.append([cat, lab, "Oil palm vs rainforest", len(t), len(c), d, v, p])
        cb = seed_counts(cat); yb = np.log((cb + 0.5) / lib_sa * 1e6)
        lu = soil.Landuse
        for contrast, tsel, csel in (("Oil palm vs unlogged forest", lu == "Oil palm", lu == "Unlogged"),
                                     ("Oil palm vs all forest (sensitivity)", lu == "Oil palm", lu != "Oil palm"),
                                     ("Once-logged vs unlogged", lu == "Once-logged", lu == "Unlogged"),
                                     ("Twice-logged vs unlogged", lu == "Twice-logged", lu == "Unlogged")):
            if "all forest" in contrast:   # 6 vs 18: exact enumeration too large; use 20,000 random permutations
                yt, yc = yb[tsel.values].values, yb[csel.values].values
                d = yt.mean() - yc.mean(); v = yt.var(ddof=1) / len(yt) + yc.var(ddof=1) / len(yc)
                allv = np.r_[yt, yc]; cnt = 0
                for _ in range(20000):
                    pm = RNG.permutation(allv); cnt += abs(pm[:len(yt)].mean() - pm[len(yt):].mean()) >= abs(d) - 1e-12
                p = (cnt + 1) / 20001
            else:
                d, v, p = effect(yb[tsel.values], yb[csel.values])
            rows.append([cat, "Sabah, Borneo", contrast, int(tsel.sum()), int(csel.sum()), d, v, p])
    eff = pd.DataFrame(rows, columns=["category", "site", "contrast", "n_treat", "n_ref", "LRR", "var", "p_exact_perm"])
    eff["SE"] = np.sqrt(eff["var"]); eff["pct_change"] = (np.exp(eff.LRR) - 1) * 100
    eff["CI_low_pct"] = (np.exp(eff.LRR - 1.96 * eff.SE) - 1) * 100; eff["CI_high_pct"] = (np.exp(eff.LRR + 1.96 * eff.SE) - 1) * 100
    eff.round(5).to_csv(f"{a.out}/X1_site_effects.csv", index=False)

    main_eff = eff[eff.contrast.isin(["Oil palm vs rainforest", "Oil palm vs unlogged forest"])]
    pooled = []
    for cat in cats:
        g = main_eff[main_eff.category == cat]
        mu, lo, hi, p, tau2, I2, Q, pQ = pool(g.LRR, g["var"])
        same = int((np.sign(g.LRR) == np.sign(mu)).sum())
        pooled.append([cat, len(g), mu, lo, hi, p, tau2, I2, Q, pQ, f"{same}/{len(g)}"])
    pooled = pd.DataFrame(pooled, columns=["category", "k_sites", "LRR", "CI_low", "CI_high", "p", "tau2", "I2_pct",
                                           "Q", "p_heterogeneity", "sites_same_direction"])
    pooled["p_BH"] = bh(pooled.p); pooled["pct_change"] = (np.exp(pooled.LRR) - 1) * 100
    pooled["CI_low_pct"] = (np.exp(pooled.CI_low) - 1) * 100; pooled["CI_high_pct"] = (np.exp(pooled.CI_high) - 1) * 100
    pooled.round(5).to_csv(f"{a.out}/X2_pooled_oilpalm_effects.csv", index=False)

    # Sabah: soil pH vs functions, and soil differences
    sab = []
    for cat in cats:
        yb = np.log((seed_counts(cat) + 0.5) / lib_sa * 1e6)
        rho, p = stats.spearmanr(soil.pH, yb); sab.append([cat, rho, p])
    sab = pd.DataFrame(sab, columns=["category", "spearman_rho_pH", "p"]); sab["p_BH"] = bh(sab.p)
    # Sumatra: same matched categories vs soil pH (pH column was not affected by the decimal errors)
    su_soil = pd.read_excel(a.s3, header=1).iloc[:32].set_index("SampleID")
    ph_su = su_soil[[c for c in su_soil.columns if str(c).startswith("pH")][0]].loc[samples].astype(float)
    rows_ph = []
    for _, r in sab.iterrows():
        cat = r.category
        ys = np.log((kegg_counts(cat) + 0.5) / lib_su * 1e6); rho_s, p_s = stats.spearmanr(ph_su, ys)
        # combine the two regions' correlations (Fisher z, weights n-3) and test their difference
        z1, z2 = np.arctanh(rho_s), np.arctanh(r.spearman_rho_pH); w1, w2 = 32 - 3, 24 - 3
        zc = (w1 * z1 + w2 * z2) / (w1 + w2); se = 1 / np.sqrt(w1 + w2)
        p_comb = 2 * stats.norm.sf(abs(zc / se)); p_diff = 2 * stats.norm.sf(abs(z1 - z2) / np.sqrt(1 / w1 + 1 / w2))
        rows_ph.append([cat, rho_s, p_s, r.spearman_rho_pH, r.p, np.tanh(zc), np.tanh(zc - 1.96 * se),
                        np.tanh(zc + 1.96 * se), p_comb, p_diff])
    php = pd.DataFrame(rows_ph, columns=["category", "rho_Sumatra", "p_Sumatra", "rho_Sabah", "p_Sabah", "rho_combined",
                                         "CI_low", "CI_high", "p_combined", "p_difference_between_regions"])
    php["p_combined_BH"] = bh(php.p_combined)
    php.round(5).to_csv(f"{a.out}/X3_pH_associations_both_regions.csv", index=False)
    sab.round(5).to_csv(f"{a.out}/X3b_sabah_pH_associations.csv", index=False)
    # Fig. 8: pH relationships in both regions (2 x 3 grid, full width; legend in the sixth cell)
    fig, axs = plt.subplots(2, 3, figsize=(W2, 118 * MM), layout="constrained")
    for ax, cat, letter in zip(axs.flat, cats, "abcde"):
        ys = np.log((kegg_counts(cat) + 0.5) / lib_su * 1e6); yb = np.log((seed_counts(cat) + 0.5) / lib_sa * 1e6)
        zs = (ys - ys.mean()) / ys.std(); zb = (yb - yb.mean()) / yb.std()
        ax.scatter(ph_su, zs, s=16, marker="o", c="#2a78d6", edgecolors=INK, linewidths=0.4, zorder=3)
        ax.scatter(soil.pH, zb, s=18, marker="^", facecolors="white", edgecolors="#eb6834", linewidths=1.0, zorder=3)
        r = php.set_index("category").loc[cat]
        ax.text(0.03, 0.97, f"Sumatra ρ = {r.rho_Sumatra:.2f}\nSabah ρ = {r.rho_Sabah:.2f}", transform=ax.transAxes,
                ha="left", va="top", fontsize=7.5, bbox=dict(boxstyle="round,pad=0.15", fc="white", ec="none", alpha=0.85))
        allz = np.r_[zs, zb]; ax.set_ylim(allz.min() - 0.3, allz.max() + 0.35 * (allz.max() - allz.min()))
        ax.set_xlabel("Soil pH"); ax.set_ylabel(f"{cat} (z-score)")
        ax.grid(color=GRID, lw=0.5); ax.set_axisbelow(True)
        ax.text(-0.02, 1.02, letter, transform=ax.transAxes, fontsize=9, fontweight="bold", ha="right", va="bottom")
    axs.flat[5].axis("off")
    axs.flat[5].legend(handles=[plt.Line2D([], [], marker="o", ls="", mfc="#2a78d6", mec=INK, mew=0.4, ms=5.5, label="Sumatra, KEGG (filled circles)"),
                                plt.Line2D([], [], marker="^", ls="", mfc="white", mec="#eb6834", mew=1.0, ms=6, label="Sabah, SEED (open triangles)")],
                       frameon=False, loc="center")
    fig.savefig(f"{a.out}/Fig8_pH_relationships_both_regions.png"); plt.close(fig)
    soil_sum = soil.groupby("Landuse").agg(["mean", "std"]).round(3); soil_sum.to_csv(f"{a.out}/X4_sabah_soil_summary.csv")
    soil_tests = []
    for vname in [c for c in soil.columns if c != "Landuse"]:
        d, v, p = effect(np.log(soil.loc[soil.Landuse == "Oil palm", vname]), np.log(soil.loc[soil.Landuse == "Unlogged", vname]))
        soil_tests.append([vname, soil.loc[soil.Landuse == "Unlogged", vname].mean(), soil.loc[soil.Landuse == "Oil palm", vname].mean(),
                           (np.exp(d) - 1) * 100, p])
    pd.DataFrame(soil_tests, columns=["variable", "unlogged_mean", "oil_palm_mean", "pct_change", "p_exact_perm"]).round(4) \
        .to_csv(f"{a.out}/X5_sabah_soil_oilpalm_vs_unlogged.csv", index=False)
    with open(f"{a.out}/X0_notes.txt", "w") as fh:
        fh.write(f"Sabah samples: {len(ids)}; SEED role columns: {len(names)}; size factors {np.round((lib_sa / lib_sa.mean()).min(), 2)}-"
                 f"{np.round((lib_sa / lib_sa.mean()).max(), 2)}\nKEGG type VI secretion KOs used: {kegg_n_t6}\n"
                 + "".join(f"SEED {c}: {int(np.isin(names, SEED[c]).sum())} role columns\n" for c in cats))

    # Fig. 7: per-site effects + pooled diamond (full width)
    fig, ax = plt.subplots(figsize=(W2, 120 * MM), layout="constrained"); ypos = 0; ticks = []
    for cat in cats:
        g = main_eff[main_eff.category == cat].set_index("site"); pr = pooled.set_index("category").loc[cat]
        for site in SITE_COL:
            r = g.loc[site]
            ax.hlines(ypos, r.LRR - 1.96 * r.SE, r.LRR + 1.96 * r.SE, color=SITE_COL[site], lw=1.1)
            ax.scatter(r.LRR, ypos, s=24, marker=SITE_SHAPE[site], c=SITE_COL[site], edgecolors=INK, linewidths=0.5, zorder=3)
            ypos += 0.8
        ax.fill([pr.CI_low, pr.LRR, pr.CI_high, pr.LRR], [ypos, ypos - 0.3, ypos, ypos + 0.3], color=INK, zorder=3)
        ticks.append((ypos - 0.8, f"{cat}\npooled {pr['pct_change']:+.0f}% (I² = {pr['I2_pct']:.0f}%)"))
        ypos += 1.6
    ax.axvline(0, color=MUTED, lw=0.7, ls="--")
    ax.set_yticks([t for t, _ in ticks], [l for _, l in ticks]); ax.invert_yaxis()
    tk = [-75, -50, 0, 100, 300, 700]
    lo, hi = ax.get_xlim(); tk = [t for t in tk if lo <= np.log(1 + t / 100) <= hi]
    ax.set_xticks([np.log(1 + t / 100) for t in tk], [f"{t:+d}%" if t else "0" for t in tk])
    ax.set_xlabel("Oil palm relative to forest (log-ratio scale)"); ax.grid(axis="x", color=GRID, lw=0.5); ax.set_axisbelow(True)
    h = [plt.Line2D([], [], marker=SITE_SHAPE[k], ls="", mfc=c, mec=INK, mew=0.5, ms=5.5, label=k) for k, c in SITE_COL.items()]
    h.append(plt.Line2D([], [], marker="D", ls="", mfc=INK, mec=INK, ms=5.5, label="Pooled, random effects"))
    ax.legend(handles=h, frameon=False, loc="upper left", bbox_to_anchor=(1.01, 1))
    fig.savefig(f"{a.out}/Fig7_crossstudy_forest_plot.png"); plt.close(fig)
    print(pooled.round(4).to_string()); print(php.round(3).to_string()); print(open(f"{a.out}/X0_notes.txt").read())


if __name__ == "__main__":
    main()