"""Objectives and diagnostics owned by the Sudoku instantiation."""

from typing import Any

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


class UniquenessLoss(nn.Module):
    """Soft uniqueness regulariser: penalises high-entropy blank-cell distributions.

    A Sudoku puzzle has a *unique* solution when every blank cell's value is
    fully determined by the declared constraints and the given clues.  A model
    that spreads probability across multiple equally-valid digit assignments
    for blank cells implicitly encodes multiple valid solutions.

    We proxy the uniqueness constraint by minimising the mean normalised
    entropy of post-reasoning distributions over blank cells, driving the
    model to commit to specific symbol assignments rather than remaining
    ambiguous.  This is a fully-differentiable continuous analog of the
    iterative nogood-based uniqueness enforcement in Mulamba et al. (2024,
    *Constraints*, doi:10.1007/s10601-024-09372-9).

    Loss::

        L_uniq = mean_{i ∈ blank} H(p_i) / log(|Σ|)

    where H(p_i) = −Σ_s p_{i,s} log p_{i,s} is per-cell entropy and |Σ| is
    the symbol-domain size.  Normalising by log(|Σ|) maps the range to [0, 1]:
    0 when every blank cell is one-hot (fully committed), 1 when every blank
    cell is uniform (maximally ambiguous).
    """

    def forward(
        self,
        probs: torch.Tensor,
        blank_mask: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            probs:      (B, 81, |Σ|) softmax probabilities — post-reasoning.
            blank_mask: (B, 81) bool tensor, True for blank (non-clue) cells.

        Returns:
            Scalar normalised entropy loss in [0, 1].
        """
        flat_blank = blank_mask.reshape(-1)  # (B*81,)
        flat_probs = probs.reshape(-1, probs.shape[-1])  # (B*81, |Σ|)
        blank_probs = flat_probs[flat_blank]  # (n_blank, |Σ|)
        if blank_probs.shape[0] == 0:
            return probs.sum() * 0.0  # differentiable zero
        entropy = -(blank_probs * (blank_probs + 1e-8).log()).sum(-1)  # (n_blank,)
        log_sigma = torch.log(torch.tensor(float(probs.shape[-1]), device=probs.device))
        return entropy.mean() / log_sigma


def compute_sudoku_objective(
    pre_logits: torch.Tensor,
    post_logits: torch.Tensor,
    labels: torch.Tensor,
    puzzle_mask: torch.Tensor,
    *,
    current_epoch: int,
    lambda_ic: float,
    lambda_blank: float,
    blank_decay_epochs: int,
    ic_warmup_epochs: int,
    lambda_uniq: float,
    constraint_loss: nn.Module,
    uniqueness_loss: nn.Module,
) -> dict[str, Any]:
    """Sudoku supervision, curriculum and diagnostics; separate from AS2 reasoning."""
    # ── 6. Loss computation ───────────────────────────────────────────
    # loss_cnn   — direct CE on the CNN concept head (pre_logits) for clue
    #              cells.  PRIMARY gradient signal for the CNN.
    #
    # loss_sup   — CE on the Transformer output (post_logits) for clue
    #              cells.  Keeps the Transformer grounded at clue positions.
    #
    # loss_blank — CE on blank-cell Transformer outputs against GT labels.
    #              Provides direct supervision for blank cells (the old
    #              model had NONE).  Weighted by a curriculum schedule
    #              that decays from full to 0.1× over training so the
    #              model gradually shifts to relying on constraint reasoning.
    #
    # loss_ic    — T_P fixed-point constraint energy on ALL 81 cells.
    #              Penalises deviation from the T_P operator's fixed point.
    #              Ramps up from 0 over ic_warmup_epochs.
    clue_flat = puzzle_mask.reshape(-1).bool()  # (B*81,)
    blank_flat = ~clue_flat
    labels_flat = labels.reshape(-1)

    # Direct CNN supervision (most important for learning digits)
    loss_cnn = F.cross_entropy(
        pre_logits.reshape(-1, 9)[clue_flat],
        labels_flat[clue_flat],
    )

    # Transformer output supervision at clue positions
    loss_sup = F.cross_entropy(
        post_logits.reshape(-1, 9)[clue_flat],
        labels_flat[clue_flat],
    )

    # Blank-cell CE supervision with curriculum decay
    if blank_flat.any():
        loss_blank = F.cross_entropy(
            post_logits.reshape(-1, 9)[blank_flat],
            labels_flat[blank_flat],
        )
    else:
        loss_blank = torch.tensor(0.0, device=post_logits.device)

    # Curriculum weight: linear decay from 1.0 to 0.1 over blank_decay_epochs
    epoch = current_epoch
    blank_scale = max(0.1, 1.0 - 0.9 * epoch / max(1, blank_decay_epochs))

    # T_P constraint loss with warmup ramp
    soft_model = F.softmax(post_logits, dim=-1)  # (B, 81, 9)
    loss_ic = constraint_loss(soft_model)
    ic_scale = min(1.0, epoch / max(1, ic_warmup_epochs))

    # Uniqueness regulariser: penalise high blank-cell entropy post-reasoning
    blank_mask_2d = ~puzzle_mask.bool()  # (B, 81) bool
    loss_uniq = uniqueness_loss(soft_model, blank_mask_2d)

    loss = (
        loss_cnn
        + loss_sup
        + lambda_blank * blank_scale * loss_blank
        + lambda_ic * ic_scale * loss_ic
        + lambda_uniq * ic_scale * loss_uniq
    )

    # ── 7. Diagnostics ────────────────────────────────────────────────
    # Blank-cell entropy: high → model is uncertain, low → confident
    if blank_flat.any():
        blank_probs = soft_model.reshape(-1, 9)[blank_flat]  # (n_blank, 9)
        blank_entropy = -(blank_probs * (blank_probs + 1e-8).log()).sum(-1).mean()
    else:
        blank_entropy = torch.tensor(0.0, device=post_logits.device)

    # ── 8. Predictions ────────────────────────────────────────────────
    predictions = post_logits.argmax(dim=-1)  # (B, 81)

    # Flat clue-cell tensors:
    # Use pre_logits (CNN output) so that clue_accuracy reflects CNN quality
    # independent of the Transformer.
    clue_logits = pre_logits.reshape(-1, 9)[clue_flat]  # (n_clue, 9)
    clue_labels = labels_flat[clue_flat]  # (n_clue,)

    # ── 9. Perception-vs-reasoning diagnostic tensors ─────────────────
    # Pre-reasoning CNN logits for blank cells (measures perception quality
    # on the cells that actually require reasoning to fill).
    blank_pre_logits = pre_logits.reshape(-1, 9)[blank_flat]  # (n_blank, 9)
    blank_labels = labels_flat[blank_flat]  # (n_blank,)

    # Post-reasoning logits split by clue / blank (measures how much
    # the Transformer improves on CNN quality, and whether it corrupts
    # the confident clue-cell predictions).
    blank_post_logits = post_logits.reshape(-1, 9)[blank_flat]  # (n_blank, 9)
    clue_post_logits = post_logits.reshape(-1, 9)[clue_flat]  # (n_clue, 9)

    return {
        "loss": loss,
        "logits": post_logits,  # (B, 81, 9) — for board-level metrics
        "predictions": predictions,  # (B, 81)
        # ── CNN (pre-reasoning) split ─────────────────────────────────
        "clue_logits": clue_logits,  # (n_clue, 9)  CNN on clue cells
        "clue_labels": clue_labels,  # (n_clue,)
        "blank_pre_logits": blank_pre_logits,  # (n_blank, 9) CNN on blank cells
        "blank_labels": blank_labels,  # (n_blank,)
        # ── Transformer (post-reasoning) split ───────────────────────
        "blank_post_logits": blank_post_logits,  # (n_blank, 9)
        "clue_post_logits": clue_post_logits,  # (n_clue, 9)
        # ── Scalar losses ─────────────────────────────────────────────
        "loss_cnn": loss_cnn.detach(),
        "loss_sup": loss_sup.detach(),
        "loss_blank": loss_blank.detach(),
        "loss_ic": loss_ic.detach(),
        "loss_uniq": loss_uniq.detach(),
        # Weighted contributions (what actually drives the total loss)
        "loss_blank_w": (lambda_blank * blank_scale * loss_blank).detach(),
        "loss_ic_w": (lambda_ic * ic_scale * loss_ic).detach(),
        "loss_uniq_w": (lambda_uniq * ic_scale * loss_uniq).detach(),
        # Curriculum scales for diagnostics
        "blank_scale": blank_scale,
        "ic_scale": ic_scale,
        "blank_entropy": blank_entropy.detach(),
    }