"""
Step 3: train the model.

    python train.py --config config/pilot.yaml

Progress goes to the console and to runs/<name>/: log.csv (metrics over time), status.json (live
state) and samples.json (what the model hears on a few validation ayahs at each evaluation).
`python dashboard.py` shows them as live charts. The best checkpoint (lowest letter error rate on
the validation clips) is runs/<name>/best.pt.

Train in as many sessions as you like: running the same command again continues exactly where
it stopped (same batch, optimizer, learning rate schedule and data order). To pause:
  - Ctrl+C, the dashboard's Pause button, or a file named "stop" in runs/<name>/: saves and exits
    within seconds;
  - --stop-after 2 (hours) or --stop-at 07:30: pauses by itself, e.g. before you need the laptop;
  - closing the window or a power cut: loses at most `checkpoint_minutes` (10) of training.
"""
import csv
import hashlib
import json
import math
import os
import shutil
import signal
import sys
import time

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

from quran_asr import config as config_mod
from quran_asr import data, model as model_mod, text

LOG_COLUMNS = ["step", "epoch", "loss", "lr", "valid_cer", "valid_wer", "hours_seen", "elapsed_min",
               "grad_norm", "gpu_peak_gb", "audio_h_per_h", "valid_skeleton", "valid_hafs_cer", "valid_warsh_cer"]


class EpochBatches:
    """
    Wraps the duration sampler so every index carries the epoch (fresh augmentation per epoch) and
    whether to augment: the first `augment_after` steps use clean audio, so the model first learns
    which sounds are which letters. Augmenting from the very first step kept it from ever learning.
    """

    def __init__(self, sampler, augment_after=0):
        self.sampler = sampler
        self.augment_after = augment_after
        self.skip = 0

    def __iter__(self):
        epoch = self.sampler.epoch
        for k, batch in enumerate(self.sampler):
            if k >= self.skip:
                on = epoch * len(self.sampler) + k >= self.augment_after
                yield [(i, epoch, on) for i in batch]
        self.skip = 0

    def __len__(self):
        return len(self.sampler) - self.skip


class EpochDataset(data.RecitationDataset):
    def __getitem__(self, key):
        i, epoch, on = key if isinstance(key, tuple) else (key, 0, True)
        self.epoch, self.augment_on = epoch, on
        return super().__getitem__(i)


def worker_init(_):
    torch.set_num_threads(1)
    # Ctrl+C (and Ctrl+Break) reach every process in the console window. Only the main process
    # should react: it saves a checkpoint, then shuts the workers down itself. A worker dying first
    # would turn the clean stop into a "worker exited unexpectedly" crash without a checkpoint.
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    if hasattr(signal, "SIGBREAK"):
        signal.signal(signal.SIGBREAK, signal.SIG_IGN)


class Pause(Exception):
    """Stop training now, keeping a checkpoint to resume from."""


def _interrupt(*_):
    raise KeyboardInterrupt


def deadline_from(args):
    """--stop-after hours / --stop-at HH:MM as a time.time() value, or None."""
    deadline = time.time() + args.stop_after * 3600 if args.stop_after else None
    if args.stop_at:
        hour, minute = (int(v) for v in args.stop_at.split(":"))
        now = time.localtime()
        at = time.mktime((now.tm_year, now.tm_mon, now.tm_mday, hour, minute, 0, 0, 0, -1))
        if at <= time.time():
            at += 24 * 3600  # that time tomorrow
        deadline = min(deadline or at, at)
    return deadline


