# Setup: dependencies, the shared helpers, non-interactive plotting.
import importlib.util
import subprocess
import sys
import warnings
from pathlib import Path

# 1. Dependencies are pinned in notebooks/requirements.txt.  Nothing is installed when the pinned
#    stack is already present (local runs, CI); a fresh Colab or Binder kernel installs it once.
#    On Colab, run from a clone of the repository so that notebooks/_shared/ is available
#    (repository URL: TODO(confirm), spec section 13 item 3).
_needed = ("mne", "scipy", "matplotlib", "pooch")
_missing = [p for p in _needed if importlib.util.find_spec(p) is None]
if _missing:
    _req = next((d / "requirements.txt" for d in (Path.cwd(), *Path.cwd().parents)
                 if (d / "requirements.txt").exists()), None)
    _cmd = [sys.executable, "-m", "pip", "install", "-q"]
    _cmd += ["-r", str(_req)] if _req else ["mne==1.10.2", "pooch>=1.8"]
    subprocess.check_call(_cmd)

# 2. Shared helpers, located relative to the working directory -- notebooks/<level>/ or notebooks/ --
#    never through an absolute path.
_shared = next((d / "_shared" for d in (Path.cwd(), *Path.cwd().parents)
                if (d / "_shared" / "helpers_l5.py").exists()), None)
if _shared is None:
    raise FileNotFoundError("start the kernel in notebooks/L5/ (or notebooks/) so that _shared/helpers_l5.py is found")
sys.path.insert(0, str(_shared))
import helpers
import helpers_l5 as L5

# 3. Plotting: Jupyter's default inline backend renders static PNGs through Agg (no windows, nothing
#    blocks); outside Jupyter the helpers select Agg.  Every MNE figure is requested with show=False
#    and each figure cell ends with plt.show().
import matplotlib.pyplot as plt
import numpy as np
import mne
import pooch

mne.set_log_level("WARNING")
pooch.get_logger().setLevel("WARNING")   # no download chatter: it would print local paths
plt.rcParams["figure.dpi"] = 72
print(f"MNE {mne.__version__}; helpers_l5 imported from notebooks/_shared")
print("mne-connectivity available:", L5.have_mne_connectivity())
print(L5.disk_line("disk at the start"))

L5.print_source_model_block()

LEMON_CHANNELS = ["Fp1", "Fp2", "F7", "F3", "Fz", "F4", "F8", "FC5", "FC1", "FC2", "FC6", "T7", "C3", "Cz",
                  "C4", "T8", "CP5", "CP1", "CP2", "CP6", "AFz", "P7", "P3", "Pz", "P4", "P8", "PO9", "O1",
                  "Oz", "O2", "PO10", "AF7", "AF3", "AF4", "AF8", "F5", "F1", "F2", "F6", "FT7", "FC3",
                  "FC4", "FT8", "C5", "C1", "C2", "C6", "TP7", "CP3", "CPz", "CP4", "TP8", "P5", "P1", "P2",
                  "P6", "PO7", "PO3", "POz", "PO4", "PO8"]

info61 = L5.sphere_info(LEMON_CHANNELS, sfreq=250.0)
models = L5.sphere_models()
for name, bem in models.items():
    print(f"  {name:11s}: {L5.SPHERE_CALLS[name]}")
    print(f"               {L5.SPHERE_NOTES[name]}")
print()

GRID_SPACING_M, GRID_MAX_FRACTION = 0.010, 0.85
rr = L5.source_grid(spacing_m=GRID_SPACING_M, max_fraction=GRID_MAX_FRACTION, min_radius_m=0.012)
nn = rr / np.linalg.norm(rr, axis=1, keepdims=True)          # radial orientation, stated and fixed
fwd1, G1 = L5.leadfield(info61, models["one-shell"], rr, fixed_normals=nn)
fwd4, G4 = L5.leadfield(info61, models["four-shell"], rr, fixed_normals=nn)
print(f"grid: {len(rr)} radial dipoles on a {GRID_SPACING_M * 1000:.0f} mm cubic lattice inside "
      f"{GRID_MAX_FRACTION * L5.HEAD_RADIUS_M * 1000:.0f} mm of the centre, i.e. from "
      f"{(1 - GRID_MAX_FRACTION) * L5.HEAD_RADIUS_M * 1000:.0f} mm below the scalp inwards")
