#!/usr/bin/env python3
"""
Minimal inference script for the Doc-to-LoRA Perceiver.
Requirements:  pip install transformers>=4.51.0 huggingface_hub torch
"""
import re, torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import hf_hub_download, login

REPO_ID  = "farpluto/doc-to-lora-niah"   # filled in automatically at packaging time
HF_TOKEN = None                 # set your token here if the base model is gated

if HF_TOKEN:
    login(token=HF_TOKEN)

ckpt   = torch.load(hf_hub_download(REPO_ID, "hypernet.pt", token=HF_TOKEN),
                    map_location="cuda", weights_only=False)
hcfg   = ckpt["hypernet_cfg"]
BASE   = ckpt["base_model"]
TGT    = ckpt["target_module"]
ALPHA  = ckpt["lora_alpha"]
EARLY  = ckpt["early_exit"]

tokenizer = AutoTokenizer.from_pretrained(BASE, token=HF_TOKEN, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

llm = AutoModelForCausalLM.from_pretrained(
    BASE, token=HF_TOKEN, torch_dtype=torch.bfloat16,
    device_map="cuda", attn_implementation="sdpa", trust_remote_code=True)
llm.eval()
for p in llm.parameters():
    p.requires_grad_(False)


class CrossAttentionBlock(nn.Module):
    def __init__(self, latent_dim, ctx_dim, n_heads=8):
        super().__init__()
        self.n_heads, self.head_dim = n_heads, latent_dim // n_heads
        self.norm_q=nn.LayerNorm(latent_dim); self.norm_ctx=nn.LayerNorm(ctx_dim)
        self.q_proj=nn.Linear(latent_dim,latent_dim,bias=False)
        self.k_proj=nn.Linear(ctx_dim,latent_dim,bias=False)
        self.v_proj=nn.Linear(ctx_dim,latent_dim,bias=False)
        self.o_proj=nn.Linear(latent_dim,latent_dim,bias=False)
        self.norm_ff=nn.LayerNorm(latent_dim)
        self.ff=nn.Sequential(nn.Linear(latent_dim,latent_dim*4,bias=False),
                              nn.GELU(),nn.Linear(latent_dim*4,latent_dim,bias=False))
    def forward(self,latents,ctx,ctx_mask=None):
        B,L,D=latents.shape; _,S,_=ctx.shape; H,Dh=self.n_heads,self.head_dim
        q=self.q_proj(self.norm_q(latents)).view(B,L,H,Dh).transpose(1,2)
        k=self.k_proj(self.norm_ctx(ctx)).view(B,S,H,Dh).transpose(1,2)
        v=self.v_proj(ctx).view(B,S,H,Dh).transpose(1,2)
        bias=None
        if ctx_mask is not None:
            bias=(1.0-ctx_mask.float()).unsqueeze(1).unsqueeze(2)*-1e4
            bias=bias.to(q.dtype)
        out=F.scaled_dot_product_attention(q,k,v,attn_mask=bias)
        latents=latents+self.o_proj(out.transpose(1,2).contiguous().view(B,L,D))
        latents=latents+self.ff(self.norm_ff(latents))
        return latents


class PerceiverHypernet(nn.Module):
    def __init__(self,ctx_dim,n_lora_layers,lora_r,target_in,target_out,
                 latent_dim=512,n_blocks=8):
        super().__init__()
        self.n_lora_layers=n_lora_layers; self.din=target_in; self.dout=target_out
        d=target_in+target_out
        self.ctx_proj=nn.Linear(ctx_dim,latent_dim,bias=False)
        self.ctx_norm=nn.LayerNorm(latent_dim)
        self.latent_q=nn.Parameter(torch.randn(lora_r,latent_dim)*latent_dim**-0.5)
        self.blocks=nn.ModuleList([CrossAttentionBlock(latent_dim,latent_dim) for _ in range(n_blocks)])
        self.head_w=nn.Parameter(torch.randn(n_lora_layers,latent_dim,d)*0.01)
        self.head_b=nn.Parameter(torch.zeros(n_lora_layers,d))
    def forward(self,ctx_acts,ctx_mask=None):
        B=ctx_acts.shape[0]
        ctx=self.ctx_norm(self.ctx_proj(ctx_acts))
        lat=self.latent_q.unsqueeze(0).expand(B,-1,-1)
        for blk in self.blocks: lat=blk(lat,ctx,ctx_mask)
        flat=torch.einsum("brd,nde->bnre",lat,self.head_w)
        flat=flat+self.head_b.unsqueeze(0).unsqueeze(2)
        return flat[...,:self.din], flat[...,self.din:].transpose(-1,-2)


hypernet=PerceiverHypernet(**hcfg).to("cuda",dtype=torch.bfloat16)
hypernet.load_state_dict(ckpt["state_dict"])
hypernet.eval()
print("Perceiver loaded OK")

_tok_open  = tokenizer.convert_tokens_to_ids("<think>")
_tok_close = tokenizer.convert_tokens_to_ids("</think>")
THINK_TOKENS = {t for t in [_tok_open,_tok_close]
                if t not in (tokenizer.unk_token_id, None)}

def _strip_think(ids):
    toks=ids.tolist()
    if not THINK_TOKENS or not any(t in THINK_TOKENS for t in toks):
        return tokenizer.decode(toks,skip_special_tokens=True).strip()
    clean,inside=[],False
    op,cl=min(THINK_TOKENS),max(THINK_TOKENS)
    for t in toks:
        if t==op: inside=True
        elif t==cl: inside=False
        elif not inside: clean.append(t)
    return tokenizer.decode(clean,skip_special_tokens=True).strip()

def _sorted_mods(model,mod_name):
    mods=[(n,m) for n,m in model.named_modules()
          if mod_name in n and isinstance(m,nn.Linear)]
    def _idx(name):
        nums=re.findall(r"\d+",name)
        return int(nums[0]) if nums else -1
    return sorted(mods,key=lambda x:_idx(x[0]))

target_mods=_sorted_mods(llm,TGT)
scale=ALPHA/hcfg["lora_r"]

@torch.no_grad()
def internalize_and_query(document,query,max_new_tokens=12):
    ctx_ids=torch.tensor(
        [tokenizer.encode(document,add_special_tokens=False)],device="cuda")
    ctx_mask=torch.ones_like(ctx_ids)
    qry_ids=torch.tensor(
        [tokenizer.encode(query+" /no_think",add_special_tokens=True)],device="cuda")
    acts=llm(input_ids=ctx_ids,attention_mask=ctx_mask,
             output_hidden_states=True,use_cache=False).hidden_states[EARLY]
    A,B=hypernet(acts,ctx_mask)
    A,B=A.squeeze(0),B.squeeze(0)
    hooks=[]
    def _mkhook(Ai,Bi):
        def h(mod,inp,out): return out+scale*(inp[0]@Ai.t())@Bi.t()
        return h
    for i,(_,mod) in enumerate(target_mods):
        hooks.append(mod.register_forward_hook(_mkhook(A[i],B[i])))
    ids=qry_ids.clone()
    for _ in range(max_new_tokens):
        out=llm(input_ids=ids,attention_mask=torch.ones_like(ids),use_cache=False)
        nxt=out.logits[:,-1,:].argmax(-1,keepdim=True)
        ids=torch.cat([ids,nxt],dim=1)
        if nxt.item()==tokenizer.eos_token_id: break
    for h in hooks: h.remove()
    return _strip_think(ids[0,qry_ids.shape[1]:])

if __name__=="__main__":
    doc = "The special magic number is 7341. The sky is blue today."
    ans = internalize_and_query(doc, "What is the special magic number?")
    print(f"Answer: {ans}")