"""The prior-fitted network: the task prior, context-only attention, in-context
inference and the Delphes-to-GamGam workflow."""

import json
from dataclasses import replace

import pytest
import torch
from conftest import GAMGAM_FIXTURES, TTH_FIXTURE, fixture_config, random_events

from gnn4colliders.cli import main
from gnn4colliders.config import ModelConfig, TaskConfig, save_config
from gnn4colliders.data import pack_events
from gnn4colliders.models import NODE_TYPE, build_model
from gnn4colliders.pfn import (
    ABSENT_LOGIT,
    TaskBatch,
    TaskLoader,
    context_ranks,
    event_observables,
    in_context_forward,
    rank_keys,
)
from gnn4colliders.sequences import ENERGY, PT, SequenceBatch, collate_sequence_batch

SMALL = {"hid_size": 16, "n_layers": 1, "n_heads": 2, "pair_bias": True}
PRIOR = {
    "name": "pfn",
    "task_classes": [2, 3],
    "context_events": [8, 16],
    "query_events": 4,
    "steps_per_epoch": 3,
    "validation_steps": 2,
}


def _events():
    events = random_events(200, classes=4)
    events.object_features[:, 6] = torch.randint(
        0, 5, (events.object_features.shape[0],)
    ).float()
    return events


def _loader(events, train=True, **prior):
    config = TaskConfig(**{**PRIOR, **prior})
    return TaskLoader(events, torch.arange(len(events)), config, 64, train=train)


#: The mean-pooled PFN, summary tokens, alternating attention and the
#: per-task scale of momenta.
VARIANTS = {
    "pooled": {},
    "summary": {"pfn_summary_tokens": 2},
    "alternating": {"pfn_summary_tokens": 2, "pfn_alternating": True},
    "normalized": {"pfn_summary_tokens": 1, "pfn_normalize": True},
    # The options of TabPFN-3.5, with and without alternating attention.
    "tabpfn": {
        "pfn_summary_tokens": 2,
        "pfn_alternating": True,
        "pfn_concat": True,
        "pfn_retrieval": True,
        "pfn_qk_norm": True,
        "pfn_query_kv": True,
        "pfn_fourier": 3,
        "pfn_ranks": True,
    },
    "retrieval": {"pfn_summary_tokens": 1, "pfn_retrieval": True, "pfn_ranks": True},
}


def _pfn(**options):
    torch.manual_seed(0)
    config = ModelConfig(name="pfn", pfn_layers=2, max_classes=4, **SMALL, **options)
    model = build_model(config, num_features=7, num_node_types=5).eval()
    # Maps that start at zero (Fourier features, ranks) are tested non-zero.
    with torch.no_grad():
        for parameter in model.parameters():
            if not parameter.any():
                parameter.normal_(std=0.3)
    return model


def _task(events, context, queries):
    """One task: the ``context`` events with their labels, then ``queries``."""
    indices = torch.cat((context, queries))
    return TaskBatch(
        events=collate_sequence_batch(events, indices),
        context_labels=events.labels[context][None].long(),
        num_classes=torch.tensor([4]),
        labels=events.labels[queries],
        weight=torch.ones(len(queries)),
        event_index=queries,
    )


def test_tasks_follow_the_prior():
    events = _events()
    for batch in _loader(events, prior_classes=[1, 2, 3]):
        tasks, context = batch.context_labels.shape
        assert 8 <= context <= 16 and tasks == 64 // (context + 4)
        assert ((batch.num_classes >= 2) & (batch.num_classes <= 3)).all()
        index = batch.events.event_index.view(tasks, -1)
        queries = batch.labels.view(tasks, -1)
        for task in range(tasks):
            k = int(batch.num_classes[task])
            labels = torch.cat((batch.context_labels[task], queries[task]))
            # Every class has a context event, and each task label is one
            # original class out of prior_classes.
            assert set(batch.context_labels[task].tolist()) == set(range(k))
            original = events.labels[index[task]]
            for label in range(k):
                assert len(set(original[labels == label].tolist())) == 1
            assert not set(original.tolist()) - {1, 2, 3}


def test_ranks_share_the_context_size_of_a_step_but_not_its_events():
    events = _events()
    config = TaskConfig(**{**PRIOR, "context_events": [8, 40]})
    first = [
        next(
            iter(TaskLoader(events, torch.arange(200), config, 64, train=True, rank=r))
        )
        for r in (0, 1)
    ]
    sizes = [batch.context_labels.shape for batch in first]
    assert sizes[0] == sizes[1]
    assert not torch.equal(first[0].events.event_index, first[1].events.event_index)


def _originals(events, batch):
    """Per task, the original labels behind each task label."""
    tasks = batch.context_labels.shape[0]
    index = batch.events.event_index.view(tasks, -1)
    labels = torch.cat((batch.context_labels, batch.labels.view(tasks, -1)), 1)
    return [
        [set(events.labels[index[t]][labels[t] == k].tolist()) for k in range(n)]
        for t, n in enumerate(batch.num_classes.tolist())
    ]