print(f"      (the four-shell model's brain compartment ends at 0.90 R = "
      f"{0.90 * L5.HEAD_RADIUS_M * 1000:.0f} mm, so the whole grid is inside it -- a source outside the "
      f"innermost shell returns a leadfield of zeros, silently)")
print(f"leadfield G: {G1.shape[0]} electrodes x {G1.shape[1]} sources, units V per A.m (MNE's own)")
print(f"      rank {np.linalg.matrix_rank(G1)} -- at most one per electrode, and one fewer once a reference "
      f"is applied")
print(f"      condition number {np.linalg.cond(G1):.3e}: the ratio between the best- and worst-seen source "
      f"directions")

# One column (a source's map) and one row (an electrode's sensitivity), drawn side by side.
depth_mm = (L5.HEAD_RADIUS_M - np.linalg.norm(rr, axis=1)) * 1000
pick_source = int(np.argmin(np.abs(depth_mm - 20) + np.abs(rr[:, 0]) * 1000 + np.abs(rr[:, 1] + 0.05) * 1000))
pick_channel = LEMON_CHANNELS.index("Oz")

fig = plt.figure(figsize=(13.5, 4.2))
ax0 = fig.add_subplot(1, 3, 1)
col = G1[:, pick_source] * 1e-9 * 1e6                       # uV per nA.m
im, _ = mne.viz.plot_topomap(col, info61, axes=ax0, show=False, contours=6, sensors=True)
cb = fig.colorbar(im, ax=ax0, shrink=0.85); cb.set_label("uV per nA.m")
ax0.set_title(f"one COLUMN of G: the map of a source\n{depth_mm[pick_source]:.0f} mm below the scalp",
              fontsize=9)

ax1 = fig.add_subplot(1, 3, 2)
row = np.abs(G1[pick_channel]) * 1e-9 * 1e6
sl = np.abs(rr[:, 0]) < GRID_SPACING_M / 2                   # the mid-sagittal slice of the grid
sc = ax1.scatter(rr[sl, 1] * 1000, rr[sl, 2] * 1000, c=row[sl], s=70, cmap="magma")
cb = fig.colorbar(sc, ax=ax1); cb.set_label("|uV| per nA.m")
ax1.set_xlabel("y, anterior + (mm)"); ax1.set_ylabel("z, superior + (mm)")
ax1.set_title(f"one ROW of G: what {LEMON_CHANNELS[pick_channel]} is sensitive to\n(mid-sagittal slice of "
              f"the grid, radial sources)", fontsize=9)
ax1.set_aspect("equal")

ax2 = fig.add_subplot(1, 3, 3)
peak = np.abs(G1).max(axis=0) * 1e-9 * 1e6
ax2.plot(depth_mm, peak, ".", ms=4, alpha=0.5, color="C0")
ax2.set_xlabel("depth below the scalp (mm)"); ax2.set_ylabel("largest |potential| (uV per nA.m)")
ax2.set_yscale("log")
ax2.set_title("every source's peak scalp amplitude\nagainst its depth", fontsize=9)
ax2.grid(alpha=0.25)
fig.suptitle("The leadfield of a 61-electrode montage on a one-shell sphere: a column, a row, and the "
             "depth dependence", y=1.02, fontsize=10)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

SHALLOW_MM, DEEP_MM = 25.0, 60.0
shallow = peak[depth_mm < SHALLOW_MM].mean()
deep = peak[depth_mm > DEEP_MM].mean()
print(f"mean peak amplitude: {shallow:.4f} uV per nA.m for the {int((depth_mm < SHALLOW_MM).sum())} sources "
      f"within {SHALLOW_MM:.0f} mm of the scalp, {deep:.4f} for the {int((depth_mm > DEEP_MM).sum())} deeper "
      f"than {DEEP_MM:.0f} mm -- a factor of {shallow / deep:.1f}")
print("That factor is the whole reason an inverse solution needs a depth weighting, and the reason a deep")
print("claim from scalp EEG needs a constraint the sensors cannot supply (L5.5).")

CASES = {
    "temporal, tangential": dict(depth=0.167, azimuth_deg=270, elevation_deg=10, orient_deg=90),
    "temporal, radial":     dict(depth=0.167, azimuth_deg=270, elevation_deg=10, orient_deg=0),
    "vertex, radial":       dict(depth=0.167, azimuth_deg=0, elevation_deg=90, orient_deg=0),
    "deep midline, radial": dict(depth=0.5, azimuth_deg=0, elevation_deg=54, orient_deg=0),
}
WIDGET = {"temporal, tangential": (0.922, 0.376), "temporal, radial": (0.943, 0.582),
          "vertex, radial": (0.959, 0.377), "deep midline, radial": (0.967, 0.650)}

