#!/usr/bin/env python3
"""Frozen public-data pilot for dendritic innovation and trial outcomes."""

from __future__ import annotations

import argparse
import csv
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import numpy as np
from scipy.io import loadmat
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import roc_auc_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler


PROTOCOL = "larkum_public_innovation_pilot_v1"
CONDITIONS = ("Sal", "DCZ")
METHODS = ("raw_branch", "task_residual", "innovation", "local_spine", "choice_only")
N_PREDICTOR_FOLDS = 5
RIDGE_ALPHA = 1.0
BOOTSTRAP_SAMPLES = 5000
BOOTSTRAP_SEED = 20270810


@dataclass(frozen=True)
class Trial:
    animal: str
    session: str
    condition: str
    trial_index: int
    instruction: int
    choice: int
    correct: int
    features: dict[str, np.ndarray]


def _trial_array(value: object) -> np.ndarray:
    array = np.asarray(value, dtype=np.float64)
    if array.ndim != 2 or array.shape[0] != 180 or array.shape[1] < 1:
        raise ValueError(f"expected a 180 x ROI trial array, found {array.shape}")
    if not np.isfinite(array).all():
        raise ValueError("non-finite neural activity")
    return array - array[:30].mean(axis=0, keepdims=True)


def _design(spines: np.ndarray, instruction: int) -> np.ndarray:
    frames = spines.shape[0]
    frame_basis = np.eye(frames, dtype=np.float64)
    interaction = frame_basis * float(instruction)
    return np.concatenate((spines, frame_basis, interaction), axis=1)


def _fit_predictor(
    spine_trials: list[np.ndarray],
    branch_trials: list[np.ndarray],
    instructions: np.ndarray,
    train_indices: np.ndarray,
) -> tuple[object, np.ndarray]:
    x = np.concatenate([_design(spine_trials[i], int(instructions[i])) for i in train_indices])
    y = np.concatenate([branch_trials[i] for i in train_indices])
    model = make_pipeline(StandardScaler(), Ridge(alpha=RIDGE_ALPHA))
    model.fit(x, y)
    return model, y


def _predict_trials(
    model: object,
    spine_trials: list[np.ndarray],
    instructions: np.ndarray,
    indices: Iterable[int],
) -> dict[int, np.ndarray]:
    return {
        int(i): np.asarray(model.predict(_design(spine_trials[i], int(instructions[i]))))
        for i in indices
    }


def _template_predict(
    branch_trials: list[np.ndarray],
    train_instructions: np.ndarray,
    train_indices: np.ndarray,
    test_instructions: np.ndarray,
    test_indices: Iterable[int],
) -> dict[int, np.ndarray]:
    templates: dict[int, np.ndarray] = {}
    for side in (0, 1):
        selected = [branch_trials[i] for i in train_indices if train_instructions[i] == side]
        if not selected:
            raise ValueError(f"training fold has no instruction-{side} trials")
        templates[side] = np.mean(selected, axis=0)
    return {int(i): templates[int(test_instructions[i])] for i in test_indices}


def _summarize(trace: np.ndarray) -> np.ndarray:
    time_axis = np.linspace(-3.0, 3.0, 180)
    indices = np.flatnonzero((time_axis >= 0.0) & (time_axis < 1.0))
    bins = np.array_split(indices, 6)
    values: list[float] = []
    for frame_indices in bins:
        block = trace[frame_indices]
        values.extend((float(block.mean()), float(np.abs(block).mean()), float(np.sqrt(np.mean(block**2)))))
    result = np.asarray(values, dtype=np.float64)
    if result.shape != (18,) or not np.isfinite(result).all():
        raise ValueError("invalid summary features")
    return result


