test_valve_engine.py

"""
Unit tests for valve_engine module.
"""

import pytest
import math
from shapely.geometry import Point, Polygon
from valve_engine import (
    calculate_pump_flow_lph,
    calculate_total_emitter_flow,
    calculate_num_zones,
    should_split_by_topography,
    choose_manifold_strategy,
    place_valves_hierarchical,
    generate_valve_zones,
    anchor_valves_to_zones,
    valve_layout_summary,
    ValveEngineError,
)


class TestPumpFlow:
    """Test pump horsepower to flow rate conversion."""

    def test_known_hp_values(self):
        """Exact lookups for known HP values."""
        assert calculate_pump_flow_lph(1.0) == 5000
        assert calculate_pump_flow_lph(2.0) == 15000
        assert calculate_pump_flow_lph(10.0) == 80000

    def test_interpolation(self):
        """Interpolate between known values."""
        flow_1_2 = calculate_pump_flow_lph(1.2)
        # Between 1.0 (5000 lph) and 1.5 (8000 lph)
        # 1.2 should be ~6200
        assert 6000 < flow_1_2 < 7000

    def test_invalid_hp(self):
        """Reject zero or negative HP."""
        with pytest.raises(ValveEngineError):
            calculate_pump_flow_lph(0)
        with pytest.raises(ValveEngineError):
            calculate_pump_flow_lph(-1)


class TestTotalFlow:
    """Test total emitter flow calculation."""

    def test_single_crop_zone(self):
        """Calculate flow for one crop zone."""
        crop_zones = [
            {"crop": "tomato", "area_m2": 1000}
        ]
        # tomato: 4.17 emitters/m2, 4 lph each
        # flow = 1000 * 4.17 * 4 = 16,680 lph
        flow = calculate_total_emitter_flow(crop_zones)
        assert abs(flow - 16680) < 5

    def test_multiple_crop_zones(self):
        """Sum flow across multiple crops."""
        crop_zones = [
            {"crop": "tomato", "area_m2": 1000},   # 1000 * 4.17 * 4 = 16,680 lph
            {"crop": "lettuce", "area_m2": 500},    # 500 * 12.5 * 2 = 12,500 lph
        ]
        flow = calculate_total_emitter_flow(crop_zones)
        assert abs(flow - 29180) < 5

    def test_polygon_area_fallback(self):
        """Use polygon area if area_m2 not provided."""
        poly = Polygon([(0, 0), (10, 0), (10, 100), (0, 100)])  # 1000 m²
        crop_zones = [
            {"crop": "tomato", "polygon": poly}
        ]
        flow = calculate_total_emitter_flow(crop_zones)
        assert abs(flow - 16680) < 5

    def test_missing_area_raises_error(self):
        """Raise error if no area_m2 or polygon provided."""
        crop_zones = [
            {"crop": "tomato"}  # No area!
        ]
        with pytest.raises(ValveEngineError):
            calculate_total_emitter_flow(crop_zones)


class TestNumZones:
    """Test zone count calculation."""

    def test_no_split_needed(self):
        """Pump capacity >= total flow -> 1 zone."""
        num_zones = calculate_num_zones(total_emitter_flow_lph=500, pump_flow_lph=1000)
        assert num_zones == 1

    def test_single_split(self):
        """Total flow = 1.5x pump capacity -> 2 zones."""
        num_zones = calculate_num_zones(total_emitter_flow_lph=1500, pump_flow_lph=1000)
        assert num_zones == 2

    def test_multiple_splits(self):
        """Total flow = 3.2x pump capacity -> 4 zones."""
        num_zones = calculate_num_zones(total_emitter_flow_lph=3200, pump_flow_lph=1000)
        assert num_zones == 4

    def test_zero_flow(self):
        """Zero emitter flow -> 1 zone (degenerate case)."""
        num_zones = calculate_num_zones(total_emitter_flow_lph=0, pump_flow_lph=1000)
        assert num_zones == 1


class TestTopographySplit:
    """Test topography-based zone splitting."""

    def test_flat_field(self):
        """No elevation data -> no split."""
        poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        should_split, delta = should_split_by_topography(poly, None)
        assert should_split is False
        assert delta == 0.0

    def test_small_elevation_delta(self):
        """Elevation delta < 5m -> no split."""
        elevation_data = {"min_elevation_m": 100, "max_elevation_m": 103}
        poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        should_split, delta = should_split_by_topography(poly, elevation_data)
        assert should_split is False
        assert delta == 3

    def test_large_elevation_delta(self):
        """Elevation delta > 5m -> split."""
        elevation_data = {"min_elevation_m": 100, "max_elevation_m": 107}
        poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        should_split, delta = should_split_by_topography(poly, elevation_data)
        assert should_split is True
        assert delta == 7


