"""
Schema-conditioned BERT candidate scorer.

One BERT encoder with a single scalar head scores (state, question+candidate)
pairs. Deterministic code groups the scalar logits per question and decodes
them into `choice`, `noul`, and `score` answers. The schema (instructions,
criteria, candidate ids) is read at inference time, never baked into weights.
"""

from __future__ import annotations

import json
import os
import random
from dataclasses import dataclass
from typing import Any

import torch
from torch import nn
from transformers import AutoModelForSequenceClassification, AutoTokenizer

SUPPORTED_PRIMITIVES = ("choice", "noul", "score")

# Serialization layout. Must match between training and inference.
# SCORER_CANDIDATE_FIRST=0 restores the trailing-candidate layout (ablation only).
CANDIDATE_FIRST = os.environ.get("SCORER_CANDIDATE_FIRST", "1") == "1"


# --------------------------------------------------------------------------- #
# 1. Request compiler
# --------------------------------------------------------------------------- #
def serialize(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=False)


def compile_question(state: Any, question: dict) -> tuple[list[tuple[str, str]], list[tuple[str, str]]]:
    """
    Turn one question into (candidates, text_pairs).

    candidates: [(candidate_id, description), ...] in decode order.
    text_pairs: [(sequence_a, sequence_b), ...], one per candidate, where
        sequence_a = the state and
        sequence_b = the full question schema plus the candidate under review.
    """
    kind = question.get("type")

    if kind == "choice":
        criteria = question.get("criteria")
        if not isinstance(criteria, dict):
            raise ValueError("choice questions need a dict of criteria {id: description}")
        candidates = [(str(k), str(v)) for k, v in criteria.items()]

    elif kind == "score":
        criteria = question.get("criteria")
        if not isinstance(criteria, list):
            raise ValueError("score questions need an ordered list of level descriptions")
        candidates = [(str(i), str(d)) for i, d in enumerate(criteria)]

    elif kind == "noul":
        criteria = question.get("criteria") or {}
        candidates = [
            ("false", str(criteria.get("false", "No. The proposition is false for this state."))),
            ("true", str(criteria.get("true", "Yes. The proposition is true for this state."))),
        ]

    else:
        raise ValueError(f"Unsupported primitive: {kind!r} (expected one of {SUPPORTED_PRIMITIVES})")

    if len(candidates) < 2:
        raise ValueError("This implementation requires at least two candidates per question.")

    state_text = state if isinstance(state, str) else serialize(state)
    schema = {
        "type": kind,
        "instructions": question.get("instructions", ""),
    }
    if question.get("criteria") is not None:
        schema["criteria"] = question["criteria"]

    # Candidate first: the only tokens that differ between a question's
    # candidates sit right after [SEP], where the encoder attends most easily.
    # Trailing placement (candidate after the full schema) trains much slower.
    if CANDIDATE_FIRST:
        pairs = [
            (state_text, serialize({"candidate": {"id": cid, "description": desc}, **schema}))
            for cid, desc in candidates
        ]
    else:
        pairs = [
            (state_text, serialize({**schema, "candidate": {"id": cid, "description": desc}}))
            for cid, desc in candidates
        ]
    return candidates, pairs


# --------------------------------------------------------------------------- #
# 2. Decoder (logits -> structured answer)
# --------------------------------------------------------------------------- #
def decode_answer(kind: str, candidates: list[tuple[str, str]], logits: torch.Tensor) -> dict:
    probabilities = logits.float().softmax(dim=0)

    if kind == "noul":
        return {"type": "noul", "noul": float(probabilities[1])}

    if kind == "choice":
        selected = int(probabilities.argmax())
        return {
            "type": "choice",
            "choice": candidates[selected][0],
            "probabilities": {cid: float(probabilities[i]) for i, (cid, _) in enumerate(candidates)},
        }

    if kind == "score":
        expected = sum(i * float(probabilities[i]) for i in range(len(candidates)))
        return {
            "type": "score",
            "score": expected,
            "probabilities": [float(p) for p in probabilities],
            "legend": dict(candidates),
        }

    raise ValueError(f"Unsupported primitive: {kind!r}")


