"""
Trains KairoGPT from scratch on backend/learn/data/corpus_v2.txt (falls back to
corpus.txt if corpus_v2.txt doesn't exist yet).
Run prepare_corpus.py then prepare_code_corpus.py first.

Usage: python train_base.py
"""
import json
import logging
import math
import os
import time
from pathlib import Path

# Must be set before the first CUDA allocation: lets the allocator grow/shrink
# its arena instead of crashing on fragmentation, which matters a lot on a 6 GB
# card running a model that fills most of VRAM.
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import numpy as np
import psutil
import torch

from model import KairoGPT, KairoGPTConfig
from tokenizer import CharTokenizer

logging.basicConfig(level=logging.INFO)
_LOG_FILE = Path(__file__).parent / "pipeline_full.err.log"
_fh = logging.FileHandler(_LOG_FILE, mode="a", encoding="utf-8")
_fh.setFormatter(logging.Formatter("%(levelname)s:%(name)s:%(message)s"))
logging.getLogger().addHandler(_fh)
logger = logging.getLogger("kairo.learn.train")

DATA_DIR = Path(__file__).parent / "data"
CKPT_DIR = Path(__file__).parent / "checkpoints"
CORPUS_V2_FILE = DATA_DIR / "corpus_v2.txt"
CORPUS_FILE = DATA_DIR / "corpus.txt"
TOKENIZER_FILE = CKPT_DIR / "tokenizer.json"
MODEL_FILE = CKPT_DIR / "base.pt"
INPROGRESS_FILE = CKPT_DIR / "base_inprogress.pt"

# ~58M model: bigger => base.pt ~230MB (vs 80MB) and smarter, still trains fine
# on a GTX 1060 in fp32 (~1.5 s/step). Reuses the finished BPE corpus.
BLOCK_SIZE = 384
BATCH_SIZE = 1
GRAD_ACCUM = 32
N_LAYER = 8
N_HEAD = 8
N_EMBD = 512
MAX_ITERS = 60000  # high on purpose: model keeps improving, user stops when happy
EVAL_INTERVAL = 250
# Save both the inference-ready base.pt AND the resume checkpoint this often, so
# stopping loses at most SAVE_INTERVAL steps ("resumable download") and a fresh
# testable base.pt exists every ~1 min. Tighter than this thrashes the HDD.
SAVE_INTERVAL = 50
LEARNING_RATE = 3e-4
# Standard GPT recipe (nanoGPT/Chinchilla style): linear warmup then cosine
# decay to 10% -- avoids early divergence and squeezes better final loss out
# of the same steps.
WARMUP_ITERS = 2000
MIN_LR = LEARNING_RATE / 10


def lr_at(it: int) -> float:
    if it < WARMUP_ITERS:
        return LEARNING_RATE * (it + 1) / WARMUP_ITERS
    progress = (it - WARMUP_ITERS) / max(1, MAX_ITERS - WARMUP_ITERS)
    return MIN_LR + 0.5 * (LEARNING_RATE - MIN_LR) * (1 + math.cos(math.pi * progress))
STEP_SLEEP = 0.02  # ponytail: throttle so training never hogs GPU/CPU, raise if PC still lags
IDS_DAT_FILE = CKPT_DIR / "corpus_ids.dat"
IDS_META_FILE = CKPT_DIR / "corpus_ids.meta.json"


def gpu_smoke_test() -> bool:
    """Small matmul+backward on CUDA -- catches CUBLAS/driver issues before a
    multi-week run commits to a device that will crash mid-training."""
    try:
        a = torch.randn(256, 256, device="cuda", requires_grad=True)
        b = torch.randn(256, 256, device="cuda")
        (a @ b).sum().backward()
        torch.cuda.synchronize()
        return True
    except Exception as exc:
        logger.warning("GPU smoke test failed (%s), falling back to CPU", exc)
        return False


def pick_device() -> str:
    if torch.cuda.is_available() and gpu_smoke_test():
        return "cuda"
    return "cpu"


DEVICE = pick_device()
CUDA_CC = torch.cuda.get_device_capability(0)[0] if DEVICE == "cuda" else 0
if DEVICE == "cuda":
    # Safe throughput win: lets cuDNN pick the fastest kernels for the fixed
    # block/batch shapes. No accuracy or stability cost.
    torch.backends.cudnn.benchmark = True
