"""Window, CAR, and label-free quality-gating diagnostic for TCH EEG."""
from __future__ import annotations

import argparse
import copy
import json
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.model_selection import StratifiedKFold

from event_stft.models import ShallowConvNet
from event_stft.support_csp import SOURCE_SUBJECTS, normalize_task
from event_stft.training import resolve_device, seed_everything, stable_seed


ROOT = Path(__file__).resolve().parent
DEFAULT_TCH_CACHE = ROOT.parent / "TCH_Sequential_8ch" / "data_cache"
WINDOWS = {
    "0.00-2.00s": (0, 400),
    "0.00-1.50s": (0, 300),
    "0.25-1.75s": (50, 350),
    "0.50-1.50s": (100, 300),
    "0.75-1.75s": (150, 350),
    "1.00-2.00s": (200, 400),
}


def car(values: np.ndarray) -> np.ndarray:
    return values - values.mean(axis=1, keepdims=True)


def task_window(epochs: np.ndarray, bounds: tuple[int, int], apply_car: bool) -> np.ndarray:
    task = np.asarray(epochs[:, :, 200:600], dtype=np.float32)
    if apply_car:
        task = car(task)
    start, stop = bounds
    return normalize_task(task[:, :, start:stop])


def quality_scores(epochs: np.ndarray, apply_car: bool) -> np.ndarray:
    """Label-free artifact score: amplitude, slope, and high-frequency excess."""
    values = np.asarray(epochs[:, :, 200:600], dtype=np.float64)
    if apply_car:
        values = car(values)
    rms = np.sqrt(np.mean(values ** 2, axis=(1, 2)))
    slope = np.mean(np.abs(np.diff(values, axis=-1)), axis=(1, 2))
    spectrum = np.abs(np.fft.rfft(values, axis=-1)) ** 2
    freqs = np.fft.rfftfreq(values.shape[-1], d=1 / 200)
    high = spectrum[:, :, (freqs >= 25) & (freqs <= 40)].mean(axis=(1, 2))
    motor = spectrum[:, :, (freqs >= 8) & (freqs <= 30)].mean(axis=(1, 2))
    features = np.log(np.c_[rms, slope, high / np.maximum(motor, 1e-12)] + 1e-12)
    center = np.median(features, axis=0)
    scale = 1.4826 * np.median(np.abs(features - center), axis=0)
    z = (features - center) / np.maximum(scale, 1e-6)
    return np.mean(np.clip(z, -5, 5), axis=1).astype(np.float32)


def augment(values: torch.Tensor, generator: torch.Generator) -> torch.Tensor:
    result = values.clone()
    batch, channels, samples = result.shape
    result *= torch.empty((batch, channels, 1), device=values.device).uniform_(.95, 1.05, generator=generator)
    result += torch.empty((batch, channels, 1), device=values.device).uniform_(-.05, .05, generator=generator)
    shifts = torch.randint(-10, 11, (batch,), generator=generator, device=values.device)
    padded = torch.nn.functional.pad(result, (10, 10), mode="reflect")
    return torch.stack([padded[i, :, 10 - int(shift): 10 - int(shift) + samples] for i, shift in enumerate(shifts)])


def train(model, x: np.ndarray, y: np.ndarray, device: torch.device, seed: int, epochs: int = 25) -> None:
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-3)
    criterion = nn.CrossEntropyLoss(label_smoothing=.05)
    generator = torch.Generator(device=device.type).manual_seed(seed)
    x_t = torch.as_tensor(x, dtype=torch.float32, device=device)
    y_t = torch.as_tensor(y, dtype=torch.long, device=device)
    for _ in range(epochs):
        order = torch.randperm(len(y_t), generator=generator, device=device)
        for indices in order.split(min(32, len(y_t))):
            optimizer.zero_grad(set_to_none=True)
            loss = criterion(model(augment(x_t[indices], generator)), y_t[indices])
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.)
            optimizer.step()


@torch.no_grad()
def accuracy(model, x: np.ndarray, y: np.ndarray, device: torch.device) -> float:
    model.eval()
    prediction = model(torch.as_tensor(x, dtype=torch.float32, device=device)).argmax(dim=-1).cpu().numpy()
    return float((prediction == y).mean())


def source_data(cache: Path, bounds: tuple[int, int], apply_car: bool) -> tuple[np.ndarray, np.ndarray]:
    x_all, y_all = [], []
    for subject in SOURCE_SUBJECTS:
        x = np.load(cache / "source" / f"{subject}_x.npy", allow_pickle=False)
        y = np.load(cache / "source" / f"{subject}_y.npy", allow_pickle=False).astype(np.int64)
        keep = y != 2
        x_all.append(task_window(x[keep], bounds, apply_car))
        y_all.append(y[keep])
    return np.concatenate(x_all), np.concatenate(y_all)