def _session_trials(path: Path) -> list[Trial]:
    session = path.stem
    animal = session.split("_")[0]
    data = loadmat(path, simplify_cells=True)["cont_data"]
    prepared: dict[str, dict[str, object]] = {}

    for condition in CONDITIONS:
        condition_data = data[condition]
        choices = np.asarray(condition_data["Choice"])[:, 0]
        instructions = np.asarray(condition_data["TrialTypes"])[:, 0].astype(np.int64)
        correct = np.asarray(condition_data["DirOut"]).astype(np.int64)
        valid_indices: list[int] = []
        spines: list[np.ndarray] = []
        branches: list[np.ndarray] = []
        for i in range(len(correct)):
            if not np.isfinite(choices[i]):
                continue
            try:
                spine = _trial_array(condition_data["spine_local"][i])
                branch = _trial_array(condition_data["branch"][i])
            except ValueError:
                continue
            valid_indices.append(i)
            spines.append(spine)
            branches.append(branch)
        prepared[condition] = {
            "source_indices": np.asarray(valid_indices, dtype=np.int64),
            "choices": choices[valid_indices].astype(np.int64),
            "instructions": instructions[valid_indices],
            "correct": correct[valid_indices],
            "spines": spines,
            "branches": branches,
        }

    saline = prepared["Sal"]
    sal_spines = saline["spines"]
    sal_branches = saline["branches"]
    sal_instructions = saline["instructions"]
    n_sal = len(sal_spines)
    if n_sal < N_PREDICTOR_FOLDS:
        raise ValueError(f"{session} has only {n_sal} usable saline trials")

    crossfit_prediction: dict[int, np.ndarray] = {}
    crossfit_template: dict[int, np.ndarray] = {}
    for test_indices in np.array_split(np.arange(n_sal), N_PREDICTOR_FOLDS):
        train_indices = np.setdiff1d(np.arange(n_sal), test_indices)
        model, _ = _fit_predictor(sal_spines, sal_branches, sal_instructions, train_indices)
        crossfit_prediction.update(_predict_trials(model, sal_spines, sal_instructions, test_indices))
        crossfit_template.update(
            _template_predict(
                sal_branches,
                sal_instructions,
                train_indices,
                sal_instructions,
                test_indices,
            )
        )

    all_sal_indices = np.arange(n_sal)
    dcz = prepared["DCZ"]
    dcz_model, _ = _fit_predictor(sal_spines, sal_branches, sal_instructions, all_sal_indices)
    dcz_prediction = _predict_trials(
        dcz_model, dcz["spines"], dcz["instructions"], range(len(dcz["spines"]))
    )
    dcz_template = _template_predict(
        sal_branches,
        sal_instructions,
        all_sal_indices,
        dcz["instructions"],
        range(len(dcz["spines"])),
    )

    output: list[Trial] = []
    for condition in CONDITIONS:
        entry = prepared[condition]
        predictions = crossfit_prediction if condition == "Sal" else dcz_prediction
        templates = crossfit_template if condition == "Sal" else dcz_template
        for i, source_index in enumerate(entry["source_indices"]):
            branch = entry["branches"][i]
            spine = entry["spines"][i]
            innovation = branch - predictions[i]
            task_residual = branch - templates[i]
            output.append(
                Trial(
                    animal=animal,
                    session=session,
                    condition=condition,
                    trial_index=int(source_index),
                    instruction=int(entry["instructions"][i]),
                    choice=int(entry["choices"][i]),
                    correct=int(entry["correct"][i]),
                    features={
                        "raw_branch": _summarize(branch),
                        "task_residual": _summarize(task_residual),
                        "innovation": _summarize(innovation),
                        "local_spine": _summarize(spine),
                        "choice_only": np.asarray([float(entry["choices"][i])]),
                    },
                )
            )
    return output


