brainEncoding/visualization/plot_predictivity_brainmaps.py

#!/usr/bin/env python3
"""Plot Huth-style pycortex flatmap predictivity heatmaps from LitCoder outputs."""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
import sys
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import nibabel.freesurfer.io as fsio
import numpy as np
from PIL import Image, ImageDraw
from scipy import sparse
from scipy.spatial import cKDTree


BRAIN_ENCODING = Path(__file__).resolve().parents[1]
DEFAULT_SCRATCH_ROOT = Path(os.environ.get("BRAINENCODING_SCRATCH_ROOT", "/storage/scratch1/8/whuang409"))
DEFAULT_RESULTS_ROOT = BRAIN_ENCODING / "results/pythia-2.8b-deduped_layer13_lookback128"
DEFAULT_ROI_DIR = BRAIN_ENCODING / "masks/requested_roi_groups_20484"
DEFAULT_PYCORTEX_DB = DEFAULT_SCRATCH_ROOT / "visionLM_data/ds003020/derivatives/pycortex-db"
DEFAULT_FREESURFER_SUBJECTS = DEFAULT_SCRATCH_ROOT / "visionLM_data/ds003020/derivatives/freesurfer_subjdir"
DEFAULT_FSAVERAGE5 = Path("/storage/project/r-aivanova7-0/shared/.env/freesurfer/subjects/fsaverage5")
HUTH_SCRIPTS = BRAIN_ENCODING.parents[1] / "huth2016" / "scripts"
VISUALIZATION_SCRIPTS = Path(__file__).resolve().parent
if str(VISUALIZATION_SCRIPTS) not in sys.path:
    sys.path.insert(0, str(VISUALIZATION_SCRIPTS))
if str(HUTH_SCRIPTS) not in sys.path:
    sys.path.append(str(HUTH_SCRIPTS))

os.environ.setdefault("MPLCONFIGDIR", str(DEFAULT_SCRATCH_ROOT / "matplotlib"))
from flatmap_plotting import (  # noqa: E402
    colorize_positive,
    crop_box_for_half,
    half_mask,
    load_font,
    render_positive_colorbar,
    split_rows,
    text_size,
)
from litcoder_style_plotting import save_litcoder_style_map  # noqa: E402
from roi_overlay import draw_roi_legend, overlay_roi_borders, roi_group_metadata  # noqa: E402


N_HEMI = 10242
N_VERTICES = 20484
HEMIS = ("lh", "rh")
HEMI_LABEL = {"lh": "Left", "rh": "Right"}
HEMI_HALF = {"lh": "left", "rh": "right"}


@dataclass(frozen=True)
class SubjectFlatmap:
    subject: str
    pixmap: sparse.csr_matrix
    mask: np.ndarray
    split: int
    nearest_fsaverage5: dict[str, np.ndarray]
    native_vertices: dict[str, int]
    projection: str


def safe_name(value: str) -> str:
    return re.sub(r"[^A-Za-z0-9_.-]+", "_", value).strip("_")


def unit_vectors(vertices: np.ndarray) -> np.ndarray:
    norms = np.linalg.norm(vertices, axis=1, keepdims=True)
    if np.any(norms == 0):
        raise RuntimeError("sphere surface contains zero-length vertices")
    return vertices / norms


def load_sparse_cache(path: Path) -> sparse.csr_matrix:
    if not path.exists():
        raise FileNotFoundError(f"Missing pycortex flatmap cache: {path}")
    with np.load(path) as npz:
        return sparse.csr_matrix((npz["data"], npz["indices"], npz["indptr"]), shape=tuple(npz["shape"]))


def fsaverage5_to_native_nearest(
    subject: str,
    hemi: str,
    freesurfer_subjects: Path,
    fsaverage5: Path,
) -> np.ndarray:
    source, _ = fsio.read_geometry(str(fsaverage5 / "surf" / f"{hemi}.sphere.reg"))
    target, _ = fsio.read_geometry(str(freesurfer_subjects / subject / "surf" / f"{hemi}.sphere.reg"))
    tree = cKDTree(unit_vectors(source))
    _dist, nearest = tree.query(unit_vectors(target), k=1)
    return nearest.astype(np.int32)


