"""Sample from a pretrained nanochat checkpoint and append one run-log line.

Purpose: close the loop on a from-scratch pretraining run: load the checkpoint you
         just trained, sample from it with a fixed prompt set so that two runs can
         be compared, read the losses and scores out of the training and evaluation
         logs, read the accelerator's power draw where the platform reports it, and
         append one JSON line to the lab notebook in the run-log format Part 11
         introduces.
Platform: all (cuda on Tracks S and N and on Track X with ROCm, mps on Track M,
          cpu anywhere as a fallback); the device is autodetected and recorded.
Minimum memory: 8 GB
Assumes: nanochat is cloned at --nanochat (default $NANOCHAT or ~/nanochat) with its
         virtual environment synced, a checkpoint exists under
         $NANOCHAT_BASE_DIR/base_checkpoints/<model-tag>, and this script is run with
         that environment's interpreter, for example
         ~/nanochat/.venv/bin/python sample-and-record.py ...

Usage: python sample-and-record.py --model-tag d6-lab [--step N]
                                   [--fields train-fields.json]
                                   [--labbook labbook.md]
                                   [--max-tokens 48] [--temperature 0.8]
                                   [--print-only]

The run log is one JSON object per line with the fields Part 11's runlog.py writes:
run_id, lab, date, config_commit, model, dataset, hyperparameters, seed, hardware,
versions, losses, scores and notes. A field that cannot be filled in is written as
null rather than omitted, so that a reader can always tell "not recorded" from "not
applicable".
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import re
import secrets
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

# Prompts that probe different things: a fact, a lookup, a small piece of
# reasoning, an antonym, a list, an open continuation and arithmetic. They are the
# set nanochat samples with during training, kept identical here so that a sample
# taken now can be compared with the ones in the training log.
FACT_PROMPTS = [
    "The capital of France is",
    "The chemical symbol of gold is",
    "If yesterday was Friday, then tomorrow will be",
    "The opposite of hot is",
    "The planets of the solar system are:",
    "My favorite color is",
    "If 5*x + 3 = 13, then x is",
]

# Longer, sampled rather than greedy, to show what the model does when it is not
# being led. This is where an undertrained model is most obvious.
FREE_PROMPTS = [
    "The history of the printing press begins",
    "In order to bake bread you will need",
]

REQUIRED_FIELDS = (
    "run_id", "lab", "date", "config_commit", "model", "dataset",
    "hyperparameters", "seed", "hardware", "versions", "losses", "scores", "notes",
)


# --------------------------------------------------------------------------- #
# Reading the logs                                                             #
# --------------------------------------------------------------------------- #

def read_text(path: str | None) -> str:
    """Return a log file's contents, or an empty string if it is missing."""
    if not path:
        return ""
    try:
        return Path(path).read_text(encoding="utf-8", errors="replace")
    except OSError:
        return ""


def last_float(pattern: str, text: str) -> float | None:
    """The last match of a single-group numeric pattern, as a float, or None.

    MULTILINE is on because several of the lines below are matched from the start
    of a line, and the logs are whole files rather than single lines.
    """
    matches = re.findall(pattern, text, flags=re.MULTILINE)
    if not matches:
        return None
    try:
        return float(str(matches[-1]).replace(",", ""))
    except ValueError:
        return None


