"""
Unit tests for pipe_network module.

Tests orthogonal routing, trunk main generation, and sub-main path calculation
for both centralized and distributed irrigation designs.
"""

import pytest
import math
from shapely.geometry import Point, LineString, Polygon
from pipe_network import (
    route_orthogonal,
    generate_pipe_network,
    calculate_pipe_lengths,
    PipeNetworkError,
)


class TestOrthogonalRouting:
    """Test orthogonal path routing with axis alignment."""

    def test_orthogonal_same_main_axis(self):
        """Points on same main axis should create straight line."""
        main_axis = (1, 0)
        lateral_axis = (0, 1)
        start = Point(0, 10)
        end = Point(100, 10)
        
        route = route_orthogonal(start, end, main_axis, lateral_axis)
        
        assert len(route.coords) == 2
        assert route.length == pytest.approx(100, rel=1e-2)

    def test_orthogonal_same_lateral_axis(self):
        """Points on same lateral axis should create straight line."""
        main_axis = (1, 0)
        lateral_axis = (0, 1)
        start = Point(50, 0)
        end = Point(50, 100)
        
        route = route_orthogonal(start, end, main_axis, lateral_axis)
        
        assert len(route.coords) == 2
        assert route.length == pytest.approx(100, rel=1e-2)

    def test_orthogonal_requires_one_bend(self):
        """Misaligned points require one 90-degree bend."""
        main_axis = (1, 0)
        lateral_axis = (0, 1)
        start = Point(0, 0)
        end = Point(100, 100)
        
        route = route_orthogonal(start, end, main_axis, lateral_axis)
        
        # Should have 3 points (start, corner, end)
        assert len(route.coords) == 3
        # Verify corner is axis-aligned
        corner = Point(route.coords[1])
        assert math.isclose(corner.x, 100) or math.isclose(corner.y, 0)

    def test_orthogonal_total_length_greater_than_direct(self):
        """Orthogonal path should be >= direct distance."""
        main_axis = (1, 0)
        lateral_axis = (0, 1)
        start = Point(0, 0)
        end = Point(100, 100)
        
        route = route_orthogonal(start, end, main_axis, lateral_axis)
        direct = start.distance(end)
        
        # Orthogonal path should be longer or equal (Manhattan distance >= Euclidean)
        assert route.length >= direct - 1e-6

    def test_orthogonal_tilted_axes(self):
        """Orthogonal routing works with tilted farm axes."""
        # 45-degree rotated axes
        sqrt2 = math.sqrt(2) / 2
        main_axis = (sqrt2, sqrt2)
        lateral_axis = (-sqrt2, sqrt2)
        
        start = Point(0, 0)
        # Use point not aligned on 45-degree diagonal to trigger bend
        end = Point(100, 50)
        
        route = route_orthogonal(start, end, main_axis, lateral_axis)
        
        # Should create a path with 3 points (requires a bend)
        assert len(route.coords) == 3


class TestPipeNetworkDistributed:
    """Test pipe network generation for distributed layouts."""

    def test_distributed_has_trunk_main(self):
        """Distributed design should generate a trunk main."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(0, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (50, 0), (50, 100), (0, 100)]),
                "area_m2": 5000,
                "valve_location": Point(25, 50),
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="distributed")
        
        assert network["trunk_main"] is not None
        assert network["trunk_main"].length > 0
        assert network["total_trunk_length_m"] > 0

    def test_distributed_sub_mains_from_trunk(self):
        """Distributed sub-mains should originate from trunk."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(0, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (50, 0), (50, 100), (0, 100)]),
                "area_m2": 5000,
                "valve_location": Point(25, 50),
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="distributed")
        
        assert "valve_000" in network["sub_mains"]
        sub_main = network["sub_mains"]["valve_000"]
        
        # Sub-main should start at a point on the trunk
        start_pt = Point(sub_main.coords[0])
        dist_to_trunk = network["trunk_main"].distance(start_pt)
        assert dist_to_trunk < 1e-6  # Should be on trunk (within numerical tolerance)

    def test_distributed_multiple_zones(self):
        """Multiple zones should each have a sub-main."""
        farm = Polygon([(0, 0), (200, 0), (200, 100), (0, 100)])
        pump = Point(0, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (100, 0), (100, 100), (0, 100)]),
                "area_m2": 5000,
                "valve_location": Point(50, 80),  # Offset from trunk to create non-zero length
            },
            {
                "valve_id": "valve_001",
                "polygon": Polygon([(100, 0), (200, 0), (200, 100), (100, 100)]),
                "area_m2": 5000,
                "valve_location": Point(150, 20),  # Offset from trunk
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="distributed")
        
        assert len(network["sub_mains"]) == 2
        assert network["total_submain_length_m"] > 0

    def test_distributed_sub_main_ends_at_valve(self):
        """Sub-main should end at the anchored valve location."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(0, 50)
        valve_loc = Point(80, 60)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (100, 0), (100, 100), (0, 100)]),
                "area_m2": 10000,
                "valve_location": valve_loc,
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="distributed")
        
        sub_main = network["sub_mains"]["valve_000"]
        end_pt = Point(sub_main.coords[-1])
        
        # End should be close to valve_location
        assert end_pt.distance(valve_loc) < 1e-6

    def test_distributed_sub_main_is_orthogonal(self):
        """Sub-main paths should be axis-aligned (orthogonal)."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(0, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (100, 0), (100, 100), (0, 100)]),
                "area_m2": 10000,
                "valve_location": Point(80, 80),
            }
        ]
        main_axis = (1, 0)
        lateral_axis = (0, 1)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="distributed")
        
        sub_main = network["sub_mains"]["valve_000"]
        coords = list(sub_main.coords)
        
        # Each segment should be either horizontal or vertical
        for i in range(len(coords) - 1):
            p1 = coords[i]
            p2 = coords[i + 1]
            # Either x or y should be the same
            is_horizontal = math.isclose(p1[1], p2[1])
            is_vertical = math.isclose(p1[0], p2[0])
            assert is_horizontal or is_vertical, f"Segment not axis-aligned: {p1} -> {p2}"