def _decode(trials: list[Trial], condition: str, method: str) -> list[dict[str, object]]:
    selected = [trial for trial in trials if trial.condition == condition]
    animals = sorted({trial.animal for trial in selected})
    predictions: list[dict[str, object]] = []
    for held_animal in animals:
        for train_instruction in (0, 1):
            test_instruction = 1 - train_instruction
            train = [
                trial
                for trial in selected
                if trial.animal != held_animal and trial.instruction == train_instruction
            ]
            test = [
                trial
                for trial in selected
                if trial.animal == held_animal and trial.instruction == test_instruction
            ]
            y_train = np.asarray([trial.correct for trial in train], dtype=np.int64)
            if len(np.unique(y_train)) != 2 or not test:
                raise ValueError(
                    f"invalid decoder fold: {condition=} {method=} {held_animal=} {train_instruction=}"
                )
            x_train = np.stack([trial.features[method] for trial in train])
            x_test = np.stack([trial.features[method] for trial in test])
            decoder = make_pipeline(
                StandardScaler(),
                LogisticRegression(C=1.0, class_weight="balanced", max_iter=5000, random_state=0),
            )
            decoder.fit(x_train, y_train)
            probabilities = decoder.predict_proba(x_test)[:, 1]
            for trial, probability in zip(test, probabilities):
                predictions.append(
                    {
                        "animal": trial.animal,
                        "session": trial.session,
                        "condition": condition,
                        "method": method,
                        "trial_index": trial.trial_index,
                        "instruction": trial.instruction,
                        "choice": trial.choice,
                        "correct": trial.correct,
                        "probability_correct": float(probability),
                    }
                )
    expected = len(selected)
    if len(predictions) != expected:
        raise RuntimeError(f"decoder produced {len(predictions)} predictions for {expected} trials")
    keys = {(p["session"], p["condition"], p["trial_index"]) for p in predictions}
    if len(keys) != expected:
        raise RuntimeError("duplicate or missing out-of-fold predictions")
    return predictions


def _auc(rows: list[dict[str, object]]) -> float:
    labels = np.asarray([row["correct"] for row in rows], dtype=np.int64)
    scores = np.asarray([row["probability_correct"] for row in rows], dtype=np.float64)
    if len(np.unique(labels)) != 2:
        return float("nan")
    return float(roc_auc_score(labels, scores))


def _bootstrap_plan(
    rows: list[dict[str, object]], rng: np.random.Generator
) -> dict[str, list[str]]:
    plan: dict[str, list[str]] = {}
    animals = sorted({str(row["animal"]) for row in rows})
    for animal in animals:
        sessions = sorted({str(row["session"]) for row in rows if row["animal"] == animal})
        plan[animal] = [str(value) for value in rng.choice(sessions, size=len(sessions), replace=True)]
    return plan


def _bootstrap_indices(
    rows: list[dict[str, object]], plan: dict[str, list[str]]
) -> np.ndarray:
    indices: list[int] = []
    for animal, sessions in plan.items():
        for session in sessions:
            indices.extend(
                i
                for i, row in enumerate(rows)
                if row["animal"] == animal and row["session"] == session
            )
    return np.asarray(indices, dtype=np.int64)


def _bootstrap_auc(
    rows: list[dict[str, object]], sample_plans: list[dict[str, list[str]]]
) -> np.ndarray:
    values: list[float] = []
    for plan in sample_plans:
        indices = _bootstrap_indices(rows, plan)
        sample = [rows[i] for i in indices]
        value = _auc(sample)
        if np.isfinite(value):
            values.append(value)
    return np.asarray(values, dtype=np.float64)


def _interval(values: np.ndarray) -> list[float]:
    return [float(v) for v in np.quantile(values, (0.025, 0.975))]


