"""Build the cortical-surface asset shown on the home page.

The figure is real data, not an illustration:

  * Geometry   FreeSurfer reconstruction (pial and inflated surfaces) of the
               subject in the MNE sample dataset.
  * Overlay    dSPM source estimate of the auditory evoked field (MEG only),
               computed here with MNE-Python from the dataset's evoked response
               and its distributed inverse operator.

Nothing is invented or embellished. The only processing beyond the standard
MNE pipeline is mesh decimation for the web, spreading the source estimate from
the source-space vertices to the display mesh (the same neighbour averaging
MNE uses for plotting), and a baked ambient-occlusion term for shading.

This script is run by hand, not in CI. It writes an uncompressed intermediate
that `pack_cortex.mjs` turns into the two committed files:

    public/data/cortex.dat          compressed geometry and frames
    src/content/cortex.meta.json    layout of the binary, time axis, colour limits

Usage (Python >= 3.10, mne, numpy, scipy, fast-simplification):

    python scripts/cortex/build_cortex.py --data /path/to/MNE-sample-data
    node scripts/cortex/pack_cortex.mjs
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import fast_simplification
import mne
import numpy as np
from scipy import ndimage, sparse

REPO = Path(__file__).resolve().parents[2]

# ---------------------------------------------------------------------------
# Parameters
# ---------------------------------------------------------------------------
CONDITION = "Right Auditory"  # right-ear tones; the left hemisphere faces the viewer
INVERSE = "sample_audvis-meg-oct-6-meg-inv.fif"
SNR = 3.0  # lambda2 = 1 / SNR^2, the MNE default for evoked data
METHOD = "dSPM"

T_START, T_STOP = -0.040, 0.300  # s; window exported for playback
FRAME_EVERY = 5  # samples between frames (8.3 ms at 600.6 Hz)
SMOOTH_STEPS = 10  # neighbour-averaging iterations when spreading to the mesh
TARGET_VERTS = 21000  # per hemisphere after decimation
ACT_MAX = 28.0  # dSPM value mapped to 255 in the stored frames
ACT_FLOOR = 3.5  # values below this are never displayed and are stored as zero
AO_DIRECTIONS = 320
AO_PIXEL_MM = 1.0
INFLATED_GAP_MM = 6.0  # space between the hemispheres in the inflated state


def adjacency(n_verts: int, tris: np.ndarray) -> sparse.csr_array:
    """Vertex adjacency including self-connections."""
    i = np.concatenate([tris[:, 0], tris[:, 1], tris[:, 2], np.arange(n_verts)])
    j = np.concatenate([tris[:, 1], tris[:, 2], tris[:, 0], np.arange(n_verts)])
    a = sparse.coo_array((np.ones(i.size, np.float32), (i, j)), shape=(n_verts, n_verts))
    a = (a + a.T).tocsr()
    a.data[:] = 1.0
    return a


def spread(adj: sparse.csr_array, vertno: np.ndarray, data: np.ndarray, steps: int) -> np.ndarray:
    """Spread values given on `vertno` to every vertex by neighbour averaging.

    Vertices without data take the mean of neighbours that have data; this is
    repeated until the surface is filled and at least `steps` times in total.
    """
    n = adj.shape[0]
    full = np.zeros((n, data.shape[1]), np.float32)
    full[vertno] = data
    have = np.zeros(n, np.float32)
    have[vertno] = 1.0
    done = 0
    while done < steps or not have.all():
        num = adj @ full
        den = adj @ have
        ok = den > 0
        full[ok] = num[ok] / den[ok, None]
        have = ok.astype(np.float32)
        done += 1
        if done > 200:
            raise RuntimeError("surface could not be filled")
    return full


def vertex_normals(rr: np.ndarray, tris: np.ndarray) -> np.ndarray:
    fn = np.cross(rr[tris[:, 1]] - rr[tris[:, 0]], rr[tris[:, 2]] - rr[tris[:, 0]])
    nn = np.zeros_like(rr)
    for k in range(3):
        np.add.at(nn, tris[:, k], fn)
    nn /= np.maximum(np.linalg.norm(nn, axis=1, keepdims=True), 1e-12)
    return nn


def fibonacci_sphere(n: int) -> np.ndarray:
    i = np.arange(n) + 0.5
    phi = np.arccos(1 - 2 * i / n)
    theta = np.pi * (1 + 5**0.5) * i
    return np.c_[np.cos(theta) * np.sin(phi), np.sin(theta) * np.sin(phi), np.cos(phi)]


def ambient_occlusion(occluders: np.ndarray, rr: np.ndarray, nn: np.ndarray) -> np.ndarray:
    """Cosine-weighted visibility of each vertex, from depth maps of the dense surface.

    For every direction the dense vertex cloud is splatted into an orthographic
    depth buffer; a display vertex is lit from that direction if nothing lies
    in front of it. Cheap, and adequate for a smooth per-vertex term.
    """
    dirs = fibonacci_sphere(AO_DIRECTIONS)
    lit = np.zeros(len(rr), np.float64)
    wsum = np.zeros(len(rr), np.float64)
    for d in dirs:
        w = nn @ d
        facing = w > 0.02
        if not facing.any():
            continue
        helper = np.array([0.0, 0.0, 1.0]) if abs(d[2]) < 0.9 else np.array([1.0, 0.0, 0.0])
        u = np.cross(d, helper)
        u /= np.linalg.norm(u)
        v = np.cross(d, u)
        ou, ov, od = occluders @ u, occluders @ v, occluders @ d
        u0, v0 = ou.min() - 3, ov.min() - 3
        nu = int((ou.max() - u0) / AO_PIXEL_MM) + 4
        nv = int((ov.max() - v0) / AO_PIXEL_MM) + 4
        zbuf = np.full((nu, nv), -np.inf, np.float32)
        np.maximum.at(zbuf, (((ou - u0) / AO_PIXEL_MM).astype(int), ((ov - v0) / AO_PIXEL_MM).astype(int)), od.astype(np.float32))
        zbuf = ndimage.maximum_filter(zbuf, size=3)
        pu = ((rr[facing] @ u - u0) / AO_PIXEL_MM).astype(int)
        pv = ((rr[facing] @ v - v0) / AO_PIXEL_MM).astype(int)
        depth = rr[facing] @ d
        # Soft test: tolerate the slope of the surface inside one dilated pixel.
        gap = zbuf[pu, pv] - depth
        vis = np.clip(1.0 - (gap - 1.6) / 2.4, 0.0, 1.0)
        lit[facing] += w[facing] * vis
        wsum[facing] += w[facing]
    return (lit / np.maximum(wsum, 1e-9)).astype(np.float32)


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--data", required=True, type=Path, help="path to MNE-sample-data")
    ap.add_argument("--cache", type=Path, default=REPO / "scripts" / "cortex" / ".cache")
    args = ap.parse_args()

    subjects_dir = args.data / "subjects"
    meg_dir = args.data / "MEG" / "sample"
    mne.set_log_level("WARNING")

    # --- Source estimate ----------------------------------------------------
    evoked = mne.read_evokeds(meg_dir / "sample_audvis-ave.fif", condition=CONDITION, baseline=(None, 0))
    inv = mne.minimum_norm.read_inverse_operator(meg_dir / INVERSE)
    stc = mne.minimum_norm.apply_inverse(evoked, inv, lambda2=1.0 / SNR**2, method=METHOD)
    # Whole samples, anchored at stimulus onset so that t = 0 is a frame.
    zero = int(np.argmin(np.abs(stc.times)))
    first = zero - FRAME_EVERY * int(round(-T_START * stc.sfreq / FRAME_EVERY))
    last = zero + FRAME_EVERY * int(round(T_STOP * stc.sfreq / FRAME_EVERY))
    frame_idx = list(range(first, last + 1, FRAME_EVERY))
    frame_times = stc.times[frame_idx]

    # Sensor-level trace for the time axis: RMS over planar gradiometers.
    grad = evoked.copy().pick("grad")
    keep = (grad.times >= T_START - 1e-9) & (grad.times <= T_STOP + 1e-9)
    rms_times = grad.times[keep]
    rms = np.sqrt((grad.data[:, keep] ** 2).mean(axis=0)) * 1e13  # T/m -> fT/cm

    hemis = []
    occluders = []
    for h, hemi in enumerate(("lh", "rh")):
        surf = subjects_dir / "sample" / "surf"
        pial, tris = mne.read_surface(surf / f"{hemi}.pial")
        infl, _ = mne.read_surface(surf / f"{hemi}.inflated")
        sulc = mne.surface.read_curvature(surf / f"{hemi}.sulc", binary=False)
        tris = tris.astype(np.int64)
        n = len(pial)
        cortex = np.zeros(n, bool)
        cortex[mne.read_label(subjects_dir / "sample" / "label" / f"{hemi}.cortex.label").vertices] = True

        adj = adjacency(n, tris)
        vertno = stc.vertices[h]
        data = (stc.lh_data if h == 0 else stc.rh_data)[:, frame_idx].astype(np.float32)
        full = spread(adj, vertno, data, SMOOTH_STEPS)
        full[~cortex] = 0.0  # no cortex on the medial wall, so no estimate either

        # Shading normals come from the dense surface, lightly smoothed; normals
        # recomputed on the decimated mesh would show its triangles.
        dense_nn = vertex_normals(pial, tris)
        for _ in range(2):
            dense_nn = adj @ dense_nn
            dense_nn /= np.maximum(np.linalg.norm(dense_nn, axis=1, keepdims=True), 1e-12)

        ratio = 1.0 - TARGET_VERTS / n
        rr32, tris32 = pial.astype(np.float32), tris.astype(np.int32)
        _, _, collapses = fast_simplification.simplify(rr32, tris32, target_reduction=ratio, agg=5, return_collapses=True)
        dec_rr, dec_tris, member_of = fast_simplification.replay_simplification(rr32, tris32, collapses)
        dec_rr = dec_rr.astype(np.float64)
        # Every display vertex stands for the group of dense vertices that were
        # collapsed into it. Its attributes come from the member closest to it.
        # (Matching by distance alone can jump to the opposite bank of a sulcus.)
        dist = np.linalg.norm(pial - dec_rr[member_of], axis=1)
        order = np.lexsort((dist, member_of))
        head = np.r_[True, member_of[order][1:] != member_of[order][:-1]]
        nearest = np.empty(len(dec_rr), np.int64)
        nearest[member_of[order][head]] = order[head]

        hemis.append(
            dict(
                hemi=hemi,
                pial=dec_rr,
                tris=dec_tris.astype(np.int64),
                infl=infl[nearest].astype(np.float64),
                normal=dense_nn[nearest],
                sulc=sulc[nearest].astype(np.float32),
                cortex=cortex[nearest],
                act=full[nearest],
                peak=float(full.max()),
            )
        )
        occluders.append(pial)
        print(f"{hemi}: {n} -> {len(dec_rr)} vertices, {len(dec_tris)} triangles")

    # --- Inflated layout: hemispheres side by side, matched in size to the pial brain
    lh, rh = hemis
    lh["infl"][:, 0] -= lh["infl"][:, 0].max() + INFLATED_GAP_MM / 2
    rh["infl"][:, 0] -= rh["infl"][:, 0].min() - INFLATED_GAP_MM / 2
    pial_all = np.vstack([lh["pial"], rh["pial"]])
    infl_all = np.vstack([lh["infl"], rh["infl"]])
    scale = np.ptp(pial_all[:, 1]) / np.ptp(infl_all[:, 1])
    pial_mid = (pial_all.max(0) + pial_all.min(0)) / 2
    infl_mid = (infl_all.max(0) + infl_all.min(0)) / 2
    for hd in hemis:
        hd["infl"] = (hd["infl"] - infl_mid) * scale + pial_mid

    # --- Shading terms --------------------------------------------------------
    occ = np.vstack(occluders)
    for hd in hemis:
        ao = ambient_occlusion(occ, hd["pial"], hd["normal"])
        adj = adjacency(len(ao), hd["tris"])
        for _ in range(2):
            ao = (adj @ ao) / (adj @ np.ones_like(ao))
        hd["ao"] = ao

    # --- Pack -------------------------------------------------------------------
    offset = len(lh["pial"])
    tris = np.vstack([lh["tris"], rh["tris"] + offset]).astype(np.uint32)
    pial = np.vstack([lh["pial"], rh["pial"]])
    infl = np.vstack([lh["infl"], rh["infl"]])
    normal = np.vstack([lh["normal"], rh["normal"]])
    ao = np.concatenate([lh["ao"], rh["ao"]])
    sulc = np.concatenate([lh["sulc"], rh["sulc"]])
    cortex = np.concatenate([lh["cortex"], rh["cortex"]])
    act = np.vstack([lh["act"], rh["act"]])  # (N, T)

    both = np.vstack([pial, infl])
    center = (both.max(0) + both.min(0)) / 2
    half = float(np.abs(both - center).max()) * 1.001

    def quant16(x: np.ndarray) -> np.ndarray:
        # 12 significant bits (0.04 mm steps) in an int16; the zeroed low bits compress away.
        return (np.round((x - center) / half * 2047).astype(np.int32) * 16).astype("<i2")

    def oct8(n: np.ndarray) -> np.ndarray:
        # Octahedral encoding of unit vectors into two signed bytes.
        n = n / np.abs(n).sum(axis=1, keepdims=True)
        xy = n[:, :2].copy()
        back = n[:, 2] < 0
        xy[back] = (1 - np.abs(xy[back][:, ::-1])) * np.where(xy[back] >= 0, 1.0, -1.0)
        return np.round(xy * 127).astype(np.int8)

    # Sulcal depth: FreeSurfer convention, positive = sulcus. Medial wall is flat.
    sulc_n = np.clip(sulc / 1.6, -1, 1)
    sulc_n[~cortex] = 0.0
    sulc_q = np.round((sulc_n * 0.5 + 0.5) * 255).astype(np.uint8)
    ao_lo, ao_hi = np.percentile(ao, [1, 99.5])
    ao_q = np.round(np.clip((ao - ao_lo) / (ao_hi - ao_lo), 0, 1) * 255).astype(np.uint8)

    act_q = np.round(np.clip(np.where(act < ACT_FLOOR, 0.0, act) / ACT_MAX, 0, 1) * 255).astype(np.uint8)

    n_verts, n_tris = len(pial), len(tris)
    parts = [
        ("index", tris.astype("<u4").tobytes()),
        ("pial", quant16(pial).tobytes()),
        ("inflated", quant16(infl).tobytes()),
        ("normal", oct8(normal).tobytes()),
        ("ao", ao_q.tobytes()),
        ("sulc", sulc_q.tobytes()),
        ("frames", np.ascontiguousarray(act_q).tobytes()),  # N x T, one row per vertex
    ]
    layout, blob, pos = {}, b"", 0
    for name, raw in parts:
        layout[name] = {"offset": pos, "length": len(raw)}
        blob += raw
        pos += len(raw)

    args.cache.mkdir(parents=True, exist_ok=True)
    (args.cache / "cortex.raw").write_bytes(blob)

    # Peak of the estimate in the hemisphere facing the viewer (left).
    lh_act = act[:offset]
    pk_v, pk_t = np.unravel_index(np.argmax(lh_act), lh_act.shape)
    peak_trace = lh_act[pk_v]

    meta = {
        "source": {
            "dataset": "MNE sample dataset (MEG/EEG auditory and visual paradigm)",
            "condition": CONDITION,
            "inverse": INVERSE,
            "method": METHOD,
            "snr": SNR,
            "software": f"MNE-Python {mne.__version__}",
            "nave": int(evoked.nave),
            "smoothSteps": SMOOTH_STEPS,
        },
        "nVerts": n_verts,
        "nTris": n_tris,
        "nFrames": int(len(frame_times)),
        "center": [round(float(c), 4) for c in center],
        "halfExtent": round(half, 4),
        "rawLayout": layout,
        "frameTimesMs": [round(float(t) * 1000, 2) for t in frame_times],
        "actMax": ACT_MAX,
        "peak": {
            "value": round(float(lh_act[pk_v, pk_t]), 2),
            "timeMs": round(float(frame_times[pk_t]) * 1000, 1),
            "position": [round(float(x), 2) for x in pial[pk_v]],
        },
        "peakTrace": [round(float(x), 2) for x in peak_trace],
        "sensorRms": {
            "unit": "fT/cm",
            "timesMs": [round(float(t) * 1000, 2) for t in rms_times[::2]],
            "values": [round(float(x), 2) for x in rms[::2]],
        },
    }
    (args.cache / "cortex.raw.json").write_text(json.dumps(meta), encoding="utf-8")

    print(f"vertices {n_verts}, triangles {n_tris}, frames {len(frame_times)}, raw {pos / 1e6:.2f} MB")
    print(f"LH peak dSPM {meta['peak']['value']} at {meta['peak']['timeMs']} ms; lh/rh peak {lh['peak']:.1f}/{rh['peak']:.1f}")
    print(f"sensor RMS peak {rms.max():.1f} fT/cm at {rms_times[rms.argmax()] * 1000:.0f} ms")
    print(f"written to {args.cache}; now run: node scripts/cortex/pack_cortex.mjs")


if __name__ == "__main__":
    main()