MONTAGE21 = ["Fp1", "Fpz", "Fp2", "F7", "F3", "Fz", "F4", "F8", "T7", "C3", "Cz", "C4", "T8",
             "P7", "P3", "Pz", "P4", "P8", "O1", "Oz", "O2"]
info21 = L5.sphere_info(MONTAGE21, sfreq=250.0)

print("Four shells against one, 21 electrodes (the w-dipole-to-scalp montage), 20 nA.m sources:")
print(f"{'source':>22s} {'peak ch (1 shell)':>18s} {'peak ch (4 shells)':>19s} {'r':>7s} {'MAG':>7s} "
      f"{'peak uV 1':>10s} {'peak uV 4':>10s} | {'widget r, MAG':>16s}")
shell_rows = []
for name, kw in CASES.items():
    p, m = L5.dipole_from_params(strength_nAm=20.0, **kw)
    v1 = L5.dipole_potentials(info21, models["one-shell"], p, m)
    v4 = L5.dipole_potentials(info21, models["four-shell"], p, m)
    r = float(np.corrcoef(v1, v4)[0, 1])
    mag = float(np.linalg.norm(v4) / np.linalg.norm(v1))
    pk1, pk4 = MONTAGE21[int(np.abs(v1).argmax())], MONTAGE21[int(np.abs(v4).argmax())]
    shell_rows.append((name, pk1, pk4, r, mag, float(np.abs(v1).max()), float(np.abs(v4).max())))
    print(f"{name:>22s} {pk1:>18s} {pk4:>19s} {r:7.3f} {mag:7.3f} {np.abs(v1).max():10.3f} "
          f"{np.abs(v4).max():10.3f} | {WIDGET[name][0]:7.3f} {WIDGET[name][1]:7.3f}")
print()
print("The peak channel is identical in every row, as it must be: both models are spherically symmetric about")
print("the same centre, so neither can move a map, only rescale and broaden it.  A real head has no such")
print("symmetry.  Whatever the skull does to a scalp map in a real head, this comparison CANNOT show it.")

# The widget reports slightly different numbers for the same four presets.  Find out why rather than tuning.
print("CROSS-CHECK against w-dipole-to-scalp's stored values.  Two rows agree, two do not:")
print(f"{'source':>22s} {'r here':>8s} {'r widget':>9s} {'MAG here':>9s} {'MAG widget':>11s} "
      f"{'MAG difference':>15s}")
for name, pk1, pk4, r, mag, a1, a4 in shell_rows:
    wr, wm = WIDGET[name]
    print(f"{name:>22s} {r:8.3f} {wr:9.3f} {mag:9.3f} {wm:11.3f} {mag - wm:15.3f}")
print()
print("The widget evaluates both models at the NEAREST POINT OF ITS STORED GRID, not at the exact preset")
print("position -- that is deliberate, and the integration notes record why (comparing the analytic model at")
print("the exact position against a stored map at a snapped one folds grid error into what is presented as")
print("physics).  The grid spacing is reported as 8.0 mm by the widget and 7.62 mm by the script that built")
print("it.  A depth sweep says that is enough to account for the gap:")
print()
print(f"{'depth (fraction of R)':>22s} {'source radius (mm)':>19s} {'r':>7s} {'MAG':>7s}   vertex, radial")
for d in np.arange(0.12, 0.361, 0.04):
    p, m = L5.dipole_from_params(float(d), 0, 90, orient_deg=0, strength_nAm=20.0)
    v1 = L5.dipole_potentials(info21, models["one-shell"], p, m)
    v4 = L5.dipole_potentials(info21, models["four-shell"], p, m)
    mark = "  <- the widget's 0.959 / 0.377 sits here" if 0.27 <= d <= 0.30 else ""
    print(f"{d:22.2f} {np.linalg.norm(p) * 1000:19.1f} {float(np.corrcoef(v1, v4)[0, 1]):7.3f} "
          f"{float(np.linalg.norm(v4) / np.linalg.norm(v1)):7.3f}{mark}")
print()
print("So the two pipelines are not disagreeing about physics: they are evaluating at positions about a")
print("centimetre apart.  BOTH sets of numbers are reported and neither is tuned to match the other.  A lesson")
print("that quotes one must say which, because the difference (up to 0.15 in MAG) is larger than the effect")
print("some exercises would test.  TODO(confirm) for the author: which position convention the lesson quotes.")

