import math
import pytest
import torch
from nexora.data import prepare
from nexora.compute import Estimate, topology, kv_cache_bytes
from nexora.evaluation import pass_at_k, wilson, percentiles, word_error_rate
from nexora.posttraining import dpo_loss, group_advantages, masked_sft_loss, rejection_sample


def record(text, **kw):
    return {"id": "x", "text": text, "source": "original", "license": "MIT", "domain": "text", **kw}


def test_pipeline_filters_and_provenance(tmp_path):
    good = "The engineering document describes reliable software with independent tests."
    held = "A hidden evaluation question asks about a particular graph algorithm and its runtime."
    rows = [record(good), record(good), record("This data has an unknown license and must not be admitted.", license="unknown"),
            record("Do not admit this private credential hf_" + "a"*30), record(held),
            record("Contact test@example.com to obtain further documentation on the system.")]
    report = prepare(rows, tmp_path, holdouts=[held])
    assert report["accepted"] == 2
    assert report["rejected"] == {"exact_duplicate": 1, "license_or_provenance_or_split": 1, "secret": 1, "contamination": 1}
    text = (tmp_path / "records.jsonl").read_text()
    assert "test@example.com" not in text and "[EMAIL]" in text


def test_compute_formulas():
    r = Estimate(120e9, 12e9, 2e12, 1024).calculate()
    assert r["weight_GB"]["bf16"] == 240
    assert r["training_flops"] == 6*12e9*2e12
    assert topology(8, 8, 2, 2, 1, 16, 8)["world"] == 64
    assert kv_cache_bytes(4, 2, 32, 256) == 262144
    with pytest.raises(ValueError):
        topology(8, 8, 8, 8, 1, 8)


def test_metrics():
    assert pass_at_k(10, 2, 1) == pytest.approx(.2)
    assert pass_at_k(10, 2, 3) == pytest.approx(1-56/120)
    assert wilson(5, 10)[0] < .5 < wilson(5, 10)[1]
    assert percentiles([1, 2, 3, 4])["p95"] == 4
    assert word_error_rate("one two three", "one four three") == pytest.approx(1/3)
    with pytest.raises(ValueError):
        pass_at_k(0, 0, 1)


def test_posttraining_gradients():
    chosen = torch.tensor([2.0, 3.0], requires_grad=True)
    loss = dpo_loss(chosen, torch.zeros(2), torch.zeros(2), torch.zeros(2))
    loss.backward()
    assert (chosen.grad < 0).all()
    assert torch.equal(group_advantages(torch.ones(2, 4)), torch.zeros(2, 4))
    logits = torch.randn(1, 3, 5, requires_grad=True)
    loss = masked_sft_loss(logits, torch.tensor([[1, 2, 3]]), torch.tensor([[False, True, True]]))
    loss.backward()
    assert torch.equal(logits.grad[0, 0], torch.zeros(5))
    assert rejection_sample(["pass", "fail", "error"], lambda s: s == "pass") == ["pass"]