import numpy as np
import torch

from worldalign.common import retrieval_metrics, sliced_wasserstein
from worldalign.energy import (
    log_sinkhorn,
    relation_field_energy,
    standardized_relation,
)
from worldalign.gw import gw_pseudo_targets
from worldalign.models import Bridge, PrefixAdapter
from worldalign.vg_diagnose import bundle_signature


def test_shapes_and_retrieval():
    bridge = Bridge(8, 12, hidden_dim=16)
    prefix = PrefixAdapter(12, 20, prefix_length=4, hidden_dim=16)
    x = torch.randn(5, 8)
    y = bridge(x)
    assert y.shape == (5, 12)
    assert prefix(y).shape == (5, 4, 20)
    metrics = retrieval_metrics(y, y)
    assert metrics["i2t_r@1"] == 1.0


def test_sliced_wasserstein_identity():
    x = torch.randn(32, 16)
    assert sliced_wasserstein(x, x).item() < 1e-10


def test_gw_recovers_structure_without_pairs():
    rng = np.random.default_rng(7)
    z = rng.normal(size=(600, 6)).astype(np.float32)
    q, _ = np.linalg.qr(rng.normal(size=(6, 6)))
    x = torch.from_numpy(z)
    y = torch.from_numpy((z @ q).astype(np.float32))[torch.randperm(len(z))]
    result = gw_pseudo_targets(x, y, clusters=12, seed=3, max_iter=50)
    assert result["coupling"].shape == (12, 12)
    assert np.isfinite(result["gw_distance"])


def test_bundle_signature_is_view_permutation_and_rotation_invariant():
    torch.manual_seed(5)
    x = torch.randn(7, 8, 12)
    q, _ = torch.linalg.qr(torch.randn(12, 12))
    permuted = (x @ q)[:, torch.randperm(8)]
    first = bundle_signature(x)
    second = bundle_signature(permuted)
    assert torch.allclose(first, second, atol=2e-5)


def test_log_sinkhorn_has_unit_marginals():
    logits = torch.randn(9, 9)
    coupling = log_sinkhorn(logits, temperature=0.7, iterations=30)
    assert torch.allclose(coupling.sum(0), torch.ones(9), atol=1e-4)
    assert torch.allclose(coupling.sum(1), torch.ones(9), atol=1e-4)


def test_transposition_delta_matches_brute_force():
    from worldalign.manifold_gate import (
        all_transposition_delta_mse,
        build_channels,
        permuted,
        relation_mse,
    )

    generator = torch.Generator().manual_seed(11)
    visual = torch.randn(24, 8, generator=generator)[:, None, :]
    text = torch.randn(24, 10, generator=generator)[:, None, :]
    text_channels, _ = build_channels(text, bundle=False)
    visual_channels, _ = build_channels(visual, bundle=False)
    delta = all_transposition_delta_mse(text_channels[0], visual_channels[0])
    base = relation_mse(text_channels[0], visual_channels[0])
    for p, q in [(0, 1), (3, 17), (5, 23), (10, 11)]:
        permutation = torch.arange(24)
        permutation[[p, q]] = permutation[[q, p]]
        brute = (
            relation_mse(permuted(text_channels[0], permutation), visual_channels[0])
            - base
        )
        assert torch.allclose(delta[p, q], brute, atol=1e-10)
        assert torch.allclose(delta[q, p], brute, atol=1e-10)


def test_gate_energy_matches_energy_module():
    from worldalign.manifold_gate import assignment_energy, build_channels

    generator = torch.Generator().manual_seed(3)
    visual = torch.randn(16, 6, generator=generator)
    text = torch.randn(16, 9, generator=generator)
    text_channels, text_relation = build_channels(text[:, None, :], bundle=False)
    visual_channels, visual_relation = build_channels(visual[:, None, :], bundle=False)
    permutation = torch.randperm(16, generator=generator)
    gate = assignment_energy(
        text_channels, visual_channels, text_relation, visual_relation, permutation
    )
    reference_relation, reference_standardized = standardized_relation(visual)
    mse, kl = relation_field_energy(
        reference_relation, reference_standardized, text[permutation]
    )
    assert abs(gate["mse"] - float(mse)) < 1e-5
    assert abs(gate["conditional_kl"] - float(kl)) < 1e-4


def test_k_derangement_moves_exactly_k():
    from worldalign.manifold_gate import k_derangement

    generator = torch.Generator().manual_seed(9)
    for k in (2, 4, 16):
        permutation = k_derangement(64, k, generator)
        moved = (permutation != torch.arange(64)).sum()
        assert int(moved) == k