# The same comparison across the whole grid, so the four presets are not the whole story.
r_grid = np.array([float(np.corrcoef(G1[:, j], G4[:, j])[0, 1]) for j in range(G1.shape[1])])
mag_grid = np.linalg.norm(G4, axis=0) / np.linalg.norm(G1, axis=0)

fig, axes = plt.subplots(1, 2, figsize=(12, 4.0))
axes[0].plot(depth_mm, r_grid, ".", ms=4, alpha=0.5)
axes[0].set_xlabel("depth below the scalp (mm)")
axes[0].set_ylabel("correlation between the two models' maps (dimensionless)")
axes[0].set_title("shape: four shells against one", fontsize=9); axes[0].grid(alpha=0.25)
axes[1].plot(depth_mm, mag_grid, ".", ms=4, alpha=0.5, color="C1")
axes[1].set_xlabel("depth below the scalp (mm)")
axes[1].set_ylabel("MAG = |four-shell map| / |one-shell map| (dimensionless)")
axes[1].set_title("amplitude: what the skull costs", fontsize=9); axes[1].grid(alpha=0.25)
fig.suptitle(f"Adding a skull: every one of the {len(rr)} radial grid sources, 61 electrodes", y=1.02,
             fontsize=10)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print(f"over all {len(rr)} grid sources, 61 electrodes:")
print(f"  shape      r: median {np.median(r_grid):.3f}, range {r_grid.min():.3f}-{r_grid.max():.3f}")
print(f"  amplitude  MAG: median {np.median(mag_grid):.3f}, range {mag_grid.min():.3f}-{mag_grid.max():.3f}")
print(f"  shallow sources (< {SHALLOW_MM:.0f} mm): MAG median {np.median(mag_grid[depth_mm < SHALLOW_MM]):.3f}")
print(f"  deep sources (> {DEEP_MM:.0f} mm):    MAG median {np.median(mag_grid[depth_mm > DEEP_MM]):.3f}")
print("The skull costs the most where it is nearest: a superficial source loses the larger share of its")
print("amplitude, which flattens the depth dependence of the leadfield and is one reason a four-shell model")
print("localises differently from a one-shell one even though neither can move a peak channel.")

SUBSETS = {
    "61 (ds-lemon)": LEMON_CHANNELS,
    "32 (a 10-20 subset)": ["Fp1", "Fp2", "F7", "F3", "Fz", "F4", "F8", "FC5", "FC1", "FC2", "FC6", "T7",
                            "C3", "Cz", "C4", "T8", "CP5", "CP1", "CP2", "CP6", "P7", "P3", "Pz", "P4",
                            "P8", "O1", "Oz", "O2", "AF3", "AF4", "PO3", "PO4"],
    "21 (10-20 + Fpz, Oz)": MONTAGE21,
    "19 (classic 10-20)": [c for c in MONTAGE21 if c not in ("Fpz", "Oz")],
}
NOISE_UV, N_TRIALS, AMBIGUITY_R = 1.0, 200, 0.99
rng = np.random.default_rng(L5.SEED)


def ambiguity_mm(G, rr, threshold=AMBIGUITY_R, cell_m=GRID_SPACING_M):
    """How big the set of sources that look alike is, as an equivalent sphere radius in mm.

    For each source, count the grid points whose scalp map correlates above ``threshold`` with it, multiply
    by the volume of one grid cell, and report the radius of a sphere of that volume.  Counting rather than
    taking a maximum distance keeps the number off the lattice: a max is quantised to the grid spacing and
    would read 0 or 10 mm and nothing in between.
    """
    Gc = G - G.mean(axis=0, keepdims=True)          # average reference: the map, not the offset
    Gc = Gc / np.linalg.norm(Gc, axis=0, keepdims=True)
    counts = (np.abs(Gc.T @ Gc) >= threshold).sum(axis=0)
    radii = (3.0 * counts * (cell_m * 1000) ** 3 / (4.0 * np.pi)) ** (1.0 / 3.0)
    return float(np.mean(radii)), float(np.percentile(radii, 90))


print(f"{'montage':>22s} {'n':>4s} {'rank':>5s} {'cond(G)':>9s} {'ambiguity radius, mean (mm)':>28s} "
      f"{'90th':>6s} {'localisation error, mean (mm)':>30s} {'90th':>7s}")
