"""SDPO-style student-top-K + pooled-tail JSD on existing student states.

The context teacher uses the actor's pre-update weights with stop-gradient.
This module is an objective primitive, not an unbiased policy-value estimator.
"""
import math
import torch
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint


def coarsened_log_probs(logits,support):
    if logits.ndim!=2 or support.ndim!=2 or support.shape[0]!=logits.shape[0]:
        raise ValueError('Expected token by vocabulary logits and token by K support')
    if not 0<support.shape[1]<logits.shape[1]:raise ValueError('K must be strictly between0 and vocabulary size')
    norm=torch.logsumexp(logits,dim=-1,keepdim=True)
    top=logits.gather(1,support)-norm
    # Compute the tail directly: 1-sum(top) is inaccurate for peaked policies.
    tail=torch.logsumexp(logits.scatter(1,support,float('-inf')),dim=-1,keepdim=True)-norm
    return torch.cat((top,tail),dim=1)


def jsd_from_logs(student,teacher):
    if student.shape!=teacher.shape:raise ValueError('Mismatched categorical partitions')
    teacher=teacher.detach()
    mixture=torch.logaddexp(student,teacher)-math.log(2.)
    return .5*((student.exp()*(student-mixture)).sum(-1)+(teacher.exp()*(teacher-mixture)).sum(-1))


def head_terms(head,hidden,targets,support=None,teacher_logs=None,*,chunk=64):
    """Summed raw-policy log score and summed coarsened JSD, bounded head memory."""
    if hidden.ndim!=2 or len(hidden)!=len(targets) or len(targets)<1:raise ValueError('Wrong target alignment')
    if (support is None)!=(teacher_logs is None):raise ValueError('Support and teacher must accompany each other')
    if support is not None and (len(support)!=len(hidden) or len(teacher_logs)!=len(hidden)):raise ValueError('Dense length mismatch')
    score=hidden.new_zeros(());dense=hidden.new_zeros(())
    for start in range(0,len(hidden),chunk):
        h=hidden[start:start+chunk];y=targets[start:start+chunk]
        idx=None if support is None else support[start:start+chunk].to(h.device)
        q=None if teacher_logs is None else teacher_logs[start:start+chunk].to(h.device)
        def compute(x,labels,indices=idx,teacher=q):
            logits=head(x)
            lp=-F.cross_entropy(logits,labels,reduction='sum')
            div=logits.new_zeros(()) if indices is None else jsd_from_logs(coarsened_log_probs(logits,indices),teacher).sum()
            return lp,div
        if torch.is_grad_enabled() and hidden.requires_grad:
            lp,div=checkpoint(compute,h,y,use_reentrant=False,preserve_rng_state=False)
        else:lp,div=compute(h,y)
        score=score+lp;dense=dense+div
    return score,dense


def trajectory_hidden(model,prompt,completion):
    if not prompt or not completion:raise ValueError('Empty trajectory')
    device=next(model.parameters()).device
    ids=torch.tensor([prompt+completion[:-1]],device=device)
    hidden,_=model.hidden(ids=ids)
    return hidden[0,len(prompt)-1:]


def forward_work(prompt,completion):
    length=len(prompt)+len(completion)-1
    return {'trunk_positions':length,'head_positions':len(completion),'attention_pairs':length*length}


@torch.no_grad()
def teacher_partition(model,native_prompt,teacher_prompt,completion,*,k=100,chunk=64):
    """Support from the student's pre-update logits; teacher is detached current model."""
    model.eval()
    student=trajectory_hidden(model,native_prompt,completion)
    teacher=trajectory_hidden(model,teacher_prompt,completion)
    supports=[];logs=[]
    for start in range(0,len(completion),chunk):
        s=model.base.out_head(student[start:start+chunk])
        idx=s.topk(k,dim=-1).indices
        t=model.base.out_head(teacher[start:start+chunk])
        q=coarsened_log_probs(t,idx)
        assert torch.isfinite(q).all()
        supports.append(idx.cpu());logs.append(q.cpu())
    return torch.cat(supports),torch.cat(logs)


def trajectory_terms(model,prompt,completion,partition=None,*,chunk=64):
    hidden=trajectory_hidden(model,prompt,completion)
    target=torch.tensor(completion,device=hidden.device)
    support,q=partition if partition is not None else (None,None)
    return head_terms(model.base.out_head,hidden,target,support,q,chunk=chunk)
