"""NCM chat fine-tune: reuse released XSum NCM (frozen multiscale VQ tokenizer + prefix-conditioned generator),
fine-tune the generator on (conversation prefix -> short reply). Needs HazyResearch/ncm on NCM_ROOT."""
import os, sys, random, torch, torch.nn.functional as F
from omegaconf import OmegaConf
NCM_ROOT = os.environ.get("NCM_ROOT", "/tmp/ncm"); sys.path.insert(0, NCM_ROOT)
from models.next_concept import NextConceptTransformer
from tokenization.vqvae_tokenizer import VQVAETokenizer

def load(ck, cfg_path=None, dev="cpu", load_gen=True):
    CFG = OmegaConf.load(cfg_path or os.path.join(ck, "config.yaml"))
    CFG.tokenizer.vqvae.checkpoint_path = os.path.join(ck, CFG.checkpoints.tokenizer)
    tk = VQVAETokenizer.from_config(CFG.tokenizer)
    tk.vqvae = tk.vqvae.to(dev).eval()
    for p in tk.vqvae.parameters(): p.requires_grad = False
    net = NextConceptTransformer(vqvae_local=tk.vqvae, config=CFG.model).to(dev)
    if load_gen:
        sd = torch.load(os.path.join(ck, CFG.checkpoints.generator), map_location="cpu", weights_only=False)
        net.load_state_dict(sd["model"], strict=True)
    return CFG, tk, net

def fmt_prompt(turns):
    """turns: list of (role, text) ending with a user turn -> prompt string"""
    s = "".join(f"{'User' if r=='user' else 'Assistant'}: {t}\n" for r, t in turns)
    return s + "Assistant:"

class Codec:
    def __init__(self, CFG, tk, net):
        self.tk, self.net, self.P = tk, net, int(CFG.model.prefix_len)
        self.pre = tk.pre_tokenizer
        self.EOS, self.PAD = self.pre.eos_token_id, self.pre.pad_token_id
        self.ppad = int(CFG.model.prefix_pad_token_id)   # pad id used in the generator prefix
        self.T = tk.vqvae.context_length                  # target length (48)
        self.hf = self.pre.tokenizer if hasattr(self.pre, "tokenizer") else None
    def enc(self, text): return self.pre.encode(text)
    def prefix_ids(self, prompt):
        ids = self.enc(prompt)[-self.P:]                   # keep most recent context
        return ids + [self.ppad] * (self.P - len(ids))
    def target_ids(self, reply):
        ids = self.enc(" " + reply.strip())[: self.T - 1] + [self.EOS]
        return ids + [self.EOS] * (self.T - len(ids))

@torch.no_grad()
def target_codes(codec, tgt):                              # tgt: LongTensor (B,T) -> idx (B, sum levels)
    return codec.tk.vqvae.encode(tgt)

def build_inputs(codec, prefix, idx, corrupt=0.1):
    net, vq = codec.net, codec.tk.vqvae.vq
    with torch.no_grad():
        idc = idx.clone()
        if corrupt > 0:                                    # robustness: corrupt conditioning codes (as in paper)
            m = torch.rand_like(idc, dtype=torch.float) < corrupt
            idc[m] = torch.randint(0, 4096, (int(m.sum()),), device=idc.device)
        nc = vq.idx_to_ncm_input(idc)                      # cumulative recon fed to levels 1..K-1
        zp = codec.tk.vqvae.encoder.wte(prefix)[:, :codec.P, :]
        q0 = net.build_l0_query_input(prefix.shape[0])
    return torch.cat([zp, q0, nc], dim=1)

def loss_fn(codec, prefix, tgt, corrupt=0.1):
    idx = target_codes(codec, tgt)
    x = build_inputs(codec, prefix, idx, corrupt)
    pad_mask = prefix == codec.ppad
    logits = codec.net(x, prefix_pad_mask=None if not hasattr(codec.net, "_prefix_pad_bias") else pad_mask)
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]).float(), idx.reshape(-1)), logits, idx

@torch.no_grad()
def reply(codec, prompts, temperature=0.8, top_k=20, top_p=0.9, dev="cpu"):
    net, tk = codec.net, codec.tk
    pre = torch.tensor([codec.prefix_ids(p) for p in prompts], device=dev)
    g = net.generate(None, pre, temperature=temperature, top_k=top_k, top_p=top_p)
    _, zl = tk.vqvae.vq.idx_to_zhats(g, return_zlevels=True)
    ids = tk.vqvae.decoder(zl[-1]).argmax(-1).tolist()
    out = []
    for r in ids:
        if codec.EOS in r: r = r[: r.index(codec.EOS)]
        out.append(codec.pre.decode([t for t in r if t not in (codec.EOS, codec.PAD)], skip_special_tokens=True).strip())
    return out