def test_bundle_mean_channel_matches_view_mean_relation():
    from worldalign.manifold_gate import view_bundle_channels

    torch.manual_seed(4)
    views = torch.nn.functional.normalize(torch.randn(6, 5, 8), dim=-1)
    channels = view_bundle_channels(views)
    assert channels.shape == (4, 6, 6)
    means = views.double().mean(1)
    assert torch.allclose(channels[0], means @ means.T, atol=1e-10)


def test_batched_transposition_delta_matches_single():
    from worldalign.manifold_gate import all_transposition_delta_mse
    from worldalign.view_gate import batched_transposition_delta, standardized_field

    torch.manual_seed(2)
    views_a = torch.nn.functional.normalize(torch.randn(3, 6, 5), dim=-1)
    views_b = torch.nn.functional.normalize(torch.randn(3, 6, 7), dim=-1)
    fields_a = standardized_field(views_a)
    fields_b = standardized_field(views_b)
    batched = batched_transposition_delta(fields_a, fields_b)
    for item in range(3):
        single = all_transposition_delta_mse(fields_a[item], fields_b[item])
        assert torch.allclose(batched[item], single, atol=1e-10)


def test_batched_descent_reaches_local_minimum():
    from worldalign.view_gate import (
        batched_descent,
        batched_transposition_delta,
        field_mse,
        standardized_field,
    )

    torch.manual_seed(6)
    visual = torch.nn.functional.normalize(torch.randn(4, 8, 5), dim=-1)
    noise = torch.nn.functional.normalize(
        visual + 0.1 * torch.randn(4, 8, 5), dim=-1
    )
    visual_fields = standardized_field(visual)
    text_fields = standardized_field(noise)
    starts = torch.stack([torch.randperm(8) for _ in range(4)])
    final_cost, permutations = batched_descent(
        text_fields.clone(), visual_fields, starts.clone(), sweeps=100
    )
    index = torch.arange(4)
    permuted = text_fields[
        index[:, None, None], permutations[:, :, None], permutations[:, None, :]
    ]
    assert torch.allclose(field_mse(permuted, visual_fields), final_cost, atol=1e-10)
    delta = batched_transposition_delta(permuted, visual_fields)
    assert (delta >= -1e-9).all()


def test_ricci_flow_returns_finite_symmetric_geometry():
    import numpy as np

    from worldalign.ricci_control import ollivier_ricci_apsp

    rng = np.random.default_rng(4)
    points = rng.normal(size=(24, 3))
    distance = np.linalg.norm(points[:, None] - points[None, :], axis=-1)

    class FlowArgs:
        neighbors = 5
        lazy_alpha = 0.5
        flow_step = 0.4

    flowed = ollivier_ricci_apsp(distance, FlowArgs(), iterations=3)
    assert np.isfinite(flowed).all()
    assert np.allclose(flowed, flowed.T, atol=1e-9)
    assert (np.diag(flowed) == 0).all()


def test_paired_relation_field_has_less_energy_than_permutation():
    generator = torch.Generator().manual_seed(7)
    visual = torch.randn(32, 12, generator=generator)
    paired = visual @ torch.randn(12, 20, generator=generator)
    visual_relation, visual_standardized = standardized_relation(visual)
    paired_energy = relation_field_energy(
        visual_relation, visual_standardized, paired
    )[0]
    permutation = torch.randperm(32, generator=generator)
    shuffled_energy = relation_field_energy(
        visual_relation, visual_standardized, paired[permutation]
    )[0]
    assert paired_energy < shuffled_energy


def test_tier0_factor_vector_is_one_hot_per_factor():
    from worldalign.tier0_pipeline import factor_vector

    vector = factor_vector(3, 2, 0, classes=10)
    assert vector.shape == (17,)
    assert vector[:10].sum() == 1.0 and vector[3] == 1.0
    assert vector[10:14].sum() == 1.0 and vector[11] == 1.0
    assert vector[14:].sum() == 1.0 and vector[14] == 1.0


