"""Future acquisition: bounded microbatches preserve equal question weights.

Not imported by the frozen pilot. No research dose or learner is chosen here.
"""
import math
import torch
from batched_latent_acquisition import answer_loss, explicit_loss
from latent_deliberation import feature_loss


def partition_rows(rows, *, max_rows=4, prompt_positions=960,
                   explicit_positions=5052, explicit_pairs=6380676):
    """Contiguous groups; bounds correspond to batch4/prompt240/trace1024."""
    if not rows:
        raise ValueError('empty optimizer batch')
    def fits(group):
        n=len(group); p=max(len(r['prompt']) for r in group)
        t=max(len(r['completion']) for r in group); length=p+t-1
        return (n<=max_rows and n*p<=prompt_positions and n*length<=explicit_positions
                and n*length*length<=explicit_pairs)
    groups=[];current=[]
    for row in rows:
        if not fits([row]):
            raise ValueError('one question exceeds the declared memory bounds')
        if current and not fits(current+[row]):
            groups.append(current);current=[]
        current.append(row)
    if current: groups.append(current)
    assert [id(r) for g in groups for r in g]==[id(r) for r in rows]
    return groups


def backward_batch(model, groups, teacher, *, mode, pad, head_chunk=128, alignment_weight=1.):
    """Accumulate one optimizer batch; caller zeros/clips/steps exactly once.

    O is ordinary full-trace loss. D/L are answer+explicit loss; F also aligns.
    Each component averages tokens within question and questions across batch.
    Model slots are configured by caller; no optimizer state is owned here.
    """
    if mode not in ('O','D','L','F') or not groups or any(not g for g in groups):
        raise ValueError('invalid arm or empty group')
    n=sum(len(g) for g in groups)
    result={'questions':n,'microbatch_sizes':[len(g) for g in groups],
            'answer_ce':0.,'feature_loss':0.,'text_ce':0.}
    for group in groups:
        weight=len(group)/n
        prompts=[r['prompt'] for r in group]
        if mode!='O':
            boundaries=[r['split']['separator'] for r in group]
            if any(b!=boundaries[0] for b in boundaries):
                raise ValueError('common boundary required')
            targets=[r['split']['canonical_answer'] for r in group]
            answer,states,_,_=answer_loss(model,prompts,targets,boundaries[0],pad=pad,chunk=head_chunk)
            target=torch.stack([teacher[r['index']] for r in group]).to(states.device)
            alignment=feature_loss(states[:,1:],target)
            objective=answer+alignment_weight*alignment if mode=='F' else answer
            if not torch.isfinite(objective):raise FloatingPointError('nonfinite answer objective')
            result['answer_ce']+=weight*float(answer.detach())
            result['feature_loss']+=weight*float(alignment.detach())
            (weight*objective).backward()
            del answer,states,target,alignment,objective
        text=explicit_loss(model,prompts,[r['completion'] for r in group],pad=pad,chunk=head_chunk)
        if not torch.isfinite(text):raise FloatingPointError('nonfinite explicit objective')
        result['text_ce']+=weight*float(text.detach())
        (weight*text).backward();del text
    assert all(math.isfinite(result[k]) for k in ('answer_ce','feature_loss','text_ce'))
    return result
