"""Strict nested-CV selection of time window, reference, and quality gate for TCH EEG."""
from __future__ import annotations

import argparse
import copy
import json
from dataclasses import dataclass, asdict
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),
}
REFERENCES = ("raw", "car", "median")


@dataclass(frozen=True)
class Candidate:
    window: str
    reference: str
    quality: str = "none"
    drop_fraction: float = 0.0

    @property
    def key(self) -> tuple[str, str]:
        return self.window, self.reference

    @property
    def name(self) -> str:
        return f"{self.window} | {self.reference} | {self.quality} | drop={self.drop_fraction:g}"


def candidates() -> list[Candidate]:
    result = [Candidate(window, reference) for window in WINDOWS for reference in REFERENCES]
    # Targeted quality-gate candidates from the first diagnostic, still selected inside nested CV.
    result += [Candidate("0.00-1.50s", "raw", "combined", drop) for drop in (.10, .20, .30)]
    result += [Candidate("0.00-1.50s", "raw", kind, .20) for kind in ("rms", "slope", "high_ratio")]
    result += [Candidate("0.00-1.50s", reference, "combined", .20) for reference in ("car", "median")]
    return result


def rereference(values: np.ndarray, reference: str) -> np.ndarray:
    if reference == "raw":
        return values
    if reference == "car":
        return values - values.mean(axis=1, keepdims=True)
    if reference == "median":
        return values - np.median(values, axis=1, keepdims=True)
    raise ValueError(reference)


def window_data(epochs: np.ndarray, candidate: Candidate) -> np.ndarray:
    values = rereference(np.asarray(epochs[:, :, 200:600], dtype=np.float32), candidate.reference)
    start, stop = WINDOWS[candidate.window]
    return normalize_task(values[:, :, start:stop])


def scores(epochs: np.ndarray, reference: str) -> dict[str, np.ndarray]:
    values = rereference(np.asarray(epochs[:, :, 200:600], dtype=np.float64), reference)
    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
    frequency = np.fft.rfftfreq(values.shape[-1], 1 / 200)
    high = spectrum[:, :, (frequency >= 25) & (frequency <= 40)].mean(axis=(1, 2))
    motor = spectrum[:, :, (frequency >= 8) & (frequency <= 30)].mean(axis=(1, 2))
    raw = {"rms": np.log(rms + 1e-12), "slope": np.log(slope + 1e-12), "high_ratio": np.log(high / np.maximum(motor, 1e-12))}
    normalized = {}
    for name, values_ in raw.items():
        center = np.median(values_); scale = 1.4826 * np.median(np.abs(values_ - center))
        normalized[name] = np.clip((values_ - center) / max(scale, 1e-6), -5, 5)
    normalized["combined"] = np.mean(np.c_[normalized["rms"], normalized["slope"], normalized["high_ratio"]], axis=1)
    normalized["none"] = np.zeros(len(values), dtype=np.float64)
    return normalized


def augmented(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 fit(model: nn.Module, x: np.ndarray, y: np.ndarray, device: torch.device, seed: int, epochs: int) -> None:
    model.train(); optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-3); loss_fn = 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):
        for index in torch.randperm(len(y_t), generator=generator, device=device).split(min(32, len(y_t))):
            optimizer.zero_grad(set_to_none=True); loss = loss_fn(model(augmented(x_t[index], generator)), y_t[index]); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); optimizer.step()


@torch.no_grad()
def predict(model: nn.Module, x: np.ndarray, device: torch.device) -> np.ndarray:
    model.eval(); return model(torch.as_tensor(x, dtype=torch.float32, device=device)).argmax(-1).cpu().numpy()


def source_epochs(cache: Path) -> tuple[np.ndarray, np.ndarray]:
    arrays, labels = [], []
    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; arrays.append(x[keep]); labels.append(y[keep])
    return np.concatenate(arrays), np.concatenate(labels)


def source_states(cache: Path, candidates_: list[Candidate], device: torch.device, seed: int) -> dict[tuple[str, str], dict]:
    source_x, source_y = source_epochs(cache); states = {}
    for window, reference in sorted({candidate.key for candidate in candidates_}):
        candidate = Candidate(window, reference)
        x = window_data(source_x, candidate)
        model = ShallowConvNet(channels=8, samples=x.shape[-1], dropout=.45).to(device)
        fit(model, x, source_y, device, stable_seed(seed, "source", window, reference), epochs=8)
        states[(window, reference)] = {key: value.detach().cpu() for key, value in model.state_dict().items()}
        print(f"[source] {window} {reference}", flush=True)
    return states


def selected_train_indices(indices: np.ndarray, score: np.ndarray, candidate: Candidate, labels: np.ndarray) -> tuple[np.ndarray, float]:
    if candidate.quality == "none": return indices, np.inf
    cutoff = float(np.quantile(score[indices], 1 - candidate.drop_fraction))
    kept = indices[score[indices] <= cutoff]
    return (kept if len(np.unique(labels[kept])) == 2 else indices), cutoff


