#!/usr/bin/env python3

from __future__ import annotations

import csv
from dataclasses import dataclass
from typing import Dict, Iterable, List, Sequence, Tuple

import numpy as np
import torch
import torch.nn as nn


# -----------------------------------------------------------------------------
# Regulatory prior
# -----------------------------------------------------------------------------
@dataclass
class TFRegulon:
    tf: str
    targets: np.ndarray          # column indices in current expression matrix
    weights: np.ndarray          # signed CollecTRI edge weights
    n_total_targets: int
    n_present_targets: int
    coverage: float
    programs: Tuple[str, ...]
    tf_weight: float
    reproducibility: float
    perturbed_cl13: bool


class RegulonScorer:
    """Simple, dependency-light signed TF activity scorer.

    Activity_r(cell) = sum_g w_rg * z_g(cell) / sqrt(sum_g w_rg^2)

    where z_g is standardized using retained TRAINING cells only.

    This transparent scoring method is the signed target-gene prior
    used in the TRACE implementation; it is not TF transcript abundance.
    """

    def __init__(
        self,
        prior_csv: str,
        var_names: Sequence[str],
        min_targets: int = 5,
        min_coverage: float = 0.15,
        candidate_source: str = "taiji",
        max_tfs: int = 0,
        edge_weight_mode: str = "weight",
    ):
        self.prior_csv = prior_csv
        self.var_names = np.asarray([str(x) for x in var_names])
        self.lookup = {g.upper(): i for i, g in enumerate(self.var_names)}
        self.min_targets = int(min_targets)
        self.min_coverage = float(min_coverage)
        self.candidate_source = candidate_source
        self.max_tfs = int(max_tfs)
        self.edge_weight_mode = str(edge_weight_mode).lower()
        if self.edge_weight_mode not in {"weight", "sign", "sign_resources"}:
            raise ValueError("edge_weight_mode must be one of: weight, sign, sign_resources")

        rows = []
        with open(prior_csv, newline="") as f:
            reader = csv.DictReader(f)
            required = {"program_id", "tf", "target_gene", "weight", "provenance"}
            if not required.issubset(set(reader.fieldnames or [])):
                raise ValueError(
                    f"{prior_csv} must contain at least {sorted(required)}; "
                    f"found {reader.fieldnames}"
                )
            for r in reader:
                rows.append(r)
        if not rows:
            raise ValueError(f"No rows found in {prior_csv}")

        # Candidate TF metadata comes from Taiji state-association rows.  This
        # selects TFs biologically linked to the T-cell programs without imposing
        # a state order on the trajectory.
        meta: Dict[str, dict] = {}
        for r in rows:
            if (r.get("provenance") or "").strip().lower() != "taiji_pagerank":
                continue
            tf = (r.get("tf") or "").strip()
            if not tf:
                continue
            m = meta.setdefault(tf, {
                "programs": set(), "tf_weight": [], "reproducibility": [], "perturbed": False
            })
            m["programs"].add((r.get("program_id") or "").strip())
            try:
                m["tf_weight"].append(abs(float(r.get("tf_weight") or 1.0)))
            except Exception:
                pass
            try:
                val = float(r.get("reproducibility") or "nan")
                if np.isfinite(val):
                    m["reproducibility"].append(val)
            except Exception:
                pass
            m["perturbed"] = m["perturbed"] or str(r.get("perturbed_cl13", "")).lower() in {"true", "1", "yes"}

        # Collect signed TF-target edges.  The source file repeats some edges in
        # multiple program blocks; aggregate duplicates by mean edge weight so a
        # TF-target relation is counted once.
        edge_vals: Dict[Tuple[str, str], List[float]] = {}
        total_targets: Dict[str, set] = {}
        for r in rows:
            if (r.get("provenance") or "").strip().lower() != "collectri":
                continue
            tf = (r.get("tf") or "").strip()
            tg = (r.get("target_gene") or "").strip()
            if not tf or not tg:
                continue
            if candidate_source == "taiji" and tf not in meta:
                continue
            try:
                if self.edge_weight_mode == "weight":
                    # Legacy BioCHIME behavior.  In the assembled T-cell prior this
                    # column may include state-specific scaling, so the paper
                    # configuration uses state-independent sign mode.
                    w = float(r.get("weight") or 0.0)
                else:
                    sg = float(r.get("sign") or 0.0)
                    if sg == 0.0:
                        raw = float(r.get("weight") or 0.0)
                        sg = float(np.sign(raw))
                    if self.edge_weight_mode == "sign":
                        w = sg
                    else:
                        nr = max(float(r.get("n_resources") or 1.0), 1.0)
                        w = sg * np.sqrt(nr)
            except Exception:
                continue
            if not np.isfinite(w) or abs(w) < 1e-15:
                continue
            edge_vals.setdefault((tf, tg), []).append(w)
            total_targets.setdefault(tf, set()).add(tg)

        # Rank candidates only if requested.  We use reproducibility first, then
        # Taiji weight.  This ranking is based on external prior metadata, not on
        # the held-out expression data.
        all_tfs = sorted(total_targets)
        if max_tfs > 0:
            def rank_key(tf: str):
                m = meta.get(tf, {})
                rep = np.nanmean(m.get("reproducibility", [np.nan])) if m.get("reproducibility") else -np.inf
                tw = np.nanmax(m.get("tf_weight", [0.0])) if m.get("tf_weight") else 0.0
                if not np.isfinite(rep):
                    rep = -1e9
                return (rep, tw)
            all_tfs = sorted(all_tfs, key=rank_key, reverse=True)[:max_tfs]

        regulons: List[TFRegulon] = []
        for tf in all_tfs:
            tg_all = sorted(total_targets[tf])
            present_idx, present_w = [], []
            for tg in tg_all:
                j = self.lookup.get(tg.upper())
                if j is None:
                    continue
                vals = edge_vals.get((tf, tg), [])
                if not vals:
                    continue
                present_idx.append(j)
                present_w.append(float(np.mean(vals)))
            nt = len(tg_all)
            npres = len(present_idx)
            cov = npres / max(nt, 1)
            if npres < min_targets or cov < min_coverage:
                continue
            m = meta.get(tf, {})
            tw = float(np.nanmax(m.get("tf_weight", [1.0]))) if m.get("tf_weight") else 1.0
            reps = m.get("reproducibility", [])
            rep = float(np.nanmax(reps)) if reps else float("nan")
            regulons.append(TFRegulon(
                tf=tf,
                targets=np.asarray(present_idx, dtype=int),
                weights=np.asarray(present_w, dtype=np.float64),
                n_total_targets=nt,
                n_present_targets=npres,
                coverage=cov,
                programs=tuple(sorted(x for x in m.get("programs", set()) if x)),
                tf_weight=tw,
                reproducibility=rep,
                perturbed_cl13=bool(m.get("perturbed", False)),
            ))

        if not regulons:
            raise ValueError(
                "No TF regulons passed coverage filters. Try --min_targets 3 "
                "and/or --min_coverage 0.05, or rebuild the expression matrix "
                "with more genes / forced inclusion of prior target genes."
            )
        self.regulons = regulons
        self.tf_names = [r.tf for r in regulons]
        self.gene_mu = None
        self.gene_sd = None

    def fit_gene_scaler(self, X_train: np.ndarray):
        X_train = np.asarray(X_train, dtype=np.float64)
        self.gene_mu = X_train.mean(0)
        self.gene_sd = X_train.std(0)
        self.gene_sd[self.gene_sd < 1e-8] = 1.0
        return self

    def score(self, X: np.ndarray) -> np.ndarray:
        if self.gene_mu is None or self.gene_sd is None:
            raise RuntimeError("Call fit_gene_scaler(X_train) before score(X)")
        X = np.asarray(X, dtype=np.float64)
        Zg = (X - self.gene_mu) / self.gene_sd
        out = np.empty((len(X), len(self.regulons)), dtype=np.float32)
        for j, r in enumerate(self.regulons):
            w = r.weights
            denom = np.sqrt(np.sum(w * w)) + 1e-12
            out[:, j] = (Zg[:, r.targets] @ w / denom).astype(np.float32)
        return out

    def coverage_rows(self) -> List[dict]:
        rows = []
        for r in self.regulons:
            rows.append({
                "tf": r.tf,
                "n_targets_total": r.n_total_targets,
                "n_targets_present": r.n_present_targets,
                "coverage": r.coverage,
                "programs": ";".join(r.programs),
                "tf_weight": r.tf_weight,
                "reproducibility": r.reproducibility,
                "perturbed_cl13": r.perturbed_cl13,
            })
        return rows