subset_rows = []
for label, chans in SUBSETS.items():
    inf = L5.sphere_info(chans, sfreq=250.0)
    _, G = L5.leadfield(inf, models["four-shell"], rr, fixed_normals=nn)
    amb_med, amb_p90 = ambiguity_mm(G, rr)
    Gc = G - G.mean(axis=0, keepdims=True)
    Gc = Gc / np.linalg.norm(Gc, axis=0, keepdims=True)
    idx = rng.choice(len(rr), size=N_TRIALS, replace=False)
    errs = []
    for j in idx:
        v = G[:, j] * 50e-9 * 1e6 + rng.standard_normal(len(chans)) * NOISE_UV
        v = v - v.mean()
        best = int(np.argmax(np.abs(Gc.T @ (v / np.linalg.norm(v)))))
        errs.append(np.linalg.norm(rr[best] - rr[j]) * 1000)
    errs = np.array(errs)
    subset_rows.append((label, len(chans), int(np.linalg.matrix_rank(G)), float(np.linalg.cond(G)),
                        amb_med, amb_p90, float(np.median(errs)), float(errs.mean()),
                        float(np.percentile(errs, 90))))
    print(f"{label:>22s} {len(chans):4d} {subset_rows[-1][2]:5d} {subset_rows[-1][3]:9.2e} "
          f"{amb_med:28.1f} {amb_p90:6.1f} {subset_rows[-1][7]:30.1f} {subset_rows[-1][8]:7.1f}")
print()
print(f"The AMBIGUITY RADIUS is the size of the set of places a source could be without its scalp map changing")
print(f"by more than a correlation of {AMBIGUITY_R:g}, expressed as the radius of a sphere of the same volume: "
      f"the resolution")
print("the sensors have before any inverse method, prior or regularization is applied.  It is the honest")
print("version of 'how many electrodes do I need', because it does not depend on a choice of estimator.")
print()
print(f"Localisation error is the distance from the true grid point to the best-matching one, over "
      f"{N_TRIALS} sources drawn")
print(f"without replacement, with {NOISE_UV:g} uV of independent sensor noise and a 50 nA.m source.  The grid "
      f"is {GRID_SPACING_M * 1000:.0f} mm, so")
print("the error is quantised to the lattice and a median of zero means 'usually the right grid point', not")
print("'exact'.  This is a best case in three ways at once: the head model is exactly right, the source really")
print("is at a grid point, and the source really is a single radial dipole.")

fig, axes = plt.subplots(1, 2, figsize=(12, 4.0))
labels = [r[0] for r in subset_rows]
axes[0].bar(range(len(labels)), [r[4] for r in subset_rows], color="C0")
axes[0].set_xticks(range(len(labels))); axes[0].set_xticklabels(labels, rotation=20, ha="right", fontsize=7)
axes[0].set_ylabel(f"mean ambiguity radius at r >= {AMBIGUITY_R:g} (mm)")
axes[0].set_title("how far a source can move before its map\nstops looking the same", fontsize=9)
axes[0].grid(alpha=0.25, axis="y")
axes[1].bar(range(len(labels)), [r[7] for r in subset_rows], color="C1",
            yerr=[[0] * len(labels), [r[8] - r[7] for r in subset_rows]], capsize=3)
axes[1].set_xticks(range(len(labels))); axes[1].set_xticklabels(labels, rotation=20, ha="right", fontsize=7)
axes[1].set_ylabel("localisation error (mm)")
axes[1].set_title(f"mean and 90th percentile, {NOISE_UV:g} uV sensor noise", fontsize=9)
axes[1].grid(alpha=0.25, axis="y")
fig.suptitle("Channel count and what it buys: four-shell sphere, radial grid, best-case dipole scan",
             y=1.02, fontsize=10)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

LOC_SUBJECT = "sub-010005"          # one of the 145 subjects with a digitised localiser
print(L5.disk_line("disk before the localiser download"))
loc_paths = []
try:
    loc = L5.lemon_localizer(LOC_SUBJECT)
    loc_paths = L5.lemon_localizer_files(LOC_SUBJECT)
    if loc is None:
        raise RuntimeError("the ds-lemon localiser could not be fetched")
    dig_names = [c for c in loc["ch_pos"] if c in LEMON_CHANNELS]
    dig_pos = np.array([loc["ch_pos"][c] for c in dig_names])