def load_subject_flatmap(
    subject: str,
    pycortex_db: Path,
    freesurfer_subjects: Path,
    fsaverage5: Path,
    height: int,
) -> SubjectFlatmap:
    cache = pycortex_db / subject / "cache"
    mask_path = cache / f"flatmask_{height}.npz"
    if not mask_path.exists():
        raise FileNotFoundError(f"Missing pycortex flatmask cache: {mask_path}")
    with np.load(mask_path) as npz:
        mask = np.asarray(npz["mask"], dtype=bool)

    pixmap = load_sparse_cache(cache / f"flatverts_{height}.npz")
    native_vertices = {}
    nearest = {}
    for hemi in HEMIS:
        nearest[hemi] = fsaverage5_to_native_nearest(subject, hemi, freesurfer_subjects, fsaverage5)
        native_vertices[hemi] = int(nearest[hemi].shape[0])

    expected_vertices = native_vertices["lh"] + native_vertices["rh"]
    if pixmap.shape[1] != expected_vertices:
        raise ValueError(
            f"{subject}: flatverts has {pixmap.shape[1]} vertex columns, "
            f"but subject surfaces have {expected_vertices} vertices."
        )
    if pixmap.shape[0] != int(mask.sum()):
        raise ValueError(f"{subject}: flatverts rows do not match flatmask pixels.")

    split, _ = split_rows(mask)
    return SubjectFlatmap(
        subject=subject,
        pixmap=pixmap,
        mask=mask,
        split=split,
        nearest_fsaverage5=nearest,
        native_vertices=native_vertices,
        projection="fsaverage5 sphere.reg -> subject native -> pycortex flatverts",
    )


def load_roi_metadata(roi_dir: Path) -> dict[str, dict[str, Any]]:
    manifest = roi_dir / "requested_roi_groups_manifest.tsv"
    metadata: dict[str, dict[str, Any]] = {}
    if not manifest.exists():
        return metadata
    with manifest.open() as f:
        for row in csv.DictReader(f, delimiter="\t"):
            key = row["roi_key"]
            if key in metadata:
                continue
            color = tuple(int(v) for v in row["color_rgb"].split(","))
            metadata[key] = {
                "label": row["label"],
                "description": row["description"],
                "color": color,
                "note": row["note"],
            }
    return metadata


def load_subject_roi_masks(roi_dir: Path, subject: str) -> dict[str, np.ndarray]:
    path = roi_dir / f"{subject}_requested_roi_groups_fsaverage5_20484.npz"
    if not path.exists():
        raise FileNotFoundError(f"Missing requested ROI group bundle for {subject}: {path}")
    with np.load(path) as npz:
        masks = {name: np.asarray(npz[name], dtype=bool) for name in npz.files if name != "roi_names"}
    for name, mask in masks.items():
        if mask.shape != (N_VERTICES,):
            raise ValueError(f"{path}:{name} has shape {mask.shape}, expected {(N_VERTICES,)}")
    return masks


def find_predictivity_files(results_root: Path, subjects: list[str] | None) -> list[Path]:
    files = sorted((results_root / "predictivity").glob("*_predictivity.npz"))
    if subjects:
        wanted = {subject.upper() for subject in subjects}
        files = [path for path in files if path.name.split("_", 1)[0].upper() in wanted]
    return files


def load_predictivity(path: Path) -> tuple[str, np.ndarray]:
    with np.load(path, allow_pickle=False) as npz:
        if "correlations" not in npz:
            raise KeyError(f"{path} does not contain a 'correlations' array")
        corr = np.asarray(npz["correlations"], dtype=np.float32)
        subject = str(npz["subject"]) if "subject" in npz else path.name.split("_", 1)[0]
    if corr.shape != (N_VERTICES,):
        raise ValueError(f"{path}: correlations shape {corr.shape}, expected {(N_VERTICES,)}")
    return subject, corr