def test_held_out_labels_make_the_validation_tasks_only():
    events = _events()
    for train, allowed in ((True, {0, 1}), (False, {2, 3})):
        for batch in _loader(events, train=train, holdout_classes=[2, 3]):
            for classes in _originals(events, batch):
                assert set().union(*classes) <= allowed


def test_mixed_classes_draw_on_disjoint_labels():
    events = _events()
    loader = _loader(events, task_classes=[2, 2], mixture_labels=[2, 2])
    for batch in loader:
        for first, second in _originals(events, batch):
            assert 1 <= len(first) <= 2 and 1 <= len(second) <= 2
            assert not first & second


def test_class_proportions_survive_a_zero_random_draw(monkeypatch):
    rand = torch.rand
    monkeypatch.setattr(torch, "rand", lambda *a, **k: rand(*a, **k) * 0)
    loader = _loader(_events(), mixture_labels=[1, 2])
    step = loader.sample(torch.Generator().manual_seed(0))
    assert (step.labels >= 0).all()


def test_training_tasks_vary_by_epoch_and_validation_tasks_do_not():
    events = _events()

    def first(loader, epoch):
        loader.set_epoch(epoch)
        return next(iter(loader)).events.event_index

    train, validation = _loader(events), _loader(events, train=False)
    assert torch.equal(first(train, 0), first(train, 0))
    assert not torch.equal(first(train, 0), first(train, 1))
    assert torch.equal(first(validation, 0), first(validation, 1))


@pytest.mark.parametrize("variant", VARIANTS)
def test_queries_see_only_the_context(variant):
    model = _pfn(**VARIANTS[variant])
    batch = next(iter(_loader(_events(), train=False)))
    tasks, context = batch.context_labels.shape
    tokens = batch.events.tokens.clone()
    first_query = context  # of the first task
    tokens[first_query, :, :6] += 1.0
    changed = replace(batch, events=replace(batch.events, tokens=tokens))
    with torch.no_grad():
        expected = model(batch).view(tasks, -1, 4)
        actual = model(changed).view(tasks, -1, 4)
    assert not torch.allclose(actual[0, 0], expected[0, 0])
    torch.testing.assert_close(actual[0, 1:], expected[0, 1:])
    torch.testing.assert_close(actual[1:], expected[1:])
    absent = (torch.arange(4) >= batch.num_classes[:, None, None]).expand_as(expected)
    assert (expected[absent] == ABSENT_LOGIT).all()


@pytest.mark.parametrize("variant", VARIANTS)
def test_in_context_inference_matches_a_task(variant):
    model, events = _pfn(**VARIANTS[variant]), _events()
    context, queries = torch.arange(0, 40), torch.arange(40, 60)
    # The context in three parts, the queries in two.
    forward = in_context_forward(
        model, events, context, num_classes=4, batch_size=16, device=torch.device("cpu")
    )
    with torch.no_grad():
        actual = torch.cat(
            [forward(collate_sequence_batch(events, q)) for q in queries.split(16)]
        )
        expected = model(_task(events, context, queries))
    torch.testing.assert_close(actual, expected)


@pytest.mark.parametrize("ranks", [False, True])
def test_normalized_pfn_ignores_the_energy_scale(ranks):
    model, events = _pfn(**VARIANTS["normalized"], pfn_ranks=ranks), _events()
    scaled = replace(events, object_features=events.object_features.clone())
    scaled.object_features[:, [0, 3]] *= 1.3
    context, queries = torch.arange(0, 40), torch.arange(40, 60)
    with torch.no_grad():
        expected = model(_task(events, context, queries))
        actual = model(_task(scaled, context, queries))
    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-4)


def test_ensemble_averages_probabilities_over_transformed_passes():
    model, events = _pfn(**VARIANTS["alternating"]), _events()
    context, queries = torch.arange(0, 40), torch.arange(40, 60)

    def logits(ensemble):
        forward = in_context_forward(
            model,
            events,
            context,
            num_classes=4,
            batch_size=64,
            device=torch.device("cpu"),
            ensemble=ensemble,
        )
        return forward(collate_sequence_batch(events, queries))

    single, ensemble = logits(1), logits(4)
    torch.testing.assert_close(ensemble.exp().sum(1), torch.ones(20))
    assert not torch.allclose(ensemble, single.log_softmax(1))
    torch.testing.assert_close(logits(4), ensemble)


def test_context_ranks_are_mid_ranks_per_task_and_object_type():
    tokens = torch.zeros(3, 2, 7)
    tokens[..., PT] = torch.tensor([[1.0, 2.0], [3.0, 2.0], [2.0, 5.0]])
    tokens[..., ENERGY] = 1.0
    tokens[2, 1, NODE_TYPE] = 1  # the only object of its type
    batch = SequenceBatch(
        tokens=tokens,
        token_mask=torch.ones(3, 2, dtype=torch.bool),
        labels=None,
        weight=None,
        global_features=None,
        event_index=None,
    )
    task = torch.tensor([0, 0, 1])
    keys = rank_keys(batch, task)
    ranks = context_ranks(keys, keys[:2].flatten().sort().values)
    # pT of task 0: 1 < 2 = 2 < 3; task 1 has no context.
    torch.testing.assert_close(
        ranks[:2, :, 0], torch.tensor([[0.125, 0.5], [0.875, 0.5]])
    )
    torch.testing.assert_close(ranks[:2, :, 1], torch.full((2, 2), 0.5))
    assert (ranks[2] == 0.5).all()


