"""Per-token KL divergence kernel.

Direction: KL(teacher || student), in nats:

    d_i = sum_v p_i(v) * (ln p_i(v) - ln q_i(v))

Inputs are FP32 logits from both sides. Both sides are masked to the real
tokenizer vocabulary (dropping padded ``lm_head`` columns) before the softmax.
Log-softmax and the per-token vocabulary sum are computed in float64 — the
protocol's FP32-log-softmax/FP64-accumulate floor, exceeded on the log-softmax
side so the kernel itself contributes < 1e-12 relative error (validated in
``tests/test_kld.py`` against an FP64 SciPy reference).

Alignment contract: teacher row r and student row r are the model outputs
after consuming tokens[0..r], predicting tokens[r+1]. Never shift.
"""
from __future__ import annotations

from dataclasses import dataclass

import numpy as np
import torch


@dataclass
class WindowScores:
    """Per-scored-position diagnostics for one window."""

    kld: np.ndarray  # float64, KL(teacher||student) per position
    teacher_entropy: np.ndarray  # float64, H(p)
    teacher_top1: np.ndarray  # int64 token id
    student_top1: np.ndarray  # int64 token id
    teacher_top1_p: np.ndarray  # float64
    student_top1_p: np.ndarray  # float64
    realized_token: np.ndarray  # int64, tokens[r+1]
    teacher_logp_realized: np.ndarray  # float64
    student_logp_realized: np.ndarray  # float64
    delta_p_realized: np.ndarray  # float64, q(x) - p(x)


def score_window(
    teacher_logits: np.ndarray,
    student_logits: np.ndarray,
    token_ids: np.ndarray,
    vocab_limit: int,
    chunk_rows: int = 256,
) -> WindowScores:
    """Compute per-token KLD and diagnostics for one aligned window."""
    if teacher_logits.ndim != 2 or student_logits.ndim != 2:
        raise ValueError("logits must be rank-2 (positions, vocab)")
    rows, teacher_vocab = teacher_logits.shape
    s_rows, student_vocab = student_logits.shape
    if s_rows != rows:
        raise ValueError(f"row mismatch: teacher {rows} vs student {s_rows}")
    if token_ids.ndim != 1 or token_ids.size != rows + 1:
        raise ValueError(
            f"token_ids must have {rows + 1} entries for {rows} scored positions"
        )
    if vocab_limit <= 0:
        raise ValueError("vocab_limit must be positive")
    if student_vocab < vocab_limit:
        raise ValueError(
            f"student vocab {student_vocab} narrower than tokenizer vocab "
            f"{vocab_limit}; refusing (exact KLD impossible)"
        )
    if teacher_vocab < vocab_limit:
        raise ValueError(
            f"teacher stored vocab {teacher_vocab} narrower than tokenizer "
            f"vocab {vocab_limit}; refusing"
        )
    realized = np.asarray(token_ids[1:], dtype=np.int64)
    if realized.min() < 0 or realized.max() >= vocab_limit:
        raise ValueError("realized token id outside tokenizer vocabulary")

    out = WindowScores(
        kld=np.empty(rows, dtype=np.float64),
        teacher_entropy=np.empty(rows, dtype=np.float64),
        teacher_top1=np.empty(rows, dtype=np.int64),
        student_top1=np.empty(rows, dtype=np.int64),
        teacher_top1_p=np.empty(rows, dtype=np.float64),
        student_top1_p=np.empty(rows, dtype=np.float64),
        realized_token=realized,
        teacher_logp_realized=np.empty(rows, dtype=np.float64),
        student_logp_realized=np.empty(rows, dtype=np.float64),
        delta_p_realized=np.empty(rows, dtype=np.float64),
    )
    for start in range(0, rows, chunk_rows):
        stop = min(start + chunk_rows, rows)
        t = torch.from_numpy(
            np.ascontiguousarray(teacher_logits[start:stop, :vocab_limit])
        ).to(torch.float64)
        s = torch.from_numpy(
            np.ascontiguousarray(student_logits[start:stop, :vocab_limit])
        ).to(torch.float64)
        t_logp = torch.log_softmax(t, dim=-1)
        s_logp = torch.log_softmax(s, dim=-1)
        p = torch.exp(t_logp)
        kld = torch.sum(p * (t_logp - s_logp), dim=-1)
        entropy = -torch.sum(p * t_logp, dim=-1)
        t_top = torch.argmax(t, dim=-1)
        s_top = torch.argmax(s, dim=-1)
        idx = torch.arange(stop - start)
        real = torch.from_numpy(realized[start:stop])
        out.kld[start:stop] = kld.numpy()
        out.teacher_entropy[start:stop] = entropy.numpy()
        out.teacher_top1[start:stop] = t_top.numpy()
        out.student_top1[start:stop] = s_top.numpy()
        out.teacher_top1_p[start:stop] = torch.exp(t_logp[idx, t_top]).numpy()
        out.student_top1_p[start:stop] = torch.exp(s_logp[idx, s_top]).numpy()
        t_lp_real = t_logp[idx, real]
        s_lp_real = s_logp[idx, real]
        out.teacher_logp_realized[start:stop] = t_lp_real.numpy()
        out.student_logp_realized[start:stop] = s_lp_real.numpy()
        out.delta_p_realized[start:stop] = (
            torch.exp(s_lp_real) - torch.exp(t_lp_real)
        ).numpy()
    return out