"""Per-subject features from preprocessed parcel time series (code/preprocess.py output).

Networks (Schaefer-2018, 7-network labels): Default (DMN), SalVentAttn (salience, SN),
Cont (control, CEN).

- Intrinsic neural timescale per parcel: the autocorrelation window ACW-50, the lag in
  seconds at which the autocorrelation function first falls below 0.5, by linear
  interpolation between lags (TR 2.5 s); ACW-0, the same for the first crossing of zero.
  Network ACW = mean over the network's parcels.
- Network-dominance states: at each volume, the network (of DMN, SN, CEN) with the
  highest mean z-scored signal. From that sequence: fractional occupancy and mean dwell
  time (s) of DMN dominance, and transition entropy, the entropy in bits of the
  empirical distribution of transitions between consecutive volumes (all nine ordered
  pairs, self-transitions included).

Usage: python code/features.py <derived dir> <labels.txt> <out.csv>
"""
import csv
import glob
import json
import os
import sys

import numpy as np

TR = 2.5
NETS = {"DMN": "Default", "SN": "SalVentAttn", "CEN": "Cont"}


def acf(x, maxlag):
    x = x - x.mean()
    d = (x * x).sum()
    return np.array([1.0] + [(x[:-k] * x[k:]).sum() / d for k in range(1, maxlag + 1)])


def crossing(a, level):
    """Lag (in TRs, interpolated) where a first falls below level; NaN if never."""
    for k in range(1, len(a)):
        if a[k] < level:
            return (k - 1) + (a[k - 1] - level) / (a[k - 1] - a[k])
    return float("nan")


def main():
    src, labels_path, out = sys.argv[1], sys.argv[2], sys.argv[3]
    names = [line.split()[1] for line in open(labels_path) if line.strip()]
    net_of = [n.split("_")[2] for n in names]
    idx = {k: [i for i, n in enumerate(net_of) if n == v] for k, v in NETS.items()}
    rows = []
    for f in sorted(glob.glob(os.path.join(src, "sub-*_parcels.npy"))):
        sub = os.path.basename(f).split("_")[0]
        ts = np.load(f)
        qc = json.load(open(f.replace("_parcels.npy", "_qc.json")))
        n = ts.shape[0]
        maxlag = n // 2
        acw50 = np.array([crossing(acf(ts[:, p], maxlag), 0.5) * TR if np.isfinite(ts[:, p]).all() else np.nan
                          for p in range(ts.shape[1])])
        acw0 = np.array([crossing(acf(ts[:, p], maxlag), 0.0) * TR if np.isfinite(ts[:, p]).all() else np.nan
                         for p in range(ts.shape[1])])
        row = {"participant_id": sub, "mean_fd_mm": qc["mean_fd_mm"], "share_fd_over_0_5": qc["share_fd_over_0_5"],
               "max_translation_mm": qc["max_translation_mm"], "registration_dice": qc["registration_dice"],
               "acw50_all": float(np.nanmean(acw50)), "acw0_all": float(np.nanmean(acw0))}
        for k, ii in idx.items():
            row[f"acw50_{k}"] = float(np.nanmean(acw50[ii]))
            row[f"acw0_{k}"] = float(np.nanmean(acw0[ii]))
        net_ts = np.column_stack([ts[:, idx[k]].mean(1) for k in NETS])
        z = (net_ts - net_ts.mean(0)) / net_ts.std(0, ddof=1)
        state = z.argmax(1)  # 0 = DMN, 1 = SN, 2 = CEN
        runs, cur, length = [], state[0], 1
        for s in state[1:]:
            if s == cur:
                length += 1
            else:
                runs.append((cur, length))
                cur, length = s, 1
        runs.append((cur, length))
        dmn_runs = [l for s, l in runs if s == 0]
        row["dmn_occupancy"] = float((state == 0).mean())
        row["dmn_dwell_s"] = float(np.mean(dmn_runs) * TR) if dmn_runs else 0.0
        T = np.zeros((3, 3))
        for a, b in zip(state[:-1], state[1:]):
            T[a, b] += 1
        p = T.ravel() / T.sum()
        p = p[p > 0]
        row["transition_entropy_bits"] = float(-(p * np.log2(p)).sum())
        rows.append(row)
    with open(out, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)
    print(len(rows), "subjects")


if __name__ == "__main__":
    main()