"""Shared, auditable current-state dense/contextual RL update."""
import collections,json,pathlib
import torch
from contextual_dense_operator import teacher_partition,trajectory_terms,forward_work
from prefix_teacher_spec import TEMPLATE


def group_advantages(rewards,finished):
    r=torch.tensor(rewards,dtype=torch.float32)
    if len(r)!=8 or len(finished)!=8:raise ValueError('Eight responses required')
    return ((r-r.mean())/(r.std(correction=1)+1e-4)*torch.tensor(finished)).tolist()


def source_roles(rewards,finished):
    donor=next((i for i,(r,e) in enumerate(zip(rewards,finished)) if r==1 and e),None)
    receivers=[] if donor is None else [i for i,(r,e) in enumerate(zip(rewards,finished)) if r==0 and e]
    return donor,receivers


def teacher_prompt(tok,prompt,context,empty=False):
    assert context[-1]==tok.eos_token_id and tok.eos_token_id not in context[:-1]
    solution='' if empty else tok.decode(context[:-1])
    return tok.encode(TEMPLATE.format(solution=solution,native_prompt=tok.decode(prompt)))


def accumulate_question(model,tok,prompt,outputs,rewards,finished,*,arm,batch_questions=4,dense_weight=.1,replay_weight=.1,topk=100,chunk=64):
    if arm not in ('G','C','E','H'):raise ValueError(arm)
    advantage=group_advantages(rewards,finished);donor,receivers=source_roles(rewards,finished)
    denominator=8*batch_questions;work=collections.Counter();parts={}
    tp=None
    if arm in ('C','E') and receivers:
        tp=teacher_prompt(tok,prompt,outputs[donor],empty=arm=='E')
        for i in receivers:
            parts[i]=teacher_partition(model,prompt,tp,outputs[i],k=topk,chunk=chunk)
            for key,value in forward_work(prompt,outputs[i]).items():work['support_'+key]+=value
            for key,value in forward_work(tp,outputs[i]).items():work['teacher_'+key]+=value
    result={'advantages':advantage,'donor':donor,'receivers':receivers,'teacher_prompt_tokens':len(tp) if tp else 0,
            'rl_objective':0.,'dense_objective':0.,'replay_objective':0.,'dense_jsd_sum':0.,'dense_tokens':0,'backward_calls':0}
    model.train()
    for i,tokens in enumerate(outputs):
        replay=arm=='H' and i==donor
        if advantage[i]==0 and i not in parts and not replay:continue
        score,div=trajectory_terms(model,prompt,tokens,parts.get(i),chunk=chunk)
        rl=-advantage[i]*score/denominator
        dense=dense_weight*div/denominator
        hard=-replay_weight*score/batch_questions if replay else score.new_zeros(())
        objective=rl+dense+hard
        if not torch.isfinite(objective):raise FloatingPointError('Nonfinite objective')
        objective.backward();result['backward_calls']+=1
        result['rl_objective']+=float(rl.detach());result['dense_objective']+=float(dense.detach());result['replay_objective']+=float(hard.detach())
        result['dense_jsd_sum']+=float(div.detach());result['dense_tokens']+=len(tokens) if i in parts else 0
        for key,value in forward_work(prompt,tokens).items():work['backward_'+key]+=value
    result['work']=dict(work);return result