class TestManifoldStrategy:
    """Test centralized vs. distributed valve choice."""

    def test_small_farm_centralized(self):
        """Farm < 1 ha -> centralized."""
        area_m2 = 5000  # 0.5 ha
        strategy = choose_manifold_strategy(area_m2)
        assert strategy == "centralized"

    def test_large_farm_distributed(self):
        """Farm >= 1 ha -> distributed."""
        area_m2 = 10000  # 1.0 ha
        strategy = choose_manifold_strategy(area_m2)
        assert strategy == "distributed"

    def test_boundary_case(self):
        """Farm at 1 ha boundary."""
        area_m2 = 10000  # Exactly 1 ha
        strategy = choose_manifold_strategy(area_m2)
        # At boundary, should be "distributed"
        assert strategy == "distributed"


class TestValvePlacement:
    """Test valve placement hierarchy (capacity + crop + area-density + topography)."""

    def test_area_density_generic(self):
        """Generic crop: ~5 valves/ha rule of thumb."""
        # 1 ha farm, generic crop, very high pump (capacity not limiting)
        farm_poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump_point = Point(0, 0)
        crop_zones = [{"crop": "generic", "area_m2": 10000}]
        pump_hp = 20.0  # 160,000 lph — well above demand

        valves = place_valves_hierarchical(
            farm_poly, pump_point, crop_zones, pump_hp, centralized=False
        )
        # 1 ha * 5/ha = 5 valves
        assert len(valves) == 5

    def test_area_density_orchard_lower(self):
        """Orchard: sparser density (2/ha) means fewer valves."""
        farm_poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])  # 1 ha
        pump_point = Point(0, 0)
        crop_zones = [{"crop": "orchard", "area_m2": 10000}]
        pump_hp = 20.0

        valves = place_valves_hierarchical(
            farm_poly, pump_point, crop_zones, pump_hp, centralized=False
        )
        # 1 ha * 2/ha = 2 valves
        assert len(valves) == 2

    def test_simple_farm_area_density_floor(self):
        """Small farm: area-density floor dominates over low capacity demand."""
        # 100m x 50m = 0.5 ha; tomato @ 6/ha -> floor = ceil(3) = 3 valves
        farm_poly = Polygon([(0, 0), (100, 0), (100, 50), (0, 50)])
        pump_point = Point(0, 0)
        crop_zones = [
            {"crop": "tomato", "area_m2": 100}  # Very small zone -> capacity = 1
        ]
        pump_hp = 10.0  # 80,000 lph — well above demand

        valves = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=True,
        )

        # Area floor (0.5 ha * 6/ha = 3) should dominate over capacity floor (1)
        assert len(valves) >= 3
        assert valves[0]["id"] == "valve_000"

    def test_high_demand_multiple_valves(self):
        """Large farm, high demand -> capped at agronomic maximum."""
        farm_poly = Polygon([(0, 0), (500, 0), (500, 500), (0, 500)])
        pump_point = Point(0, 0)
        crop_zones = [
            {"crop": "tomato", "area_m2": 50000}  # 50,000 * 4.17 * 4 = 834,000 lph
        ]
        pump_hp = 10.0  # 80,000 lph

        valves = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=False,
        )

        # 500x500m = 25 ha, tomato @ 6/ha -> floor = 150;
        # cap (>= 10ha) = 100 → capped at 100
        assert len(valves) <= 100
        assert len(valves) >= 10

    def test_multiple_crops_separate_valves(self):
        """Multiple crops -> at least as many valves as crops."""
        farm_poly = Polygon([(0, 0), (200, 0), (200, 200), (0, 200)])
        pump_point = Point(0, 0)
        crop_zones = [
            {"crop": "tomato", "area_m2": 5000},
            {"crop": "lettuce", "area_m2": 5000},
            {"crop": "orchard", "area_m2": 5000},
        ]
        pump_hp = 100.0  # Very high capacity

        valves = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=False,
        )

        # 200x200m = 4 ha, tomato @ 6/ha -> area floor = 24;
        # cap (< 5ha) = 35 → expect 24 valves
        assert len(valves) >= 3   # at minimum one per crop
        assert len(valves) <= 35  # cap for 4 ha

    def test_topography_forces_extra_valve(self):
        """Elevation split adds a valve."""
        farm_poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump_point = Point(0, 0)
        crop_zones = [
            {"crop": "tomato", "area_m2": 1000}
        ]
        pump_hp = 10.0
        elevation_data = {"min_elevation_m": 100, "max_elevation_m": 110}

        valves_flat = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=True,
            elevation_data=None,
        )

        valves_sloped = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=True,
            elevation_data=elevation_data,
        )

        # Sloped should have at least one more valve
        assert len(valves_sloped) >= len(valves_flat)

    def test_centralized_vs_distributed_locations(self):
        """Centralized places near pump; distributed places at zone boundary."""
        farm_poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump_point = Point(10, 10)
        crop_zones = [
            {"crop": "tomato", "area_m2": 1000}
        ]
        pump_hp = 10.0

        valves_centralized = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=True,
        )

        valves_distributed = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=False,
        )

        # Both should have at least one valve
        assert len(valves_centralized) >= 1
        assert len(valves_distributed) >= 1

        # Centralized valve should have 'centralized' strategy
        assert valves_centralized[0]["strategy"] == "centralized"
        assert valves_distributed[0]["strategy"] == "distributed"