def model_from_state(state: dict, samples: int, device: torch.device) -> ShallowConvNet:
    model = ShallowConvNet(channels=8, samples=samples, dropout=.45).to(device)
    model.load_state_dict(copy.deepcopy(state)); return model


def candidate_validation(candidate: Candidate, x: np.ndarray, y: np.ndarray, score_map: dict[str, np.ndarray], state: dict, outer_train: np.ndarray, device: torch.device, seed: int) -> float:
    splitter = StratifiedKFold(n_splits=3, shuffle=True, random_state=stable_seed(seed, candidate.name, "inner"))
    scores_ = []
    for fold, (relative_train, relative_val) in enumerate(splitter.split(outer_train, y[outer_train]), 1):
        train_index, val_index = outer_train[relative_train], outer_train[relative_val]
        kept, _ = selected_train_indices(train_index, score_map[candidate.quality], candidate, y)
        model = model_from_state(state, x.shape[-1], device)
        fit(model, x[kept], y[kept], device, stable_seed(seed, candidate.name, fold), epochs=12)
        scores_.append(float((predict(model, x[val_index], device) == y[val_index]).mean()))
    return float(np.mean(scores_))


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_nested_pipeline_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); options = candidates(); states = source_states(args.source_cache, options, device, args.seed)
    outer_rows, inner_rows = [], []
    for epoch_path in sorted((args.tch_cache / "sessions").glob("*_x.npy")):
        session = epoch_path.stem[:-2]; epochs = np.load(epoch_path, allow_pickle=False); y = np.load(epoch_path.with_name(f"{session}_y.npy"), allow_pickle=False).astype(np.int64)
        x_cache = {option.key: window_data(epochs, option) for option in options}
        score_cache = {reference: scores(epochs, reference) for reference in REFERENCES}
        splitter = StratifiedKFold(n_splits=5, shuffle=True, random_state=stable_seed(args.seed, session, "outer"))
        for outer_fold, (train_index, test_index) in enumerate(splitter.split(epochs, y), 1):
            validation = []
            for option in options:
                result = candidate_validation(option, x_cache[option.key], y, score_cache[option.reference], states[option.key], train_index, device, stable_seed(args.seed, session, outer_fold))
                validation.append((result, option))
                inner_rows.append({"session_id":session,"outer_fold":outer_fold,"candidate":option.name,"inner_accuracy":result,**asdict(option)})
            # Deterministic tie-break: simpler candidate order wins if mean inner accuracy is equal.
            best_score, best = max(validation, key=lambda item: item[0])
            x = x_cache[best.key]; kept, cutoff = selected_train_indices(train_index, score_cache[best.reference][best.quality], best, y)
            model = model_from_state(states[best.key], x.shape[-1], device)
            fit(model, x[kept], y[kept], device, stable_seed(args.seed, session, outer_fold, "outer"), epochs=25)
            pred = predict(model, x[test_index], device); test_score = score_cache[best.reference][best.quality][test_index]
            high = test_score <= cutoff
            outer_rows.append({"session_id":session,"outer_fold":outer_fold,"outer_accuracy":float((pred == y[test_index]).mean()),"inner_selected_accuracy":best_score,"train_total":len(train_index),"train_kept":len(kept),"test_total":len(test_index),"test_high_quality":int(high.sum()),"test_high_quality_accuracy":float((pred[high] == y[test_index][high]).mean()) if high.any() else np.nan,**asdict(best)})
            print(f"[outer] {session} fold={outer_fold} selected={best.name} inner={best_score:.3f}", flush=True)
    outer = pd.DataFrame(outer_rows); inner = pd.DataFrame(inner_rows)
    outer.to_csv(args.output / "outer_fold_results.csv", index=False); inner.to_csv(args.output / "inner_candidate_results.csv", index=False)
    outer.groupby("session_id", as_index=False).agg(mean_accuracy=("outer_accuracy","mean"), std_accuracy=("outer_accuracy","std")).to_csv(args.output / "session_summary.csv", index=False)
    selection = outer.groupby(["window","reference","quality","drop_fraction"], as_index=False).size().rename(columns={"size":"outer_folds_selected"})
    selection.to_csv(args.output / "selection_frequency.csv", index=False)
    (args.output / "run_config.json").write_text(json.dumps({"design":"Per-session nested 5-fold outer CV. Window/reference/quality candidates are selected by 3-fold inner CV using outer-train only. Outer test data are never used for selection; quality gate only excludes train trials.","candidates":[asdict(candidate) for candidate in options],"seed":args.seed},ensure_ascii=False,indent=2),encoding="utf-8")
    print("[complete]", args.output, flush=True)


if __name__ == "__main__":
    main()