def parse_train_log(text: str) -> dict:
    """Pull the figures nanochat prints during and after training out of its log."""
    step_rates = [float(m.replace(",", "")) for m in re.findall(r"tok/sec: ([0-9,]+)", text)]
    return {
        "min_val_bpb": last_float(r"Minimum validation bpb: ([0-9.]+)", text),
        "last_val_bpb": last_float(r"Validation bpb: ([0-9.]+)", text),
        "last_train_loss": last_float(r"\| loss: ([0-9.]+)", text),
        "total_training_minutes": last_float(r"Total training time: ([0-9.]+)m", text),
        "peak_memory_mib": last_float(r"Peak memory usage: ([0-9.]+)MiB", text),
        "parameters_total": last_float(r"^total\s+: ([0-9,]+)", text),
        "flops_per_token": last_float(r"Estimated FLOPs per token: ([0-9.e+]+)", text),
        "tokens_per_scaling_param": last_float(r"Tokens : Scaling params ratio: ([0-9.]+)", text),
        "training_tokens": last_float(r"Total number of training tokens: ([0-9,]+)", text),
        "median_tokens_per_second": (
            sorted(step_rates)[len(step_rates) // 2] if step_rates else None
        ),
        "last_mfu_percent": last_float(r"bf16_mfu: ([0-9.]+)", text),
    }


def parse_eval_log(text: str) -> dict:
    """Pull the evaluation figures out of the base_eval log."""
    return {
        "train_bpb": last_float(r"^train bpb: ([0-9.]+)", text),
        "val_bpb": last_float(r"^val bpb: ([0-9.]+)", text),
        "core_metric": last_float(r"CORE metric: ([0-9.]+)", text),
    }


# --------------------------------------------------------------------------- #
# The machine                                                                  #
# --------------------------------------------------------------------------- #

def run_command(args: list[str]) -> str | None:
    """Run a command and return its stripped output, or None if it is unavailable."""
    if not shutil.which(args[0]):
        return None
    try:
        out = subprocess.run(args, capture_output=True, text=True, timeout=20, check=True)
    except (subprocess.SubprocessError, OSError):
        return None
    return out.stdout.strip() or None


def read_power() -> dict:
    """Accelerator power draw, where the platform has a documented way to report it.

    NVIDIA and AMD both ship a query tool. macOS has no equivalent one-line query
    that this course has been able to cite, so on Track M the power field is null
    and the notebook records wall-clock time instead.
    """
    nvidia = run_command(
        ["nvidia-smi", "--query-gpu=power.draw", "--format=csv,noheader,nounits"]
    )
    if nvidia:
        first = nvidia.splitlines()[0].strip()
        try:
            return {"watts": float(first), "source": "nvidia-smi --query-gpu=power.draw"}
        except ValueError:
            return {"watts": None, "source": f"nvidia-smi returned {first!r}"}

    amd = run_command(["rocm-smi", "--showpower"])
    if amd:
        match = re.search(r"([0-9]+\.?[0-9]*)\s*W", amd)
        return {
            "watts": float(match.group(1)) if match else None,
            "source": "rocm-smi --showpower",
        }

    return {"watts": None, "source": "not reported on this platform"}


def describe_hardware(device) -> dict:
    import torch

    accelerator = "cpu"
    name = platform.processor() or platform.machine()
    if device.type == "cuda":
        accelerator = "cuda"
        name = torch.cuda.get_device_name(0)
    elif device.type == "mps":
        accelerator = "mps"
        name = "Apple silicon GPU (Metal Performance Shaders)"
    return {
        "os": f"{platform.system()} {platform.release()}",
        "arch": platform.machine(),
        "accelerator": accelerator,
        "device_name": name,
        "power": read_power(),
    }


def describe_versions() -> dict:
    import torch

    versions = {
        "python": platform.python_version(),
        "torch": torch.__version__,
        "cuda": getattr(torch.version, "cuda", None),
        "hip": getattr(torch.version, "hip", None),
    }
    for name in ("tiktoken", "rustbpe"):
        try:
            module = __import__(name)
        except ImportError:
            versions[name] = None
        else:
            versions[name] = getattr(module, "__version__", "installed")
    return versions


def new_run_id() -> str:
    """A short identifier that sorts by time and does not collide between runs."""
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    return f"{stamp}-{secrets.token_hex(3)}"


# --------------------------------------------------------------------------- #
# Sampling                                                                     #
# --------------------------------------------------------------------------- #

def sample_from_model(engine, tokenizer, prompts, *, max_tokens, temperature, top_k, seed):
    """Return one continuation per prompt, decoded back to text."""
    out = []
    for prompt in prompts:
        tokens = tokenizer(prompt, prepend="<|bos|>")
        kwargs = {"num_samples": 1, "max_tokens": max_tokens, "temperature": temperature, "seed": seed}
        if temperature > 0 and top_k:
            kwargs["top_k"] = top_k
        completion, _ = engine.generate_batch(tokens, **kwargs)
        out.append({"prompt": prompt, "completion": tokenizer.decode(completion[0])})
    return out


# --------------------------------------------------------------------------- #

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model-tag", required=True, help="checkpoint directory name, e.g. d6-lab")
    parser.add_argument("--step", type=int, default=None, help="checkpoint step (default: the last one)")
    parser.add_argument("--nanochat", default=os.environ.get("NANOCHAT", str(Path.home() / "nanochat")))
    parser.add_argument("--fields", default=None, help="train-fields.json written by train-small.sh")
    parser.add_argument("--labbook", default=None, help="append one JSON line to this file")
    parser.add_argument("--max-tokens", type=int, default=48)
    parser.add_argument("--temperature", type=float, default=0.8, help="for the free-continuation prompts")
    parser.add_argument("--top-k", type=int, default=50)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device-type", default="", help="cuda|mps|cpu (empty = autodetect)")
    parser.add_argument("--print-only", action="store_true", help="show the record without writing it")
    args = parser.parse_args()

    # nanochat is imported from the clone rather than installed, so put it on the path.
    repo = Path(args.nanochat).expanduser()
    if not (repo / "nanochat").is_dir():
        sys.exit(f"No nanochat package under {repo}. Pass --nanochat or set NANOCHAT.")
    sys.path.insert(0, str(repo))

    try:
        from nanochat.checkpoint_manager import load_model
        from nanochat.common import autodetect_device_type, compute_cleanup, compute_init
        from nanochat.engine import Engine
    except ImportError as exc:
        sys.exit(
            f"Could not import nanochat ({exc}). Run this script with the repository's own\n"
            f"interpreter, for example {repo}/.venv/bin/python sample-and-record.py ..."
        )

    device_type = autodetect_device_type() if args.device_type == "" else args.device_type
    _, _, _, _, device = compute_init(device_type)
    model, tokenizer, meta = load_model(
        "base", device, phase="eval", model_tag=args.model_tag, step=args.step
    )
    engine = Engine(model, tokenizer)

    # Greedy on the fact prompts so the answer is the model's most likely one, and
    # sampled on the free prompts so the text is representative rather than flat.
    greedy = sample_from_model(
        engine, tokenizer, FACT_PROMPTS,
        max_tokens=16, temperature=0.0, top_k=None, seed=args.seed,
    )
    free = sample_from_model(
        engine, tokenizer, FREE_PROMPTS,
        max_tokens=args.max_tokens, temperature=args.temperature, top_k=args.top_k, seed=args.seed,
    )

    print("\nGreedy continuations (temperature 0)")
    for item in greedy:
        print("-" * 78)
        print(item["completion"])
    print("\nSampled continuations (temperature %.2f, top-k %d)" % (args.temperature, args.top_k))
    for item in free:
        print("-" * 78)
        print(item["completion"])

    fields = {}
    if args.fields:
        try:
            fields = json.loads(Path(args.fields).read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            print(f"\nwarning: could not read {args.fields} ({exc}); recording without it")

    train_stats = parse_train_log(read_text(fields.get("train_log")))
    eval_stats = parse_eval_log(read_text(fields.get("eval_log")))

    record = {
        "run_id": new_run_id(),
        "lab": fields.get("lab", "part-12/train-a-small-model"),
        "date": datetime.now(timezone.utc).strftime("%Y-%m-%d"),
        "config_commit": fields.get("commit"),
        "model": {
            "name": f"nanochat base, tag {args.model_tag}",
            "trained_from": "random initialisation",
            "step": meta.get("step"),
            "config": meta.get("model_config"),
            "parameters_total": train_stats["parameters_total"],
        },
        "dataset": {
            "name": fields.get("dataset"),
            "training_tokens": train_stats["training_tokens"],
            "tokens_per_scaling_param": train_stats["tokens_per_scaling_param"],
        },
        "hyperparameters": {
            key: fields.get(key)
            for key in (
                "depth", "head_dim", "window_pattern", "max_seq_len",
                "device_batch_size", "total_batch_size", "num_iterations",
                "save_every", "target_minutes", "dtype_override",
            )
        },
        "seed": args.seed,
        "hardware": describe_hardware(device),
        "versions": describe_versions(),
        "losses": {
            "final_train_loss": train_stats["last_train_loss"],
            "min_val_bpb": train_stats["min_val_bpb"],
            "train_bpb": eval_stats["train_bpb"],
            "val_bpb": eval_stats["val_bpb"],
        },
        "scores": {
            "core_metric": eval_stats["core_metric"],
            "median_tokens_per_second": train_stats["median_tokens_per_second"],
            "last_mfu_percent": train_stats["last_mfu_percent"],
            "flops_per_token": train_stats["flops_per_token"],
            "total_training_minutes": train_stats["total_training_minutes"],
            "peak_memory_mib": train_stats["peak_memory_mib"],
        },
        "notes": {
            "greedy_samples": greedy,
            "sampled_completions": free,
            "track": fields.get("track"),
            "calibrated_tokens_per_second": fields.get("calibrated_tokens_per_second"),
        },
    }
    for key in REQUIRED_FIELDS:
        record.setdefault(key, None)

    line = json.dumps(record, ensure_ascii=False, sort_keys=False)
    if args.print_only or not args.labbook:
        print("\n" + line)
    else:
        notebook = Path(args.labbook)
        if not notebook.exists():
            notebook.write_text("# Lab notebook\n\n## Results\n\n", encoding="utf-8")
        with notebook.open("a", encoding="utf-8") as handle:
            handle.write(line + "\n")
        print(f"\nrecorded in {notebook}")

    compute_cleanup()


if __name__ == "__main__":
    main()