"""
Pretraining loop for Wisp on Apple Silicon via MLX.

Precision policy: parameters are held in float32 as master weights for the
optimizer, and cast down to the compute dtype (bfloat16 by default) for the
forward and backward pass. Gradients come back in the compute dtype and are
promoted to float32 before the optimizer step. This is the standard mixed
precision recipe and it matters here: bfloat16 second moments in AdamW will
quietly stall a from-scratch run.

Usage:
    python train.py --config config/run1.json
    python train.py --config config/run1.json --resume out/run1/ckpt_latest
    python train.py --config config/run1.json \
        --save-initialized out/run1-untrained/ckpt
"""

import argparse
import json
import math
import os
import shutil
import sys
import time

import mlx.core as mx
import mlx.nn as nn
import mlx.optimizers as optim
import numpy as np
from mlx.utils import tree_flatten, tree_map, tree_unflatten

from checkpoint_fs import (
    install_checkpoint,
    prepare_snapshot_boundary_recovery,
    recover_checkpoint,
)
from data import (
    ShardDataset,
    sampler_batches_since_reset,
    sampler_reset_steps,
    validate_resume_sampling_contract,
    validate_data_contract,
)
from model import Wisp, ModelArgs

DTYPES = {"float32": mx.float32, "bfloat16": mx.bfloat16, "float16": mx.float16}


def peak_memory_gb() -> float:
    """MLX moved this call between versions, so probe both spellings."""
    for getter in (getattr(mx, "get_peak_memory", None),
                   getattr(getattr(mx, "metal", None), "get_peak_memory", None)):
        if callable(getter):
            try:
                return getter() / 1e9
            except Exception:
                continue
    return 0.0


def load_config(path: str) -> dict:
    with open(path) as f:
        return json.load(f)


def build_schedule(cfg: dict):
    # cosine_decay divides by decay_steps, so a config where max_steps is at or
    # below warmup_steps crashes with a bare ZeroDivisionError from inside MLX.
    # That happens for real in short probe runs, where the message points at the
    # scheduler rather than at the config that caused it.
    decay_steps = max(1, cfg["max_steps"] - cfg["warmup_steps"])
    warmup = optim.linear_schedule(0.0, cfg["lr"], cfg["warmup_steps"])
    decay = optim.cosine_decay(cfg["lr"], decay_steps, cfg["lr_min"])
    return optim.join_schedules([warmup, decay], [cfg["warmup_steps"]])


def cast_tree(tree, dtype):
    return tree_map(lambda a: a.astype(dtype) if isinstance(a, mx.array) else a, tree)


def flat_arrays(tree) -> dict:
    return {k: v for k, v in tree_flatten(tree) if isinstance(v, mx.array)}


def prune_empty(tree):
    """
    Rebuild a parameter tree with its array-free nodes removed.

    `nn.RoPE` holds no learnable parameters, but it still occupies a node in
    `model.parameters()` as an empty dict. safetensors stores arrays and nothing
    else, so that node does not survive a save and load. A resumed run then has a
    master tree without it while the live gradient tree still has it, and the
    `tree_map` inside `apply_gradients` raises `KeyError: 'rope'` on the first
    optimizer step. A fresh run never sees this, because there the gradient tree
    drives the traversal and the extra node in master is simply ignored.

    Normalising every tree the same way makes a resumed run structurally identical
    to a fresh one.
    """
    return tree_unflatten(list(flat_arrays(tree).items()))