class TestValveZoneGeneration:
    """Test zone polygon generation."""

    def test_zones_generated_for_valves(self):
        """Each valve should have a zone polygon."""
        farm_poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump_point = Point(50, 50)
        crop_zones = [
            {"crop": "tomato", "area_m2": 5000}
        ]
        pump_hp = 50.0  # Very high capacity

        valves = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=False,
        )

        zones = generate_valve_zones(farm_poly, len(valves))

        # Should have zones generated (not necessarily 1:1 with valves due to merging)
        assert len(zones) > 0

        # Each zone should have a valid polygon
        for zone in zones:
            assert zone["polygon"].is_valid
            assert zone["area_m2"] > 0

    def test_total_zone_area_equals_farm_area(self):
        """Sum of all zone areas should equal farm area (approximately)."""
        farm_poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump_point = Point(50, 50)
        crop_zones = [
            {"crop": "tomato", "area_m2": 5000}
        ]
        pump_hp = 1.0  # 5000 lph

        valves = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=False,
        )

        zones = generate_valve_zones(farm_poly, len(valves))
        total_zone_area = sum(z["area_m2"] for z in zones)

        # Should be close to farm area (within 5% tolerance)
        farm_area = farm_poly.area
        assert abs(total_zone_area - farm_area) < 0.05 * farm_area


class TestSummary:
    """Test human-readable summary generation."""

    def test_summary_format(self):
        """Summary should contain valve info."""
        farm_poly = Polygon([(0, 0), (100, 0), (100, 100), (0, 100)])
        pump_point = Point(50, 50)
        crop_zones = [
            {"crop": "tomato", "area_m2": 5000}
        ]
        pump_hp = 10.0  # 80,000 lph

        valves = place_valves_hierarchical(
            farm_poly,
            pump_point,
            crop_zones,
            pump_hp,
            centralized=True,
        )

        zones = generate_valve_zones(farm_poly, len(valves))
        summary = valve_layout_summary(valves, zones)

        assert "Valve Placement Summary" in summary
        assert "Total Valves" in summary
        assert "Valve Details" in summary