class TestPipeNetworkCentralized:
    """Test pipe network generation for centralized layouts."""

    def test_centralized_no_trunk_main(self):
        """Centralized design should NOT generate a trunk main."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(50, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (50, 0), (50, 100), (0, 100)]),
                "area_m2": 5000,
                "valve_location": Point(25, 50),
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="centralized")
        
        assert network["trunk_main"] is None
        assert network["total_trunk_length_m"] == 0

    def test_centralized_sub_mains_from_pump(self):
        """Centralized sub-mains should originate from pump."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(50, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (50, 0), (50, 100), (0, 100)]),
                "area_m2": 5000,
                "valve_location": Point(25, 50),
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="centralized")
        
        assert "valve_000" in network["sub_mains"]
        sub_main = network["sub_mains"]["valve_000"]
        start_pt = Point(sub_main.coords[0])
        
        # Should start at pump (within tolerance)
        assert start_pt.distance(pump) < 1e-6

    def test_centralized_multiple_zones(self):
        """Centralized layout with multiple zones."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(50, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (50, 0), (50, 100), (0, 100)]),
                "area_m2": 5000,
                "valve_location": Point(25, 75),
            },
            {
                "valve_id": "valve_001",
                "polygon": Polygon([(50, 0), (100, 0), (100, 100), (50, 100)]),
                "area_m2": 5000,
                "valve_location": Point(75, 25),
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="centralized")
        
        assert len(network["sub_mains"]) == 2
        assert network["total_submain_length_m"] > 0
        # Centralized should not have trunk
        assert network["total_trunk_length_m"] == 0

    def test_centralized_sub_main_is_orthogonal(self):
        """Centralized sub-mains should also be orthogonal."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(50, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (100, 0), (100, 100), (0, 100)]),
                "area_m2": 10000,
                "valve_location": Point(10, 10),
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="centralized")
        
        sub_main = network["sub_mains"]["valve_000"]
        coords = list(sub_main.coords)
        
        # Each segment should be axis-aligned
        for i in range(len(coords) - 1):
            p1 = coords[i]
            p2 = coords[i + 1]
            is_horizontal = math.isclose(p1[1], p2[1])
            is_vertical = math.isclose(p1[0], p2[0])
            assert is_horizontal or is_vertical


class TestPipeNetworkGeneral:
    """General pipe network tests."""

    def test_empty_zones_returns_empty_network(self):
        """Empty zone list should return empty network."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(50, 50)
        zones = []
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis)
        
        assert network["trunk_main"] is None
        assert len(network["sub_mains"]) == 0
        assert network["total_trunk_length_m"] == 0
        assert network["total_submain_length_m"] == 0

    def test_fallback_to_centroid_if_no_valve_location(self):
        """Zone without valve_location should use centroid as fallback."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(50, 50)
        zone_poly = Polygon([(0, 0), (50, 0), (50, 100), (0, 100)])
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": zone_poly,
                "area_m2": 5000,
                # No valve_location provided
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="distributed")
        
        assert "valve_000" in network["sub_mains"]
        sub_main = network["sub_mains"]["valve_000"]
        
        # Should end at zone centroid (fallback)
        end_pt = Point(sub_main.coords[-1])
        assert end_pt.distance(zone_poly.centroid) < 1

    def test_pipe_lengths_calculation(self):
        """Pipe length calculation should match geometry."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(0, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (100, 0), (100, 100), (0, 100)]),
                "area_m2": 10000,
                "valve_location": Point(100, 50),
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis, design_type="distributed")
        
        # Sum should match totals
        total = (
            network["total_trunk_length_m"] +
            network["total_submain_length_m"]
        )
        assert total > 0
        
        # Trunk should be 100m (from pump at x=0 to x=100)
        assert network["total_trunk_length_m"] == pytest.approx(100, rel=1e-2)

    def test_no_negative_lengths(self):
        """All pipe lengths should be non-negative."""
        farm = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump = Point(50, 50)
        zones = [
            {
                "valve_id": "valve_000",
                "polygon": Polygon([(0, 0), (100, 0), (100, 100), (0, 100)]),
                "area_m2": 10000,
                "valve_location": Point(75, 75),
            }
        ]
        main_axis = (1, 0)
        
        network = generate_pipe_network(farm, pump, zones, main_axis)
        
        assert network["total_trunk_length_m"] >= 0
        assert network["total_submain_length_m"] >= 0
        for sub_main in network["sub_mains"].values():
            assert sub_main.length >= 0


if __name__ == "__main__":
    pytest.main([__file__, "-v"])