"""Patient-grouped k-fold cross-validation driver (single process, sequential).

Runs the ordinary train->evaluate pipeline once per fold, each on a different
patient-grouped partition, writing per-fold artifacts under
``<work_dir>/folds/fold_<k>/``. It then concatenates every fold's held-out
evaluation along the ``sample`` axis into out-of-fold (OOF) files at
``<work_dir>/evaluations/oof_*.nc`` -- a single evaluation covering EVERY sample,
each predicted by the fold-ensemble that never trained on it. The existing
analyses then run unchanged on the OOF files (via ``scen_analyze`` or the
``run_analysis`` step appended to ``scen_kfold``).

Reuses the existing task functions verbatim: each fold just gets its own
dataloaders in ``state`` and a per-fold ``work_dir`` in a shallow config copy.
The full dataset is loaded ONCE and every fold is a set of Subset indices over
it, so memory use matches a single run.
"""

import pathlib as pl
from typing import Any, Dict, List

import numpy as np
import xarray as xr

import model.dataset as ds
from evaluation import CLASS_NAMES, load_evaluation, save_evaluation
from util.control import check_control_events
from util.progress import ProgressTracker
from util.ui_logger import PipelineLogger

def _inner_tasks(include_noise: bool):
    """Per-fold pipeline (load_data is handled by the driver; noise optional).

    Sibling task modules are imported lazily so this module never needs the
    ``tasks`` package to be fully initialized at import time (``tasks/__init__``
    imports this module in order to register the k-fold task).
    """
    from . import evaluate, load_models, train_bayesian, train_normal

    tasks = [
        ("Train Regular", train_normal.train_normal_task),
        ("Train Bayesian", train_bayesian.train_bayesian_task),
        ("Load Models", load_models.load_models_task),
        ("Evaluate Regular", evaluate.evaluate_normal_task),
        ("Evaluate Bayesian", evaluate.evaluate_bayesian_task),
    ]
    if include_noise:
        tasks.append(("Evaluate Noisy", evaluate.evaluate_noisy_task))
    return tasks


_FOLDS_SUBDIR = "folds"
_EVAL_SUBDIR = "evaluations"


def _kfold_cfg(config: Dict[str, Any]):
    cfg = config.get("kfold", {})
    k_folds = int(cfg.get("k_folds", 4))
    val_fraction = float(cfg.get("val_fraction", 0.15))
    include_noise = bool(cfg.get("include_noise", False))
    return k_folds, val_fraction, include_noise


def _aggregate_oof(
    work_dir: pl.Path, fold_dirs: List[pl.Path], log: PipelineLogger
) -> None:
    """Concatenate each fold's held-out evaluations into OOF files."""
    out_dir = work_dir / _EVAL_SUBDIR
    out_dir.mkdir(parents=True, exist_ok=True)

    # Evaluation file names are identical across folds (normal.nc, bayesian.nc,
    # and noisy_* when enabled); discover them from the folds that produced them.
    stems = sorted(
        {p.stem for d in fold_dirs for p in (d / _EVAL_SUBDIR).glob("*.nc")}
    )
    for stem in stems:
        parts = [
            load_evaluation(str(d / _EVAL_SUBDIR / f"{stem}.nc"))
            for d in fold_dirs
            if (d / _EVAL_SUBDIR / f"{stem}.nc").exists()
        ]
        if not parts:
            continue
        # Folds hold disjoint samples; concatenating along `sample` yields full,
        # non-overlapping out-of-fold coverage. Reset the (per-fold) sample index.
        oof = xr.concat(parts, dim="sample")
        oof = oof.assign_coords(sample=np.arange(oof.sizes["sample"]))
        oof.attrs["n_folds"] = len(parts)
        oof.attrs["kfold"] = 1
        save_evaluation(oof, str(out_dir / f"oof_{stem}.nc"))
        for part in parts:
            part.close()
        log.info(
            f"Aggregated OOF: oof_{stem}.nc "
            f"({oof.sizes['sample']} samples from {len(parts)} folds)."
        )


def run_kfold_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
) -> None:
    events = state.get("events")
    stop_event = events.get("stop") if isinstance(events, dict) else None
    pause_event = events.get("pause") if isinstance(events, dict) else None

    from . import load_data  # lazy, as in _inner_tasks

    base_seed = load_data.seed_pipeline(config, log)
    k_folds, val_fraction, include_noise = _kfold_cfg(config)

    dataset, ptids, image_to_ptid = load_data.load_full_dataset(config, log)
    state["image_to_ptid"] = image_to_ptid

    inner_tasks = _inner_tasks(include_noise)

    work_dir = pl.Path(config["work_dir"])
    fold_dirs: List[pl.Path] = []

    log.info(
        f"Starting {k_folds}-fold cross-validation "
        f"(val_fraction={val_fraction}, include_noise={include_noise})."
    )
    track.set_title("K-Fold Cross-Validation")
    track.update(total=k_folds, advance=0)

    for fold in range(k_folds):
        check_control_events(stop_event, pause_event)

        subsets = ds.kfold_split_by_patient_id(
            dataset, ptids, k_folds, fold, base_seed, val_fraction
        )
        train_loader, val_loader, test_loader = ds.initalize_dataloaders(
            subsets, batch_size=config["training"]["batch_size"], seed=base_seed
        )
        state["train_loader"] = train_loader
        state["val_loader"] = val_loader
        state["test_loader"] = test_loader

        fold_dir = work_dir / _FOLDS_SUBDIR / f"fold_{fold}"
        fold_dirs.append(fold_dir)
        # Per-fold config: same everything, but outputs go under the fold dir.
        fold_config = {**config, "work_dir": str(fold_dir)}

        log.info(
            f"===== Fold {fold + 1}/{k_folds} "
            f"(train={len(subsets[0])}, val={len(subsets[1])}, "
            f"test={len(subsets[2])} scans) ====="
        )
        for task_label, task_func in inner_tasks:
            check_control_events(stop_event, pause_event)
            fold_track = track.get_sub_tracker(
                f"Fold {fold + 1}/{k_folds}: {task_label}"
            )
            task_func(fold_track, log, fold_config, state)

        track.update(advance=1)

    log.info("All folds complete; aggregating out-of-fold evaluations.")
    _aggregate_oof(work_dir, fold_dirs, log)