"""
Phase L - standalone val-set perplexity for Abyss-1.

Walks sequentially through val.bin with non-overlapping windows (no
randomness), ao the reported perplexity is a true full-dataset figure.

Usage (from the_final_product/ directory):
    python -m ml.eval
    python -m ml.eval --checkpoint ml/checkpoints/ckpt.pt --batch_size 32
"""

import argparse
import math
import os
import numpy as np
import torch
from contextlib import nullcontext

from ml.model import TheAbyssModel

VAL_BIN = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                       "data", "tinystories", "val.bin")


@torch.no_grad()
def compute_perplexity(model, device, block_size, batch_size, ctx):
    model.eval()
    data = np.memmap(VAL_BIN, dtype=np.uint16, mode="r")
    n = len(data)
 
    total_loss   = 0.0
    total_tokens = 0
    batches_done = 0
    pos = 0

    while pos + block_size + 1 <= n:
        starts = []
        while len(starts) < batch_size and pos + block_size + 1 <= n:
            starts.append(pos)
            pos += block_size

        xs = torch.stack([
            torch.from_numpy(data[i : i + block_size].astype(np.int64))
            for i in starts
        ])
        ys = torch.stack([
            torch.from_numpy(data[i + 1 : i + 1 + block_size].astype(np.int64))
            for i in starts
        ])

        if device.startswith("cuda"):
            xs = xs.pin_memory().to(device, non_blocking=True)
            ys = ys.pin_memory().to(device, non_blocking=True)
        else:
            xs, ys = xs.to(device), ys.to(device)

        with ctx:
            _, loss = model(xs, ys)

        n_tok = xs.numel()
        total_loss   += loss.item() * n_tok
        total_tokens += n_tok
        batches_done += 1

        if batches_done % 50 == 0:
            running_ppl = math.exp(total_loss / total_tokens)
            pct = 100.0 * pos / n
            print(f"  {pct:5.1f}% | {total_tokens:>10,} tok | running ppl {running_ppl:.2f}",
                  flush=True)
            
    avg_loss = total_loss / total_tokens
    return avg_loss, math.exp(avg_loss),  total_tokens


def main():
    parser = argparse.ArgumentParser(description="Val-set perplexity for Abyss-1.")
    parser.add_argument("--checkpoint", default="ml/checkpoints/ckpt.pt")
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--block_size", type=int, default=512)
    parser.add_argument("--device",
                        default="cuda" if torch.cuda.is_available() else "cpu")
    args = parser.parse_args()

    # Load Checkpoint
    print(f"loading {args.checkpoint} ...", flush=True)
    ck = torch.load(args.checkpoint, map_location=args.device, weights_only=False)
    config = ck["config"]
    model  = TheAbyssModel(config).to(args.device)
    model.load_state_dict(ck["model"])
    print(f"  iter {ck['iter']} | best_val {ck['best_val']:.4f} "
                    f" | params {model.num_params():,}", flush=True)
    
    ptdtype = (torch.bfloat16 if args.device == "cuda" and torch.cuda.is_bf16_supported()
               else torch.float16 if args.device == "cuda" else torch.float32)
    ctx = (nullcontext() if args.device == "cpu"
           else torch.autocast(device_type="cuda", dtype=ptdtype))
    
    print(f"evaluating full val set "
          f"(block={args.block_size}, batch={args.batch_size}, {args.device}) ...\n",
          flush=True)
    loss, ppl, n_tok = compute_perplexity(
        model, args.device, args.block_size, args.batch_size, ctx
    )

    bar = "=" * 42
    print(f"\n{bar}")
    print(f"  checkpoint : {args.checkpoint}")
    print(f"  iter       : {ck['iter']}")
    print(f"  val loss   : {loss:.4f}")
    print(f"  perplexity : {ppl:.2f}")
    print(f"  tokens     : {n_tok:,}")
    print(bar)

if __name__ == "__main__":
    main()