def source_state(cache: Path, bounds: tuple[int, int], apply_car: bool, device: torch.device, seed: int) -> dict:
    x, y = source_data(cache, bounds, apply_car)
    model = ShallowConvNet(channels=8, samples=x.shape[-1], dropout=.45).to(device)
    train(model, x, y, device, seed, epochs=8)
    return {name: value.detach().cpu() for name, value in model.state_dict().items()}


def evaluate_condition(cache: Path, tch_cache: Path, window_name: str, bounds: tuple[int, int], apply_car: bool, drop_fraction: float, device: torch.device, seed: int) -> list[dict]:
    state = source_state(cache, bounds, apply_car, device, stable_seed(seed, "source", window_name, apply_car))
    rows = []
    for path in sorted((tch_cache / "sessions").glob("*_x.npy")):
        session = path.stem[:-2]
        epochs = np.load(path, allow_pickle=False)
        y = np.load(path.with_name(f"{session}_y.npy"), allow_pickle=False).astype(np.int64)
        x = task_window(epochs, bounds, apply_car)
        scores = quality_scores(epochs, apply_car)
        splitter = StratifiedKFold(n_splits=5, shuffle=True, random_state=stable_seed(seed, session, "fold"))
        for fold, (train_indices, test_indices) in enumerate(splitter.split(x, y), 1):
            cutoff = np.quantile(scores[train_indices], 1 - drop_fraction) if drop_fraction else np.inf
            kept_train = train_indices[scores[train_indices] <= cutoff]
            # Keep both labels represented even for very small sessions.
            if len(np.unique(y[kept_train])) < 2:
                kept_train = train_indices
            model = ShallowConvNet(channels=8, samples=x.shape[-1], dropout=.45).to(device)
            model.load_state_dict(copy.deepcopy(state))
            train(model, x[kept_train], y[kept_train], device, stable_seed(seed, window_name, apply_car, drop_fraction, session, fold))
            rows.append({"session_id": session, "fold": fold, "window": window_name, "car": apply_car, "drop_fraction": drop_fraction, "test_accuracy": accuracy(model, x[test_indices], y[test_indices], device), "train_total": len(train_indices), "train_kept": len(kept_train), "test_total": len(test_indices), "test_high_quality": int(np.sum(scores[test_indices] <= cutoff))})
    return rows


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--source-cache", type=Path, default=ROOT / "support_csp_v3_source_cache")
    parser.add_argument("--tch-cache", type=Path, default=DEFAULT_TCH_CACHE)
    parser.add_argument("--output", type=Path, default=ROOT / "tch_window_quality_runs")
    parser.add_argument("--seed", type=int, default=20260721)
    args = parser.parse_args()
    seed_everything(args.seed)
    device = resolve_device("cuda")
    args.output.mkdir(parents=True, exist_ok=True)
    rows = []
    for name, bounds in WINDOWS.items():
        print(f"[window] {name}", flush=True)
        rows.extend(evaluate_condition(args.source_cache, args.tch_cache, name, bounds, False, 0., device, args.seed))
    window_frame = pd.DataFrame(rows)
    best_window = window_frame.groupby("window").test_accuracy.mean().sort_values(ascending=False).index[0]
    bounds = WINDOWS[best_window]
    print(f"[best window] {best_window}", flush=True)
    for apply_car, drop_fraction in ((False, .10), (False, .20), (True, 0.), (True, .10), (True, .20)):
        print(f"[quality] car={apply_car} drop={drop_fraction}", flush=True)
        rows.extend(evaluate_condition(args.source_cache, args.tch_cache, best_window, bounds, apply_car, drop_fraction, device, args.seed))
    result = pd.DataFrame(rows)
    result.to_csv(args.output / "fold_results.csv", index=False)
    summary = result.groupby(["window", "car", "drop_fraction"], as_index=False).agg(mean_accuracy=("test_accuracy", "mean"), std_accuracy=("test_accuracy", "std"), sessions=("session_id", "nunique"), mean_train_kept=("train_kept", "mean"))
    summary.to_csv(args.output / "condition_summary.csv", index=False)
    (args.output / "run_config.json").write_text(json.dumps({"pipeline": "1-40 Hz zero-phase IIR cache; ShallowConvNet source pretraining per tested window; current-session independent 5-fold adaptation", "quality_gate": "label-free robust score from RMS, absolute first derivative, and 25-40/8-30 Hz power ratio; applied only to fold train trials; test remains all trials", "windows": WINDOWS, "best_window_for_quality": best_window, "seed": args.seed}, ensure_ascii=False, indent=2), encoding="utf-8")
    print("[complete]", args.output, flush=True)


if __name__ == "__main__":
    main()