"""Deterministic train/validation-only training for the DeepWeeds comparison."""

from __future__ import annotations

import argparse
import json
import math
import random
import sys
import time
from collections.abc import Mapping, Sequence
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch
from safetensors.torch import save_file
from torch import nn
from torch.optim import Adam, Optimizer
from torch.utils.data import DataLoader

try:  # Package import in the organizer workspace.
    from .data import build_comparison_transforms, load_fold_splits, make_loader
    from .evaluate import compute_classification_metrics, package_environment, resolve_device, write_json
    from .model import SmallDeepWeedsCNN
except ImportError:  # Flat-file import after copying into Hub staging.
    from data import build_comparison_transforms, load_fold_splits, make_loader
    from evaluate import compute_classification_metrics, package_environment, resolve_device, write_json
    from model import SmallDeepWeedsCNN


@dataclass(frozen=True)
class TrainingConfig:
    """The immutable optimization contract shared by good and bad runs."""

    seed: int = 2026
    batch_size: int = 64
    epochs: int = 12
    learning_rate: float = 1e-3
    optimizer: str = "Adam"
    scheduler: None = None
    loss: str = "CrossEntropyLoss"
    class_weights: None = None
    checkpoint_selection: str = "minimum_validation_loss"
    fold: int = 0


def training_config_for_run(run_kind: str) -> TrainingConfig:
    """Return the same declared training configuration for either arm."""

    if run_kind not in {"good", "bad"}:
        raise ValueError("run_kind must be 'good' or 'bad'")
    return TrainingConfig()


def build_criterion() -> nn.CrossEntropyLoss:
    """Build ordinary, unweighted multiclass cross-entropy."""

    return nn.CrossEntropyLoss(weight=None)


def build_optimizer(model: nn.Module, config: TrainingConfig) -> Adam:
    """Build the sole optimizer declared by the comparison contract."""

    if config.optimizer != "Adam" or config.learning_rate != 1e-3:
        raise ValueError("The comparison requires Adam with learning rate 1e-3")
    return Adam(model.parameters(), lr=config.learning_rate)


def seed_everything(seed: int) -> None:
    """Seed supported random sources and deterministic CUDA behavior."""

    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    torch.use_deterministic_algorithms(True)


def select_checkpoint_by_validation_loss(
    validation_losses: Mapping[str, float],
) -> str:
    """Return the key with the lowest finite validation loss and no other input."""

    if not validation_losses:
        raise ValueError("At least one validation loss is required")
    if any(not math.isfinite(float(loss)) for loss in validation_losses.values()):
        raise ValueError("Validation losses must be finite")
    return min(validation_losses, key=validation_losses.__getitem__)


def _run_epoch(
    model: nn.Module,
    loader: DataLoader,
    criterion: nn.CrossEntropyLoss,
    device: torch.device,
    optimizer: Optimizer | None,
) -> dict[str, float]:
    training = optimizer is not None
    model.train(training)
    loss_sum = 0.0
    sample_count = 0
    true_labels: list[int] = []
    predicted_labels: list[int] = []
    class_count: int | None = None

    with torch.set_grad_enabled(training):
        for images, labels, _filenames in loader:
            images = images.to(device, non_blocking=device.type == "cuda")
            labels = labels.to(device, non_blocking=device.type == "cuda")
            if optimizer is not None:
                optimizer.zero_grad(set_to_none=True)
            logits = model(images)
            if logits.ndim != 2:
                raise ValueError(f"Expected two-dimensional logits, got {tuple(logits.shape)}")
            class_count = logits.shape[1]
            loss = criterion(logits, labels)
            if not torch.isfinite(loss):
                raise FloatingPointError(f"Encountered non-finite {'train' if training else 'validation'} loss")
            if optimizer is not None:
                loss.backward()
                optimizer.step()

            batch_size = labels.shape[0]
            sample_count += batch_size
            loss_sum += float(loss.detach().item()) * batch_size
            true_labels.extend(int(label) for label in labels.detach().cpu().tolist())
            predicted_labels.extend(int(label) for label in logits.detach().argmax(dim=1).cpu().tolist())

    if sample_count == 0 or class_count is None:
        raise ValueError("Training and validation loaders must not be empty")
    metrics = compute_classification_metrics(
        true_labels,
        predicted_labels,
        class_names=[str(class_id) for class_id in range(class_count)],
    )
    return {
        "loss": loss_sum / sample_count,
        "accuracy": metrics["accuracy"],
        "macro_f1": metrics["macro_f1"],
        "balanced_accuracy": metrics["balanced_accuracy"],
    }


def _save_safetensors(model: nn.Module, path: Path) -> None:
    state = {
        name: tensor.detach().cpu().contiguous()
        for name, tensor in model.state_dict().items()
    }
    save_file(state, str(path))


