"""Create a conservative ASR20 + frontal-proxy ICA cache and visual report for TCH EEG."""
from __future__ import annotations

import csv
import json
import re
import shutil
from datetime import datetime
from pathlib import Path

import mne
import matplotlib.pyplot as plt
import numpy as np
from asrpy import ASR
from scipy.signal import butter, resample_poly, sosfiltfilt

from prepare_tch_data import CHANNELS, EXCLUDED_TRIALS, LABELS, SKIPPED_SESSIONS, _event_rows, session_id


ROOT = Path(__file__).resolve().parent.parent
SOURCE = ROOT / "中榮data"
RAW_CACHE = Path(__file__).resolve().parent / "data_cache"
OUTPUT_CACHE = Path(__file__).resolve().parent / "data_cache_asr20_ica"
REPORT = ROOT / "輸出呈現用" / "20260721_中榮data_ASR20_ICA視覺清理"
COLORS = {0: "#2c7a7b", 1: "#d46a6a"}
NAMES = {0: "Left (-1)", 1: "Right (1)"}


def date_from_name(path: Path) -> str:
    matched = re.search(r"_(\d{4})_data", path.stem)
    return datetime.strptime(matched.group(1), "%m%d").replace(year=2026).strftime("%Y-%m-%d")


def eog_proxy_components(ica: mne.preprocessing.ICA, raw: mne.io.BaseRaw) -> tuple[list[int], dict[int, tuple[float, float, float]]]:
    component_scores = {index: [0.0, 0.0] for index in range(ica.n_components_)}
    for channel, position in (("Fp1", 0), ("Fp2", 1)):
        _, scores = ica.find_bads_eog(raw, ch_name=channel, threshold=.8, measure="correlation", verbose=False)
        for index, value in enumerate(scores):
            component_scores[index][position] = float(value)
    diagnostics = {index: (abs(values[0]), abs(values[1]), max(abs(values[0]), abs(values[1]))) for index, values in component_scores.items()}
    rejected = [index for index, values in sorted(diagnostics.items(), key=lambda item: item[1][2], reverse=True) if values[2] >= .8][:2]
    return rejected, diagnostics