finally:
    print()
    L5.delete_files(loc_paths)
    print(L5.disk_line("disk after the localiser was deleted"))

template = mne.channels.make_standard_montage("standard_1005").get_positions()["ch_pos"]
tpl_pos = np.array([template[c] for c in dig_names])
c_dig, r_dig = L5.fit_sphere(dig_pos)
c_tpl, r_tpl = L5.fit_sphere(tpl_pos)
print()
print(f"{len(dig_names)} electrodes matched by name between the localiser and the template montage")
print(f"  this subject's fitted head radius : {r_dig * 1000:6.1f} mm")
print(f"  MNE standard_1005 fitted radius   : {r_tpl * 1000:6.1f} mm")
print(f"  difference                        : {(r_dig - r_tpl) * 1000:+6.1f} mm "
      f"({100 * (r_dig / r_tpl - 1):+.1f} %)")
print("  (fiducials in the localiser file: nasion " + str(np.round(loc["nasion"] * 1000, 1)) +
      " mm, LPA " + str(np.round(loc["lpa"] * 1000, 1)) + ", RPA " + str(np.round(loc["rpa"] * 1000, 1)) +
      "; the file is in Brainstorm's SCS frame and helpers_l5 converts it to MNE head coordinates)")

u_dig = (dig_pos - c_dig) / np.linalg.norm(dig_pos - c_dig, axis=1, keepdims=True)
u_tpl = (tpl_pos - c_tpl) / np.linalg.norm(tpl_pos - c_tpl, axis=1, keepdims=True)
sep_mm = np.arccos(np.clip((u_dig * u_tpl).sum(axis=1), -1, 1)) * L5.HEAD_RADIUS_M * 1000
order = np.argsort(-sep_mm)
print()
print(f"with the scale removed (both projected onto the same {L5.HEAD_RADIUS_M * 100:.0f} cm sphere), each "
      f"electrode moves by")
print(f"  median {np.median(sep_mm):.1f} mm, 90th percentile {np.percentile(sep_mm, 90):.1f} mm, "
      f"largest {sep_mm.max():.1f} mm ({dig_names[order[0]]})")
print("  the five that move most: " + ", ".join(f"{dig_names[i]} {sep_mm[i]:.0f} mm" for i in order[:5]))

# What does using the template instead of this subject's own positions cost a localisation?
info_dig = L5.sphere_info(dig_names, u_dig * L5.HEAD_RADIUS_M, sfreq=250.0)
info_tpl = L5.sphere_info(dig_names, u_tpl * L5.HEAD_RADIUS_M, sfreq=250.0)
_, G_dig = L5.leadfield(info_dig, models["four-shell"], rr, fixed_normals=nn)
_, G_tpl = L5.leadfield(info_tpl, models["four-shell"], rr, fixed_normals=nn)

r_pos = np.array([float(np.corrcoef(G_dig[:, j], G_tpl[:, j])[0, 1]) for j in range(G_dig.shape[1])])
mag_pos = np.linalg.norm(G_tpl, axis=0) / np.linalg.norm(G_dig, axis=0)
print(f"the same source seen through the two electrode sets, over all {len(rr)} grid sources:")
print(f"  map correlation r : median {np.median(r_pos):.4f}, worst {r_pos.min():.4f}")
print(f"  amplitude ratio   : median {np.median(mag_pos):.4f}, range {mag_pos.min():.4f}-{mag_pos.max():.4f}")
print()


def scan(G_model, v):
    Gc = G_model - G_model.mean(axis=0, keepdims=True)
    Gc = Gc / np.linalg.norm(Gc, axis=0, keepdims=True)
    vv = v - v.mean()
    return int(np.argmax(np.abs(Gc.T @ (vv / np.linalg.norm(vv)))))


rng2 = np.random.default_rng(L5.SEED + 3)
idx = rng2.choice(len(rr), size=N_TRIALS, replace=False)
err_right, err_wrong = [], []
for j in idx:
    truth = G_dig[:, j] * 50e-9 * 1e6
    noise = rng2.standard_normal(len(dig_names)) * NOISE_UV
    err_right.append(np.linalg.norm(rr[scan(G_dig, truth + noise)] - rr[j]) * 1000)
    err_wrong.append(np.linalg.norm(rr[scan(G_tpl, truth + noise)] - rr[j]) * 1000)
err_right, err_wrong = np.array(err_right), np.array(err_wrong)
print(f"dipole scan over {len(idx)} sources, {NOISE_UV:g} uV sensor noise, four-shell model, "
      f"{len(dig_names)} electrodes:")