# Pascal (sm_6x, e.g. GTX 1060) has no usable fp16/bf16 hardware path --
# autocast there runs emulated and is far SLOWER than plain fp32. AMP only
# pays off on Volta+ (compute capability >= 7).
USE_AMP = DEVICE == "cuda" and CUDA_CC >= 7
AMP_DTYPE = torch.bfloat16 if (USE_AMP and torch.cuda.is_bf16_supported()) else torch.float16
SCALER = torch.cuda.amp.GradScaler(enabled=USE_AMP and AMP_DTYPE == torch.float16)


def get_batch(data):
    # data is a numpy memmap: only the sampled BLOCK_SIZE windows are ever
    # copied into RAM, so the multi-GB corpus stays on disk.
    max_start = len(data) - BLOCK_SIZE - 1
    # dtype=int64: corpus is ~3.9B tokens, past int32 max, and numpy defaults to
    # int32 on Windows -> "high is out of bounds for int32" without this.
    ix = np.random.randint(0, max_start, size=BATCH_SIZE, dtype=np.int64)
    x = np.stack([data[i:i + BLOCK_SIZE] for i in ix]).astype(np.int64)
    y = np.stack([data[i + 1:i + BLOCK_SIZE + 1] for i in ix]).astype(np.int64)
    x = torch.from_numpy(x).to(DEVICE, non_blocking=True)
    y = torch.from_numpy(y).to(DEVICE, non_blocking=True)
    return x, y


@torch.no_grad()
def estimate_loss(model, train_data, val_data):
    model.eval()
    losses = {}
    for name, data in (("train", train_data), ("val", val_data)):
        total = 0.0
        for _ in range(20):
            x, y = get_batch(data)
            with torch.autocast(device_type=DEVICE, dtype=AMP_DTYPE, enabled=USE_AMP):
                _, loss = model(x, y)
            total += loss.item()
        losses[name] = total / 20
    model.train()
    return losses


def _save_atomic(obj, dest):
    # Write to a temp file then rename: a test process reading base.pt never
    # sees a half-written file, and a crash mid-save can't corrupt the good one.
    tmp = dest.with_suffix(dest.suffix + ".tmp")
    torch.save(obj, tmp)
    os.replace(tmp, dest)