def prepare() -> tuple[list[dict], list[dict], list[dict]]:
    source_rate, target_rate = 1000, 200
    pre_samples, post_samples = 200, 400
    OUTPUT_CACHE.mkdir(parents=True, exist_ok=True)
    sessions_dir = OUTPUT_CACHE / "sessions"; sessions_dir.mkdir(exist_ok=True)
    sos = butter(4, (1., 40.), btype="bandpass", fs=source_rate, output="sos")
    info = mne.create_info(list(CHANNELS), source_rate, "eeg")
    sessions, quality, components = [], [], []
    for path in sorted(SOURCE.glob("*.npz")):
        session = session_id(path)
        if session in SKIPPED_SESSIONS:
            continue
        raw_values, events = _event_rows(path, source_rate, target_rate, pre_samples, post_samples)
        filtered = sosfiltfilt(sos, raw_values.astype(np.float64), axis=1)
        filtered_raw = mne.io.RawArray(filtered * 1e-9, info.copy(), verbose=False)
        asr = ASR(sfreq=source_rate, cutoff=20, win_len=.5, win_overlap=.66, max_bad_chans=.1)
        _, clean_window = asr.fit(filtered_raw, return_clean_window=True)
        asr_raw = asr.transform(filtered_raw)
        ica = mne.preprocessing.ICA(n_components=len(CHANNELS), method="fastica", random_state=20260721, max_iter=512, verbose=False)
        ica.fit(asr_raw, verbose=False)
        rejected, scores = eog_proxy_components(ica, asr_raw)
        ica_raw = asr_raw.copy(); ica.apply(ica_raw, exclude=rejected, verbose=False)
        cleaned = ica_raw.get_data() * 1e9
        continuous = resample_poly(cleaned, up=target_rate, down=source_rate, axis=1).astype(np.float32)
        kept, labels, numbers = [], [], []
        for number, code, index, timestamp in events:
            start, stop = index - pre_samples, index + post_samples
            if number in EXCLUDED_TRIALS.get(session, set()) or start < 0 or stop > continuous.shape[1]:
                continue
            kept.append(continuous[:, start:stop]); labels.append(LABELS[code]); numbers.append(number)
        x = np.stack(kept).astype(np.float32); y = np.asarray(labels, dtype=np.int64)
        np.save(sessions_dir / f"{session}_x.npy", x, allow_pickle=False)
        np.save(sessions_dir / f"{session}_y.npy", y, allow_pickle=False)
        np.save(sessions_dir / f"{session}_trial_numbers.npy", np.asarray(numbers, dtype=np.int64), allow_pickle=False)
        asr_values = asr_raw.get_data() * 1e9
        quality.append({"session_id": session, "date": date_from_name(path), "asr_calibration_fraction": float(clean_window.mean()), "filtered_sd_nV": float(np.std(filtered)), "after_asr_sd_nV": float(np.std(asr_values)), "after_ica_sd_nV": float(np.std(cleaned)), "asr_correction_rms_nV": float(np.sqrt(np.mean((asr_values-filtered)**2)),), "ica_correction_rms_nV": float(np.sqrt(np.mean((cleaned-asr_values)**2)),), "ica_rejected_components": ",".join(map(str, rejected)) if rejected else "none", "trials": len(y)})
        for component, (fp1, fp2, maximum) in scores.items():
            components.append({"session_id":session, "component":component, "fp1_abs_correlation":fp1, "fp2_abs_correlation":fp2, "max_abs_correlation":maximum, "rejected":component in rejected})
        sessions.append({"session_id":session, "date":date_from_name(path), "left":int((y==0).sum()), "right":int((y==1).sum()), "trials":len(y), "ica_rejected_components":",".join(map(str,rejected)) if rejected else "none"})
        print(f"[ASR20+ICA] {session}: rejected={rejected or 'none'}, trials={len(y)}", flush=True)
    for name, rows in (("session_summary.csv", sessions), ("quality_summary.csv", quality), ("ica_components.csv", components)):
        with (OUTPUT_CACHE / name).open("w", encoding="utf-8-sig", newline="") as handle:
            writer = csv.DictWriter(handle, fieldnames=sorted({key for row in rows for key in row})); writer.writeheader(); writer.writerows(rows)
    (OUTPUT_CACHE / "config.json").write_text(json.dumps({"pipeline":"continuous 1-40 Hz fourth-order zero-phase SOS IIR -> ASR cutoff 20 -> FastICA -> reject up to two components with |correlation to Fp1 or Fp2| >= 0.8 -> resample to 200 Hz -> epoch[-1,2] s", "warning":"Fp1/Fp2 are frontal EEG proxies, not dedicated EOG. ICA rejection is intentionally conservative and requires review."}, ensure_ascii=False, indent=2), encoding="utf-8")
    return sessions, quality, components


def plot_session(session: dict, output: Path) -> str:
    time = np.arange(600) / 200 - 1
    datasets = (("IIR only", RAW_CACHE / "sessions" / f"{session['session_id']}_x.npy", RAW_CACHE / "sessions" / f"{session['session_id']}_y.npy"), ("ASR20 + ICA", OUTPUT_CACHE / "sessions" / f"{session['session_id']}_x.npy", OUTPUT_CACHE / "sessions" / f"{session['session_id']}_y.npy"))
    figure, axes = plt.subplots(2, 2, figsize=(13, 6.5), sharex=True, constrained_layout=True)
    for column, (title, x_path, y_path) in enumerate(datasets):
        values = np.load(x_path, allow_pickle=False); labels = np.load(y_path, allow_pickle=False)
        corrected = values - values[:, :, 100:200].mean(axis=-1, keepdims=True)
        for row, (channel, channel_index) in enumerate((("C3", 3), ("C4", 4))):
            axis = axes[row, column]
            for label in (0, 1):
                trials = corrected[labels == label, channel_index]
                mean = trials.mean(axis=0); sem = trials.std(axis=0, ddof=1) / np.sqrt(len(trials))
                axis.plot(time, mean, color=COLORS[label], lw=1.6, label=f"{NAMES[label]} (n={len(trials)})")
                axis.fill_between(time, mean-sem, mean+sem, color=COLORS[label], alpha=.16, linewidth=0)
            axis.axvspan(-1, 0, color="#e8edf2", alpha=.6); axis.axvspan(0, 2, color="#edf6f0", alpha=.55); axis.axvline(0, color="#526273", lw=1, ls="--")
            axis.set_ylabel(f"{channel}\nbaseline-corrected nV"); axis.grid(axis="y", alpha=.2); axis.spines[["top","right"]].set_visible(False)
            if row == 0: axis.set_title(title)
            if row == 0 and column == 1: axis.legend(frameon=False, ncol=2, loc="upper right")
            if row == 1: axis.set_xlabel("Time relative to event (s)")
    figure.suptitle(f"{session['session_id']} | {session['date']} | ICA excluded: {session['ica_rejected_components']}", fontsize=14)
    filename = f"{session['session_id']}_{session['date'].replace('-', '')}_IIR_vs_ASR20ICA.png"
    figure.savefig(output / "img" / filename, dpi=160, bbox_inches="tight"); plt.close(figure)
    return f"img/{filename}"


