"""Qualitative generation demo: full-attn vs dense block-attn vs sparse (router top-k).

Runs several diverse RAG examples through the block-distilled model and prints actual generations so we
can eyeball quality (not just NLL). Blocks = documents (+4 sink tokens each); landmark = last content
token; router selects top-k blocks per decode step.
"""
from __future__ import annotations

import argparse
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer

NEG = -1e9
N_SINKS = 4
SYSTEM = "You are a helpful assistant. Answer the question using ONLY the documents. Be concise."

EXAMPLES = [
    # (docs, question) — mix of counterfactual (must read), multi-hop, real, numeric, distractor-heavy
    ([ "The Zorban Reactor at Vale Station reached full output in the year 3471.",
       "Mount Kilimanjaro, in Tanzania, is the highest mountain in Africa.",
       "The Great Wall of China stretches across northern China.",
       "Photosynthesis converts carbon dioxide and water into glucose and oxygen." ],
     "In what year did the Zorban Reactor reach full output?"),
    ([ "The CEO of Vantacorp is Marisa Quen.",
       "Marisa Quen was born in the city of Drennholm.",
       "Drennholm is famous for its glass bridges.",
       "The CEO of Bexil Industries is Tomas Ray." ],
     "In which city was the CEO of Vantacorp born?"),
    ([ "Python is a programming language created by Guido van Rossum, first released in 1991.",
       "Rust is a systems language emphasizing memory safety.",
       "The mitochondrion is the powerhouse of the cell.",
       "Java was developed by James Gosling at Sun Microsystems." ],
     "Who created Python and in what year was it first released?"),
    ([ "Order #4471 shipped on March 3 and contains 2 laptops.",
       "Order #4472 shipped on March 5 and contains 1 monitor.",
       "Order #4473 is delayed and contains 3 keyboards.",
       "Order #4474 shipped on March 6 and contains 1 laptop." ],
     "Which orders shipped in March and what did order #4473 contain?"),
    ([ "The Treaty of Kessel was signed in 1804 between Aldoria and Brenne.",
       "Aldoria's capital is Feldspar City.",
       "Brenne is known for its copper mines.",
       "The Kessel treaty ended the Seven Rivers War.",
       "Feldspar City sits on the river Onn.",
       "Copper was Brenne's main export in the 1800s." ],
     "What war did the Treaty of Kessel end, and in what year was it signed?"),
]


def build(tok, docs, question):
    sink = tok("\n", add_special_tokens=False)["input_ids"] * N_SINKS
    ids, blk, snk = [], [], []
    def add(t, b, special=False):
        for x in tok(t, add_special_tokens=special)["input_ids"]:
            ids.append(x); blk.append(b); snk.append(False)
    add(SYSTEM + "\n", -1, special=True)
    for bi, d in enumerate(docs):
        for s in sink:
            ids.append(s); blk.append(bi); snk.append(True)
        add(f"[Document {bi+1}] {d}\n", bi)
    q0 = len(ids); add(f"Question: {question}\nAnswer:", -2)
    return ids, blk, snk, q0


def landmarks(blk, snk, nb):
    lm, isc = {}, [False]*len(blk)
    for j,(b,s) in enumerate(zip(blk,snk)):
        if b>=0 and not s: isc[j]=True; lm[b]=j
    return [lm[b] for b in range(nb)], isc


def mask(blk, snk, q0, nb, active_rows, n, dev, full=False):
    b = torch.tensor(blk+[-2]*(n-len(blk)), device=dev)
    causal = torch.tril(torch.ones(n,n,dtype=torch.bool,device=dev))
    if full:
        return torch.where(causal,0.0,NEG).view(1,1,n,n).float()
    bi,bj=b.view(n,1),b.view(1,n)
    if active_rows is None:
        allowed=((bj==-1)|(bi==bj)|(bi==-2))&causal
        return torch.where(allowed,0.0,NEG).view(1,1,n,n).float()
    lm,isc=landmarks(blk,snk,nb)
    is_lm=torch.zeros(n,dtype=torch.bool,device=dev); is_lm[torch.tensor(lm,device=dev)]=True
    isc_t=torch.zeros(n,dtype=torch.bool,device=dev); isc_t[:len(isc)]=torch.tensor(isc,device=dev)
    static=(b==-1)|(b==-2)|is_lm
    allowed=torch.zeros(n,n,dtype=torch.bool,device=dev)
    allowed[:q0]=(((bj==-1)|(bi==bj)|(bi==-2))&causal)[:q0]
    for i in range(q0,n):
        act=active_rows.get(i-q0,set(range(nb))); vis=static.clone()
        if act: vis=vis|(isc_t&torch.isin(bj.view(n),torch.tensor(sorted(act),device=dev)))
        allowed[i]=vis
    allowed&=causal
    return torch.where(allowed,0.0,NEG).view(1,1,n,n).float()


