"""Regression tests for data isolation, flow-map times, gradients and population metrics."""

from pathlib import Path
import json
import numpy as np
import pandas as pd
import pytest
import torch
from pivot.data.preprocess import assign_splits, prepare
from pivot.data.perturb_data import PerturbData
from pivot.models.flow_map import FlowMap
from pivot.models.encoders import PerturbationEncoder
from pivot.training.losses import compute_losses
from pivot.evaluation.rewards import Reward, rbf_mmd2
from pivot.evaluation.metrics import mmd2, retrieval_metrics


def test_diagonal_identity():
    m = FlowMap(3, 2, hidden=8, depth=1)
    torch.nn.init.normal_(m.net[-1].weight)
    t = torch.rand(5)
    c = torch.randn(5, 3)
    e = torch.randn(5, 2)
    torch.testing.assert_close(m(t, t, c, e), c, rtol=0, atol=0)


def test_pooling_permutation():
    m = PerturbationEncoder(5, 2, 10, emb_dim=3)
    g = torch.tensor([[1, 3]])
    o = torch.ones_like(g)
    mask = torch.ones_like(g)
    torch.testing.assert_close(m(g, o, mask), m(g.flip(1), o, mask))


def test_semigroup_uses_corresponding_source_time():
    class RecordedFlow:
        def __init__(self):
            self.calls = []

        def __call__(self, s, t, c, e):
            self.calls.append((s.clone(), t.clone(), c.clone()))
            return c + (t - s)[:, None] * 2

        def velocity(self, s, t, c, e):
            return torch.ones_like(c) * 2

    f = RecordedFlow()
    c0 = torch.zeros(12, 3)
    c1 = c0 + 2
    e = torch.zeros(12, 2)
    loss, parts = compute_losses(f, e, c0, c1, {})
    # Call 0 is map supervision; calls 1 and 2 start the composition comparison.
    for s, t, c in [f.calls[1], f.calls[2]]:
        torch.testing.assert_close(c, 2 * s[:, None].expand_as(c))
    assert parts["semi"] < 1e-10


def test_reward_gradient_matches_finite_difference():
    torch.manual_seed(1)
    a = torch.randn(3, 2, dtype=torch.double)
    e = torch.randn(2, dtype=torch.double, requires_grad=True)
    target = torch.randn(3, dtype=torch.double)
    r = lambda z: -((a @ z - target) ** 2).sum()
    (g,) = torch.autograd.grad(r(e), e)
    h = 1e-6
    numeric = torch.stack(
        [
            (
                r(e.detach() + h * torch.eye(2, dtype=torch.double)[i])
                - r(e.detach() - h * torch.eye(2, dtype=torch.double)[i])
            )
            / (2 * h)
            for i in range(2)
        ]
    )
    torch.testing.assert_close(g, numeric, rtol=1e-6, atol=1e-6)


def test_numpy_torch_population_statistic():
    rng = np.random.default_rng(2)
    x = rng.normal(size=(7, 3))
    y = rng.normal(size=(11, 3))
    assert np.isclose(
        mmd2(x, y, 0.2), rbf_mmd2(torch.tensor(x), torch.tensor(y), 0.2).item()
    )
    assert abs(mmd2(x, x, 0.2)) < 1e-12


def test_retrieval_censoring():
    assert retrieval_metrics(["A", "B"], "C") == dict(
        top1=0.0, top5=0.0, ndcg10=0.0, rank=None
    )


def test_held_out_genes_do_not_occur_in_training():
    labels = [
        "control",
        "A",
        "B",
        "C",
        "D",
        "E",
        "F",
        "G",
        "H",
        "A_B",
        "C_D",
        "E_F",
        "G_H",
    ]
    o = pd.DataFrame({"perturbation": np.repeat(labels, 20)})
    o["is_control"] = o.perturbation.eq("control")
    split = assign_splits(o, "gene", 5)
    genes = lambda part: {
        g
        for p in o.loc[(split == part) & ~o.is_control, "perturbation"]
        for g in p.split("_")
    }
    held_single = set(
        o.loc[
            (split == "test") & ~o.is_control & ~o.perturbation.str.contains("_"),
            "perturbation",
        ]
    )
    assert held_single and not genes("train") & held_single
    assert not genes("val") & held_single


@pytest.fixture(scope="module")
def prepared(tmp_path_factory):
    raw = Path(__file__).parents[1] / "fixtures/norman_small.h5ad"
    d = tmp_path_factory.mktemp("cache")
    prepare(str(raw), str(d), "norman", n_hvg=100, n_pca=5, seed=0)
    return raw, PerturbData(str(d))


def test_partition_ids_are_disjoint(prepared):
    _, d = prepared
    parts = [set(d.indices(p)) for p in ["train", "val", "test", "reference"]]
    assert len(set.union(*parts)) == len(d.obs)
    for i, a in enumerate(parts):
        for b in parts[i + 1 :]:
            assert not a & b
    assert all(len(d.indices(p, True)) for p in ["train", "val", "test"])
    assert all(
        len(np.intersect1d(d.indices("reference", False), ids))
        for ids in d.pert_to_idx.values()
    )


def test_pca_fit_uses_only_training_cells(prepared):
    _, d = prepared
    train = d.indices("train")
    mean = np.asarray(d.Xhvg[train].mean(0)).ravel()
    np.testing.assert_allclose(d.pca_mean, mean, rtol=1e-5, atol=1e-5)


def test_test_expression_cannot_change_fitted_features(prepared, tmp_path):
    import anndata as ad

    raw, d = prepared
    a = ad.read_h5ad(raw)
    held = d.obs.loc[d.obs.split.ne("train"), "cell_id"]
    ids = a.obs_names.get_indexer(held)
    x = a.X.tocsr()
    x[ids] = x[ids] * 7
    a.X = x
    changed = tmp_path / "changed.h5ad"
    a.write_h5ad(changed)
    out = tmp_path / "second"
    prepare(str(changed), str(out), "norman", n_hvg=100, n_pca=5, seed=0)
    e = PerturbData(str(out))
    assert e.genes == d.genes
    np.testing.assert_allclose(e.pca_mean, d.pca_mean, atol=1e-6)
    np.testing.assert_allclose(e.pca_components, d.pca_components, atol=1e-5)


def test_reference_outcomes_never_change_model_ranking(prepared):
    from pivot.evaluation.runner import evaluate

    _, d = prepared
    # The predictor accesses source populations and action labels only.
    predict = lambda c, p: c + len(d.parse(p)) * 0.1
    root = Path(d.dir)
    a = evaluate(d, predict, root / "a.json", catalog="all", n_cells=8)
    ref = d.indices("reference", False)
    saved = d.emb[ref].copy()
    d.emb[ref] += 5
    try:
        b = evaluate(d, predict, root / "b.json", catalog="all", n_cells=8)
    finally:
        d.emb[ref] = saved
    assert [r["selected"] for r in a["inverse"]] == [
        r["selected"] for r in b["inverse"]
    ]
    assert [r["measured_reward"] for r in a["inverse"]] != [
        r["measured_reward"] for r in b["inverse"]
    ]