def test_synthetic_tasks_rank_events_of_few_labels():
    events = _events()
    loader = _loader(events, synthetic_tasks=1.0, prior_classes=None)
    for batch in loader:
        for task, classes in enumerate(_originals(events, batch)):
            k = int(batch.num_classes[task])
            assert set(batch.context_labels[task].tolist()) == set(range(k))
            assert len(set().union(*classes)) <= 3
    # Without synthetic tasks the prior draws exactly as before.
    plain, zero = _loader(events), _loader(events, synthetic_tasks=0.0)
    first = plain.sample(torch.Generator().manual_seed(1))
    assert torch.equal(
        first.indices, zero.sample(torch.Generator().manual_seed(1)).indices
    )


def test_event_observables_per_object_type():
    # Two photons back to back (pT 3 and 4, eta 0), then MET of pT 2.
    objects = [
        torch.tensor(
            [
                [3.0, 0.0, 0.0, 3.0, 0.0, 0.0, 3.0],
                [4.0, 0.0, torch.pi, 4.0, 1.0, 0.0, 3.0],
                [2.0, 0.0, 1.0, 2.0, 0.0, 0.0, 4.0],
            ]
        )
    ]
    observables = event_observables(pack_events(objects, torch.tensor([0])))[0]
    counts, total, leading, mass = observables[:20].view(4, 5)
    assert counts.tolist() == [0, 0, 0, 2, 1]
    torch.testing.assert_close(total[3], torch.log1p(torch.tensor(7.0)))
    torch.testing.assert_close(leading[3], torch.log1p(torch.tensor(4.0)))
    # m^2 = (3 + 4)^2 - (3 - 4)^2 = 48
    torch.testing.assert_close(mass[3], torch.log1p(torch.tensor(48.0).sqrt()))
    assert observables[20] == 1  # b-tagged objects


def test_encoder_starts_from_a_jepa_checkpoint(tmp_path):
    jepa = build_model(
        ModelConfig(name="jepa", **SMALL), num_features=7, num_node_types=5
    )
    path = tmp_path / "jepa.pt"
    torch.save({"model_state_dict": jepa.state_dict()}, path)
    config = ModelConfig(name="pfn", pretrained=str(path), **SMALL)
    pfn = build_model(config, num_features=7, num_node_types=5)
    expected = jepa.context_encoder.state_dict()
    for name, value in pfn.encoder.state_dict().items():
        assert torch.equal(value, expected[name])


def test_meta_train_on_delphes_then_classify_gamgam_in_context(tmp_path):
    small = [f"model.{key}={value}" for key, value in SMALL.items()]
    prior = [
        "task.context_events=[4,16]",
        "task.query_events=4",
        "task.steps_per_epoch=2",
        "task.validation_steps=1",
    ]
    meta = fixture_config(
        "delphes_pfn",
        [TTH_FIXTURE, TTH_FIXTURE],
        [0, 1],
        tmp_path / "meta",
        "train.max_epochs=2",
        "model.pfn_layers=1",
        *small,
        *prior,
    )
    path = str(save_config(meta, tmp_path / "meta.yaml"))
    assert main(["prepare", path]) == 0
    assert main(["train", path]) == 0
    history = (meta.path / "history.jsonl").read_text().splitlines()
    assert "accuracy" in json.loads(history[-1])["validation"]

    gamgam = fixture_config(
        "gamgam_pfn",
        GAMGAM_FIXTURES,
        list(range(5)),
        tmp_path / "gamgam",
        f"checkpoint={meta.path}/best.pt",
        "data.train_events=20",
        "model.pfn_layers=1",
        "split=validation",
        *small,
    )
    path = str(save_config(gamgam, tmp_path / "gamgam.yaml"))
    assert main(["prepare", path]) == 0
    assert main(["train", path]) == 2  # nothing to train on GamGam
    assert main(["evaluate", path]) == 0
    metrics = json.loads((gamgam.path / "validation_metrics.json").read_text())
    assert 0 <= metrics["metrics"]["auc"] <= 1
    assert main(["predict", path]) == 0
    assert main(["export", path]) == 2


def test_in_context_models_need_a_context(tmp_path):
    config = fixture_config(
        "gamgam_tabpfn",
        GAMGAM_FIXTURES,
        list(range(5)),
        tmp_path,
        "data.train_events=null",
    )
    path = str(save_config(config, tmp_path / "tabpfn.yaml"))
    assert main(["prepare", path]) == 0
    assert main(["train", path]) == 2
    assert main(["evaluate", path]) == 2