print(f"  model built on the subject's own digitised positions : median {np.median(err_right):5.1f} mm, "
      f"mean {err_right.mean():5.1f} mm, 90th pct {np.percentile(err_right, 90):5.1f} mm")
print(f"  model built on the TEMPLATE positions instead        : median {np.median(err_wrong):5.1f} mm, "
      f"mean {err_wrong.mean():5.1f} mm, 90th pct {np.percentile(err_wrong, 90):5.1f} mm")
print(f"  the electrode-position choice alone costs {err_wrong.mean() - err_right.mean():+.1f} mm at the mean "
      f"and {np.percentile(err_wrong, 90) - np.percentile(err_right, 90):+.1f} mm at the 90th percentile")
print()
print("Both rows are still best cases -- the conductor model is exactly right, the source is a single radial")
print("dipole, and it sits exactly on the search grid.  The template-anatomy error this notebook CANNOT")
print("measure sits on top of this one, not instead of it.")

fig, axes = plt.subplots(1, 3, figsize=(15, 4.2))
sc = axes[0].scatter(u_tpl[:, 0] * 1000, u_tpl[:, 1] * 1000, c=sep_mm, s=60, cmap="magma")
for i in order[:5]:
    axes[0].annotate(dig_names[i], (u_tpl[i, 0] * 1000, u_tpl[i, 1] * 1000), fontsize=7)
cb = fig.colorbar(sc, ax=axes[0]); cb.set_label("displacement (mm)")
axes[0].set_xlabel("x, right + (mm, projected)"); axes[0].set_ylabel("y, anterior + (mm, projected)")
axes[0].set_title("how far each electrode moves\n(scale removed)", fontsize=9); axes[0].set_aspect("equal")
axes[1].hist(sep_mm, bins=20, color="C0")
axes[1].set_xlabel("displacement (mm)"); axes[1].set_ylabel("number of electrodes")
axes[1].set_title(f"median {np.median(sep_mm):.1f} mm", fontsize=9)
axes[2].hist(err_right, bins=np.arange(0, 60, 4), alpha=0.6, label="own positions")
axes[2].hist(err_wrong, bins=np.arange(0, 60, 4), alpha=0.6, label="template positions")
axes[2].set_xlabel("localisation error (mm)"); axes[2].set_ylabel("number of sources")
axes[2].set_title("dipole scan error", fontsize=9); axes[2].legend(fontsize=8)
fig.suptitle(f"ds-lemon {LOC_SUBJECT}: digitised electrode positions against MNE's standard_1005 template",
             y=1.02, fontsize=10)
fig.tight_layout()
plt.show()   # render the static figure(s) of this cell inline

print("nb-5-4-forward -- L5.4 numbers (draft; TODO(confirm) at author review)")
print(f"Models: {L5.SPHERE_CALLS['one-shell']}")
print(f"        {L5.SPHERE_CALLS['four-shell']}")
print(f"        NO boundary-element model and NO template or individual anatomy -- see the licence block in "
      f"section 0.")
print(f"Grid: {len(rr)} radial dipoles, {GRID_SPACING_M * 1000:.0f} mm lattice, inside "
      f"{GRID_MAX_FRACTION * L5.HEAD_RADIUS_M * 1000:.0f} mm of the centre of a "
      f"{L5.HEAD_RADIUS_M * 100:.0f} cm sphere")
print(f"Electrodes: {len(LEMON_CHANNELS)} ds-lemon channel names at MNE standard_1005 template positions, "
      f"plus one subject's digitised set")
print(f"Data downloaded: one 3.3 kB ds-lemon localiser file, deleted.  {L5.licence_line('ds-lemon')}")
print()
print("ANSWER KEY -- ex-5-4, reworded: 'one shell versus four -- what does the skull do?'")
print("  On the 21-electrode montage, 20 nA.m sources, evaluated at the exact preset positions:")
print(f"{'':6s}{'source':>22s} {'peak channel':>14s} {'r (shape)':>10s} {'MAG (amplitude)':>16s}")
for name, pk1, pk4, r, mag, a1, a4 in shell_rows:
    print(f"{'':6s}{name:>22s} {pk1 + ' (both)':>14s} {r:10.3f} {mag:16.3f}")