def standardize_activity(A_train: np.ndarray, A: np.ndarray | None = None):
    A_train = np.asarray(A_train, dtype=np.float64)
    mu = A_train.mean(0)
    sd = A_train.std(0)
    sd[sd < 1e-8] = 1.0
    if A is None:
        return ((A_train - mu) / sd).astype(np.float32), mu.astype(np.float32), sd.astype(np.float32)
    return ((np.asarray(A) - mu) / sd).astype(np.float32), mu.astype(np.float32), sd.astype(np.float32)


def fit_ridge_map(A: np.ndarray, Z: np.ndarray, lam: float = 1.0) -> Tuple[np.ndarray, np.ndarray]:
    """Fit Z ~= A @ W + b; derivatives map as dZ/dt ~= dA/dt @ W."""
    A = np.asarray(A, dtype=np.float64)
    Z = np.asarray(Z, dtype=np.float64)
    am, zm = A.mean(0), Z.mean(0)
    Ac, Zc = A - am, Z - zm
    G = Ac.T @ Ac + float(lam) * np.eye(A.shape[1])
    W = np.linalg.solve(G, Ac.T @ Zc)
    b = zm - am @ W
    return W.astype(np.float32), b.astype(np.float32)


# -----------------------------------------------------------------------------
#Hermite helpers
# -----------------------------------------------------------------------------
def pchip_slopes(x: np.ndarray, Y: np.ndarray) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    K = len(x)
    h = np.diff(x)
    d = (Y[1:] - Y[:-1]) / h[:, None, None]
    m = np.zeros_like(Y)
    if K == 2:
        m[0] = d[0]
        m[1] = d[0]
        return m
    for k in range(1, K - 1):
        d0, d1 = d[k - 1], d[k]
        h0, h1 = h[k - 1], h[k]
        w0, w1 = 2 * h1 + h0, h1 + 2 * h0
        mono = (d0 * d1) > 0
        with np.errstate(divide="ignore", invalid="ignore"):
            mk = (w0 + w1) / (w0 / d0 + w1 / d1)
        m[k] = np.where(mono, mk, 0.0)

    def endpoint(h0, h1, d0, d1):
        me = ((2 * h0 + h1) * d0 - h0 * d1) / (h0 + h1)
        me = np.where(np.sign(me) != np.sign(d0), 0.0, me)
        me = np.where((np.sign(d0) != np.sign(d1)) & (np.abs(me) > 3 * np.abs(d0)), 3 * d0, me)
        return me

    m[0] = endpoint(h[0], h[1], d[0], d[1])
    m[-1] = endpoint(h[-1], h[-2], d[-1], d[-2])
    return m