# --------------------------------------------------------------------------- #
# 3. Model wrapper: inference
# --------------------------------------------------------------------------- #
def pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


class LocalSystemOne:
    """TypeSafe-shaped adapter over a trained one-logit candidate scorer."""

    def __init__(self, checkpoint: str, max_length: int | None = None, device: torch.device | None = None):
        self.device = device or pick_device()
        self.tokenizer = AutoTokenizer.from_pretrained(checkpoint)
        self.model = AutoModelForSequenceClassification.from_pretrained(checkpoint, dtype=torch.float32).to(self.device).eval()
        if self.model.config.num_labels != 1:
            raise ValueError("Expected a trained one-logit candidate scorer (num_labels == 1).")
        # Default to the architecture's positional limit (512 for BERT/DeBERTa, 8192 for ModernBERT).
        self.max_length = max_length or getattr(self.model.config, "max_position_embeddings", 512)

    @torch.inference_mode()
    def score_pairs(self, pairs: list[tuple[str, str]], batch_size: int = 32) -> torch.Tensor:
        chunks = []
        for start in range(0, len(pairs), batch_size):
            batch = pairs[start : start + batch_size]
            encoded = self.tokenizer(
                [a for a, _ in batch],
                [b for _, b in batch],
                padding=True,
                truncation=False,
                return_tensors="pt",
            )
            if encoded["input_ids"].shape[1] > self.max_length:
                raise ValueError(
                    f"State + question + criteria exceed the input limit "
                    f"({encoded['input_ids'].shape[1]} > {self.max_length} tokens)."
                )
            encoded = {k: v.to(self.device) for k, v in encoded.items()}
            chunks.append(self.model(**encoded).logits.squeeze(-1).float().cpu())
        return torch.cat(chunks) if chunks else torch.empty(0)

    def system_one(self, state: Any, questions: dict[str, dict], batch_size: int = 32) -> dict:
        all_pairs: list[tuple[str, str]] = []
        groups = []
        for name, question in questions.items():
            candidates, pairs = compile_question(state, question)
            start = len(all_pairs)
            all_pairs.extend(pairs)
            groups.append((name, question["type"], candidates, start, len(all_pairs)))

        if not all_pairs:
            return {"model": "local-bert-scorer", "answers": {}}

        logits = self.score_pairs(all_pairs, batch_size=batch_size)
        answers = {
            name: decode_answer(kind, candidates, logits[start:end])
            for name, kind, candidates, start, end in groups
        }
        return {"model": "local-bert-scorer", "answers": answers}


# --------------------------------------------------------------------------- #
# 4. Training: grouped distribution loss
# --------------------------------------------------------------------------- #
@dataclass
class Example:
    state: Any
    question: dict
    target: list[float]  # distribution over the compiled candidates, sums to 1


def target_for(kind: str, label: Any, n_candidates: int) -> list[float]:
    """Helper to build a target distribution from a plain label."""
    if kind == "noul":
        p = float(label)
        return [1.0 - p, p]
    if kind == "choice":
        # label is the index of the correct candidate
        t = [0.0] * n_candidates
        t[int(label)] = 1.0
        return t
    if kind == "score":
        # label is either a level index or a full distribution
        if isinstance(label, (list, tuple)):
            return [float(x) for x in label]
        t = [0.0] * n_candidates
        t[int(label)] = 1.0
        return t
    raise ValueError(kind)


