"""Unfrozen development: batching and bounded vocabulary-head memory.

Not used by the current pilot. Preserves mean token loss within each question,
then equal question weight. Right-padded target positions carry zero loss;
they cannot affect valid earlier positions under causal attention.
"""
from functools import partial

import torch
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint


def collate(prompts, targets, pad, device):
    pwidth=max(map(len,prompts));twidth=max(map(len,targets))
    assert all(prompts) and all(targets) and len(prompts)==len(targets)
    p=torch.tensor([[pad]*(pwidth-len(x))+x for x in prompts],device=device)
    mask=torch.tensor([[False]*(pwidth-len(x))+[True]*len(x) for x in prompts],
                      dtype=torch.bool,device=device)
    t=torch.tensor([x+[pad]*(twidth-len(x)) for x in targets],device=device)
    valid=torch.tensor([[True]*len(x)+[False]*(twidth-len(x)) for x in targets],
                       dtype=torch.bool,device=device)
    return p,mask,t,valid


def bounded_head_ce(model, hidden, targets, valid, chunk=128):
    """Per-token CE, recomputing bounded head chunks only when differentiating."""
    assert hidden.shape[:2]==targets.shape==valid.shape and chunk>0
    labels=targets.masked_fill(~valid,-100).flatten()
    flat=hidden.flatten(0,1)
    losses=[]
    def project(h,y):
        return F.cross_entropy(model.base.out_head(h.to(model.cfg['dtype'])),y,
                               ignore_index=-100,reduction='none')
    for start in range(0,len(labels),chunk):
        h=flat[start:start+chunk];y=labels[start:start+chunk]
        if torch.is_grad_enabled():
            loss=checkpoint(project,h,y,use_reentrant=False,preserve_rng_state=False)
        else:
            loss=project(h,y)
        losses.append(loss)
    tokens=torch.cat(losses).reshape(targets.shape)
    per_question=tokens.sum(1)/valid.sum(1)
    return per_question.mean(),tokens


def answer_loss(model,prompts,targets,boundary,*,pad,chunk=128):
    device=next(model.parameters()).device
    p,mask,t,valid=collate(prompts,targets,pad,device)
    delim=torch.tensor([boundary]*len(prompts),device=device)
    h,features,cache,mask=model.prefill(p,delim,mask,features=True)
    if t.shape[1]>1:
        prefix=t[:,:-1]
        mask=torch.cat((mask,torch.ones_like(prefix,dtype=torch.bool)),1)
        tail,_=model.hidden(ids=prefix,cache=cache,mask=mask)
        h=torch.cat((h,tail),1)
    loss,tokens=bounded_head_ce(model,h,t,valid,chunk)
    return loss,features,tokens,valid


def explicit_loss(model,prompts,targets,*,pad,chunk=128):
    device=next(model.parameters()).device
    p,mask,t,valid=collate(prompts,targets,pad,device)
    ids=torch.cat((p,t),1)[:,:-1]
    mask=torch.cat((mask,torch.ones_like(t[:,:-1],dtype=torch.bool)),1)
    h,_=model.hidden(ids=ids,mask=mask)
    h=h[:,p.shape[1]-1:]
    return bounded_head_ce(model,h,t,valid,chunk)[0]