def project_shape_safe_slopes(x: np.ndarray, Y: np.ndarray, M: np.ndarray) -> np.ndarray:
    """Conservative coordinate-wise projection that prevents PCHIP-like overshoot."""
    x = np.asarray(x, dtype=float)
    h = np.diff(x)
    d = (Y[1:] - Y[:-1]) / h[:, None, None]
    out = np.asarray(M).copy()
    K = len(x)
    if K == 2:
        sec = d[0]
        for k in [0, 1]:
            ok = np.sign(out[k]) == np.sign(sec)
            out[k] = np.where(ok, out[k], 0.0)
            out[k] = np.sign(sec) * np.minimum(np.abs(out[k]), 3.0 * np.abs(sec))
        return out

    for k, sec in [(0, d[0]), (K - 1, d[-1])]:
        ok = np.sign(out[k]) == np.sign(sec)
        out[k] = np.where(ok, out[k], 0.0)
        out[k] = np.sign(sec) * np.minimum(np.abs(out[k]), 3.0 * np.abs(sec))

    for k in range(1, K - 1):
        dl, dr = d[k - 1], d[k]
        mono = (dl * dr) > 0
        sgn = np.sign(dl)
        cap = 3.0 * np.minimum(np.abs(dl), np.abs(dr))
        cand = np.where(np.sign(out[k]) == sgn, out[k], 0.0)
        cand = sgn * np.minimum(np.abs(cand), cap)
        out[k] = np.where(mono, cand, 0.0)
    return out


