"""
Phase K - text generation from an Abyss-1 checkpoint.

Apllies temoerature scaling, top-k filtering, and top-p (nucleus)
filtering in that order, then samples from the resulting distribution.

Usage (from the_final_product/ directory)
    python -m ml.generate
    python -m ml.generate --prompt "The dragon" --max_tokens 300
    python -m ml.generate --checkpoint ml/checkpoints/ckpt.pt --temperature 1.0
"""

import argparse
import sys
import torch
import torch.nn.functional as F

from ml.model import AbyssConfig, TheAbyssModel
from ml.tokenizer.tokenizer import AbyssTokenizer

TOKENIZER_DIR = "ml/tokenizer/vocab"


@torch.no_grad()
def sample(model, idx, max_new_tokens, temperature, top_k, top_p):
    """Autoregressively yield one token id at a time."""
    block_size = model.config.block_size
    for _ in range(max_new_tokens):
        idx_cond = idx if idx.size(1) <= block_size else idx[:, -block_size:]
        logits, _ = model(idx_cond)
        logits = logits[:, -1, :]

        # 1. Temperature
        logits = logits / temperature

        #2. top-k: mask everything bellow the k-th highest logit
        if top_k > 0:
            k = min(top_k, logits.size(-1))
            kth_val = torch.topk(logits, k).values[:, -1:]
            logits = logits.masked_fill(logits < kth_val, float("-inf"))

        # 3. top-p (nucleus): keep smallest set whose cumulative prob >= top_p
        if 0.0 < top_p <1.0:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            probs_sorted = F.softmax(sorted_logits, dim=-1)
            cum_probs = torch.cumsum(probs_sorted, dim=-1)
            # Shift right so the token that first pushes cumsum over top_p is kept
            remove = (cum_probs - probs_sorted) > top_p
            sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
            logits = torch.zeros_like(logits).scatter_(1, sorted_idx, sorted_logits)

        probs = F.softmax(logits, dim=-1)
        next_tok = torch.multinomial(probs, num_samples=1)
        yield next_tok.item()
        idx = torch.cat([idx, next_tok], dim=1)


def main():
    parser = argparse.ArgumentParser(description="Generate text from Abyss-1")
    parser.add_argument("--checkpoint",  default="ml/checkpoints/ckpt.pt")
    parser.add_argument("--prompt",      default="Once upon a time")
    parser.add_argument("--max_tokens",  type=int,   default=200)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top_k",       type=int,   default=75)
    parser.add_argument("--top_p",       type=float, default=0.95)
    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"])
    model.eval()
    print(f"  iter {ck['iter']} | best_val {ck['best_val']:.4f} "
          f"| params {model.num_params():,}", flush=True)
    
    # Load Tokenizer
    tok = AbyssTokenizer.load(TOKENIZER_DIR)

    # Encode prompt and Generate
    ids = tok.encode(args.prompt)
    idx = torch.tensor([ids], dtype=torch.long, device=args.device)

    print(f"\n[temp={args.temperature} top_k={args.top_k} top_p={args.top_p}]\n")
    sys.stdout.write(args.prompt)
    sys.stdout.flush()

    for tok_id in sample(model, idx, args.max_tokens,
                         args.temperature, args.top_k,args.top_p):
        sys.stdout.write(tok.decode([tok_id]))
        sys.stdout.flush()

    print()  # Trailing Newline

if __name__ == "__main__":
    main()