def build_report(sessions: list[dict], quality: list[dict]) -> None:
    (REPORT / "img").mkdir(parents=True, exist_ok=True); (REPORT / "code").mkdir(exist_ok=True)
    sessions = sorted(sessions, key=lambda row:(row["date"], row["session_id"]))
    for row in sessions: row["image"] = plot_session(row, REPORT)
    for source in [Path(__file__), OUTPUT_CACHE / "quality_summary.csv", OUTPUT_CACHE / "ica_components.csv", OUTPUT_CACHE / "session_summary.csv", OUTPUT_CACHE / "config.json"]:
        destination = REPORT / "code" / source.name if source.suffix == ".py" else REPORT / source.name
        shutil.copy2(source, destination)
    rows = "".join(f"<tr><td>{r['session_id']}</td><td>{r['date']}</td><td>{r['ica_rejected_components']}</td><td>{r['asr_calibration_fraction']*100:.1f}%</td><td>{r['asr_correction_rms_nV']:.2f} nV</td><td>{r['ica_correction_rms_nV']:.2f} nV</td></tr>" for r in sorted(quality, key=lambda row:(row["date"],row["session_id"])))
    cards = "".join(f"<section><h2>{s['session_id']} | {s['date']}</h2><p>有效 trials：Left {s['left']}、Right {s['right']}；ICA excluded: <code>{s['ica_rejected_components']}</code></p><img src=\"{s['image']}\" alt=\"{s['session_id']} IIR vs ASR ICA\"></section>" for s in sessions)
    page = f"""<!doctype html><html lang=\"zh-Hant\"><head><meta charset=\"utf-8\"><meta name=\"viewport\" content=\"width=device-width,initial-scale=1\"><title>中榮資料 ASR20 + ICA 視覺清理</title><style>
body{{font-family:Arial,'Microsoft JhengHei',sans-serif;background:#f4f6f8;color:#18232e;margin:0;line-height:1.6}}main{{max-width:1360px;margin:auto;padding:32px}}section{{background:#fff;border:1px solid #d9e0e6;border-radius:8px;padding:20px;margin:18px 0}}h1,h2{{margin:0 0 6px}}p{{margin:0 0 12px;color:#536274}}img{{width:100%;display:block}}table{{border-collapse:collapse;width:100%}}th,td{{border:1px solid #d9e0e6;padding:8px;text-align:left}}th{{background:#edf2f5}}code{{background:#edf1f4;padding:2px 4px}}.note{{background:#fff7dd;border-left:4px solid #d69e2e;padding:12px 14px}}
</style></head><body><main><h1>中榮資料：ASR20 + ICA 視覺清理</h1><p>左欄為原始 <code>1–40 Hz IIR</code>，右欄為 <code>ASR cutoff 20 + ICA</code>。每一條線為該日有效 trial 的平均、陰影為 SEM；均做 <code>-0.5–0 s</code> baseline correction。</p><div class=\"note\"><strong>保守規則：</strong>ICA 於每段獨立 fit，只移除與 Fp1 或 Fp2 的 absolute correlation ≥ 0.8 的前兩個 component。Fp1/Fp2 是 frontal EEG proxy 而非專用 EOG，這不是生理上確定的 artifact 標記；目標是便於視覺檢查，而不是以「好看」為標準強行刪除腦訊號。</div><section><h2>每段處理摘要</h2><table><thead><tr><th>Session</th><th>Date</th><th>ICA excluded</th><th>ASR calibration</th><th>ASR correction RMS</th><th>ICA correction RMS</th></tr></thead><tbody>{rows}</tbody></table></section>{cards}<section><h2>檔案</h2><p><a href=\"quality_summary.csv\">quality summary</a> · <a href=\"ica_components.csv\">ICA component diagnostics</a> · <a href=\"config.json\">pipeline config</a> · <a href=\"code/prepare_tch_asr20_ica.py\">code</a></p></section></main></body></html>"""
    (REPORT / "index.html").write_text(page, encoding="utf-8")


if __name__ == "__main__":
    sessions, quality, _ = prepare()
    build_report(sessions, quality)
    print(f"Wrote {REPORT}")