"""Tests for solver.datagen.generator -- synthetic assembly generation."""

from __future__ import annotations

from typing import ClassVar

import numpy as np
import pytest

from solver.datagen.generator import COMPLEXITY_RANGES, SyntheticAssemblyGenerator
from solver.datagen.types import JointType

# ---------------------------------------------------------------------------
# Original generators (chain / rigid / overconstrained)
# ---------------------------------------------------------------------------


class TestChainAssembly:
    """generate_chain_assembly produces valid underconstrained chains."""

    def test_returns_three_tuple(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        bodies, joints, _analysis = gen.generate_chain_assembly(4)
        assert len(bodies) == 4
        assert len(joints) == 3

    def test_chain_underconstrained(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        _, _, analysis = gen.generate_chain_assembly(4)
        assert analysis.combinatorial_classification == "underconstrained"

    def test_chain_body_ids(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        bodies, _, _ = gen.generate_chain_assembly(5)
        ids = [b.body_id for b in bodies]
        assert ids == [0, 1, 2, 3, 4]

    def test_chain_joint_connectivity(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        _, joints, _ = gen.generate_chain_assembly(4)
        for i, j in enumerate(joints):
            assert j.body_a == i
            assert j.body_b == i + 1

    def test_chain_custom_joint_type(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        _, joints, _ = gen.generate_chain_assembly(
            3,
            joint_type=JointType.BALL,
        )
        assert all(j.joint_type is JointType.BALL for j in joints)


class TestRigidAssembly:
    """generate_rigid_assembly produces rigid assemblies."""

    def test_rigid(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, _, analysis = gen.generate_rigid_assembly(4)
        assert analysis.is_rigid

    def test_spanning_tree_structure(self) -> None:
        """n bodies should have at least n-1 joints (spanning tree)."""
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, joints, _ = gen.generate_rigid_assembly(5)
        assert len(joints) >= len(bodies) - 1

    def test_deterministic(self) -> None:
        """Same seed produces same results."""
        g1 = SyntheticAssemblyGenerator(seed=99)
        g2 = SyntheticAssemblyGenerator(seed=99)
        _, j1, a1 = g1.generate_rigid_assembly(4)
        _, j2, a2 = g2.generate_rigid_assembly(4)
        assert a1.jacobian_rank == a2.jacobian_rank
        assert len(j1) == len(j2)


class TestOverconstrainedAssembly:
    """generate_overconstrained_assembly adds redundant constraints."""

    def test_has_redundant(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, _, analysis = gen.generate_overconstrained_assembly(
            4,
            extra_joints=2,
        )
        assert analysis.combinatorial_redundant > 0

    def test_extra_joints_added(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints_base, _ = gen.generate_rigid_assembly(4)

        gen2 = SyntheticAssemblyGenerator(seed=42)
        _, joints_over, _ = gen2.generate_overconstrained_assembly(
            4,
            extra_joints=3,
        )
        # Overconstrained has base joints + extra
        assert len(joints_over) > len(joints_base)


# ---------------------------------------------------------------------------
# New topology generators
# ---------------------------------------------------------------------------


class TestTreeAssembly:
    """generate_tree_assembly produces tree-structured assemblies."""

    def test_body_and_joint_counts(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, joints, _ = gen.generate_tree_assembly(6)
        assert len(bodies) == 6
        assert len(joints) == 5  # n - 1

    def test_underconstrained(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, _, analysis = gen.generate_tree_assembly(6)
        assert analysis.combinatorial_classification == "underconstrained"

    def test_branching_factor(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, joints, _ = gen.generate_tree_assembly(
            10,
            branching_factor=2,
        )
        assert len(bodies) == 10
        assert len(joints) == 9

    def test_mixed_joint_types(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        types = [JointType.REVOLUTE, JointType.BALL, JointType.FIXED]
        _, joints, _ = gen.generate_tree_assembly(10, joint_types=types)
        used = {j.joint_type for j in joints}
        # With 9 joints and 3 types, very likely to use at least 2
        assert len(used) >= 2

    def test_single_joint_type(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints, _ = gen.generate_tree_assembly(
            5,
            joint_types=JointType.BALL,
        )
        assert all(j.joint_type is JointType.BALL for j in joints)

    def test_sequential_body_ids(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, _, _ = gen.generate_tree_assembly(7)
        assert [b.body_id for b in bodies] == list(range(7))


class TestLoopAssembly:
    """generate_loop_assembly produces closed-loop assemblies."""

    def test_body_and_joint_counts(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, joints, _ = gen.generate_loop_assembly(5)
        assert len(bodies) == 5
        assert len(joints) == 5  # n joints for n bodies

    def test_has_redundancy(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, _, analysis = gen.generate_loop_assembly(5)
        assert analysis.combinatorial_redundant > 0

    def test_wrap_around_connectivity(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints, _ = gen.generate_loop_assembly(4)
        edges = {(j.body_a, j.body_b) for j in joints}
        assert (0, 1) in edges
        assert (1, 2) in edges
        assert (2, 3) in edges
        assert (3, 0) in edges  # wrap-around

    def test_minimum_bodies_error(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        with pytest.raises(ValueError, match="at least 3"):
            gen.generate_loop_assembly(2)

    def test_mixed_joint_types(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        types = [JointType.REVOLUTE, JointType.FIXED]
        _, joints, _ = gen.generate_loop_assembly(8, joint_types=types)
        used = {j.joint_type for j in joints}
        assert len(used) >= 2


class TestStarAssembly:
    """generate_star_assembly produces hub-and-spoke assemblies."""

    def test_body_and_joint_counts(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, joints, _ = gen.generate_star_assembly(6)
        assert len(bodies) == 6
        assert len(joints) == 5  # n - 1

    def test_all_joints_connect_to_hub(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints, _ = gen.generate_star_assembly(6)
        for j in joints:
            assert j.body_a == 0 or j.body_b == 0

    def test_underconstrained(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, _, analysis = gen.generate_star_assembly(5)
        assert analysis.combinatorial_classification == "underconstrained"

    def test_minimum_bodies_error(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        with pytest.raises(ValueError, match="at least 2"):
            gen.generate_star_assembly(1)

    def test_hub_at_origin(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, _, _ = gen.generate_star_assembly(4)
        np.testing.assert_array_equal(bodies[0].position, np.zeros(3))


class TestMixedAssembly:
    """generate_mixed_assembly produces tree+loop hybrid assemblies."""

    def test_more_joints_than_tree(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, joints, _ = gen.generate_mixed_assembly(
            8,
            edge_density=0.3,
        )
        assert len(joints) > len(bodies) - 1

    def test_density_zero_is_tree(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _bodies, joints, _ = gen.generate_mixed_assembly(
            5,
            edge_density=0.0,
        )
        assert len(joints) == 4  # spanning tree only

    def test_density_validation(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        with pytest.raises(ValueError, match="must be in"):
            gen.generate_mixed_assembly(5, edge_density=1.5)
        with pytest.raises(ValueError, match="must be in"):
            gen.generate_mixed_assembly(5, edge_density=-0.1)

    def test_no_duplicate_edges(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints, _ = gen.generate_mixed_assembly(6, edge_density=0.5)
        edges = [frozenset([j.body_a, j.body_b]) for j in joints]
        assert len(edges) == len(set(edges))

    def test_high_density(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _bodies, joints, _ = gen.generate_mixed_assembly(
            5,
            edge_density=1.0,
        )
        # Fully connected: 5*(5-1)/2 = 10 edges
        assert len(joints) == 10


# ---------------------------------------------------------------------------
# Axis sampling strategies
# ---------------------------------------------------------------------------


class TestAxisStrategy:
    """Axis sampling strategies produce valid unit vectors."""

    def test_cardinal_axis_from_six(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        axes = {tuple(gen._cardinal_axis()) for _ in range(200)}
        expected = {
            (1, 0, 0),
            (-1, 0, 0),
            (0, 1, 0),
            (0, -1, 0),
            (0, 0, 1),
            (0, 0, -1),
        }
        assert axes == expected

    def test_random_axis_unit_norm(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        for _ in range(50):
            axis = gen._sample_axis("random")
            assert abs(np.linalg.norm(axis) - 1.0) < 1e-10

    def test_near_parallel_close_to_base(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        base = np.array([0.0, 0.0, 1.0])
        for _ in range(50):
            axis = gen._near_parallel_axis(base)
            assert abs(np.linalg.norm(axis) - 1.0) < 1e-10
            assert np.dot(axis, base) > 0.95

    def test_sample_axis_cardinal(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        axis = gen._sample_axis("cardinal")
        cardinals = [
            np.array(v, dtype=float)
            for v in [(1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1), (0, 0, -1)]
        ]
        assert any(np.allclose(axis, c) for c in cardinals)

    def test_sample_axis_near_parallel(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=0)
        axis = gen._sample_axis("near_parallel")
        z = np.array([0.0, 0.0, 1.0])
        assert np.dot(axis, z) > 0.95


# ---------------------------------------------------------------------------
# Geometric diversity: orientations
# ---------------------------------------------------------------------------


class TestRandomOrientations:
    """Bodies should have non-identity orientations."""

    def test_bodies_have_orientations(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, _, _ = gen.generate_tree_assembly(5)
        non_identity = sum(1 for b in bodies if not np.allclose(b.orientation, np.eye(3)))
        assert non_identity >= 3

    def test_orientations_are_valid_rotations(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        bodies, _, _ = gen.generate_star_assembly(6)
        for b in bodies:
            r = b.orientation
            # R^T R == I
            np.testing.assert_allclose(r.T @ r, np.eye(3), atol=1e-10)
            # det(R) == 1
            assert abs(np.linalg.det(r) - 1.0) < 1e-10

    def test_all_generators_set_orientations(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        # Chain
        bodies, _, _ = gen.generate_chain_assembly(3)
        assert not np.allclose(bodies[1].orientation, np.eye(3))
        # Loop
        bodies, _, _ = gen.generate_loop_assembly(4)
        assert not np.allclose(bodies[1].orientation, np.eye(3))
        # Mixed
        bodies, _, _ = gen.generate_mixed_assembly(4)
        assert not np.allclose(bodies[1].orientation, np.eye(3))


# ---------------------------------------------------------------------------
# Geometric diversity: grounded parameter
# ---------------------------------------------------------------------------


class TestGroundedParameter:
    """Grounded parameter controls ground_body in analysis."""

    def test_chain_grounded_default(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, _, analysis = gen.generate_chain_assembly(4)
        assert analysis.combinatorial_dof >= 0

    def test_chain_floating(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, _, analysis = gen.generate_chain_assembly(
            4,
            grounded=False,
        )
        # Floating: 6 trivial DOF not subtracted by ground
        assert analysis.combinatorial_dof >= 6

    def test_floating_vs_grounded_dof_difference(self) -> None:
        gen1 = SyntheticAssemblyGenerator(seed=42)
        _, _, a_grounded = gen1.generate_chain_assembly(4, grounded=True)
        gen2 = SyntheticAssemblyGenerator(seed=42)
        _, _, a_floating = gen2.generate_chain_assembly(4, grounded=False)
        # Floating should have higher DOF due to missing ground constraint
        assert a_floating.combinatorial_dof > a_grounded.combinatorial_dof

    @pytest.mark.parametrize(
        "gen_method",
        [
            "generate_chain_assembly",
            "generate_rigid_assembly",
            "generate_tree_assembly",
            "generate_loop_assembly",
            "generate_star_assembly",
            "generate_mixed_assembly",
        ],
    )
    def test_all_generators_accept_grounded(
        self,
        gen_method: str,
    ) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        method = getattr(gen, gen_method)
        n = 4
        # Should not raise
        if gen_method in ("generate_chain_assembly", "generate_rigid_assembly"):
            method(n, grounded=False)
        else:
            method(n, grounded=False)


# ---------------------------------------------------------------------------
# Geometric diversity: parallel axis injection
# ---------------------------------------------------------------------------


class TestParallelAxisInjection:
    """parallel_axis_prob causes shared axis direction."""

    def test_parallel_axes_similar(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints, _ = gen.generate_chain_assembly(
            6,
            parallel_axis_prob=1.0,
        )
        base = joints[0].axis
        for j in joints[1:]:
            # Near-parallel: |dot| close to 1
            assert abs(np.dot(j.axis, base)) > 0.9

    def test_zero_prob_no_forced_parallel(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints, _ = gen.generate_chain_assembly(
            6,
            parallel_axis_prob=0.0,
        )
        base = joints[0].axis
        dots = [abs(np.dot(j.axis, base)) for j in joints[1:]]
        # With 5 random axes, extremely unlikely all are parallel
        assert min(dots) < 0.95

    def test_parallel_on_loop(self) -> None:
        """Parallel axes on a loop assembly."""
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints, _ = gen.generate_loop_assembly(
            5,
            parallel_axis_prob=1.0,
        )
        base = joints[0].axis
        for j in joints[1:]:
            assert abs(np.dot(j.axis, base)) > 0.9

    def test_parallel_on_star(self) -> None:
        """Parallel axes on a star assembly."""
        gen = SyntheticAssemblyGenerator(seed=42)
        _, joints, _ = gen.generate_star_assembly(
            5,
            parallel_axis_prob=1.0,
        )
        base = joints[0].axis
        for j in joints[1:]:
            assert abs(np.dot(j.axis, base)) > 0.9


# ---------------------------------------------------------------------------
# Complexity tiers
# ---------------------------------------------------------------------------


class TestComplexityTiers:
    """Complexity tier parameter on batch generation."""

    def test_simple_range(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(20, complexity_tier="simple")
        lo, hi = COMPLEXITY_RANGES["simple"]
        for ex in batch:
            assert lo <= ex["n_bodies"] < hi

    def test_medium_range(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(20, complexity_tier="medium")
        lo, hi = COMPLEXITY_RANGES["medium"]
        for ex in batch:
            assert lo <= ex["n_bodies"] < hi

    def test_complex_range(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(3, complexity_tier="complex")
        lo, hi = COMPLEXITY_RANGES["complex"]
        for ex in batch:
            assert lo <= ex["n_bodies"] < hi

    def test_tier_overrides_range(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(
            10,
            n_bodies_range=(2, 3),
            complexity_tier="medium",
        )
        lo, hi = COMPLEXITY_RANGES["medium"]
        for ex in batch:
            assert lo <= ex["n_bodies"] < hi


# ---------------------------------------------------------------------------
# Training batch
# ---------------------------------------------------------------------------


class TestTrainingBatch:
    """generate_training_batch produces well-structured examples."""

    EXPECTED_KEYS: ClassVar[set[str]] = {
        "example_id",
        "generator_type",
        "grounded",
        "n_bodies",
        "n_joints",
        "body_positions",
        "body_orientations",
        "joints",
        "joint_labels",
        "labels",
        "assembly_classification",
        "is_rigid",
        "is_minimally_rigid",
        "internal_dof",
        "geometric_degeneracies",
    }

    VALID_GEN_TYPES: ClassVar[set[str]] = {
        "chain",
        "rigid",
        "overconstrained",
        "tree",
        "loop",
        "star",
        "mixed",
    }

    def test_batch_size(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(20)
        assert len(batch) == 20

    def test_example_keys(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(10)
        for ex in batch:
            assert set(ex.keys()) == self.EXPECTED_KEYS

    def test_example_ids_sequential(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(15)
        assert [ex["example_id"] for ex in batch] == list(range(15))

    def test_generator_type_valid(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(50)
        for ex in batch:
            assert ex["generator_type"] in self.VALID_GEN_TYPES

    def test_generator_type_diversity(self) -> None:
        """100-sample batch should use at least 5 of 7 generator types."""
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(100)
        types = {ex["generator_type"] for ex in batch}
        assert len(types) >= 5

    def test_default_body_range(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(30)
        for ex in batch:
            # default (3, 8), but loop/star may clamp
            assert 2 <= ex["n_bodies"] <= 7

    def test_joint_label_consistency(self) -> None:
        """independent + redundant == total for every joint."""
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(30)
        for ex in batch:
            for label in ex["joint_labels"].values():
                total = label["independent_constraints"] + label["redundant_constraints"]
                assert total == label["total_constraints"]

    def test_body_orientations_present(self) -> None:
        """Each example includes body_orientations as 3x3 lists."""
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(10)
        for ex in batch:
            orients = ex["body_orientations"]
            assert len(orients) == ex["n_bodies"]
            for o in orients:
                assert len(o) == 3
                assert len(o[0]) == 3

    def test_labels_structure(self) -> None:
        """Each example has labels dict with expected sub-keys."""
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(10)
        for ex in batch:
            labels = ex["labels"]
            assert "per_constraint" in labels
            assert "per_joint" in labels
            assert "per_body" in labels
            assert "assembly" in labels

    def test_grounded_field_present(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(10)
        for ex in batch:
            assert isinstance(ex["grounded"], bool)


# ---------------------------------------------------------------------------
# Batch grounded ratio
# ---------------------------------------------------------------------------


class TestBatchGroundedRatio:
    """grounded_ratio controls the mix in batch generation."""

    def test_all_grounded(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(20, grounded_ratio=1.0)
        assert all(ex["grounded"] for ex in batch)

    def test_none_grounded(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(20, grounded_ratio=0.0)
        assert not any(ex["grounded"] for ex in batch)

    def test_mixed_ratio(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(100, grounded_ratio=0.5)
        grounded_count = sum(1 for ex in batch if ex["grounded"])
        # With 100 samples and p=0.5, should be roughly 50 +/- 20
        assert 20 < grounded_count < 80

    def test_batch_axis_strategy_cardinal(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(
            10,
            axis_strategy="cardinal",
        )
        assert len(batch) == 10

    def test_batch_parallel_axis_prob(self) -> None:
        gen = SyntheticAssemblyGenerator(seed=42)
        batch = gen.generate_training_batch(
            10,
            parallel_axis_prob=0.5,
        )
        assert len(batch) == 10


# ---------------------------------------------------------------------------
# Seed reproducibility
# ---------------------------------------------------------------------------


class TestSeedReproducibility:
    """Different seeds produce different results."""

    def test_different_seeds_differ(self) -> None:
        g1 = SyntheticAssemblyGenerator(seed=1)
        g2 = SyntheticAssemblyGenerator(seed=2)
        b1 = g1.generate_training_batch(
            batch_size=5,
            n_bodies_range=(3, 6),
        )
        b2 = g2.generate_training_batch(
            batch_size=5,
            n_bodies_range=(3, 6),
        )
        c1 = [ex["assembly_classification"] for ex in b1]
        c2 = [ex["assembly_classification"] for ex in b2]
        r1 = [ex["is_rigid"] for ex in b1]
        r2 = [ex["is_rigid"] for ex in b2]
        assert c1 != c2 or r1 != r2

    def test_same_seed_identical(self) -> None:
        g1 = SyntheticAssemblyGenerator(seed=123)
        g2 = SyntheticAssemblyGenerator(seed=123)
        b1, j1, _ = g1.generate_tree_assembly(5)
        b2, j2, _ = g2.generate_tree_assembly(5)
        for a, b in zip(b1, b2, strict=True):
            np.testing.assert_array_almost_equal(a.position, b.position)
        assert len(j1) == len(j2)
