"""Seeding, device selection, the training loop, checkpoints and the shared command line.

Every lecture script ends in the same four lines: parse the arguments, build the model, call
``train`` with a step function, and print a report. Only the step function differs between
frameworks, and each lecture is about that function.
"""

from __future__ import annotations

import argparse
import json
import random
import time
from pathlib import Path
from typing import Callable

import torch


# %% 01. Reproducibility and hardware


def set_seed(seed: int) -> None:
    """Seed Python, NumPy through torch, and torch itself, so a run repeats exactly."""
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def pick_device(requested: str = "auto") -> torch.device:
    """Choose CUDA, then Apple MPS, then CPU, unless the reader names one."""
    if requested != "auto":
        return torch.device(requested)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


# %% 02. The command line every script shares
#
# A reader who has learned the flags once knows them for all fifty scripts. A script adds its own
# flags on top of these with `parser.add_argument`.


def common_parser(description: str) -> argparse.ArgumentParser:
    """The flags shared by every script in the course."""
    parser = argparse.ArgumentParser(
        description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--steps", type=int, default=600, help="optimizer steps")
    parser.add_argument("--batch-size", type=int, default=64, help="examples per step")
    parser.add_argument("--lr", type=float, default=1e-3, help="learning rate")
    parser.add_argument("--width", type=int, default=128, help="hidden width of the network")
    parser.add_argument("--length", type=int, default=16, help="sequence length L")
    parser.add_argument("--samples", type=int, default=16, help="sequences or points to generate")
    parser.add_argument("--sample-steps", type=int, default=64, help="integration steps")
    parser.add_argument("--seed", type=int, default=0, help="random seed")
    parser.add_argument("--device", type=str, default="auto", help="auto, cpu, cuda or mps")
    parser.add_argument("--out", type=str, default="", help="directory for the run's output")
    parser.add_argument("--quiet", action="store_true", help="print only the final report")
    return parser


# %% 03. The training loop
#
# The caller passes a function of the step index that returns a scalar loss with a graph attached;
# everything else, the optimizer, the logging and the timing, is the same in every lecture.


def train(
    model: torch.nn.Module,
    loss_fn: Callable[[int], torch.Tensor],
    steps: int = 600,
    lr: float = 1e-3,
    log_every: int = 100,
    quiet: bool = False,
    grad_clip: float = 1.0,
    extra_parameters: list[torch.nn.Parameter] | None = None,
) -> list[float]:
    """Run ``steps`` optimizer steps on ``loss_fn`` and return the loss at every step.

    ``loss_fn(step)`` draws its own batch, so a lecture with fresh noise at every step runs without
    a DataLoader. ``extra_parameters`` covers the rare case of a second trained object, such as the
    bias force of Chapter 7.
    """
    parameters = list(model.parameters()) + list(extra_parameters or [])
    optimizer = torch.optim.AdamW(parameters, lr=lr, weight_decay=0.0)
    losses: list[float] = []
    started = time.time()
    model.train()
    for step in range(steps):
        loss = loss_fn(step)  # The lecture's own objective.
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(parameters, grad_clip)  # Keep a bad batch from derailing.
        optimizer.step()
        losses.append(float(loss.detach()))
        if not quiet and log_every and (step + 1) % log_every == 0:
            window = losses[-log_every:]  # Average over the last window, a less jumpy number.
            print(f"  step {step + 1:5d}/{steps}  loss {sum(window) / len(window):.4f}")
    if not quiet:
        print(f"  trained {steps} steps in {time.time() - started:.1f}s")
    model.eval()
    return losses


# %% 04. Saving a run
#
# A run writes three files: the weights, the configuration that produced them, and a short report of
# what came out. No optimizer state, because none of these examples resumes training.


def save_run(
    out: str | Path,
    model: torch.nn.Module | None = None,
    config: dict | None = None,
    report: dict | None = None,
    losses: list[float] | None = None,
    samples: list[str] | None = None,
) -> Path | None:
    """Write the weights, the configuration, the report, the loss curve and any samples."""
    if not out:
        return None
    directory = Path(out)
    directory.mkdir(parents=True, exist_ok=True)
    if model is not None:
        torch.save({"model": model.state_dict(), "config": config or {}},
                   directory / "checkpoint.pt")
    if config is not None:
        (directory / "config.json").write_text(json.dumps(config, indent=2) + "\n")
    if report is not None:
        (directory / "report.json").write_text(json.dumps(report, indent=2) + "\n")
    if losses is not None:
        # Thin the curve to at most 200 points so the file stays small and readable.
        stride = max(1, len(losses) // 200)
        thinned = [round(v, 5) for v in losses[::stride]]
        (directory / "losses.json").write_text(json.dumps(thinned) + "\n")
    if samples is not None:
        (directory / "samples.txt").write_text("\n".join(samples) + "\n")
    return directory


def load_run(out: str | Path, model: torch.nn.Module) -> dict:
    """Read the weights back into ``model`` and return the configuration that trained them."""
    payload = torch.load(Path(out) / "checkpoint.pt", map_location="cpu", weights_only=False)
    model.load_state_dict(payload["model"])
    model.eval()
    return payload.get("config", {})


# %% 05. Printing a report
#
# Every script ends with the same two-column block, so we can compare two reports side by side.


def print_report(title: str, report: dict, samples: list[str] | None = None) -> None:
    """Print a titled block of name and value pairs, and a few samples."""
    print(f"\n{title}")
    print("-" * len(title))
    for key, value in report.items():
        if isinstance(value, float):
            print(f"  {key:<34s} {value:.4f}")
        else:
            print(f"  {key:<34s} {value}")
    if samples:
        print("\n  samples")
        for sample in samples[:8]:
            print(f"    {sample}")