def _summarize_predictions(all_predictions: list[dict[str, object]]) -> dict[str, object]:
    grouped: dict[tuple[str, str], list[dict[str, object]]] = {}
    for row in all_predictions:
        grouped.setdefault((str(row["condition"]), str(row["method"])), []).append(row)

    rng = np.random.default_rng(BOOTSTRAP_SEED)
    reference = grouped[("Sal", "raw_branch")]
    sample_plans = [_bootstrap_plan(reference, rng) for _ in range(BOOTSTRAP_SAMPLES)]
    results: dict[str, object] = {}
    bootstrap_values: dict[tuple[str, str], np.ndarray] = {}
    for condition in CONDITIONS:
        results[condition] = {}
        for method in METHODS:
            rows = grouped[(condition, method)]
            values = _bootstrap_auc(rows, sample_plans)
            bootstrap_values[(condition, method)] = values
            by_animal = {}
            for animal in sorted({str(row["animal"]) for row in rows}):
                animal_rows = [row for row in rows if row["animal"] == animal]
                by_animal[animal] = {
                    "auc": _auc(animal_rows),
                    "trials": len(animal_rows),
                    "errors": int(sum(1 - int(row["correct"]) for row in animal_rows)),
                }
            results[condition][method] = {
                "auc": _auc(rows),
                "session_block_bootstrap_95ci": _interval(values),
                "trials": len(rows),
                "errors": int(sum(1 - int(row["correct"]) for row in rows)),
                "by_animal": by_animal,
            }

    comparisons: dict[str, object] = {}
    for condition in CONDITIONS:
        difference = (
            bootstrap_values[(condition, "innovation")]
            - bootstrap_values[(condition, "raw_branch")]
        )
        comparisons[f"innovation_minus_raw_{condition.lower()}"] = {
            "difference": float(
                results[condition]["innovation"]["auc"]
                - results[condition]["raw_branch"]["auc"]
            ),
            "session_block_bootstrap_95ci": _interval(difference),
            "bootstrap_probability_positive": float(np.mean(difference > 0.0)),
        }
    interaction = (
        bootstrap_values[("Sal", "innovation")]
        - bootstrap_values[("DCZ", "innovation")]
    )
    comparisons["innovation_sal_minus_dcz"] = {
        "difference": float(
            results["Sal"]["innovation"]["auc"]
            - results["DCZ"]["innovation"]["auc"]
        ),
        "session_block_bootstrap_95ci": _interval(interaction),
        "bootstrap_probability_positive": float(np.mean(interaction > 0.0)),
    }
    return {"endpoints": results, "comparisons": comparisons}


def _write_predictions(path: Path, rows: list[dict[str, object]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = [
        "animal",
        "session",
        "condition",
        "method",
        "trial_index",
        "instruction",
        "choice",
        "correct",
        "probability_correct",
    ]
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-root", type=Path, required=True, help="Extracted public archive root")
    parser.add_argument("--output", type=Path, default=Path("results/larkum_public_pilot.json"))
    parser.add_argument(
        "--predictions",
        type=Path,
        default=Path("results/larkum_public_pilot_predictions.csv"),
    )
    parser.add_argument("--audit-only", action="store_true", help="Report eligible files and counts only")
    args = parser.parse_args()

    files = sorted((args.data_root / "Figure2" / "Data").glob("*_dff.mat"))
    if len(files) != 11:
        raise SystemExit(f"expected 11 Figure 2 session files, found {len(files)}")
    if args.audit_only:
        print(json.dumps({"protocol": PROTOCOL, "files": [path.name for path in files]}, indent=2))
        return

    trials: list[Trial] = []
    for path in files:
        trials.extend(_session_trials(path))

    predictions: list[dict[str, object]] = []
    for condition in CONDITIONS:
        for method in METHODS:
            predictions.extend(_decode(trials, condition, method))
    summary = _summarize_predictions(predictions)

    counts = {}
    for condition in CONDITIONS:
        subset = [trial for trial in trials if trial.condition == condition]
        counts[condition] = {
            "trials": len(subset),
            "correct": int(sum(trial.correct for trial in subset)),
            "errors": int(sum(1 - trial.correct for trial in subset)),
            "animals": sorted({trial.animal for trial in subset}),
            "sessions": len({trial.session for trial in subset}),
        }
    result = {
        "protocol": PROTOCOL,
        "source": {
            "paper_doi": "10.1126/science.adx4358",
            "dataset_doi": "10.12751/g-node.etlk5k",
            "files": [path.name for path in files],
        },
        "fixed_configuration": {
            "ridge_alpha": RIDGE_ALPHA,
            "predictor_folds": N_PREDICTOR_FOLDS,
            "decoder": "StandardScaler + balanced logistic regression, C=1.0",
            "evaluation": "cross-instruction leave-one-animal-out",
            "outcome_window_seconds": [0.0, 1.0],
            "bootstrap_samples": BOOTSTRAP_SAMPLES,
            "bootstrap_seed": BOOTSTRAP_SEED,
        },
        "counts": counts,
        **summary,
        "claim_boundary": (
            "Exploratory biological consistency test only; the release has three animals, "
            "no paired soma signal, and no trial-linked relearning endpoint."
        ),
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    _write_predictions(args.predictions, predictions)
    print(json.dumps(result, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()