"""Detached contextual teaching with an explicit, possibly frozen, teacher.

An unfrozen experimental primitive. It changes neither the completed experiment
nor its operator. The vocabulary partition still follows the current student;
fixed teacher parameters do not mean fixed coarsened targets or visited states.
"""
import collections
import torch

from contextual_dense_operator import (
    coarsened_log_probs, trajectory_hidden, trajectory_terms, forward_work,
)
from contextual_rl_common import group_advantages, source_roles, teacher_prompt


@torch.no_grad()
def reference_partition(student, teacher, native_prompt, feedback_prompt,
                        completion, *, k=100, chunk=64):
    student.eval()
    teacher.eval()
    student_hidden = trajectory_hidden(student, native_prompt, completion)
    teacher_hidden = trajectory_hidden(teacher, feedback_prompt, completion)
    supports, distributions = [], []
    for start in range(0, len(completion), chunk):
        logits = student.base.out_head(student_hidden[start:start + chunk])
        support = logits.topk(k, dim=-1).indices
        teacher_logits = teacher.base.out_head(teacher_hidden[start:start + chunk])
        target = coarsened_log_probs(teacher_logits, support)
        if not torch.isfinite(target).all():
            raise FloatingPointError('Nonfinite detached teacher distribution')
        supports.append(support.cpu())
        distributions.append(target.cpu())
    return torch.cat(supports), torch.cat(distributions)


def accumulate_reference_question(student, teacher, tok, prompt, outputs,
                                  rewards, finished, *, context_mode,
                                  batch_questions=4, dense_weight=.1,
                                  topk=100, chunk=64):
    if context_mode not in ('contextual', 'empty', 'native'):
        raise ValueError(context_mode)
    advantages = group_advantages(rewards, finished)
    donor, receivers = source_roles(rewards, finished)
    denominator = 8 * batch_questions
    work, partitions = collections.Counter(), {}
    feedback = None
    if receivers:
        feedback = (prompt if context_mode == 'native' else
                    teacher_prompt(tok, prompt, outputs[donor],
                                   empty=context_mode == 'empty'))
        for i in receivers:
            partitions[i] = reference_partition(
                student, teacher, prompt, feedback, 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(feedback, outputs[i]).items():
                work['teacher_' + key] += value
    result = {
        'advantages': advantages, 'donor': donor, 'receivers': receivers,
        'teacher_prompt_tokens': len(feedback) if feedback else 0,
        'rl_objective': 0., 'dense_objective': 0., 'dense_jsd_sum': 0.,
        'dense_tokens': 0, 'backward_calls': 0,
    }
    student.train()
    for i, tokens in enumerate(outputs):
        if advantages[i] == 0 and i not in partitions:
            continue
        score, divergence = trajectory_terms(
            student, prompt, tokens, partitions.get(i), chunk=chunk)
        rl = -advantages[i] * score / denominator
        dense = dense_weight * divergence / denominator
        objective = rl + dense
        if not torch.isfinite(objective):
            raise FloatingPointError('Nonfinite reference-teaching objective')
        objective.backward()
        result['backward_calls'] += 1
        result['rl_objective'] += float(rl.detach())
        result['dense_objective'] += float(dense.detach())
        result['dense_jsd_sum'] += float(divergence.detach())
        result['dense_tokens'] += len(tokens) if i in partitions else 0
        for key, value in forward_work(prompt, tokens).items():
            work['backward_' + key] += value
    result['work'] = dict(work)
    return result
