"""Paired intervention-value targets: g(s,a;y) = loss(D(s),y) - loss(D(T_a(s)),y).

An intervention on the model's own computation (same checkpoint, same target, cloned state). Positive = the
action reduced the teacher-forced next-token loss. These labels are used ONLY offline (controller training and
oracle audits); the deployed controller never sees y.
"""
from __future__ import annotations

import torch

from .losses import token_nll
from .model import EncodedContext, PrismeModel


@torch.no_grad()
def state_loss(model: PrismeModel, S: torch.Tensor, ctx: EncodedContext, targets: torch.Tensor,
               edited: torch.Tensor) -> torch.Tensor:
    logits, _ = model.logits_from_state(S, ctx, edited)
    return token_nll(logits, targets)  # [B,T]


@torch.no_grad()
def paired_gains(model: PrismeModel, S: torch.Tensor, ctx: EncodedContext, targets: torch.Tensor,
                 edited: torch.Tensor, actions: torch.Tensor, step, remaining) -> torch.Tensor:
    """actions: [B,T,n] candidate actions evaluated from the SAME state S (cloned implicitly: apply_action is
    functional). Returns gains [B,T,n] = base_loss - loss_after_action (STOP gives exactly 0)."""
    base = state_loss(model, S, ctx, targets, edited)
    gains = []
    for i in range(actions.shape[-1]):
        a = actions[:, :, i]
        S_a, _ = model.apply_action(S, ctx, a, step, remaining)
        ed = edited | (a != model.cfg.stop_action)
        la = state_loss(model, S_a, ctx, targets, ed)
        gains.append(base - la)
    return torch.stack(gains, -1)


@torch.no_grad()
def all_action_gains(model: PrismeModel, S, ctx, targets, edited, step, remaining) -> torch.Tensor:
    """Exhaustive audit over all nonterminal actions -> [B,T,A]."""
    A = model.cfg.n_actions
    B, T = targets.shape
    acts = torch.arange(A, device=targets.device).view(1, 1, A).expand(B, T, A)
    return paired_gains(model, S, ctx, targets, edited, acts, step, remaining)


def sample_logged_actions(cfg, B, T, policy_pref: torch.Tensor | None, generator=None, device=None):
    """Four candidates per state: current-policy preference (or random if none), a competitive alternative
    (random single/pair of the other group), and two uniformly random actions, balanced singles/pairs.
    Returns actions [B,T,4] and log-propensities [B,T,4] under the mixture used here."""
    A, M = cfg.n_actions, cfg.n_slots
    g = generator
    def rand_single():
        return torch.randint(0, M, (B, T), generator=g, device=device)
    def rand_pair():
        return torch.randint(M, A, (B, T), generator=g, device=device)
    a0 = policy_pref if policy_pref is not None else rand_single()
    a1 = torch.where(a0 < M, rand_pair(), rand_single())
    a2, a3 = rand_single(), rand_pair()
    acts = torch.stack([a0, a1, a2, a3], -1)
    # propensity of each candidate under the logging mixture (uniform over the 4 lanes, lanes 2/3 uniform in group)
    p = torch.zeros(B, T, 4, device=device)
    p[..., 0] = 1.0 if policy_pref is not None else 1.0 / M
    p[..., 1] = torch.where(a0 < M, torch.full_like(p[..., 1], 1.0 / (A - M)), torch.full_like(p[..., 1], 1.0 / M))
    p[..., 2] = 1.0 / M
    p[..., 3] = 1.0 / (A - M)
    return acts, (p / 4).log()