"""Development implementation of continuous own-model feedback.

No research run is registered by this module. All original Parameters remain
untied and trainable. Training caches retain their autograd graph. This follows
Coconut/CODI and rStation paper8's continuous-state primitive, not a new layer.
"""
from functools import partial

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

from run_baseline import KVCache


class LatentDeliberation(nn.Module):
    def __init__(self, base, slots, checkpoint_blocks=False):
        super().__init__()
        if not isinstance(slots, int) or slots < 0:
            raise ValueError('slots must be a nonnegative integer')
        self.base = base
        self.slots = slots
        self.checkpoint_blocks = checkpoint_blocks
        self.cfg = dict(base.cfg)

    def hidden(self, *, ids=None, embeds=None, cache=None, mask=None,
               boundary_features=False):
        if (ids is None) == (embeds is None):
            raise ValueError('supply exactly one of ids or embeds')
        x = self.base.tok_emb(ids) if embeds is None else embeds
        b, t, d = x.shape
        if t < 1 or d != self.cfg['emb_dim']:
            raise ValueError('empty or wrong-dimensional input')
        start = cache.get(0)[0].shape[2] if cache is not None and cache.get(0) is not None else 0
        end = start + t
        if end > self.cfg['context_length']:
            raise ValueError('context limit exceeded, including latent positions')
        if mask is None:
            mask = torch.ones((b, end), dtype=torch.bool, device=x.device)
        if mask.dtype != torch.bool or mask.shape != (b, end):
            raise ValueError('mask must cover all cached and new positions')
        if not mask[:, -1].all() or (mask[:, :-1] & ~mask[:, 1:]).any():
            raise ValueError('nonempty left-padded sequences required')
        x = x * mask[:, start:end, None].to(x.dtype)
        positions = torch.arange(start, end, device=x.device)
        allowed = (torch.arange(end, device=x.device)[None, :] <= positions[:, None])[None, None]
        allowed = allowed & mask[:, None, None, :]
        positions = positions[None].expand(b, -1)
        features = [x[:, -1]] if boundary_features else None
        for layer, block in enumerate(self.base.trf_blocks):
            old = cache.get(layer) if cache is not None else None
            if old is not None and old[0].shape[2] != start:
                raise ValueError('inconsistent cache length')
            # Bind all per-call inputs; recomputation must not see the mutable
            # live cache or a later loop iteration's block/positions.
            call = partial(block, mask=allowed, cos=self.base.cos, sin=self.base.sin,
                           cache=old, pos_ids=positions)
            if self.checkpoint_blocks and self.training and torch.is_grad_enabled():
                x, new = checkpoint(call, x, use_reentrant=False, preserve_rng_state=False)
            else:
                x, new = call(x)
            if cache is not None:
                cache.update(layer, new)
            if boundary_features:
                features.append(x[:, -1])
        x = self.base.final_norm(x)
        if boundary_features:
            # Match the usual hidden-state tuple: embedding, intermediate
            # residual outputs, then final normalized output.
            features[-1] = x[:, -1]
        return x, torch.stack(features, 1) if boundary_features else None

    def prefill(self, prompt_ids, boundary_ids, mask=None, *, features=False):
        """Question -> K feedback states -> common textual answer delimiter.

        Each feedback input is the preceding final normalized hidden state.
        Nothing is detached, including the prompt KV entries during training.
        """
        b, width = prompt_ids.shape
        if mask is None:
            mask = torch.ones_like(prompt_ids, dtype=torch.bool)
        if boundary_ids.ndim != 2 or boundary_ids.shape[0] != b or boundary_ids.shape[1] < 1:
            raise ValueError('one nonempty answer boundary per sequence required')
        cache = KVCache(self.cfg['n_layers'])
        h, _ = self.hidden(ids=prompt_ids, cache=cache, mask=mask)
        for _ in range(self.slots):
            mask = torch.cat((mask, torch.ones((b, 1), dtype=torch.bool, device=mask.device)), 1)
            h, _ = self.hidden(embeds=h[:, -1:], cache=cache, mask=mask)
        mask = torch.cat((mask, torch.ones_like(boundary_ids, dtype=torch.bool)), 1)
        h, feature = self.hidden(ids=boundary_ids, cache=cache, mask=mask,
                                 boundary_features=features)
        return h[:, -1:], feature, cache, mask

    def answer_loss(self, prompt_ids, boundary_ids, target_ids, *, features=False):
        if target_ids.ndim != 2 or target_ids.shape[0] != prompt_ids.shape[0] or target_ids.shape[1] < 1:
            raise ValueError('one nonempty target per sequence required')
        h, feature, cache, mask = self.prefill(prompt_ids, boundary_ids, features=features)
        if target_ids.shape[1] > 1:
            prefix = target_ids[:, :-1]
            mask = torch.cat((mask, torch.ones_like(prefix, dtype=torch.bool)), 1)
            tail, _ = self.hidden(ids=prefix, cache=cache, mask=mask)
            h = torch.cat((h, tail), 1)
        logits = self.base.out_head(h.to(self.cfg['dtype']))
        assert logits.shape[:2] == target_ids.shape
        loss = F.cross_entropy(logits.flatten(0, 1), target_ids.flatten(), reduction='mean')
        return loss, feature, logits

    def text_loss(self, prompt_ids, completion_ids):
        if prompt_ids.shape[0] != 1 or completion_ids.shape[0] != 1:
            raise ValueError('text fitting accumulates unpadded individual rows')
        if not prompt_ids.shape[1] or not completion_ids.shape[1]:
            raise ValueError('empty prompt or completion')
        ids = torch.cat((prompt_ids, completion_ids), 1)
        h, _ = self.hidden(ids=ids[:, :-1])
        h = h[:, prompt_ids.shape[1] - 1:]
        logits = self.base.out_head(h.to(self.cfg['dtype']))
        assert logits.shape[:2] == completion_ids.shape
        return F.cross_entropy(logits.flatten(0, 1), completion_ids.flatten())


def feature_loss(student, teacher):
    """Per-layer SmoothL1 scaled by teacher feature std, with teacher detached.

    The layerwise scale differs explicitly from CODI's batch-global std. No
    minibatch-size-dependent normalization or teacher gradient is introduced.
    """
    if student.shape != teacher.shape or student.ndim != 3:
        raise ValueError('features must match [batch, layer, width]')
    teacher = teacher.detach()
    scale = teacher.std(dim=-1, correction=0, keepdim=True).clamp_min(1e-6)
    return (F.smooth_l1_loss(student, teacher, reduction='none') / scale).mean()