def grouped_distribution_loss(logits: torch.Tensor, targets: torch.Tensor, group_sizes: list[int]) -> torch.Tensor:
    """
    logits, targets: flat 1-D tensors, concatenation of per-question groups.
    Loss = mean over questions of cross-entropy(target_dist, softmax(group logits)).
    """
    losses = []
    offset = 0
    for size in group_sizes:
        g_logits = logits[offset : offset + size]
        g_target = targets[offset : offset + size]
        losses.append(-(g_target * g_logits.log_softmax(dim=0)).sum())
        offset += size
    return torch.stack(losses).mean()


def train_scorer(
    base_model: str,
    examples: list[Example],
    output_dir: str,
    *,
    epochs: int = 2,
    questions_per_batch: int = 8,
    lr: float = 3e-5,
    max_length: int = 256,
    seed: int = 0,
    device: torch.device | None = None,
    log_every: int = 25,
    eval_fn=None,
    bf16: bool | None = None,
    gradient_checkpointing: bool = False,
) -> str:
    """
    Fine-tune an encoder as a one-logit candidate scorer and save it.

    bf16: autocast matmuls to bfloat16. Default False; see note below.
    gradient_checkpointing: trade compute for activation memory (large models
          on 24 GB cards).
    """
    random.seed(seed)
    torch.manual_seed(seed)
    device = device or pick_device()
    # Default fp32. Candidates of one question differ by a handful of tokens, so
    # their logits differ by tiny amounts early in training; bf16 rounds those to
    # identical values and the grouped loss gets no gradient (observed: bert-base
    # stuck at chance and DeBERTa-v3-large NaN under bf16 autocast).
    if bf16 is None:
        bf16 = False

    tokenizer = AutoTokenizer.from_pretrained(base_model)
    # dtype=float32: transformers 5 otherwise keeps the checkpoint's stored dtype;
    # deberta-v3-large ships in fp16, and fp16 master weights train to NaN.
    model = AutoModelForSequenceClassification.from_pretrained(base_model, num_labels=1, dtype=torch.float32)
    if gradient_checkpointing:
        model.gradient_checkpointing_enable()
    model.to(device).train()
    print(f"  base={base_model} params={sum(p.numel() for p in model.parameters()) / 1e6:.0f}M "
          f"device={device} bf16={bf16} grad_ckpt={gradient_checkpointing}")

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    total_steps = epochs * ((len(examples) + questions_per_batch - 1) // questions_per_batch)
    warmup = max(1, int(0.06 * total_steps))

    def lr_lambda(step: int) -> float:
        if step < warmup:
            return (step + 1) / warmup
        return max(0.0, (total_steps - step) / max(1, total_steps - warmup))

    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

    step = 0
    for epoch in range(epochs):
        order = list(range(len(examples)))
        random.shuffle(order)
        running = 0.0
        for b in range(0, len(order), questions_per_batch):
            batch = [examples[i] for i in order[b : b + questions_per_batch]]
            pairs, targets, sizes = [], [], []
            for ex in batch:
                candidates, ex_pairs = compile_question(ex.state, ex.question)
                if len(ex.target) != len(candidates):
                    raise ValueError("target distribution length must equal candidate count")
                pairs.extend(ex_pairs)
                targets.extend(ex.target)
                sizes.append(len(candidates))

            encoded = tokenizer(
                [a for a, _ in pairs],
                [b_ for _, b_ in pairs],
                padding=True,
                truncation="only_first",  # truncate the state, never the schema
                max_length=max_length,
                return_tensors="pt",
            )
            encoded = {k: v.to(device) for k, v in encoded.items()}
            with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=bf16):
                logits = model(**encoded).logits.squeeze(-1).float()
            loss = grouped_distribution_loss(logits, torch.tensor(targets, device=device), sizes)

            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()

            running += loss.item()
            step += 1
            if step % log_every == 0:
                print(f"  epoch {epoch + 1}/{epochs} step {step}/{total_steps} loss {running / log_every:.4f}")
                running = 0.0

        if eval_fn is not None:
            model.eval()
            eval_fn(model, tokenizer, epoch)
            model.train()

    model.eval()
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    return output_dir