def positive_limit(correlations_by_subject: dict[str, np.ndarray], percentile: float) -> float:
    values = []
    for corr in correlations_by_subject.values():
        pos = corr[np.isfinite(corr) & (corr > 0)]
        if pos.size:
            values.append(pos)
    if not values:
        return 1.0
    limit = float(np.percentile(np.concatenate(values), percentile))
    return limit if np.isfinite(limit) and limit > 0 else 1.0


def hemi_values(values: np.ndarray, hemi: str) -> np.ndarray:
    if hemi == "lh":
        return values[:N_HEMI]
    if hemi == "rh":
        return values[N_HEMI:]
    raise ValueError(f"Unknown hemisphere: {hemi}")


def project_fsaverage5_to_native(values: np.ndarray, flatmap: SubjectFlatmap) -> np.ndarray:
    left = hemi_values(values, "lh")[flatmap.nearest_fsaverage5["lh"]]
    right = hemi_values(values, "rh")[flatmap.nearest_fsaverage5["rh"]]
    return np.concatenate([left, right]).astype(np.float32, copy=False)


def native_to_flatmap(native_values: np.ndarray, flatmap: SubjectFlatmap) -> np.ndarray:
    flat_values = np.asarray(flatmap.pixmap @ native_values, dtype=np.float32)
    image = np.full(flatmap.mask.shape, np.nan, dtype=np.float32)
    image[flatmap.mask] = flat_values
    return image


def flatmap_roi_masks(roi_masks: dict[str, np.ndarray], flatmap: SubjectFlatmap) -> dict[str, np.ndarray]:
    out: dict[str, np.ndarray] = {}
    for key, mask in roi_masks.items():
        native = project_fsaverage5_to_native(mask.astype(np.float32), flatmap)
        flat_values = np.asarray(flatmap.pixmap @ native, dtype=np.float32)
        image = np.zeros(flatmap.mask.shape, dtype=bool)
        image[flatmap.mask] = flat_values > 0.5
        out[key] = image
    return out


def flatmap_roi_groups(
    flat_roi_masks: dict[str, np.ndarray],
    roi_metadata: dict[str, dict[str, Any]],
    panel_mask: np.ndarray,
) -> list[dict[str, object]]:
    groups: list[dict[str, object]] = []
    for key, mask in flat_roi_masks.items():
        if not np.any(mask & panel_mask):
            continue
        meta = roi_metadata.get(key, {"label": key, "description": "", "color": (80, 80, 80)})
        groups.append(
            {
                "label": str(meta["label"]),
                "description": str(meta.get("description", "")),
                "color": tuple(int(value) for value in meta["color"]),
                "rois": [key],
                "mask": mask,
            }
        )
    return groups


def composite_rgba_on_white(image: Image.Image) -> Image.Image:
    rgba = image.convert("RGBA")
    background = Image.new("RGBA", rgba.size, (255, 255, 255, 255))
    background.alpha_composite(rgba)
    return background.convert("RGB")


def estimate_roi_legend_height(
    draw: ImageDraw.ImageDraw,
    roi_groups: list[dict[str, object]],
    max_width: int,
    font,
    title_font,
) -> int:
    if not roi_groups:
        return 0
    title = "ROI borders:"
    title_w, _ = text_size(draw, title, title_font)
    cursor_x = title_w + 18
    cursor_y = 1
    line_h = 24
    for group in roi_groups:
        label = f"{group['label']}: {group['description']}"
        label_w, _ = text_size(draw, label, font)
        item_w = 34 + label_w + 22
        if cursor_x + item_w > max_width and cursor_x > title_w + 18:
            cursor_x = 0
            cursor_y += line_h
        cursor_x += item_w
    return cursor_y + line_h + 6