class TestValveAnchoring:
    """Test valve anchoring to zones (Phase 3)."""

    def test_anchor_adds_valve_location_to_zones(self):
        """anchor_valves_to_zones should add 'valve_location' to each zone."""
        zones = [
            {"polygon": Polygon([(0, 0), (50, 0), (50, 50), (0, 50)]), "area_m2": 2500},
            {"polygon": Polygon([(50, 0), (100, 0), (100, 50), (50, 50)]), "area_m2": 2500},
        ]
        pump_location = Point(25, 25)
        
        anchored = anchor_valves_to_zones(zones, pump_location, "distributed")
        
        assert len(anchored) == 2
        for zone in anchored:
            assert "valve_location" in zone
            assert isinstance(zone["valve_location"], Point)
            assert zone["polygon"].is_valid
            assert zone["area_m2"] > 0

    def test_centralized_valves_cluster_near_pump(self):
        """Centralized design should place all valve_locations near pump."""
        zones = [
            {"polygon": Polygon([(0, 0), (50, 0), (50, 50), (0, 50)]), "area_m2": 2500},
            {"polygon": Polygon([(50, 0), (100, 0), (100, 50), (50, 50)]), "area_m2": 2500},
            {"polygon": Polygon([(0, 50), (50, 50), (50, 100), (0, 100)]), "area_m2": 2500},
        ]
        pump_location = Point(25, 25)
        
        anchored = anchor_valves_to_zones(zones, pump_location, "centralized")
        
        assert len(anchored) == 3
        # All valves should be within ~20m of pump (10m offset + some margin)
        for zone in anchored:
            dist = zone["valve_location"].distance(pump_location)
            assert dist <= 15, f"Centralized valve too far from pump: {dist}m"

    def test_distributed_valves_on_zone_edges(self):
        """Distributed design should place valve_locations on zone boundaries."""
        zones = [
            {"polygon": Polygon([(0, 0), (50, 0), (50, 50), (0, 50)]), "area_m2": 2500},
            {"polygon": Polygon([(50, 0), (100, 0), (100, 50), (50, 50)]), "area_m2": 2500},
        ]
        pump_location = Point(25, 25)
        
        anchored = anchor_valves_to_zones(zones, pump_location, "distributed")
        
        assert len(anchored) == 2
        # Each valve_location should be on the zone boundary
        for zone in anchored:
            valve_loc = zone["valve_location"]
            zone_boundary = zone["polygon"].boundary
            # Distance from valve to boundary should be ~0 (allow small numerical error)
            dist_to_boundary = valve_loc.distance(zone_boundary)
            assert dist_to_boundary < 0.1, (
                f"Distributed valve not on zone boundary: {dist_to_boundary}m"
            )

    def test_preserves_zone_metadata(self):
        """Anchoring should preserve existing zone properties (polygon, area, crop)."""
        zones = [
            {
                "polygon": Polygon([(0, 0), (50, 0), (50, 50), (0, 50)]),
                "area_m2": 2500,
                "crop": "tomato",
            },
        ]
        pump_location = Point(25, 25)
        
        anchored = anchor_valves_to_zones(zones, pump_location, "distributed")
        
        assert len(anchored) == 1
        zone = anchored[0]
        assert zone["polygon"] is not None
        assert zone["area_m2"] == 2500
        assert zone["crop"] == "tomato"
        assert "valve_location" in zone

    def test_empty_zones_list(self):
        """Anchoring empty zone list should return empty list."""
        zones = []
        pump_location = Point(50, 50)
        
        anchored = anchor_valves_to_zones(zones, pump_location, "distributed")
        
        assert anchored == []

    def test_valve_count_unchanged(self):
        """Anchoring should not change the number of zones."""
        zones = [
            {"polygon": Polygon([(0, 0), (40, 0), (40, 40), (0, 40)]), "area_m2": 1600},
            {"polygon": Polygon([(40, 0), (80, 0), (80, 40), (40, 40)]), "area_m2": 1600},
            {"polygon": Polygon([(0, 40), (40, 40), (40, 80), (0, 80)]), "area_m2": 1600},
            {"polygon": Polygon([(40, 40), (80, 40), (80, 80), (40, 80)]), "area_m2": 1600},
        ]
        pump_location = Point(40, 40)
        
        anchored_cent = anchor_valves_to_zones(zones, pump_location, "centralized")
        anchored_dist = anchor_valves_to_zones(zones, pump_location, "distributed")
        
        # Both should have same number of zones as input
        assert len(anchored_cent) == len(zones)
        assert len(anchored_dist) == len(zones)

    def test_design_type_drives_placement_strategy(self):
        """Design type should determine valve placement approach."""
        zones = [
            {"polygon": Polygon([(0, 0), (100, 0), (100, 100), (0, 100)]), "area_m2": 10000},
        ]
        pump_location = Point(50, 50)  # Center of zone
        
        # Centralized: valve near pump
        anchored_cent = anchor_valves_to_zones(zones, pump_location, "centralized")
        valve_cent = anchored_cent[0]["valve_location"]
        dist_cent = valve_cent.distance(pump_location)
        
        # Distributed: valve on boundary (further from pump at corner)
        anchored_dist = anchor_valves_to_zones(zones, pump_location, "distributed")
        valve_dist = anchored_dist[0]["valve_location"]
        dist_dist = valve_dist.distance(pump_location)
        
        # Centralized should be closer to pump
        assert dist_cent < dist_dist


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