def hermite_eval(x: np.ndarray, Y: np.ndarray, M: np.ndarray, t: np.ndarray):
    x = np.asarray(x, dtype=float)
    t = np.asarray(t, dtype=float)
    K, B, _ = Y.shape
    idx = np.clip(np.searchsorted(x, t, side="right") - 1, 0, K - 2)
    xi, xi1 = x[idx], x[idx + 1]
    hi = xi1 - xi
    s = ((t - xi) / hi)[:, None]
    hb = hi[:, None]
    ar = np.arange(B)
    yi, yi1 = Y[idx, ar], Y[idx + 1, ar]
    mi, mi1 = M[idx, ar], M[idx + 1, ar]
    mu = ((2*s**3 - 3*s**2 + 1)*yi + (s**3 - 2*s**2 + s)*hb*mi
          + (-2*s**3 + 3*s**2)*yi1 + (s**3 - s**2)*hb*mi1)
    mup = ((6*s**2 - 6*s)*yi + (3*s**2 - 4*s + 1)*hb*mi
           + (-6*s**2 + 6*s)*yi1 + (3*s**2 - 2*s)*hb*mi1) / hb
    return mu, mup


def align_joint(
    clouds_z: Sequence[np.ndarray],
    clouds_a: Sequence[np.ndarray],
    rng: np.random.Generator,
    batch: int = 128,
):
    """Use the existing Euclidean OT plan, carrying TF activities with cells."""
    B = min(batch, min(len(c) for c in clouds_z))
    i0 = rng.choice(len(clouds_z[0]), B, replace=False)
    z = [clouds_z[0][i0]]
    a = [clouds_a[0][i0]]
    for k in range(1, len(clouds_z)):
        cand = rng.choice(len(clouds_z[k]), B, replace=False)
        zc, ac = clouds_z[k][cand], clouds_a[k][cand]
        import benchmark_table as BT
        P = BT.ot_plan(z[-1], zc)
        Pn = P / (P.sum(1, keepdims=True) + 1e-12)
        jj = np.array([rng.choice(B, p=Pn[i]) for i in range(B)])
        z.append(zc[jj])
        a.append(ac[jj])
    return np.stack(z, 0), np.stack(a, 0)


def _match_bio_vector_norm(m_geom: np.ndarray, m_bio: np.ndarray, clip=(0.25, 4.0)) -> np.ndarray:
    """Make biology contribute direction rather than arbitrary numerical scale."""
    ng = np.linalg.norm(m_geom, axis=-1, keepdims=True)
    nb = np.linalg.norm(m_bio, axis=-1, keepdims=True)
    scale = ng / (nb + 1e-12)
    scale = np.clip(scale, clip[0], clip[1])
    scale = np.where((ng > 1e-8) & (nb > 1e-8), scale, 0.0)
    return m_bio * scale


@dataclass
class BioFitInfo:
    rho: float
    ridge_W: np.ndarray
    ridge_b: np.ndarray