def feature_stats(entries, count=400):
    """Global mean/std of the features over a sample of clips (stored inside the model)."""
    ds = data.RecitationDataset(entries[:: max(1, len(entries) // count)][:count])
    total, total_sq, n = 0, 0, 0
    for k in range(len(ds)):
        f, _, _ = ds[k]
        total = total + f.sum(0)
        total_sq = total_sq + (f ** 2).sum(0)
        n += f.shape[0]
    mean = total / n
    std = (total_sq / n - mean ** 2).clamp(min=1e-4).sqrt()
    return mean, std


def evaluate(model, loader, device, amp_dtype, sample_ids=()):
    """Error rates on the validation clips (overall and per riwaya) and the text heard on sample_ids."""
    model.eval()
    totals, heard, sample_ids = {}, {}, set(sample_ids)
    with torch.no_grad():
        for x, lens, _tokens, _tlens, idx in loader:
            with torch.autocast(device.type, dtype=amp_dtype, enabled=amp_dtype is not None):
                log_probs, out_lens = model(x.to(device), lens.to(device))
            for ids, i in zip(model_mod.greedy_decode(log_probs, out_lens), idx):
                entry = loader.dataset.entries[i]
                hyp = text.decode(ids)
                counts = text.error_counts(hyp, entry["text"])
                for key in ("all", entry["riwaya"]):
                    acc = totals.setdefault(key, [0] * 6)
                    for j, c in enumerate(counts):
                        acc[j] += c
                if i in sample_ids:
                    heard[i] = hyp
    model.train()

    def rate(key, j):
        return totals[key][j] / max(1, totals[key][j + 1]) if key in totals else None

    scores = {"cer": rate("all", 0), "wer": rate("all", 2), "skeleton": rate("all", 4),
              "hafs_cer": rate("hafs", 0), "warsh_cer": rate("warsh", 0)}
    return scores, heard


def pick_samples(entries, count=12):
    """A fixed handful of validation clips for samples.json: Warsh first, then one reciter at a time."""
    by_reciter = {}
    for i, e in enumerate(entries):
        if 3 <= e["duration"] <= 20:
            by_reciter.setdefault((e["riwaya"] != "warsh", e["reciter"]), []).append(i)
    queues = [by_reciter[key] for key in sorted(by_reciter)]
    picked = []
    while len(picked) < count and any(queues):
        for q in queues:
            if q and len(picked) < count:
                picked.append(q.pop(len(q) // 2))
    return picked


def write_json(path, value):
    """Atomic write. Skipped (the next update retries) if a reader holds the file open on Windows."""
    tmp = path + ".tmp"
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(value, f, ensure_ascii=False)
        os.replace(tmp, path)
    except OSError:
        pass


def lr_at(step, total, cfg):
    if step < cfg.warmup_steps:
        return cfg.lr * (step + 1) / cfg.warmup_steps
    progress = min(1.0, (step - cfg.warmup_steps) / max(1, total - cfg.warmup_steps))
    return cfg.lr * (cfg.min_lr_ratio + (1 - cfg.min_lr_ratio) * 0.5 * (1 + math.cos(math.pi * progress)))


def main():
    parser = config_mod.from_args(__doc__)
    parser.add_argument("--max-steps", type=int, default=0, help="stop early (for quick tests)")
    parser.add_argument("--stop-after", type=float, default=0, help="pause after this many hours")
    parser.add_argument("--stop-at", default="", help="pause at this time of day, e.g. 07:30")
    args = parser.parse_args()
    if os.name == "nt":
        # Keep Windows from sleeping while training (it counts only keyboard/mouse as activity);
        # the screen may still turn off. Ends by itself when train.py exits.
        import ctypes
        ctypes.windll.kernel32.SetThreadExecutionState(0x80000000 | 0x00000001)  # CONTINUOUS | SYSTEM_REQUIRED
    if hasattr(signal, "SIGBREAK"):
        signal.signal(signal.SIGBREAK, _interrupt)  # Ctrl+Break saves too, like Ctrl+C
    cfg = config_mod.load(args.config)
    t = cfg.train
    torch.manual_seed(t.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type == "cpu":
        print("WARNING: no CUDA GPU found, training on the CPU will be very slow")
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    amp_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16}.get(t.precision) if device.type == "cuda" else None

    mdir = config_mod.manifest_dir(cfg)
    train_entries = data.read_manifest(os.path.join(mdir, "train.jsonl"))
    valid_entries = data.read_manifest(os.path.join(mdir, "valid.jsonl"))
    print(f"train {len(train_entries)} clips ({data.hours(train_entries):.1f} h), "
          f"valid {len(valid_entries)} clips ({data.hours(valid_entries):.1f} h)")

    os.makedirs(cfg.paths.run_dir, exist_ok=True)
    shutil.copy(cfg.config_path, os.path.join(cfg.paths.run_dir, "config.yaml"))
    shutil.copy(os.path.join(mdir, "tokens.txt"), os.path.join(cfg.paths.run_dir, "tokens.txt"))

    model = model_mod.build(cfg.model, len(text.VOCAB)).to(device)
    params = model_mod.count_parameters(model)
    print(f"model: {params / 1e6:.1f}M parameters (~{params / 1e6:.0f} MB as int8 ONNX)")
    decay = [p for n, p in model.named_parameters() if p.ndim >= 2]
    no_decay = [p for n, p in model.named_parameters() if p.ndim < 2]
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": t.weight_decay},
                             {"params": no_decay, "weight_decay": 0.0}], lr=t.lr, betas=(0.9, 0.98))
    scaler = torch.amp.GradScaler(enabled=amp_dtype == torch.float16)

    # Ordinary people (Iqra'Eval, tlog) are sampled more or less often so that they make up
    # `ordinary_share` of each epoch's audio, the professional reciters the rest.
    ordinary = ("iqraeval", "tlog")
    source_weights = {}
    share = getattr(cfg.data, "ordinary_share", None)
    if share is None:  # older configs: only Iqra'Eval
        share, ordinary = getattr(cfg.data, "iqraeval_share", 0) or 0, ("iqraeval",)
    hours = {s: sum(e["duration"] for e in train_entries if e.get("source") == s) / 3600 for s in ordinary}
    if share and sum(hours.values()):
        rest = sum(e["duration"] * (cfg.data.warsh_weight if e["riwaya"] == "warsh" else 1.0)
                   for e in train_entries if e.get("source") not in ordinary) / 3600
        weight = share * rest / ((1 - share) * sum(hours.values()))
        source_weights = {s: weight for s in ordinary}
        print(f"ordinary people: {', '.join(f'{s} {h:.0f} h' for s, h in hours.items())}, sampled {weight:.2f}x "
              f"per epoch = {share:.0%} of each epoch's audio; professional reciters {rest:.0f} h = {1 - share:.0%}")
    sampler = data.DurationBatchSampler(train_entries, t.max_batch_seconds, cfg.data.warsh_weight, True, t.seed,
                                        source_weights)
    batches = EpochBatches(sampler, getattr(cfg.augment, "start_after_steps", 0) or 0)
    train_ds = EpochDataset(train_entries, cfg.augment, cfg.spec_augment, seed=t.seed)
    loader = DataLoader(train_ds, batch_sampler=batches, collate_fn=data.collate, num_workers=t.num_workers,
                        persistent_workers=t.num_workers > 0, pin_memory=device.type == "cuda",
                        worker_init_fn=worker_init, prefetch_factor=4 if t.num_workers > 0 else None)
    valid_loader = DataLoader(data.RecitationDataset(valid_entries), batch_sampler=data.DurationBatchSampler(
        valid_entries, t.max_batch_seconds, 1.0, False), collate_fn=data.collate,
        num_workers=min(4, t.num_workers), worker_init_fn=worker_init)

    total_steps = t.epochs * len(sampler)
    # Identifies the training data, so a resumed run notices new data (e.g. Iqra'Eval added).
    data_id = hashlib.md5("\n".join(e["path"] for e in train_entries).encode("utf-8")).hexdigest()
    valid_id = hashlib.md5("\n".join(e["path"] for e in valid_entries).encode("utf-8")).hexdigest()
    inter_layers = tuple(getattr(t, "inter_ctc_layers", None) or ())
    inter_weight = getattr(t, "inter_ctc_weight", 0.3)

    # Live state for the dashboard, rewritten at every log line.
    status_path = os.path.join(cfg.paths.run_dir, "status.json")
    status = dict(name=cfg.name, phase="starting", device=torch.cuda.get_device_name(0) if device.type == "cuda"
                  else "CPU", params_m=round(params / 1e6, 2), total_steps=total_steps, epochs=t.epochs,
                  steps_per_epoch=len(sampler), train_clips=len(train_entries), valid_clips=len(valid_entries),
                  train_hours=round(data.hours(train_entries), 2), valid_hours=round(data.hours(valid_entries), 2),
                  hours_by_reciter=data.summarize(train_entries), model=vars(cfg.model), train=vars(t),
                  session_started_at=time.time())

    def publish(**changes):
        status.update(changes, updated_at=time.time())
        write_json(status_path, status)

    step, epoch, best_cer, best_step = 0, 0, float("inf"), None
    rewarm_from = None
    seen_seconds, elapsed_before = 0.0, 0.0
    last_path = os.path.join(cfg.paths.run_dir, "last.pt")
    if os.path.exists(last_path):
        ckpt = torch.load(last_path, map_location=device, weights_only=False)
        model.load_state_dict(ckpt["model"])
        opt.load_state_dict(ckpt["opt"])
        scaler.load_state_dict(ckpt["scaler"])
        step, epoch, best_cer = ckpt["step"], ckpt["epoch"], ckpt["best_cer"]
        best_step = ckpt.get("best_step")
        if ckpt.get("valid_id") != valid_id:
            # Scores on a different validation set aren't comparable (e.g. ordinary readers added):
            # the best checkpoint is chosen afresh from the next validation on.
            best_cer, best_step = float("inf"), None
            print("the validation set changed: best.pt will be replaced at the next validation")
        seen_seconds, elapsed_before = ckpt.get("seen_seconds", 0.0), ckpt.get("elapsed", 0.0)
        batches.skip = ckpt["batch_in_epoch"]
        rewarm_from = ckpt.get("rewarm_from")
        if "data_id" in ckpt:
            changed = ckpt["data_id"] != data_id
        else:  # a checkpoint from before data_id: the data changed if prepare.py ran after it was saved
            changed = os.path.getmtime(os.path.join(mdir, "train.jsonl")) > os.path.getmtime(last_path)
        if changed:
            # New training data: the batch position in the old epoch means nothing any more, so a
            # new epoch starts, and the learning rate warms up again (Adam's statistics and the
            # model meet different data).
            if batches.skip:
                epoch += 1
            batches.skip = 0
            rewarm_from = step
            print(f"the training data changed: continuing from step {step} with the new data "
                  f"(epoch {epoch + 1}, learning rate warming up again)")
        sampler.set_epoch(epoch)
        # The remaining epochs at the current data's length (it changes when data or epochs do).
        total_steps = step - batches.skip + (t.epochs - epoch) * len(sampler)
        planned = ckpt.get("total_steps")
        if not changed and planned and abs(total_steps - planned) > 0.05 * planned:
            rewarm_from = step  # more (or fewer) epochs: the schedule stretched, ease into it
            print(f"the number of epochs changed ({planned} -> {total_steps} steps): learning rate warming up again")
        status.update(total_steps=total_steps, steps_per_epoch=len(sampler))
        print(f"resumed at step {step} (epoch {epoch + 1}/{t.epochs}, batch {batches.skip}); "
              f"{total_steps} steps in total")
    else:
        publish(phase="preparing")
        print("computing feature statistics...")
        mean, std = feature_stats(train_entries)
        model.feat_mean.copy_(mean.to(device))
        model.feat_std.copy_(std.to(device))

    session_start = time.time()

    def elapsed():
        """Training time over all sessions, in seconds."""
        return elapsed_before + time.time() - session_start

    log_path = os.path.join(cfg.paths.run_dir, "log.csv")
    columns = LOG_COLUMNS
    if os.path.exists(log_path) and os.path.getsize(log_path) > 0:
        with open(log_path, newline="", encoding="utf-8") as f:
            columns = next(csv.reader(f))  # a log started by an older train.py keeps its columns
    new_log = columns is LOG_COLUMNS
    log_file = open(log_path, "a", newline="", encoding="utf-8")
    log = csv.DictWriter(log_file, fieldnames=columns, extrasaction="ignore")
    if new_log:
        log.writeheader()

    sample_ids = pick_samples(valid_entries)
    samples_path = os.path.join(cfg.paths.run_dir, "samples.json")
    clips = [dict({k: valid_entries[i][k] for k in ("reciter", "riwaya", "sura", "aya", "duration")},
                  ref=valid_entries[i]["text"]) for i in sample_ids]
    samples = {"clips": clips, "evals": []}
    if os.path.exists(samples_path):
        try:
            with open(samples_path, encoding="utf-8") as f:
                old = json.load(f)
            if old.get("clips") == clips:
                samples["evals"] = [e for e in old["evals"] if e["step"] <= step]
        except (OSError, ValueError, KeyError):
            pass
    # What the model hears on every validation clip, one file per validation, so the dashboard can
    # show any of them ("Show other ayahs") with their whole history.
    heard_dir = os.path.join(cfg.paths.run_dir, "heard")
    os.makedirs(heard_dir, exist_ok=True)
    write_json(os.path.join(heard_dir, "clips.json"), [
        dict({k: e[k] for k in ("reciter", "riwaya", "sura", "aya", "duration")}, ref=e["text"]) for e in valid_entries])
    for name in os.listdir(heard_dir):
        if name.startswith("step_") and int(name[5:12]) > step:
            os.remove(os.path.join(heard_dir, name))  # validations a resumed run will redo
    publish(phase="training", step=step, epoch=min(epoch + 1, t.epochs), hours_seen=round(seen_seconds / 3600, 2),
            elapsed_min=round(elapsed() / 60, 1), best_cer=None if best_step is None else best_cer,
            best_step=best_step)

    def save(batch_in_epoch, keep=True):
        """last.pt to resume from; keep=True also keeps a step_*.pt for export.py to average."""
        state = dict(model=model.state_dict(), opt=opt.state_dict(), scaler=scaler.state_dict(), step=step,
                     epoch=epoch, batch_in_epoch=batch_in_epoch, best_cer=best_cer, best_step=best_step,
                     seen_seconds=seen_seconds, elapsed=elapsed(), model_config=vars(cfg.model), vocab=text.VOCAB,
                     data_id=data_id, valid_id=valid_id, rewarm_from=rewarm_from, total_steps=total_steps)
        torch.save(state, last_path + ".tmp")
        os.replace(last_path + ".tmp", last_path)
        if not keep:
            return
        step_path = os.path.join(cfg.paths.run_dir, f"step_{step:07d}.pt")
        torch.save({"model": state["model"], "model_config": state["model_config"], "vocab": text.VOCAB, "step": step},
                   step_path + ".tmp")
        os.replace(step_path + ".tmp", step_path)
        steps = sorted(f for f in os.listdir(cfg.paths.run_dir) if f.startswith("step_") and f.endswith(".pt"))
        for old in steps[: -t.keep_checkpoints]:
            os.remove(os.path.join(cfg.paths.run_dir, old))

    def validate():
        nonlocal best_cer, best_step
        publish(phase="validating", step=step)
        scores, heard = evaluate(model, valid_loader, device, amp_dtype, range(len(valid_entries)))
        write_json(os.path.join(heard_dir, f"step_{step:07d}.json"),
                   {"step": step, "hyps": [heard.get(i, "") for i in range(len(valid_entries))]})
        print(f"  valid: letter error {scores['cer'] * 100:.2f}%  word error {scores['wer'] * 100:.2f}%", flush=True)
        log.writerow({"step": step, "epoch": min(epoch + 1, t.epochs), "hours_seen": round(seen_seconds / 3600, 2),
                      "elapsed_min": round(elapsed() / 60, 1),
                      **{f"valid_{k}": "" if v is None else round(v, 5) for k, v in scores.items()}})
        log_file.flush()
        samples["evals"].append({"step": step, "hyps": [heard.get(i, "") for i in sample_ids]})
        write_json(samples_path, samples)
        if scores["cer"] < best_cer:
            best_cer, best_step = scores["cer"], step
            torch.save({"model": model.state_dict(), "model_config": vars(cfg.model), "vocab": text.VOCAB,
                        "step": step, "valid_cer": scores["cer"]}, os.path.join(cfg.paths.run_dir, "best.pt"))
            print("  new best checkpoint", flush=True)
        publish(phase="training", valid=scores, best_cer=best_cer, best_step=best_step)

    stop_path = os.path.join(cfg.paths.run_dir, "stop")
    if os.path.exists(stop_path):
        os.remove(stop_path)  # left from the previous pause
    deadline = deadline_from(args)
    if deadline:
        print(f"will pause at {time.strftime('%H:%M', time.localtime(deadline))} "
              f"(in {(deadline - time.time()) / 3600:.1f} h)")
    publish(stop_at=deadline)
    every = getattr(t, "checkpoint_minutes", 10) * 60
    last_checkpoint = time.time()

    losses, norms, session_steps, skipped, overall_peak = [], [], 0, 0, 0.0
    last_log_time, last_log_seen = time.time(), seen_seconds
    done = batches.skip if step else 0  # batches of the current epoch already trained
    model.train()
    try:
        while epoch < t.epochs:
            for k, (x, lens, targets, target_lens, idx) in enumerate(loader, start=done):
                for g in opt.param_groups:
                    g["lr"] = lr_at(step, total_steps, t)
                    if rewarm_from is not None:  # after new data: warm up again from the resume point
                        g["lr"] *= min(1.0, (step - rewarm_from + 1) / t.warmup_steps)
                x, lens = x.to(device, non_blocking=True), lens.to(device, non_blocking=True)
                opt.zero_grad(set_to_none=True)
                try:
                    with torch.autocast(device.type, dtype=amp_dtype, enabled=amp_dtype is not None):
                        log_probs, out_lens, *rest = model(x, lens, inter_layers)
                    inter = rest[0] if rest else []
                    tgt, tlen = targets.to(device), target_lens.to(device)
                    letters = target_lens.sum().clamp(min=1)

                    def ctc(lp):  # per letter: a several-minute ayah weighs no more per letter than a short one
                        return F.ctc_loss(lp.transpose(0, 1), tgt, out_lens, tlen, blank=0, reduction="sum",
                                          zero_infinity=True) / letters

                    loss = ctc(log_probs)
                    if inter:
                        # Intermediate CTC: the same head also reads a few middle layers, so the early
                        # layers get a short, direct learning signal. Without it the 16-layer model can
                        # stay stuck predicting the same letters for every clip.
                        total = (1 - inter_weight) * loss + inter_weight * sum(ctc(lp) for lp in inter) / len(inter)
                    else:
                        total = loss
                    scaler.scale(total).backward()
                except torch.OutOfMemoryError:
                    # A rare oversized batch: skip it rather than stop the run.
                    del x, lens
                    log_probs = out_lens = loss = None
                    opt.zero_grad(set_to_none=True)
                    torch.cuda.empty_cache()
                    skipped += 1
                    print(f"  skipped a batch that ran out of GPU memory ({len(idx)} clips); "
                          "lower max_batch_seconds if this happens often", flush=True)
                    continue
                scaler.unscale_(opt)
                norms.append(float(torch.nn.utils.clip_grad_norm_(model.parameters(), t.grad_clip)))
                scaler.step(opt)
                scaler.update()
                step += 1
                done = k + 1
                session_steps += 1
                seen_seconds += float(lens.sum()) / 100.0
                losses.append(loss.item())

                if step % t.log_every_steps == 0:
                    now = time.time()
                    rate = session_steps / max(1e-6, now - session_start)  # steps per second in this session
                    eta_h = (total_steps - step) / max(1e-6, rate) / 3600
                    avg = sum(losses) / len(losses)
                    finite = [g for g in norms if math.isfinite(g)]
                    grad = sum(finite) / len(finite) if finite else float("nan")
                    speed = (seen_seconds - last_log_seen) / max(1e-6, now - last_log_time)  # audio h per hour
                    losses, norms, last_log_time, last_log_seen = [], [], now, seen_seconds
                    peak = None
                    if device.type == "cuda":
                        peak = torch.cuda.max_memory_allocated() / 2**30  # since the previous log line
                        torch.cuda.reset_peak_memory_stats()
                        overall_peak = max(overall_peak, peak)
                    lr = float(f"{opt.param_groups[0]['lr']:.4g}")
                    gpu = f"  GPU peak {overall_peak:.1f} GB" if peak is not None else ""
                    print(f"step {step}/{total_steps}  epoch {epoch + 1}/{t.epochs}  loss/letter {avg:.3f}  "
                          f"lr {lr:.2e}  {seen_seconds / 3600:.1f} h of audio  ETA {eta_h:.1f} h{gpu}", flush=True)
                    row = {"step": step, "epoch": epoch + 1, "loss": round(avg, 4), "lr": lr,
                           "hours_seen": round(seen_seconds / 3600, 2), "elapsed_min": round(elapsed() / 60, 1),
                           "grad_norm": round(grad, 3), "audio_h_per_h": round(speed),
                           "gpu_peak_gb": "" if peak is None else round(peak, 2)}
                    log.writerow(row)
                    log_file.flush()
                    publish(phase="training", **{k: v for k, v in row.items() if k != "gpu_peak_gb"},
                            eta_h=round(eta_h, 3), gpu_peak_gb=round(overall_peak, 2), skipped_batches=skipped)

                if step % t.eval_every_steps == 0:
                    validate()
                    save(done)
                    last_checkpoint = time.time()
                elif time.time() - last_checkpoint >= every:
                    save(done, keep=False)  # a crash or a closed window loses at most this much
                    last_checkpoint = time.time()

                if deadline and time.time() >= deadline:
                    raise Pause(f"--stop-at {args.stop_at}" if args.stop_at else f"--stop-after {args.stop_after:g}")
                if step % 10 == 0 and os.path.exists(stop_path):
                    raise Pause("pause requested")
                if args.max_steps and step >= args.max_steps:
                    raise Pause("--max-steps")
            epoch += 1
            done = 0
            sampler.set_epoch(epoch)
            save(0)
            last_checkpoint = time.time()
    except (KeyboardInterrupt, Pause) as e:
        reason = str(e) or "Ctrl+C"
        print(f"pausing ({reason}): saving a checkpoint...", flush=True)
        save(done, keep=False)
        if os.path.exists(stop_path):
            os.remove(stop_path)
        log_file.close()
        publish(phase="stopped", stop_reason=reason)
        print(f"saved at step {step}. Run the same command to continue.")
        return 0
    except BaseException as e:
        log_file.close()
        try:
            message = f"{type(e).__name__}: {e}"
        except Exception:  # some exceptions can't format themselves after crossing from a worker
            message = type(e).__name__
        publish(phase="crashed", error=message)
        raise

    try:
        print("final validation...")
        validate()
        print(f"final valid: best letter error {best_cer * 100:.2f}% (step {best_step})")
    finally:
        log_file.close()
        publish(phase="finished", eta_h=0)


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