def main():
    psutil.Process().nice(psutil.BELOW_NORMAL_PRIORITY_CLASS)
    if DEVICE == "cpu":
        torch.set_num_threads(psutil.cpu_count(logical=True) or 4)

    CKPT_DIR.mkdir(parents=True, exist_ok=True)
    corpus_path = CORPUS_V2_FILE if CORPUS_V2_FILE.exists() else CORPUS_FILE
    logger.info("Corpus: %s (%d bytes), device=%s, amp=%s", corpus_path.name, corpus_path.stat().st_size, DEVICE, AMP_DTYPE if USE_AMP else "off")

    if TOKENIZER_FILE.exists() and IDS_DAT_FILE.exists() and IDS_META_FILE.exists():
        tokenizer = CharTokenizer.load(TOKENIZER_FILE)
        meta = json.loads(IDS_META_FILE.read_text())
        data = np.memmap(IDS_DAT_FILE, dtype=np.int32, mode="r", shape=(meta["length"],))
        # The corpus lives on an HDD: random batch sampling from a memmap means
        # constant disk seeks that starve the GPU. Pull it fully into RAM when
        # there's comfortable headroom (needs corpus + 8 GB spare).
        need = data.nbytes + 8 * 1024**3
        avail = psutil.virtual_memory().available
        if avail > need:
            logger.info("Loading %d MB of ids into RAM (avail %d MB)...", data.nbytes // 2**20, avail // 2**20)
            data = np.asarray(data)
            logger.info("Corpus in RAM, HDD seeks eliminated")
        else:
            logger.info("Keeping ids on-disk memmap (avail RAM %d MB too small)", avail // 2**20)
        logger.info("Loaded tokenizer + ids (%d tokens)", len(data))
    else:
        raise SystemExit(
            f"Missing encoded corpus. Run run_full_pipeline.py first "
            f"(need {TOKENIZER_FILE.name}, {IDS_DAT_FILE.name}, {IDS_META_FILE.name})."
        )
    logger.info("Vocab size: %d, tokens: %d", tokenizer.vocab_size, len(data))

    split = int(0.9 * len(data))
    train_data, val_data = data[:split], data[split:]

    cfg = KairoGPTConfig(
        vocab_size=tokenizer.vocab_size, block_size=BLOCK_SIZE,
        n_layer=N_LAYER, n_head=N_HEAD, n_embd=N_EMBD,
    )
    model = KairoGPT(cfg).to(DEVICE)
    logger.info("Params: %.2fM", sum(p.numel() for p in model.parameters()) / 1e6)

    try:
        optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, fused=True)
    except (RuntimeError, ValueError, TypeError):
        optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

    start_it = 0
    if INPROGRESS_FILE.exists():
        ckpt = torch.load(INPROGRESS_FILE, map_location=DEVICE)
        model.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        start_it = ckpt["it"] + 1
        logger.info("Resuming from %s at step %d", INPROGRESS_FILE, start_it)

    run_start = time.time()
    win_start, win_steps = run_start, 0
    last_val = float("nan")
    for it in range(start_it, MAX_ITERS):
        # last_val != last_val is a NaN check: forces an eval on the very first
        # iteration after a resume so the LIVE status shows a real loss, not NaN.
        if it % EVAL_INTERVAL == 0 or it == MAX_ITERS - 1 or last_val != last_val:
            losses = estimate_loss(model, train_data, val_data)
            last_val = losses["val"]
            logger.info("step %d: train %.4f val %.4f", it, losses["train"], losses["val"])

        cur_lr = lr_at(it)
        for g in optimizer.param_groups:
            g["lr"] = cur_lr

        optimizer.zero_grad(set_to_none=True)
        try:
            for _ in range(GRAD_ACCUM):
                x, y = get_batch(train_data)
                with torch.autocast(device_type=DEVICE, dtype=AMP_DTYPE, enabled=USE_AMP):
                    _, loss = model(x, y)
                SCALER.scale(loss / GRAD_ACCUM).backward()
            if SCALER.is_enabled():
                SCALER.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            SCALER.step(optimizer)
            SCALER.update()
        except RuntimeError as exc:
            # VRAM spike (desktop shares this GPU). Try to recover in-process;
            # if the CUDA context is already corrupt even empty_cache() throws --
            # then exit(3) and let supervisor.py restart from the checkpoint.
            if "out of memory" in str(exc).lower():
                try:
                    optimizer.zero_grad(set_to_none=True)
                    if DEVICE == "cuda":
                        torch.cuda.empty_cache()
                    time.sleep(5)
                    logger.warning("CUDA OOM at step %d -> cleared cache, skipped step", it)
                    continue
                except RuntimeError:
                    logger.error("CUDA context dead after OOM at step %d, exiting for supervisor restart", it)
                    raise SystemExit(3)
            raise
        if STEP_SLEEP:
            time.sleep(STEP_SLEEP)
        win_steps += 1

        # Early throughput probe: report real s/step in the first ~2 minutes so
        # the ETA is measured, not guessed.
        if start_it < it <= start_it + 60 and it % 20 == 0:
            dt = time.time() - win_start
            sps = win_steps / dt if dt > 0 else 0.0
            tok_s = sps * GRAD_ACCUM * BATCH_SIZE * BLOCK_SIZE
            logger.info("throughput: %.2f s/step, %.0f tok/s", (1 / sps if sps else 0), tok_s)
            win_start, win_steps = time.time(), 0

        # Heartbeat every 250 steps with ETA so the window keeps updating.
        if it % 250 == 0 and it > 0:
            elapsed = time.time() - run_start
            done = max(1, it - start_it)
            sps = done / elapsed
            eta_h = (MAX_ITERS - it) / sps / 3600 if sps > 0 else 0
            logger.info("HEARTBEAT: %d/%d (%.1f%%), %.2f s/step, ETA %.1f h",
                        it, MAX_ITERS, 100.0 * it / MAX_ITERS, 1 / sps, eta_h)

        # Save often so base.pt is always fresh/testable and a stop loses <=50
        # steps. base.pt = inference-ready (model + cfg); base_inprogress.pt =
        # full resume state. The status line prints every ~minute so the user
        # sees it's alive and getting smarter (loss falling).
        if it % SAVE_INTERVAL == 0 and it > start_it:
            _save_atomic({"model": model.state_dict(), "cfg": vars(cfg)}, MODEL_FILE)
            _save_atomic(
                {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "it": it},
                INPROGRESS_FILE,
            )
            mb = MODEL_FILE.stat().st_size / 2**20
            logger.info("[LIVE] step %d/%d (%.1f%%) | loss %.3f | base.pt %.0f MB gespeichert -- laeuft",
                        it, MAX_ITERS, 100.0 * it / MAX_ITERS, last_val, mb)

    _save_atomic({"model": model.state_dict(), "cfg": vars(cfg)}, MODEL_FILE)
    logger.info("Training reached MAX_ITERS, final base.pt saved -> %s", MODEL_FILE)


if __name__ == "__main__":
    main()