import argparse, os
import numpy as np
import torch
import torch.nn.functional as F
from torch_predictive_attn import ppmi_embed, learn_map, doc_index, apply_stack, features, log


def build(args, device):
    import sentencepiece as spm
    sp = spm.SentencePieceProcessor(model_file=args.spm_model)
    eos = sp.eos_id(); V = sp.get_piece_size()
    train = np.fromfile(args.train_bin, dtype=np.uint16)
    E = ppmi_embed(train, V, args.d, args.window, args.cooc_tokens, device)
    Ps, Bs = [], []
    for _ in range(args.layers):
        P, B = learn_map(train, E, Ps, Bs, eos, args, device)
        Ps.append(P); Bs.append(B)
    return sp, eos, E, Ps, Bs


def next_logits(ids, E, Ps, Bs, W, b, eos, args, device):
    x = torch.tensor(ids, device=device, dtype=torch.long)
    _, within = doc_index(x, eos)
    H, phis = apply_stack(x, E, Ps, Bs, within, args)
    Phi = features(H, within, phis, args.extra_context)[-1:]
    return (Phi @ W + b).squeeze(0)


def sample(logits, temp=0.8, top_k=40):
    logits = logits / max(temp, 1e-6)
    vals, idx = torch.topk(logits, min(top_k, logits.numel()))
    probs = F.softmax(vals, dim=-1)
    return int(idx[torch.multinomial(probs, 1)])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="/workspace/oneshot/logs/glm_d896_readout/stream_ce_lam10.pt")
    ap.add_argument("--spm_model", default="/workspace/glm/glm16k.model")
    ap.add_argument("--train_bin", default="/workspace/glm/glm_train.bin")
    ap.add_argument("--d", type=int, default=896)
    ap.add_argument("--r", type=int, default=320)
    ap.add_argument("--layers", type=int, default=10)
    ap.add_argument("--att_window", type=int, default=10)
    ap.add_argument("--temp", type=float, default=0.28)
    ap.add_argument("--window", type=int, default=10)
    ap.add_argument("--extra_context", type=int, default=1)
    ap.add_argument("--res_scale", type=float, default=0.07)
    ap.add_argument("--pred_scale", type=float, default=0.035)
    ap.add_argument("--pred_schedule", default="late")
    ap.add_argument("--orth_delta", type=int, default=1)
    ap.add_argument("--pred_norm", type=int, default=1)
    ap.add_argument("--pred_features", type=int, default=1)
    ap.add_argument("--map_lam", type=float, default=0.001)
    ap.add_argument("--cooc_tokens", type=int, default=3_600_000)
    ap.add_argument("--proj_tokens", type=int, default=3_600_000)
    ap.add_argument("--chunk_docs", type=int, default=8)
    ap.add_argument("--value_mode", default="dual_ridge_delta")
    ap.add_argument("--max_new", type=int, default=80)
    ap.add_argument("--sample_temp", type=float, default=0.8)
    ap.add_argument("--top_k", type=int, default=40)
    args = ap.parse_args()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    sp, eos, E, Ps, Bs = build(args, device)
    ck = torch.load(args.ckpt, map_location=device)
    W = ck["W"].to(device); b = ck["b"].to(device)
    prompts = [
        "Question: What is gravity?\nAnswer:",
        "Question: Why is the sky blue?\nAnswer:",
        "Question: Explain photosynthesis in simple words.\nAnswer:",
        "Question: If John has 3 apples and buys 2 more, how many apples does he have?\nAnswer:",
        "Question: Write a short friendly story about a robot learning to read.\nAnswer:",
    ]
    for p in prompts:
        ids = sp.encode(p)
        for _ in range(args.max_new):
            tok = sample(next_logits(ids, E, Ps, Bs, W, b, eos, args, device), args.sample_temp, args.top_k)
            ids.append(tok)
            if tok == eos:
                break
        print("=== PROMPT ===")
        print(p)
        print("=== OUTPUT ===")
        print(sp.decode(ids))


if __name__ == "__main__":
    main()