def save_checkpoint(
    path: str,
    model,
    master,
    optimizer,
    step: int,
    cfg: dict,
    args: ModelArgs,
    include_optimizer: bool = True,
    train_sampler: dict | None = None,
):
    """
    Write a checkpoint atomically.

    Files are staged and flushed, then macOS atomically exchanges the complete
    staging and live directories. The canonical path therefore always names
    either the old checkpoint or the new checkpoint. The prior version remains
    at `.prev`. Initialization-only controls omit optimizer state because no
    optimizer step exists and they are evaluation artifacts, not resume points.
    """
    parent = os.path.dirname(os.path.abspath(path)) or "."
    os.makedirs(parent, exist_ok=True)
    staging = path + ".partial"
    if os.path.exists(staging):
        shutil.rmtree(staging)
    os.makedirs(staging)

    mx.save_safetensors(os.path.join(staging, "master.safetensors"), flat_arrays(master))
    if include_optimizer:
        mx.save_safetensors(
            os.path.join(staging, "optimizer.safetensors"),
            flat_arrays(optimizer.state),
        )
    meta = {
        "step": step,
        "config": cfg,
        "model_args": args.to_dict(),
        "optimizer_state_included": include_optimizer,
    }
    if train_sampler is not None:
        meta["train_sampler"] = train_sampler
    with open(os.path.join(staging, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)
        f.flush()
        os.fsync(f.fileno())

    install_checkpoint(staging, path)


def load_checkpoint(path: str, model, optimizer):
    master = tree_unflatten(list(mx.load(os.path.join(path, "master.safetensors")).items()))
    opt_state = tree_unflatten(list(mx.load(os.path.join(path, "optimizer.safetensors")).items()))
    with open(os.path.join(path, "meta.json")) as f:
        meta = json.load(f)
    optimizer.state.update(opt_state)
    return master, meta


def checkpoint_sampler_state(
    dataset,
    batch_size: int,
    reset_steps: list[int],
) -> dict:
    state = dataset.sampler_state(batch_size)
    state["reset_steps"] = reset_steps
    return state


def evaluate(model, dataset, cfg, compute_dtype, n_batches: int = 20):
    model.eval()
    totals = np.zeros(3 + cfg["mtp_depth"] - 1, dtype=np.float64)
    main_sum, mtp_sums, count = 0.0, [0.0] * cfg["mtp_depth"], 0
    for batch in dataset.iter_eval(cfg["micro_batch"], n_batches):
        total, main, mtps = model.loss(mx.array(batch), cfg["mtp_weight"])
        mx.eval(total, main, *mtps)
        main_sum += float(main)
        for i, m in enumerate(mtps):
            mtp_sums[i] += float(m)
        count += 1
    model.train()
    del totals
    return main_sum / count, [s / count for s in mtp_sums]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", required=True)
    ap.add_argument("--resume", default=None)
    ap.add_argument("--smoke", action="store_true", help="tiny run to shake out the pipeline")
    ap.add_argument("--set", action="append", default=[], metavar="KEY=VALUE",
                    help="override a config key, value parsed as JSON, repeatable")
    ap.add_argument("--compile", action="store_true",
                    help="mx.compile the forward and backward micro step")
    ap.add_argument(
        "--save-initialized",
        default=None,
        metavar="PATH",
        help="save an exact-geometry, seed-matched untrained control and exit",
    )
    cli = ap.parse_args()

    cfg = load_config(cli.config)
    if cli.smoke:
        cfg.update({
            "max_steps": 20, "warmup_steps": 5, "eval_interval": 10,
            "ckpt_interval": 20, "micro_batch": 2, "grad_accum": 2, "seq_len": 256,
        })
    for item in cli.set:
        key, _, raw = item.partition("=")
        try:
            cfg[key] = json.loads(raw)
        except json.JSONDecodeError:
            cfg[key] = raw
        print(f"override {key} = {cfg[key]!r}")
    if cli.save_initialized:
        if cli.resume:
            raise ValueError("--save-initialized and --resume are mutually exclusive")
        if os.path.exists(cli.save_initialized):
            raise FileExistsError(
                f"refusing to replace initialized control: {cli.save_initialized}"
            )
        cfg["run_name"] = f"{cfg['run_name']}-untrained-control"
        cfg["lr"] = 0.0
        cfg["lr_min"] = 0.0
        cfg["initialization_only"] = True
    else:
        validate_data_contract(cfg, cfg["data_index"])

    if not cli.save_initialized:
        default_checkpoint = os.path.join(cfg["out_dir"], "ckpt_latest")
        if cli.resume:
            if recover_checkpoint(cli.resume):
                print(f"recovered legacy checkpoint rename gap at {cli.resume}")
        elif recover_checkpoint(default_checkpoint):
            cli.resume = default_checkpoint
            print(
                "recovered legacy checkpoint rename gap and enabled resume at "
                f"{default_checkpoint}"
            )

    mx.random.seed(cfg.get("seed", 1337))
    compute_dtype = DTYPES[cfg.get("compute_dtype", "bfloat16")]

    args = ModelArgs.from_dict(cfg)
    args.max_seq_len = cfg["seq_len"]
    model = Wisp(args)
    mx.eval(model.parameters())

    total_params = model.n_params()
    print(f"parameters: {total_params / 1e6:.1f}M total, "
          f"{model.n_params(trunk_only=True) / 1e6:.1f}M trunk")

    master = prune_empty(cast_tree(model.parameters(), mx.float32))
    model.update(cast_tree(master, compute_dtype))
    mx.eval(model.parameters())

    if cli.save_initialized:
        save_checkpoint(
            cli.save_initialized,
            model,
            master,
            None,
            0,
            cfg,
            args,
            include_optimizer=False,
        )
        print(
            f"saved seed-matched untrained control at {cli.save_initialized}; "
            "no data was opened and no optimizer step ran"
        )
        return

    schedule = build_schedule(cfg)
    # Weight decay is applied by hand below, not by the optimizer, so that it
    # reaches matrices only. A single AdamW over the whole tree also decays every
    # RMSNorm scale, and over this schedule the decay-only multiplier on those
    # one-dimensional parameters is about exp(-0.63), roughly halving them. Norm
    # scales are not a capacity knob and shrinking them is not regularisation,
    # it is a slow drift in the function being learned.
    weight_decay = cfg.get("weight_decay", 0.1)
    # bias_correction defaults to False in MLX 0.32, which is not the Adam most
    # references describe. Without it the second moment is initialised at zero and
    # stays underestimated for roughly 1/(1-beta2) steps, which at beta2 0.95 is
    # about 20 steps, so the effective step size is inflated exactly where a from
    # scratch run is least stable. Turned on deliberately rather than inherited.
    # The 1000 step warmup would mask most of it either way, but "masked by the
    # warmup" is not a reason to run a different optimizer than the one written
    # down in the config.
    optimizer = optim.AdamW(
        learning_rate=schedule,
        betas=[cfg.get("beta1", 0.9), cfg.get("beta2", 0.95)],
        eps=cfg.get("eps", 1e-8),
        weight_decay=0.0,
        bias_correction=cfg.get("bias_correction", True),
    )

    start_step = 0
    resume_meta = None
    if cli.resume:
        master, resume_meta = load_checkpoint(cli.resume, model, optimizer)
        start_step = resume_meta["step"]
        validate_resume_sampling_contract(resume_meta, cfg)
        model.update(cast_tree(master, compute_dtype))
        mx.eval(model.parameters())
        print(f"resumed from {cli.resume} at step {start_step}")

    span = cfg["seq_len"] + 1 + cfg["mtp_depth"]
    train_ds = ShardDataset(cfg["data_index"], "train", span, seed=cfg.get("seed", 1337))
    val_ds = ShardDataset(cfg["data_index"], "val", span, seed=cfg.get("seed", 1337) + 1)
    reset_steps = sampler_reset_steps(cfg)
    if resume_meta is not None:
        expected_batches = sampler_batches_since_reset(
            start_step,
            cfg["grad_accum"],
            reset_steps,
        )
        saved_sampler = resume_meta.get("train_sampler")
        if saved_sampler is None:
            train_ds.advance_batches(cfg["micro_batch"], expected_batches)
            print(
                "reconstructed legacy training sampler at "
                f"{expected_batches} batches since its last reset"
            )
        else:
            if saved_sampler.get("reset_steps") != reset_steps:
                raise ValueError(
                    "checkpoint sampler reset schedule differs from config"
                )
            train_ds.restore_sampler_state(
                saved_sampler,
                cfg["micro_batch"],
                expected_batches,
            )
            print(
                "restored exact training sampler at "
                f"{expected_batches} batches since its last reset"
            )

    tokens_per_step = cfg["micro_batch"] * cfg["grad_accum"] * cfg["seq_len"]
    print(f"tokens/step: {tokens_per_step:,}  "
          f"total: {tokens_per_step * cfg['max_steps'] / 1e9:.2f}B over {cfg['max_steps']:,} steps")

    out_dir = cfg["out_dir"]
    os.makedirs(out_dir, exist_ok=True)
    log_path = os.path.join(out_dir, "log.jsonl")
    snapshot_interval = cfg.get("snapshot_interval", 0)
    if (
        resume_meta is not None
        and isinstance(snapshot_interval, int)
        and not isinstance(snapshot_interval, bool)
        and snapshot_interval > 0
        and start_step > 0
        and start_step % snapshot_interval == 0
        and os.path.abspath(cli.resume) == os.path.abspath(default_checkpoint)
    ):
        # The legacy writer installs ckpt_latest and then the immortal snapshot
        # before printing its snapshot JSON. A crash in that narrow window leaves
        # a valid resume point whose evidence marker is absent. Recreate or verify
        # the exact snapshot first, then print the legacy marker only when its
        # active log lineage lacks one. This runs before the resumed loop, so a
        # repeated crash at the same boundary converges without duplicate markers.
        sys.stdout.flush()
        snapshot_marker = prepare_snapshot_boundary_recovery(
            cli.resume,
            os.path.join(out_dir, f"ckpt_step{start_step:06d}"),
            os.path.join(out_dir, "train.log"),
            cli.resume,
            start_step,
            resume_meta,
        )
        if snapshot_marker is not None:
            print(snapshot_marker.decode("ascii"), flush=True)

    def loss_fn(m, batch):
        total, _, _ = m.loss(batch, cfg["mtp_weight"])
        return total

    grad_fn = nn.value_and_grad(model, loss_fn)

    def micro_step(batch):
        loss, grads = grad_fn(model, batch)
        return loss, prune_empty(cast_tree(grads, mx.float32))

    if cli.compile:
        # Shapes are constant across micro steps, so there is exactly one trace to
        # build. Compilation remains opt-in so benchmark comparisons stay explicit.
        micro_step = mx.compile(micro_step, inputs=model.state, outputs=model.state)
        print("compiled the micro step")

    t0 = time.time()
    window_tokens = 0
    for step in range(start_step, cfg["max_steps"]):
        if step in reset_steps:
            train_ds.reset_sampler()
            record = {
                "step": step,
                "sampler_reset": True,
                "seed": cfg.get("seed", 1337),
            }
            print(json.dumps(record))
            with open(log_path, "a") as f:
                f.write(json.dumps(record) + "\n")
        accum_grads = None
        loss_acc = 0.0

        for _ in range(cfg["grad_accum"]):
            batch = mx.array(train_ds.batch(cfg["micro_batch"]))
            loss, grads = micro_step(batch)
            accum_grads = grads if accum_grads is None else tree_map(
                lambda a, b: a + b, accum_grads, grads
            )
            # Force the accumulation graph before the next micro batch. MLX is lazy,
            # so without this the whole grad_accum loop stays unevaluated and every
            # micro batch's activations are held live at once. At grad_accum 16 that
            # is an out of memory kill, not a slowdown.
            mx.eval(accum_grads, loss)
            loss_acc += float(loss)

        accum_grads = tree_map(lambda g: g / cfg["grad_accum"], accum_grads)
        accum_grads, grad_norm = optim.clip_grad_norm(accum_grads, cfg.get("grad_clip", 1.0))

        master = optimizer.apply_gradients(accum_grads, master)
        if weight_decay:
            # Decoupled AdamW decay, matrices only. ndim > 1 selects the
            # embeddings and projections and skips the norm scales.
            decay_now = float(schedule(optimizer.step)) * weight_decay
            master = tree_map(
                lambda p: p * (1.0 - decay_now) if p.ndim > 1 else p, master
            )
        model.update(cast_tree(master, compute_dtype))
        mx.eval(master, model.parameters(), optimizer.state)

        window_tokens += tokens_per_step
        loss_val = loss_acc / cfg["grad_accum"]

        if (step + 1) % cfg.get("log_interval", 10) == 0:
            dt = time.time() - t0
            tps = window_tokens / dt
            lr_now = float(schedule(optimizer.step)) if callable(schedule) else cfg["lr"]
            record = {
                "step": step + 1,
                "loss": round(loss_val, 4),
                "grad_norm": round(float(grad_norm), 3),
                "lr": lr_now,
                "tok_per_sec": round(tps),
                "eta_hours": round((cfg["max_steps"] - step - 1) * tokens_per_step / tps / 3600, 2),
                "peak_gb": round(peak_memory_gb(), 2),
            }
            print(json.dumps(record))
            with open(log_path, "a") as f:
                f.write(json.dumps(record) + "\n")
            t0, window_tokens = time.time(), 0

        if (step + 1) % cfg["eval_interval"] == 0:
            val_main, val_mtp = evaluate(model, val_ds, cfg, compute_dtype)
            record = {
                "step": step + 1,
                "val_main": round(val_main, 4),
                "val_ppl": round(math.exp(min(val_main, 20)), 2),
                "val_mtp": [round(v, 4) for v in val_mtp],
            }
            print(json.dumps(record))
            with open(log_path, "a") as f:
                f.write(json.dumps(record) + "\n")
            t0, window_tokens = time.time(), 0

        if (step + 1) % cfg["ckpt_interval"] == 0 or step + 1 == cfg["max_steps"]:
            save_checkpoint(
                os.path.join(out_dir, "ckpt_latest"),
                model,
                master,
                optimizer,
                step + 1,
                cfg,
                args,
                train_sampler=checkpoint_sampler_state(
                    train_ds,
                    cfg["micro_batch"],
                    reset_steps,
                ),
            )

        # Periodic immortal snapshots, separate from ckpt_latest which is
        # overwritten. Acceptance measured against tokens seen is a curve the
        # research needs and cannot reconstruct afterwards from a single final
        # checkpoint, so the snapshots have to be taken while the run is going.
        snap = cfg.get("snapshot_interval", 0)
        if snap and (step + 1) % snap == 0:
            save_checkpoint(
                os.path.join(out_dir, f"ckpt_step{step + 1:06d}"),
                model,
                master,
                optimizer,
                step + 1,
                cfg,
                args,
                train_sampler=checkpoint_sampler_state(
                    train_ds,
                    cfg["micro_batch"],
                    reset_steps,
                ),
            )
            print(json.dumps({"step": step + 1, "snapshot": True}))

    print("done")


if __name__ == "__main__":
    main()