import numpy as np

from heapr.scoring import (
    InMemoryCovarianceStore,
    accumulate_covariance_from_trace,
    compute_atomic_scores_for_expert,
    compute_down_covariance_quadratic,
    score_atomic_outputs_np,
    score_delta_vectors_np,
    score_group_outputs_np,
)


def test_score_delta_vectors_identity_covariance():
    delta = np.array([[1.0, 2.0], [3.0, 4.0]])
    cov = np.eye(2)

    np.testing.assert_allclose(score_delta_vectors_np(delta, cov), [2.5, 12.5])


def test_atomic_scores_average_over_tokens():
    outputs = np.array(
        [
            [[1.0, 0.0], [0.0, 2.0]],
            [[3.0, 0.0], [0.0, 4.0]],
        ]
    )
    cov = np.eye(2)

    np.testing.assert_allclose(score_atomic_outputs_np(outputs, cov), [2.5, 5.0])


def test_exact_group_score_includes_cross_terms():
    outputs = np.array([[[1.0, 0.0], [1.0, 0.0]]])
    cov = np.eye(2)
    group_indices = np.array([[0, 1]])

    # Exact grouped delta is [2, 0], so 0.5 * 4 = 2. A naive sum of atomic scores
    # would be 1, which is not the desired exact grouped-output score.
    np.testing.assert_allclose(score_group_outputs_np(outputs, cov, group_indices), [2.0])


def test_compute_atomic_scores_matches_expanded_delta_formula():
    import torch

    hidden_states = torch.tensor(
        [[0.2, -0.4, 0.8], [1.2, 0.5, -0.3], [-0.7, 0.1, 0.4]],
        dtype=torch.float32,
    )
    routing_weights = torch.tensor([0.9, 0.4, 0.7], dtype=torch.float32)
    gate = torch.tensor([[0.5, -0.2, 0.1], [-0.3, 0.7, 0.2]], dtype=torch.float32)
    up = torch.tensor([[0.1, 0.4, -0.6], [0.8, -0.5, 0.3]], dtype=torch.float32)
    gate_up = torch.cat([gate, up], dim=0)
    down = torch.tensor([[0.2, -0.1], [0.5, 0.7], [-0.4, 0.3]], dtype=torch.float32)
    covariance = np.array(
        [[1.2, 0.1, -0.2], [0.1, 0.8, 0.3], [-0.2, 0.3, 1.5]],
        dtype=np.float32,
    )
    moe_scale = 1.3

    actual = compute_atomic_scores_for_expert(
        hidden_states,
        routing_weights,
        gate_up,
        down,
        covariance,
        moe_scale=moe_scale,
        atom_chunk=1,
    )

    gate_values = torch.nn.functional.silu(hidden_states @ gate.T)
    up_values = hidden_states @ up.T
    atomic_values = gate_values * up_values
    expected = []
    cov = torch.as_tensor(covariance)
    for atom in range(down.shape[1]):
        delta = (
            atomic_values[:, atom, None]
            * down[:, atom][None, :]
            * routing_weights[:, None]
            * moe_scale
        )
        expected.append(0.5 * torch.sum((delta @ cov) * delta, dim=-1).mean())

    torch.testing.assert_close(actual, torch.stack(expected))


def test_device_covariance_accumulation_matches_cpu_path():
    import types

    import torch

    grad = torch.tensor(
        [[[1.0, 2.0], [3.0, 4.0], [-1.0, 0.5]]],
        dtype=torch.float32,
    )
    selected = torch.tensor([[0, 1], [1, 1], [0, 0]])
    output = types.SimpleNamespace(grad=grad)
    record = types.SimpleNamespace(layer_sparse_idx=0, output=output, selected_experts=selected)
    trace = types.SimpleNamespace(routes=[record])

    cpu_store = InMemoryCovarianceStore(hidden_size=2)
    device_store = InMemoryCovarianceStore(hidden_size=2)
    accumulate_covariance_from_trace(trace, cpu_store, device_accumulation=False)
    accumulate_covariance_from_trace(trace, device_store, device_accumulation=True)

    assert cpu_store.counts == device_store.counts
    np.testing.assert_allclose(cpu_store.normalized(0, 0), device_store.normalized(0, 0))
    np.testing.assert_allclose(cpu_store.normalized(0, 1), device_store.normalized(0, 1))


def test_compute_down_covariance_quadratic_matches_manual_formula():
    import torch

    down = torch.tensor([[0.2, -0.1], [0.5, 0.7], [-0.4, 0.3]], dtype=torch.float32)
    covariance = np.array(
        [[1.2, 0.1, -0.2], [0.1, 0.8, 0.3], [-0.2, 0.3, 1.5]],
        dtype=np.float32,
    )

    actual = compute_down_covariance_quadratic(down, covariance, atom_chunk=1)
    cov = torch.as_tensor(covariance)
    expected = torch.stack([down[:, atom] @ cov @ down[:, atom] for atom in range(down.shape[1])])
    torch.testing.assert_close(actual, expected)