"""Diagnostic oracles (use the target; NEVER deployable):
  * greedy one-step oracle: at each step pick argmax_a g(s,a), STOP when max gain <= 0;
  * two-action lookahead oracle on audit subsets: best pair (a1, a2) by total gain;
  * uniform / random / fixed-budget reference policies for the opportunity analysis.
"""
from __future__ import annotations

import torch

from .counterfactual import all_action_gains, state_loss
from .model import PrismeModel


@torch.no_grad()
def greedy_oracle_rollout(model: PrismeModel, tokens, targets, max_budget: int, ctx=None, threshold: float = 0.0):
    """Returns dict(loss_per_step [R+1,B,T], actions [B,T,R], n_edits [B,T])."""
    cfg = model.cfg
    ctx = model.encode(tokens) if ctx is None else ctx
    S = ctx.S0
    B, T = targets.shape
    edited = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)
    stopped = torch.zeros_like(edited)
    losses = [state_loss(model, S, ctx, targets, edited)]
    trace = torch.full((B, T, max_budget), cfg.stop_action, dtype=torch.long, device=tokens.device)
    for r in range(max_budget):
        gains = all_action_gains(model, S, ctx, targets, edited, r, max_budget - r - 1)
        best_g, best_a = gains.max(-1)
        a = torch.where((best_g > threshold) & ~stopped, best_a, torch.full_like(best_a, cfg.stop_action))
        stopped = stopped | (a == cfg.stop_action)
        trace[:, :, r] = a
        S, _ = model.apply_action(S, ctx, a, r, max_budget - r - 1)
        edited = edited | (a != cfg.stop_action)
        losses.append(state_loss(model, S, ctx, targets, edited))
    return {"loss_per_step": torch.stack(losses), "actions": trace, "n_edits": (trace != cfg.stop_action).sum(-1)}


@torch.no_grad()
def two_step_oracle(model: PrismeModel, tokens, targets, ctx=None):
    """Exhaustive best pair of actions (a1 then a2, a2 may be STOP) on a (small) batch -> total gain [B,T],
    and the best greedy-then-greedy total for comparison. Cost: A + A*(A+1) editor passes."""
    cfg = model.cfg
    A = cfg.n_actions
    ctx = model.encode(tokens) if ctx is None else ctx
    S0 = ctx.S0
    B, T = targets.shape
    no = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)
    yes = ~no
    base = state_loss(model, S0, ctx, targets, no)
    g1 = all_action_gains(model, S0, ctx, targets, no, 0, 1)             # [B,T,A]
    best_total = g1.clone()                                               # a2 = STOP
    best_pair_a2 = torch.full((B, T, A), cfg.stop_action, dtype=torch.long, device=tokens.device)
    for a1 in range(A):
        act1 = torch.full((B, T), a1, device=tokens.device)
        S1, _ = model.apply_action(S0, ctx, act1, 0, 1)
        g2 = all_action_gains(model, S1, ctx, targets, yes, 1, 0)        # [B,T,A]
        tot = g1[..., a1:a1 + 1] + g2
        m, am = tot.max(-1)
        better = m > best_total[..., a1]
        best_total[..., a1] = torch.where(better, m, best_total[..., a1])
        best_pair_a2[..., a1] = torch.where(better, am, best_pair_a2[..., a1])
    two_step_best, a1_star = best_total.max(-1)
    # greedy baseline: a1 = argmax g1, then best a2 given a1 (already in best_total[a1_greedy])
    g1_best, a1_greedy = g1.max(-1)
    greedy_then_best = best_total.gather(-1, a1_greedy.unsqueeze(-1)).squeeze(-1)
    return {"base_loss": base, "one_step_best": g1_best.clamp(min=0), "two_step_best": two_step_best.clamp(min=0),
            "greedy_then_best": greedy_then_best.clamp(min=0), "a1_star": a1_star, "a1_greedy": a1_greedy}


@torch.no_grad()
def fixed_policy_rollout(model: PrismeModel, tokens, targets, actions: torch.Tensor):
    """Loss trajectory under a forced action tensor [B,T,R]."""
    cfg = model.cfg
    ctx = model.encode(tokens)
    S = ctx.S0
    B, T, R = actions.shape
    edited = torch.zeros(B, T, dtype=torch.bool, device=tokens.device)
    losses = [state_loss(model, S, ctx, targets, edited)]
    for r in range(R):
        a = actions[:, :, r]
        S, _ = model.apply_action(S, ctx, a, r, R - r - 1)
        edited = edited | (a != cfg.stop_action)
        losses.append(state_loss(model, S, ctx, targets, edited))
    return torch.stack(losses)