"""Shared held-out loss evaluator: one definition of "mean NLL" for every arm of a comparison.

Standard library only, so experiment scripts can copy this file next to themselves and `import loss_eval`.

Masking: a token counts when its loss weight is > 0. Weights are only a mask; their size is ignored, so a recipe that
normalises weights to sum to 1 per example gets the same numbers as one that uses 0/1 weights.

Aggregations (pick one before training and report it):
- token_mean:   sum of NLL over every counted token / number of counted tokens. Long responses weigh more.
- example_mean: mean over examples of (that example's NLL sum / its counted tokens). Every example weighs the same.
Examples with no counted tokens are left out of both and reported as n_empty.

Totals are kept per example, so results do not depend on how examples were split into batches. Never average
per-batch `loss:mean` values: that is neither of the two.
"""

import math

AGGREGATIONS = ("token_mean", "example_mean")
DEFAULT_AGGREGATION = "token_mean"
REL_TOLERANCE = 1e-6


def _floats(values):
    if hasattr(values, "tolist"):
        values = values.tolist()
    elif hasattr(values, "data"):
        values = values.data
    return [float(v) for v in values]


def example_nll(logprobs, weights):
    """(NLL summed over tokens with weight > 0, number of those tokens) for one example."""
    logprobs, weights = _floats(logprobs), _floats(weights)
    if len(logprobs) != len(weights):
        raise ValueError(f"{len(logprobs)} logprobs but {len(weights)} weights")
    counted = [lp for lp, w in zip(logprobs, weights) if w > 0]
    return -sum(counted), len(counted)


class LossEval:
    """Accumulates per-example NLL totals for one arm and reports both aggregations."""

    def __init__(self, aggregation=DEFAULT_AGGREGATION):
        if aggregation not in AGGREGATIONS:
            raise ValueError(f"aggregation must be one of {AGGREGATIONS}, got {aggregation!r}")
        self.aggregation = aggregation
        self.examples = []

    def add(self, logprobs, weights):
        self.examples.append(example_nll(logprobs, weights))

    def add_output(self, output, weights_per_example):
        """Add a Tinker forward/forward_backward output ("cross_entropy" loss) for the batch it scored."""
        for item, weights in zip(output.loss_fn_outputs, weights_per_example, strict=True):
            self.add(item["logprobs"], weights)

    def record(self):
        counted = [(s, n) for s, n in self.examples if n > 0]
        nll_sum = sum(s for s, _ in counted)
        n_tokens = sum(n for _, n in counted)
        token_mean = nll_sum / n_tokens if n_tokens else None
        example_mean = sum(s / n for s, n in counted) / len(counted) if counted else None
        return {
            "aggregation": self.aggregation,
            "value": token_mean if self.aggregation == "token_mean" else example_mean,
            "token_mean": token_mean,
            "example_mean": example_mean,
            "nll_sum": nll_sum,
            "n_tokens": n_tokens,
            "n_examples": len(counted),
            "n_empty": len(self.examples) - len(counted),
        }


def report(aggregation, arms):
    """The `loss_eval` record a run reports: {"aggregation", "arms": {arm name: LossEval.record()}}."""
    return {"aggregation": aggregation, "arms": {name: arm.record() for name, arm in arms.items()}}


def requested_aggregation(spec):
    """The aggregation a spec asks for: eval.aggregation when set, else read from the eval metric text, else None."""
    eval_spec = (spec or {}).get("eval") or {}
    explicit = str(eval_spec.get("aggregation") or "").strip().lower().replace("-", "_").replace(" ", "_")
    if explicit in AGGREGATIONS:
        return explicit
    text = " ".join(str(eval_spec.get(k) or "") for k in ("metric", "benchmark")).lower().replace("-", " ")
    if any(p in text for p in ("per example", "example mean", "example weighted", "mean of example", "per sequence")):
        return "example_mean"
    if any(p in text for p in ("per token", "token mean", "token weighted", "per assistant token", "nats/token")):
        return "token_mean"
    return None


def is_loss_metric(spec):
    eval_spec = (spec or {}).get("eval") or {}
    text = " ".join(str(eval_spec.get(k) or "") for k in ("metric", "benchmark")).lower()
    return any(word in text for word in ("nll", "loss", "perplexity", "cross entropy", "cross-entropy", "log likelihood"))


def check(spec, loss_eval):
    """Problems that make a run's loss comparison unusable for the spec's verdict; [] when it is consistent.

    Checks that the reported aggregation is the one the spec asked for, that every arm used it on the same tokens and
    examples, and that each arm's value really is that aggregation of its own totals."""
    requested = requested_aggregation(spec)
    if not isinstance(loss_eval, dict) or not isinstance(loss_eval.get("arms"), dict) or not loss_eval["arms"]:
        if requested or is_loss_metric(spec):
            return ["no loss_eval record: the loss aggregation behind the reported numbers is unknown"]
        return []
    problems = []
    aggregation = loss_eval.get("aggregation")
    if aggregation not in AGGREGATIONS:
        problems.append(f"loss_eval aggregation {aggregation!r} is not one of {', '.join(AGGREGATIONS)}")
    elif requested and aggregation != requested:
        problems.append(f"spec asks for {requested} but the run reported {aggregation}")
    counts = {}
    for name, arm in loss_eval["arms"].items():
        if not isinstance(arm, dict):
            problems.append(f"arm {name}: no loss record")
            continue
        if arm.get("aggregation", aggregation) != aggregation:
            problems.append(f"arm {name} used {arm.get('aggregation')}, not the run's {aggregation}")
        problems += _arm_problems(name, arm, aggregation)
        counts[name] = (arm.get("n_tokens"), arm.get("n_examples"))
    if len(set(counts.values())) > 1:
        problems.append("arms were scored on different tokens or examples: "
                        + ", ".join(f"{n} {t} tokens / {e} examples" for n, (t, e) in counts.items()))
    return problems


def _arm_problems(name, arm, aggregation):
    nll_sum, n_tokens, value = arm.get("nll_sum"), arm.get("n_tokens"), arm.get("value")
    if not all(_number(v) for v in (nll_sum, n_tokens, value)) or n_tokens <= 0:
        return [f"arm {name}: needs numeric value, nll_sum and n_tokens > 0"]
    token_mean = nll_sum / n_tokens
    reported_token_mean = arm.get("token_mean", token_mean)
    if not _number(reported_token_mean) or not _close(reported_token_mean, token_mean):
        return [f"arm {name}: token_mean {reported_token_mean} is not nll_sum / n_tokens = {token_mean:.6g}"]
    if aggregation == "token_mean" and not _close(value, token_mean):
        return [f"arm {name}: value {value} is not the token mean {token_mean:.6g} (nll_sum / n_tokens)"]
    if aggregation == "example_mean" and _number(arm.get("example_mean")) and not _close(value, arm["example_mean"]):
        return [f"arm {name}: value {value} is not its example_mean {arm['example_mean']}"]
    return []


def _number(value):
    return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value)


def _close(a, b):
    return math.isclose(a, b, rel_tol=REL_TOLERANCE, abs_tol=1e-12)