def block_content_pos(blk, snk, nb):
    bp={b:[] for b in range(nb)}
    for j,(b,s) in enumerate(zip(blk,snk)):
        if b>=0 and not s: bp[b].append(j)
    return bp


@torch.no_grad()
def gen(model, tok, ids, blk, snk, q0, mode, router=None, k=3, max_new=24, sticky=0, summ=1):
    """mode: full | dense | sparse. summ = block summary size (last-N content tokens) for router feats.
    sticky>0 -> keep a block active for `sticky` more steps after last selected."""
    dev=model.device; nb=max(blk)+1
    seq,bl,sk=list(ids),list(blk),list(snk)
    ar={} if mode=="sparse" else None
    active=set(range(nb)); last_seen={}
    for step in range(max_new):
        n=len(seq)
        if mode=="sparse": ar[n-q0]=set(range(nb)) if step==0 else set(active)
        m=mask(bl,sk,q0,nb,ar,n,dev,full=(mode=="full"))
        out=model(input_ids=torch.tensor([seq],device=dev),attention_mask=m,
                  output_attentions=(mode=="sparse"))
        nxt=int(out.logits[0,-1].argmax())
        if mode=="sparse":
            att=torch.stack(out.attentions,0)[:,0,:,-1,:]              # [L,H,n] attn of last pos
            bp=block_content_pos(bl,sk,nb)
            feat=torch.zeros(nb, att.shape[0]*att.shape[1], device=dev)
            for b in range(nb):
                cols=torch.tensor(bp[b][-summ:],device=dev)
                feat[b]=att[:,:,cols].mean(-1).reshape(-1).float()     # attn to block b's last-summ tokens
            top=router(feat).squeeze(-1).argsort(descending=True)[:k].tolist()
            for b in top: last_seen[b]=step
            active=set(b for b,s in last_seen.items() if step-s<=sticky) if sticky else set(top)
        seq.append(nxt); bl.append(-2); sk.append(False)
        if nxt==tok.eos_token_id: break
    return tok.decode(seq[q0:],skip_special_tokens=True).strip().replace("\n"," ")


def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--model",required=True); ap.add_argument("--router",default=None)
    ap.add_argument("--k",type=int,default=2); ap.add_argument("--max-new",type=int,default=24)
    ap.add_argument("--sticky",type=int,default=0,help="keep a block active this many steps after last selected")
    args=ap.parse_args()
    tok=AutoTokenizer.from_pretrained(args.model)
    model=AutoModelForCausalLM.from_pretrained(args.model,dtype=torch.bfloat16,
            attn_implementation="eager",device_map="cuda").eval()
    router=None; summ=1
    if args.router:
        ck=torch.load(args.router,map_location="cuda")
        router=nn.Linear(ck["in_dim"],1).to("cuda"); router.load_state_dict(ck["state_dict"]); router.eval()
        summ=ck.get("summary_tokens",1)
    for docs,q in EXAMPLES:
        ids,blk,snk,q0=build(tok,docs,q)
        print("\n"+"="*100); print(f"Q: {q}  ({max(blk)+1} blocks, {q0} ctx tok)")
        print(f"  full-attn      : {gen(model,tok,ids,blk,snk,q0,'full',max_new=args.max_new)!r}")
        print(f"  block dense    : {gen(model,tok,ids,blk,snk,q0,'dense',max_new=args.max_new)!r}")
        if router is not None:
            print(f"  sparse k={args.k} (summ{summ}): {gen(model,tok,ids,blk,snk,q0,'sparse',router,args.k,args.max_new,0,summ)!r}")
            print(f"  sparse k={args.k} sticky{args.sticky}: {gen(model,tok,ids,blk,snk,q0,'sparse',router,args.k,args.max_new,args.sticky,summ)!r}")


if __name__=="__main__":
    main()