print(f"  Over the whole {len(rr)}-source grid with 61 electrodes: r median {np.median(r_grid):.3f} "
      f"(range {r_grid.min():.3f}-{r_grid.max():.3f}),")
print(f"  MAG median {np.median(mag_grid):.3f} (range {mag_grid.min():.3f}-{mag_grid.max():.3f}); "
      f"shallow sources lose more than deep ones")
print(f"  ({np.median(mag_grid[depth_mm < SHALLOW_MM]):.3f} within {SHALLOW_MM:.0f} mm of the scalp against "
      f"{np.median(mag_grid[depth_mm > DEEP_MM]):.3f} deeper than {DEEP_MM:.0f} mm).")
print("  MODEL ANSWER: the skull attenuates by roughly a factor of two to three and smears the map, but on two")
print("  concentric spheres it CANNOT move the peak channel -- both models are spherically symmetric about the")
print("  same centre.  A real head is not symmetric, so this is a floor on what a realistic model would change,")
print("  not an estimate of it.  No boundary-element model ships, for the licence reason stated in section 0.")
print()
print("  WIDGET CROSS-CHECK (this is a finding, not an agreement):")
for name, pk1, pk4, r, mag, a1, a4 in shell_rows:
    wr, wm = WIDGET[name]
    print(f"      {name:>22s}: here r {r:.3f} MAG {mag:.3f};  w-dipole-to-scalp r {wr:.3f} MAG {wm:.3f}")
print("      Two of four agree to about 1 %.  The other two differ by up to 0.15 in MAG because the widget")
print("      evaluates at the nearest point of its stored grid (spacing 8.0 mm by the widget's own report,")
print("      7.62 mm by the script that built it) and this notebook evaluates the analytic model at the exact")
print("      position.  The depth sweep in section 2 reproduces the widget's vertex value about a centimetre")
print("      deeper, which is the size of that snap.  Neither number is tuned; the lesson must say which")
print("      convention it quotes.")
print()
print("ANSWER KEY -- channel count, four-shell model, best-case dipole scan:")
print(f"{'':6s}{'montage':>22s} {'n':>4s} {'rank':>5s} {'ambiguity radius (mm)':>22s} "
      f"{'90th':>6s} {'error mean (mm)':>16s} {'90th':>7s}")
for label, n, rank, cond, amb, amb90, med, mean_e, p90 in subset_rows:
    print(f"{'':6s}{label:>22s} {n:4d} {rank:5d} {amb:22.1f} {amb90:6.1f} {mean_e:16.1f} {p90:7.1f}")
print()
print("ANSWER KEY -- template electrode positions against one subject's digitised set "
      f"(ds-lemon {LOC_SUBJECT}, {len(dig_names)} electrodes):")
print(f"      head radius {r_dig * 1000:.1f} mm against the template's {r_tpl * 1000:.1f} mm "
      f"({100 * (r_dig / r_tpl - 1):+.1f} %)")
print(f"      with the scale removed, each electrode moves a median {np.median(sep_mm):.1f} mm "
      f"(90th pct {np.percentile(sep_mm, 90):.1f} mm, largest {sep_mm.max():.1f} mm at {dig_names[order[0]]})")
print(f"      map correlation between the two electrode sets: median {np.median(r_pos):.4f}, "
      f"worst {r_pos.min():.4f}")
print(f"      localisation error (mean over {N_TRIALS} sources, {NOISE_UV:g} uV noise, "
      f"{GRID_SPACING_M * 1000:.0f} mm grid): {err_right.mean():.1f} mm with the")
print(f"      subject's own positions, {err_wrong.mean():.1f} mm with the template's -- the "
      f"electrode-position choice alone costs {err_wrong.mean() - err_right.mean():+.1f} mm")
print()
print("WHAT THIS NOTEBOOK CANNOT ANSWER, AND WHY:")
print("  * a boundary-element head model of any kind (ds-fsaverage: FreeSurfer licence, settled, fails 10.7)")
print("  * a template brain against an individual MRI (ds-mne-sample: contested, a new decision for the author)")
print("  * therefore also: coregistration error between electrodes and an anatomy, cortical-surface source")
print("    spaces, and every localisation error figure that would be quoted in millimetres of cortex.")
print("  The numbers above are the electrode-geometry and conductivity halves of the same question.  They are")
print("  lower bounds on the error a template costs, not estimates of it.")
print()
print("Pitfall: pf-deep-source-claims.  Widget: w-dipole-to-scalp (mode forward).")
print(L5.disk_line("disk at the end"))