"""Twin codes for the rank: matched on total cost, opposite in how the rank moves.

The manuscript's claim is that the standard code's rank among randomized
alternatives is not a property of the code. This is the sharpest form of that
claim: among the randomized codes themselves, pairs whose total error-minimizing
cost is indistinguishable have ranks that move in opposite directions along the
same gradient of transition bias.

Writes data/twin_ranks.json and data/twin_rank_pairs.csv.

    CORAL_XLSX=/path/to/media-1.xlsx python3 code/twin_ranks.py
"""
import csv
import json
import os
import sys

import numpy as np
import pandas as pd

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
DATA = os.path.join(ROOT, "data")
sys.path.insert(0, HERE)

COST_TOL = 0.005    # pairs match to within 0.5% of total cost
SLOPE_MIN = 0.05    # and each must have a slope large enough to have a sign
NDRAW = int(os.environ.get("NDRAW", 2000))


def main():
    import coral
    from codecore import (AA20, POLAR_REQUIREMENT, STANDARD_CODE,
                          scalar_metric_matrix)

    ai = {a: i for i, a in enumerate(AA20 + ["*"])}
    codons = sorted(STANDARD_CODE)
    ci = {c: i for i, c in enumerate(codons)}
    sgc = np.array([ai[STANDARD_CODE[c]] for c in codons])
    dm = scalar_metric_matrix(POLAR_REQUIREMENT, True)
    D = np.zeros((21, 21))
    for a in AA20:
        for b in AA20:
            D[ai[a], ai[b]] = dm[(a, b)]

    ens = json.load(open(os.path.join(DATA, "confirmatory_ensemble.json")))
    perm = [ci[c] for c in ens["codon_order"]]
    codes = np.zeros((ens["n"], 64), dtype=np.int8)
    for i, s in enumerate(ens["codes"]):
        for k, aa in enumerate(s):
            codes[i, perm[k]] = ai[aa]

    fr, to = np.load(os.path.join(DATA, "edge_meta.npy"))
    pq = np.load(os.path.join(DATA, "edge_pq.npy"))
    w = pq / pq.sum()
    cost = (D[codes[:, fr], codes[:, to]] * w[None, :]).sum(axis=1)
    print(f"  {ens['n']} codes, total cost {cost.min():.3f} to {cost.max():.3f}")

    # --- the spectra, aligned to the 4,972 of the study -----------------
    counts = pd.read_pickle(os.path.join(DATA, "coral_counts.pkl"))
    abund = pd.read_pickle(os.path.join(DATA, "coral_abund.pkl"))
    spectra = coral.spectra_from_sums(counts, abund)
    keys = [tuple(k) for k in json.load(open(os.path.join(DATA,
                                                         "edge_keys.json")))]
    cats = sorted(set(keys))
    col = np.array([cats.index(k) for k in keys])
    al = pd.read_csv(os.path.join(DATA, "confirmatory_aligned.csv"))
    want = set(al["species"].astype(str))
    names = [n for n in spectra if n in want]
    assert len(names) == 4972, f"{len(names)} species aligned, expected 4,972"
    P = np.array([[spectra[n].get(k, 0.0) for k in cats] for n in names],
                 dtype=np.float32)
    P = P / P.sum(axis=1, keepdims=True)
    W = P[:, col] * pq[None, :]
    W = (W / W.sum(axis=1, keepdims=True)).astype(np.float32)
    print(f"  {len(names)} species aligned")

    def total(t):
        idx = [i for i, k in enumerate(keys)
               if k[0][1] == t[0] and k[1] == t[1]]
        return P[:, [cats.index(keys[i]) for i in idx]].sum(axis=1)

    kappa = np.log((total(("C", "U")) + total(("U", "C")))
                   / (total(("C", "A")) + total(("C", "G"))
                      + total(("U", "A")) + total(("U", "G"))))
    kz = ((kappa - kappa.mean()) / kappa.std()).astype(np.float32)

    # --- the rank of every code in every species ------------------------
    CE = D[codes[:, fr], codes[:, to]].astype(np.float32)
    PHI = W @ CE.T
    print(f"  cost matrix {PHI.shape}")

    A = np.vstack([np.ones_like(kz), kz]).T
    AtAi = np.linalg.pinv(A.T @ A) @ A.T

    def slope(phi_col):
        """How log(1 + rank) moves with standardized log kappa."""
        r = (PHI < phi_col[:, None]).sum(axis=1)
        return float((AtAi @ np.log1p(r).astype(np.float64))[1])

    rng = np.random.default_rng(20260727)
    pick = rng.choice(ens["n"], NDRAW, replace=False)
    print(f"  slopes for {NDRAW} codes")
    sl = np.array([slope(PHI[:, j]) for j in pick])
    s_sgc = slope(W @ D[sgc[fr], sgc[to]].astype(np.float32))

    print(f"\n  slope of the rank against log kappa, {NDRAW} randomized codes:")
    print(f"    {sl.min():+.4f} to {sl.max():+.4f}, median {np.median(sl):+.4f}")
    print(f"    {100 * (sl > 0).mean():.1f}% positive, "
          f"{100 * (sl < 0).mean():.1f}% negative")
    print(f"  the standard code: {s_sgc:+.4f}, "
          f"percentile {100 * (sl < s_sgc).mean():.1f}")

    r_cost = float(np.corrcoef(cost[pick], sl)[0, 1])
    print(f"\n  corr(total cost, slope) = {r_cost:+.3f}")

    # --- pairs a scalar cannot separate ---------------------------------
    # Every pair within the cost tolerance is counted, whether or not its
    # slopes agree: counting only the discordant ones would make the
    # "opposite" fraction one by construction and measure nothing.
    c = cost[pick]
    order = np.argsort(c)
    all_pairs, filt_pairs = [], []
    for a in range(len(order)):
        i = order[a]
        for b in range(a + 1, min(a + 300, len(order))):
            j = order[b]
            if abs(c[j] - c[i]) / c[i] > COST_TOL:
                break
            all_pairs.append((int(i), int(j)))
            if abs(sl[i]) > SLOPE_MIN and abs(sl[j]) > SLOPE_MIN:
                filt_pairs.append((int(i), int(j)))
    n_all = len(all_pairs)
    opp_all = sum(1 for i, j in all_pairs if sl[i] * sl[j] < 0)
    n_filt = len(filt_pairs)
    opp_filt = sum(1 for i, j in filt_pairs if sl[i] * sl[j] < 0)
    print(f"\n  {n_all} pairs within {100 * COST_TOL:.1f}% of cost; "
          f"{opp_all} discordant in sign ({100 * opp_all / n_all:.1f}%)")
    print(f"  of those, {n_filt} have both slopes above {SLOPE_MIN}; "
          f"{opp_filt} discordant ({100 * opp_filt / n_filt:.1f}%)")
    print(f"  under independence the discordant fraction would be ~50%")
    # the discordant pairs are what the manuscript shows
    pairs = [(i, j) for i, j in filt_pairs if sl[i] * sl[j] < 0]
    d = np.array([abs(c[j] - c[i]) / c[i] for i, j in pairs])
    opp = len(pairs)
    print(f"\n  {len(pairs)} pairs within {100 * COST_TOL:.1f}% of cost with "
          f"opposite slopes; {opp} of them opposite in sign")
    print(f"  median within-pair cost difference {100 * np.median(d):.3f}%")
    k = int(np.argmin(d))
    i, j = pairs[k]
    print(f"  closest pair: cost {c[i]:.4f} / {c[j]:.4f}"
          f"  slope {sl[i]:+.4f} / {sl[j]:+.4f}")

    with open(os.path.join(DATA, "twin_rank_pairs.csv"), "w",
              newline="") as fh:
        wr = csv.writer(fh)
        wr.writerow(["code_a", "code_b", "cost_a", "cost_b", "cost_diff_pct",
                     "slope_a", "slope_b", "opposite_sign"])
        for i2, j2 in pairs:
            wr.writerow([int(pick[i2]), int(pick[j2]),
                         f"{c[i2]:.6f}", f"{c[j2]:.6f}",
                         f"{100 * abs(c[j2] - c[i2]) / c[i2]:.4f}",
                         f"{sl[i2]:.6f}", f"{sl[j2]:.6f}",
                         "yes" if sl[i2] * sl[j2] < 0 else "no"])

    json.dump(dict(n_codes=NDRAW, n_all_pairs=n_all, opp_all=opp_all,
                   frac_opp_all=round(opp_all / n_all, 4),
                   n_filtered=n_filt, opp_filtered=opp_filt,
                   frac_opp_filtered=round(opp_filt / n_filt, 4), frac_positive=round(float((sl > 0).mean()), 4),
                   slope_min=round(float(sl.min()), 4),
                   slope_max=round(float(sl.max()), 4),
                   slope_median=round(float(np.median(sl)), 4),
                   sgc_slope=round(s_sgc, 4),
                   sgc_percentile=round(float(100 * (sl < s_sgc).mean()), 1),
                   corr_cost_slope=round(r_cost, 4),
                   n_pairs=len(pairs), n_opposite=opp,
                   med_cost_diff_pct=round(float(100 * np.median(d)), 4),
                   example=dict(cost_a=float(c[i]), cost_b=float(c[j]),
                                slope_a=float(sl[i]), slope_b=float(sl[j]),
                                diff_pct=float(100 * d[k]))),
              open(os.path.join(DATA, "twin_ranks.json"), "w"), indent=1)
    print("\n  data/twin_ranks.json and data/twin_rank_pairs.csv written\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