def fit_regulon_bio_flow(
    clouds_z_by_day: Dict[float, np.ndarray],
    clouds_a_by_day: Dict[float, np.ndarray],
    dn: Sequence[float],
    dev: str,
    steps: int,
    seed: int,
    rho: float = 1.0,
    batch: int = 128,
    lr: float = 1e-3,
    ridge: float = 1.0,
    ridge_map: tuple[np.ndarray, np.ndarray] | None = None,
    shape_project: bool = True,
    log_every: int = 0,
):
    """Train TRACE with TF-regulon-guided Hermite knot tangents."""
    if not (0.0 <= rho <= 1.0):
        raise ValueError("rho must be in [0,1]")
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    x = np.asarray(sorted(dn), dtype=float)
    zclouds = [clouds_z_by_day[t] for t in x]
    aclouds = [clouds_a_by_day[t] for t in x]
    d = zclouds[0].shape[1]

    Apool = np.vstack(aclouds)
    Zpool = np.vstack(zclouds)
    if ridge_map is None:
        W, b = fit_ridge_map(Apool, Zpool, lam=ridge)
    else:
        # Negative controls such as TF-channel permutation must use the
        # ridge map fitted on the TRUE TF identities.  Re-fitting after a
        # global column permutation would simply permute the rows of W and
        # make the control algebraically equivalent to the unpermuted model.
        W, b = ridge_map
        W = np.asarray(W, dtype=np.float32)
        b = np.asarray(b, dtype=np.float32)
        if W.shape != (Apool.shape[1], Zpool.shape[1]):
            raise ValueError(
                f"ridge_map W has shape {W.shape}, expected "
                f"{(Apool.shape[1], Zpool.shape[1])}"
            )
        if b.shape != (Zpool.shape[1],):
            raise ValueError(
                f"ridge_map b has shape {b.shape}, expected {(Zpool.shape[1],)}"
            )

    import benchmark_table as BT
    field = BT.Field(d).to(dev)
    opt = torch.optim.Adam(field.parameters(), lr=lr)
    ni = len(x) - 1

    for step in range(steps):
        Z, A = align_joint(zclouds, aclouds, rng, batch=batch)
        B = Z.shape[1]
        per = B // ni
        ts = []
        for i in range(ni):
            n = per + (1 if i < B - per * ni else 0)
            ts.append(rng.uniform(x[i], x[i + 1], size=n))
        tq = np.concatenate(ts)[:B]
        rng.shuffle(tq)

        m_geom = pchip_slopes(x, Z)
        if rho > 0:
            m_tf = pchip_slopes(x, A)
            m_bio = np.einsum("kbr,rd->kbd", m_tf, W)
            m_bio = _match_bio_vector_norm(m_geom, m_bio)
            M = (1.0 - rho) * m_geom + rho * m_bio
            if shape_project:
                M = project_shape_safe_slopes(x, Z, M)
        else:
            M = m_geom

        pos, vel = hermite_eval(x, Z, M, tq)
        tt = torch.tensor(tq[:, None], dtype=torch.float32, device=dev)
        pred = field(torch.tensor(pos, dtype=torch.float32, device=dev), tt)
        target = torch.tensor(vel, dtype=torch.float32, device=dev)
        loss = ((pred - target) ** 2).mean()
        opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(field.parameters(), 1.0)
        opt.step()
        if log_every and (step + 1) % log_every == 0:
            print(f"[BioCHIME rho={rho:g}] {step+1}/{steps} loss={float(loss):.6f}", flush=True)

    field.eval()
    return field, BioFitInfo(rho=rho, ridge_W=W, ridge_b=b)


def permute_within_day(clouds_a_by_day: Dict[float, np.ndarray], seed: int):
    """Negative control preserving per-day TF distributions but breaking cell correspondence."""
    rng = np.random.default_rng(seed)
    return {t: A[rng.permutation(len(A))].copy() for t, A in clouds_a_by_day.items()}


def permute_tf_channels(clouds_a_by_day: Dict[float, np.ndarray], seed: int):
    """Wrong-biology control: keep values but scramble TF identities globally."""
    rng = np.random.default_rng(seed)
    r = next(iter(clouds_a_by_day.values())).shape[1]
    p = rng.permutation(r)
    return {t: A[:, p].copy() for t, A in clouds_a_by_day.items()}