def test_watershed_splits_touching_objects():
    from worldalign.tier0_pipeline import extract_objects, group_objects

    image = torch.full((3, 64, 64), 0.1)
    # Two same-coloured discs whose supports touch at a single column.
    ys, xs = torch.meshgrid(torch.arange(64), torch.arange(64), indexing="ij")
    for centre in (22, 38):
        disc = ((ys - 32) ** 2 + (xs - centre) ** 2) <= 8**2
        image[:, disc] = torch.tensor([0.9, 0.2, 0.2])[:, None]
    objects = extract_objects(image, peak_distance=4)
    assert len(objects) == 2
    groups = group_objects(objects, threshold=0.35)
    assert len(groups) == 1 and groups[0]["members"] == 2


def test_caption_encoding_matches_declared_factors():
    from worldalign.tier0_pipeline import encode_caption

    ranks = {"red": 0, "blue": 1}
    states = encode_caption("there are three small red circles, a blue square.", ranks, 2)
    assert states.shape == (2, 9)
    assert states[0][0] > 0 and states[0][2 + 2] > 0 and states[0][2 + 4 + 0] > 0
    assert states[1][1] > 0 and states[1][2 + 0] > 0 and states[1][2 + 4 + 1] > 0


def test_shared_rank_separates_aligned_from_shuffled():
    """The shared-direction count must sit at the null when nothing is shared."""
    import numpy as np
    from worldalign.shared_rank import spectrum

    size, width = 192, 16
    generator = np.random.default_rng(0)
    common = generator.normal(size=(size, 8))
    shared_field = common @ common.T
    visual = shared_field + 0.05 * generator.normal(size=(size, size))
    text = shared_field + 0.05 * generator.normal(size=(size, size))
    visual, text = (visual + visual.T) / 2, (text + text.T) / 2

    def count(first, second):
        _, first_vectors = spectrum(first)
        _, second_vectors = spectrum(second)
        cosines = np.linalg.svd(
            first_vectors[:, :width].T @ second_vectors[:, :width], compute_uv=False
        )
        return int((np.clip(cosines, 0, 1) > 0.7).sum())

    order = generator.permutation(size)
    aligned = count(visual, text)
    shuffled = count(visual, text[np.ix_(order, order)])
    assert aligned >= 8, aligned
    assert shuffled <= 2, shuffled


def test_degree_decomposition_is_exact_and_orthogonal():
    """Fitted plus residual must reconstruct the field off the diagonal."""
    import numpy as np
    from worldalign.field_anatomy import degree_part

    generator = np.random.default_rng(1)
    size = 32
    rows = generator.normal(size=(size, 1))
    field = rows + rows.T + 0.1 * generator.normal(size=(size, size))
    np.fill_diagonal(field, 0.0)
    fitted, residual = degree_part(field)
    mask = ~np.eye(size, dtype=bool)
    assert np.allclose(fitted[mask] + residual[mask], field[mask])
    # a field that is purely additive leaves almost nothing in the residual
    pure = rows + rows.T
    np.fill_diagonal(pure, 0.0)
    _, pure_residual = degree_part(pure)
    assert np.abs(pure_residual[mask]).max() < 1e-9


def test_third_moment_kernel_raises_field_rank():
    """A degree-3 set kernel must span more directions than a degree-2 one."""
    import numpy as np
    from worldalign.natural_pipeline import moment_field_degree

    generator = torch.Generator().manual_seed(0)
    sets = [torch.randn(4, 8, generator=generator) for _ in range(24)]
    second = moment_field_degree(sets, 2, 8).double().numpy()
    third = moment_field_degree(sets, 3, 8).double().numpy()

    def effective_rank(matrix):
        values = np.abs(np.linalg.eigvalsh((matrix + matrix.T) / 2))
        weights = values / values.sum()
        weights = weights[weights > 1e-15]
        return float(np.exp(-(weights * np.log(weights)).sum()))

    assert effective_rank(third) > effective_rank(second)


def test_structured_phrase_encoding_separates_head_from_modifier():
    """Word order must change the encoding, which averaging cannot express."""
    import numpy as np

    vectors = {"red": np.array([1.0, 0.0]), "bus": np.array([0.0, 1.0])}

    def structured(tokens):
        head = vectors[tokens[-1]]
        rest = (
            np.mean([vectors[t] for t in tokens[:-1]], axis=0)
            if len(tokens) > 1
            else np.zeros(2)
        )
        return np.concatenate([head, rest])

    assert not np.allclose(structured(["red", "bus"]), structured(["bus", "red"]))
    assert np.allclose(
        np.mean([vectors[t] for t in ["red", "bus"]], axis=0),
        np.mean([vectors[t] for t in ["bus", "red"]], axis=0),
    )