import pytest
import numpy as np
from pathlib import Path

torch = pytest.importorskip("torch")

from cil.models import CTTConfig, CausalTangentTransport, ChartEncoder, TangentNormalizer
from cil.models.ctt import chamfer_to_target_set, diversity_loss, negative_boundary_loss
from scripts.eval_ctt_generated_rollout import ChartItem, _source_pool_for_target


def test_ctt_variants_preserve_tangent_shape():
    z_source = torch.randn(3, 8)
    z_target = torch.randn(3, 8)
    xi_source = torch.randn(3, 5)

    for variant in ("residual", "gated_residual"):
        model = CausalTangentTransport(
            CTTConfig(chart_feature_dim=10, chart_dim=8, tangent_dim=5, hidden_dim=16, variant=variant)
        )
        output = model(z_source, z_target, xi_source)
        assert output.shape == xi_source.shape
        assert torch.isfinite(output).all()


def test_chart_encoder_and_losses_are_finite():
    encoder = ChartEncoder(input_dim=6, hidden_dim=12, output_dim=8)
    z = encoder(torch.randn(4, 6))
    assert z.shape == (4, 8)

    predicted = torch.randn(4, 5)
    positives = predicted + 0.1
    negatives = predicted + 10.0

    assert torch.isfinite(chamfer_to_target_set(predicted, positives))
    assert negative_boundary_loss(predicted, negatives, margin=0.2).item() == pytest.approx(0.0)
    assert diversity_loss(predicted).item() >= 0.0


def test_gated_residual_matches_documented_formula():
    model = CausalTangentTransport(
        CTTConfig(chart_feature_dim=10, chart_dim=2, tangent_dim=2, hidden_dim=4, variant="gated_residual")
    )
    for parameter in model.delta.parameters():
        parameter.data.zero_()
    for parameter in model.gate.parameters():
        parameter.data.zero_()
    model.delta[-1].bias.data[:] = torch.tensor([4.0, -2.0])
    model.gate[2].bias.data[:] = torch.tensor([0.0, 0.0])

    xi = torch.tensor([[2.0, 6.0]])
    output = model(torch.zeros(1, 2), torch.zeros(1, 2), xi)

    assert torch.allclose(output, torch.tensor([[3.0, 2.0]]), atol=1.0e-6)


def test_tangent_normalizer_round_trips():
    values = torch.randn(8, 5)
    normalizer = TangentNormalizer.fit(values)
    restored = normalizer.inverse_transform(normalizer.transform(values))
    assert torch.allclose(values, restored, atol=1.0e-5)


def test_rollout_source_pool_excludes_target_chart_and_state_hash():
    def chart(chart_id: str, state_hash: str) -> ChartItem:
        return ChartItem(
            chart_id=chart_id,
            task_id="task",
            seed="0",
            state_hash=state_hash,
            instruction="",
            source_dataset=Path("."),
            base_action=np.zeros((2, 7), dtype=np.float32),
            feature=np.zeros(4, dtype=np.float32),
            positive_tangents=np.zeros((1, 21), dtype=np.float32),
            negative_tangents=np.zeros((0, 21), dtype=np.float32),
            hidden_utilities=[],
            hidden_candidate_types=[],
            stored_base_utility=None,
        )

    target = chart("chart_a", "state_a")
    sources = [chart("chart_a", "state_a"), chart("chart_b", "state_a"), chart("chart_c", "state_c")]

    pool = _source_pool_for_target(
        target,
        task_pool=sources,
        source_charts=sources,
        exclude_self_source=True,
    )

    assert [item.chart_id for item in pool] == ["chart_c"]