def train_model(
    model: nn.Module,
    train_loader: DataLoader,
    validation_loader: DataLoader,
    output_dir: str | Path,
    device: torch.device,
    config: TrainingConfig,
) -> dict[str, Any]:
    """Fit from train data and select from validation loss; no test input exists."""

    if config != TrainingConfig():
        raise ValueError("TrainingConfig must match the fixed comparison contract")
    destination = Path(output_dir).resolve()
    destination.mkdir(parents=True, exist_ok=True)
    checkpoint_path = destination / "model.safetensors"
    seed_everything(config.seed)
    model.to(device)
    criterion = build_criterion()
    optimizer = build_optimizer(model, config)
    history: list[dict[str, Any]] = []
    validation_losses: dict[str, float] = {}
    started = time.perf_counter()

    for epoch in range(1, config.epochs + 1):
        train_metrics = _run_epoch(model, train_loader, criterion, device, optimizer)
        validation_metrics = _run_epoch(model, validation_loader, criterion, device, None)
        history.append(
            {
                "epoch": epoch,
                "train": train_metrics,
                "validation": validation_metrics,
            }
        )
        epoch_key = f"epoch_{epoch}"
        validation_losses[epoch_key] = validation_metrics["loss"]
        if select_checkpoint_by_validation_loss(validation_losses) == epoch_key:
            _save_safetensors(model, checkpoint_path)

    elapsed = time.perf_counter() - started
    best_key = select_checkpoint_by_validation_loss(validation_losses)
    best_epoch = int(best_key.removeprefix("epoch_"))
    result = {
        "history": history,
        "best_epoch": best_epoch,
        "best_validation_loss": validation_losses[best_key],
        "checkpoint_path": str(checkpoint_path),
        "checkpoint_selection": config.checkpoint_selection,
        "runtime": {"training_seconds": elapsed},
        "environment": package_environment(device),
    }
    write_json(destination / "history.json", result)
    return result


def run_training(
    run_kind: str,
    labels_dir: str | Path,
    images_dir: str | Path,
    output_dir: str | Path,
    device: str = "auto",
    num_workers: int = 4,
) -> dict[str, Any]:
    """Create official fold-0 train/validation loaders and train one declared arm."""

    config = training_config_for_run(run_kind)
    if num_workers < 0:
        raise ValueError("num_workers must be non-negative")
    labels_path = Path(labels_dir).resolve()
    images_path = Path(images_dir).resolve()
    output_path = Path(output_dir).resolve()
    selected_device = resolve_device(device)
    splits = load_fold_splits(labels_dir=labels_path, images_dir=images_path, fold=config.fold)
    transforms = build_comparison_transforms()
    pin_memory = selected_device.type == "cuda"
    train_loader = make_loader(
        splits["train"],
        transforms[f"{run_kind}_train"],
        shuffle=True,
        seed=config.seed,
        batch_size=config.batch_size,
        num_workers=num_workers,
        pin_memory=pin_memory,
    )
    validation_loader = make_loader(
        splits["val"],
        transforms[f"{run_kind}_val"],
        shuffle=False,
        seed=config.seed,
        batch_size=config.batch_size,
        num_workers=num_workers,
        pin_memory=pin_memory,
    )
    seed_everything(config.seed)
    model = SmallDeepWeedsCNN(
        num_classes=9,
        in_channels=3 if run_kind == "good" else 1,
    )
    training_args = {
        "training": asdict(config),
        "run": {
            "kind": run_kind,
            "device": str(selected_device),
            "num_workers": num_workers,
            "pin_memory": pin_memory,
            "fold_files": {
                "train": str((labels_path / "train_subset0.csv").resolve()),
                "validation": str((labels_path / "val_subset0.csv").resolve()),
            },
            "images_dir": str(images_path),
        },
    }
    write_json(output_path / "training_args.json", training_args)
    result = train_model(
        model,
        train_loader,
        validation_loader,
        output_path,
        selected_device,
        config,
    )
    result["run"] = training_args["run"]
    return result


def _parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--run-kind", choices=("good", "bad"), required=True)
    parser.add_argument("--labels-dir", type=Path, required=True)
    parser.add_argument("--images-dir", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--device", default="auto")
    parser.add_argument("--num-workers", type=int, default=4)
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    args = _parser().parse_args(argv)
    result = run_training(
        run_kind=args.run_kind,
        labels_dir=args.labels_dir,
        images_dir=args.images_dir,
        output_dir=args.output_dir,
        device=args.device,
        num_workers=args.num_workers,
    )
    print(
        json.dumps(
            {
                "run_kind": args.run_kind,
                "best_epoch": result["best_epoch"],
                "best_validation_loss": result["best_validation_loss"],
                "checkpoint_path": result["checkpoint_path"],
                "training_seconds": result["runtime"]["training_seconds"],
            },
            indent=2,
            allow_nan=False,
        )
    )
    return 0


if __name__ == "__main__":
    sys.exit(main())