def write_flatmap_heatmap(
    *,
    flat_values: np.ndarray,
    output_path: Path,
    subject: str,
    hemi: str,
    flatmap: SubjectFlatmap,
    limit: float,
    percentile: float,
    roi_groups: list[dict[str, object]],
    cmap: str,
) -> dict[str, object]:
    half = HEMI_HALF[hemi]
    panel_mask = half_mask(flatmap.mask, flatmap.split, half)
    crop = crop_box_for_half(flatmap.mask, flatmap.split, half, margin=0)
    crop_mask = panel_mask[crop]

    panel = colorize_positive(flat_values[crop], crop_mask, limit, cmap_name=cmap)
    panel = composite_rgba_on_white(overlay_roi_borders(panel, roi_groups, crop, crop_mask, width=1, smooth_iterations=1, dot_step=7))

    pad = 18
    title_h = 44
    label_h = 34
    colorbar_h = 52
    title_font = load_font(24, bold=True)
    label_font = load_font(18, bold=True)
    small_font = load_font(15)
    roi_title_font = load_font(15, bold=True)

    probe = Image.new("RGB", (1, 1), "white")
    probe_draw = ImageDraw.Draw(probe)
    roi_legend_h = estimate_roi_legend_height(probe_draw, roi_groups, panel.width, small_font, roi_title_font)
    legend_h = colorbar_h + (roi_legend_h + 10 if roi_groups else 0)
    canvas_w = pad * 2 + panel.width
    canvas_h = pad * 2 + title_h + label_h + panel.height + legend_h

    canvas = Image.new("RGB", (canvas_w, canvas_h), "white")
    draw = ImageDraw.Draw(canvas)

    title = f"{subject} {HEMI_LABEL[hemi]} hemisphere voxel predictivity"
    draw.text((pad, pad), title, fill=(20, 20, 20), font=title_font)
    label = "Pythia-2.8B layer 13, lookback 128"
    label_w, _ = text_size(draw, label, label_font)
    y0 = pad + title_h
    draw.text((pad + max((panel.width - label_w) // 2, 0), y0), label, fill=(20, 20, 20), font=label_font)
    canvas.paste(panel, (pad, y0 + label_h))

    legend_y = y0 + label_h + panel.height + 14
    bar_w = min(720, canvas_w - pad * 2)
    bar = render_positive_colorbar(bar_w, 18, cmap_name=cmap)
    canvas.paste(bar.convert("RGB"), (pad, legend_y))
    draw.rectangle((pad, legend_y, pad + bar_w - 1, legend_y + 17), outline=(40, 40, 40), width=1)
    draw.text((pad, legend_y + 22), "0", fill=(20, 20, 20), font=small_font)
    right = f"+{limit:.3g}"
    right_w, _ = text_size(draw, right, small_font)
    draw.text((pad + bar_w - right_w, legend_y + 22), right, fill=(20, 20, 20), font=small_font)
    mid = f"prediction corr r, negatives clipped to 0, shared {percentile:g}th percentile scale"
    mid_w, _ = text_size(draw, mid, small_font)
    draw.text((pad + max((bar_w - mid_w) // 2, 0), legend_y + 22), mid, fill=(20, 20, 20), font=small_font)
    if roi_groups:
        draw_roi_legend(draw, roi_groups, pad, legend_y + 54, canvas_w - pad * 2, small_font, roi_title_font)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    canvas.save(output_path)
    canvas.save(output_path.with_suffix(".pdf"), "PDF")
    return {
        "subject": subject,
        "hemisphere": hemi,
        "path": str(output_path),
        "pdf_path": str(output_path.with_suffix(".pdf")),
        "projection": flatmap.projection,
        "pycortex_height": int(flatmap.mask.shape[1]),
        "vmax": float(limit),
        "color_percentile": float(percentile),
        "cmap": cmap,
        "roi_border_groups": roi_group_metadata(roi_groups),
        "image_size": list(canvas.size),
        "crop_rows": [int(crop[0].start), int(crop[0].stop)],
        "crop_cols": [int(crop[1].start), int(crop[1].stop)],
    }


def plot_subject_hemi(
    *,
    subject: str,
    correlations: np.ndarray,
    flat_roi_masks: dict[str, np.ndarray],
    roi_metadata: dict[str, dict[str, Any]],
    flatmap: SubjectFlatmap,
    hemi: str,
    output_path: Path,
    vmax: float,
    percentile: float,
    cmap: str,
) -> dict[str, Any]:
    native_values = project_fsaverage5_to_native(np.clip(correlations, 0.0, None), flatmap)
    flat_values = native_to_flatmap(native_values, flatmap)
    panel_mask = half_mask(flatmap.mask, flatmap.split, HEMI_HALF[hemi])
    roi_groups = flatmap_roi_groups(flat_roi_masks, roi_metadata, panel_mask)
    record = write_flatmap_heatmap(
        flat_values=flat_values,
        output_path=output_path,
        subject=subject,
        hemi=hemi,
        flatmap=flatmap,
        limit=vmax,
        percentile=percentile,
        roi_groups=roi_groups,
        cmap=cmap,
    )
    finite = correlations[np.isfinite(correlations)]
    record.update(
        {
            "finite_vertices": int(finite.size),
            "mean_correlation": float(np.nanmean(correlations)),
            "median_correlation": float(np.nanmedian(correlations)),
            "max_correlation": float(np.nanmax(correlations)),
        }
    )
    return record


def roi_summary_rows(subject: str, correlations: np.ndarray, roi_masks: dict[str, np.ndarray], roi_metadata: dict[str, dict[str, Any]]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for key, mask in roi_masks.items():
        meta = roi_metadata.get(key, {"label": key, "description": ""})
        for hemi in ("both", "lh", "rh"):
            if hemi == "both":
                hemi_mask = mask
                values = correlations[hemi_mask]
            else:
                hemi_mask = hemi_values(mask, hemi)
                values = hemi_values(correlations, hemi)[hemi_mask]
            finite = values[np.isfinite(values)]
            rows.append(
                {
                    "subject": subject,
                    "roi_key": key,
                    "label": meta["label"],
                    "description": meta.get("description", ""),
                    "hemisphere": hemi,
                    "vertices": int(np.sum(hemi_mask)),
                    "finite_vertices": int(finite.size),
                    "mean_r": float(np.mean(finite)) if finite.size else float("nan"),
                    "median_r": float(np.median(finite)) if finite.size else float("nan"),
                    "max_r": float(np.max(finite)) if finite.size else float("nan"),
                }
            )
    return rows


def write_tsv(path: Path, rows: list[dict[str, Any]]) -> None:
    if not rows:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(rows[0].keys()), delimiter="\t")
        writer.writeheader()
        writer.writerows(rows)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--results-root", type=Path, default=DEFAULT_RESULTS_ROOT)
    parser.add_argument("--roi-dir", type=Path, default=DEFAULT_ROI_DIR)
    parser.add_argument("--pycortex-db", type=Path, default=DEFAULT_PYCORTEX_DB)
    parser.add_argument("--freesurfer-subjects", type=Path, default=DEFAULT_FREESURFER_SUBJECTS)
    parser.add_argument("--fsaverage5", type=Path, default=DEFAULT_FSAVERAGE5)
    parser.add_argument("--subjects", nargs="+", default=None)
    parser.add_argument("--hemispheres", nargs="+", choices=HEMIS, default=list(HEMIS))
    parser.add_argument("--color-percentile", type=float, default=99.0)
    parser.add_argument("--flatmap-height", type=int, default=1024)
    parser.add_argument("--cmap", default="coolwarm", help="Matplotlib colormap for positive predictivity values.")
    parser.add_argument("--litcoder-style", action=argparse.BooleanOptionalAction, default=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    subjects = [subject.upper() for subject in args.subjects] if args.subjects else None
    pred_files = find_predictivity_files(args.results_root, subjects)
    if not pred_files:
        raise FileNotFoundError(f"No predictivity files found under {args.results_root / 'predictivity'}")

    correlations_by_subject: dict[str, np.ndarray] = {}
    for path in pred_files:
        subject, corr = load_predictivity(path)
        correlations_by_subject[subject] = corr

    vmax = positive_limit(correlations_by_subject, args.color_percentile)
    roi_metadata = load_roi_metadata(args.roi_dir)
    output_dir = args.results_root / "brainmaps"

    records: list[dict[str, Any]] = []
    roi_rows: list[dict[str, Any]] = []
    for subject, corr in sorted(correlations_by_subject.items()):
        flatmap = load_subject_flatmap(subject, args.pycortex_db, args.freesurfer_subjects, args.fsaverage5, args.flatmap_height)
        roi_masks = load_subject_roi_masks(args.roi_dir, subject)
        flat_roi_masks = flatmap_roi_masks(roi_masks, flatmap)
        roi_rows.extend(roi_summary_rows(subject, corr, roi_masks, roi_metadata))
        for hemi in args.hemispheres:
            path = output_dir / f"{subject}_{hemi}_predictivity_roi_overlay.png"
            records.append(
                plot_subject_hemi(
                    subject=subject,
                    correlations=corr,
                    flat_roi_masks=flat_roi_masks,
                    roi_metadata=roi_metadata,
                    flatmap=flatmap,
                    hemi=hemi,
                    output_path=path,
                    vmax=vmax,
                    percentile=args.color_percentile,
                    cmap=args.cmap,
                )
            )
            print(f"Wrote {path}")
        if args.litcoder_style:
            path = output_dir / f"{subject}_predictivity_litcoder_style.png"
            clipped = np.clip(corr, 0.0, None)
            save_litcoder_style_map(
                clipped,
                path,
                title=f"{subject} voxel predictivity",
                fsaverage5=args.fsaverage5,
                positive=True,
                vmax=vmax,
                cmap=args.cmap,
            )
            finite = corr[np.isfinite(corr)]
            records.append(
                {
                    "subject": subject,
                    "style": "litcoder_surface",
                    "hemisphere": "both",
                    "path": str(path),
                    "finite_vertices": int(finite.size),
                    "mean_correlation": float(np.nanmean(corr)),
                    "median_correlation": float(np.nanmedian(corr)),
                    "max_correlation": float(np.nanmax(corr)),
                    "vmax": float(vmax),
                    "color_percentile": float(args.color_percentile),
                    "cmap": args.cmap,
                }
            )
            print(f"Wrote {path}")

    write_tsv(args.results_root / "tables/roi_predictivity_summary.tsv", roi_rows)
    manifest = {
        "created_utc": datetime.now(timezone.utc).isoformat(),
        "results_root": str(args.results_root),
        "roi_dir": str(args.roi_dir),
        "pycortex_db": str(args.pycortex_db),
        "freesurfer_subjects": str(args.freesurfer_subjects),
        "fsaverage5": str(args.fsaverage5),
        "flatmap_height": int(args.flatmap_height),
        "flatmap_projection": "fsaverage5 sphere.reg -> subject native -> pycortex flatverts cache",
        "huth_scripts": str(HUTH_SCRIPTS),
        "color_percentile": float(args.color_percentile),
        "color_limit": float(vmax),
        "cmap": args.cmap,
        "litcoder_style": bool(args.litcoder_style),
        "records": records,
    }
    manifest_path = output_dir / "brainmap_manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    print(f"Wrote {args.results_root / 'tables/roi_predictivity_summary.tsv'}")
    print(f"Wrote {manifest_path}")


if __name__ == "__main__":
    main()