Ai/crop_simulation/tests.py

import importlib.util
import os
import sqlite3
import tempfile
from collections import namedtuple
from datetime import date, timedelta
from unittest.mock import patch
from unittest import skipUnless

from django.test import TestCase
from rest_framework.test import APIRequestFactory

from .models import SimulationRun, SimulationScenario
from farm_data.models import PlantCatalogSnapshot, SensorData
from farm_data.services import assign_farm_plants_from_backend_ids
from irrigation.models import IrrigationMethod
from location_data.models import SoilLocation
from weather.models import WeatherForecast
from .services import CropSimulationService, CropSimulationError, PcseSimulationManager
from .views import PlantGrowthSimulationView


def build_weather(days: int = 5) -> list[dict]:
    start = date(2026, 4, 1)
    return [
        {
            "DAY": start + timedelta(days=index),
            "LAT": 35.7,
            "LON": 51.4,
            "ELEV": 1200,
            "IRRAD": 16_000_000 + (index * 100_000),
            "TMIN": 11 + index,
            "TMAX": 22 + index,
            "VAP": 12,
            "WIND": 2.4,
            "RAIN": 0.8,
            "E0": 0.35,
            "ES0": 0.3,
            "ET0": 0.32,
        }
        for index in range(days)
    ]


def build_agromanagement(n_amount: float = 30.0) -> list[dict]:
    return [
        {
            date(2026, 4, 1): {
                "CropCalendar": {
                    "crop_name": "wheat",
                    "variety_name": "winter-wheat",
                    "crop_start_date": date(2026, 4, 5),
                    "crop_start_type": "sowing",
                    "crop_end_date": date(2026, 9, 1),
                    "crop_end_type": "harvest",
                    "max_duration": 180,
                },
                "TimedEvents": [
                    {
                        "event_signal": "apply_n",
                        "name": "N strategy",
                        "events_table": [
                            {
                                date(2026, 4, 20): {
                                    "N_amount": n_amount,
                                    "N_recovery": 0.7,
                                }
                            }
                        ],
                    }
                ],
                "StateEvents": [],
            }
        },
        {},
    ]


class CropSimulationServiceTests(TestCase):
    def setUp(self):
        self.service = CropSimulationService()
        self.weather = build_weather()
        self.soil = {"SMFCF": 0.34, "SMW": 0.12, "RDMSOL": 120.0}
        self.site = {"WAV": 40.0}
        self.crop = {"crop_name": "wheat", "TSUM1": 800, "YIELD_SCALE": 1.0}

    def test_failure_marks_scenario_and_run_failed(self):
        with patch.object(
            self.service.manager,
            "run_simulation",
            side_effect=CropSimulationError("pcse failed"),
        ):
            with self.assertRaises(CropSimulationError):
                self.service.run_single_simulation(
                    weather=self.weather,
                    soil=self.soil,
                    crop_parameters=self.crop,
                    agromanagement=build_agromanagement(),
                    site_parameters=self.site,
                    name="broken run",
                )

        scenario = SimulationScenario.objects.get()
        run = SimulationRun.objects.get()

        self.assertEqual(scenario.status, SimulationScenario.Status.FAILURE)
        self.assertEqual(run.status, SimulationScenario.Status.FAILURE)
        self.assertEqual(scenario.error_message, "pcse failed")

    def test_requires_at_least_two_fertilization_strategies(self):
        with self.assertRaises(CropSimulationError):
            self.service.compare_fertilization_strategies(
                weather=self.weather,
                soil=self.soil,
                crop_parameters=self.crop,
                strategies=[{"label": "only", "agromanagement": build_agromanagement()}],
                site_parameters=self.site,
                    )


class CropSimulationViewContractTests(TestCase):
    def setUp(self):
        self.factory = APIRequestFactory()

    @patch("crop_simulation.views.run_growth_simulation_task.delay")
    def test_growth_queue_response_includes_live_ai_metadata(self, mock_delay):
        mock_delay.return_value.id = "task-123"
        request = self.factory.post(
            "/api/crop-simulation/growth/",
            {
                "plant_name": "wheat",
                "dynamic_parameters": ["DVS"],
                "weather": [
                    {
                        "DAY": "2026-04-01",
                        "LAT": 35.7,
                        "LON": 51.4,
                        "TMIN": 12,
                        "TMAX": 24,
                        "RAIN": 0.0,
                        "ET0": 0.32,
                    }
                ],
                "soil_parameters": {"SMFCF": 0.34, "SMW": 0.14, "RDMSOL": 120.0},
                "site_parameters": {"WAV": 40.0},
                "agromanagement": [
                    {
                        "2026-04-01": {
                            "CropCalendar": {
                                "crop_name": "wheat",
                                "variety_name": "winter-wheat",
                                "crop_start_date": "2026-04-05",
                                "crop_start_type": "sowing",
                                "crop_end_date": "2026-09-01",
                                "crop_end_type": "harvest",
                                "max_duration": 180,
                            },
                            "TimedEvents": [],
                            "StateEvents": [],
                        }
                    },
                    {},
                ],
            },
            format="json",
        )

        response = PlantGrowthSimulationView.as_view()(request)

        self.assertEqual(response.status_code, 202)
        self.assertEqual(response.data["meta"]["flow_type"], "live_ai_inference")
        self.assertEqual(response.data["meta"]["source_service"], "ai_crop_simulation")

    def test_recommend_best_crop_returns_best_candidate(self):
        with patch.object(
            self.service.manager,
            "run_simulation",
            side_effect=[
                {
                    "engine": "pcse",
                    "model_name": "Wofost81_NWLP_CWB_CNB",
                    "metrics": {
                        "yield_estimate": 5200.0,
                        "biomass": 9800.0,
                        "max_lai": 4.1,
                    },
                    "daily_output": [],
                    "summary_output": [],
                    "terminal_output": [],
                },
                {
                    "engine": "pcse",
                    "model_name": "Wofost81_NWLP_CWB_CNB",
                    "metrics": {
                        "yield_estimate": 6100.0,
                        "biomass": 11000.0,
                        "max_lai": 4.4,
                    },
                    "daily_output": [],
                    "summary_output": [],
                    "terminal_output": [],
                },
            ],
        ):
            result = self.service.recommend_best_crop(
                weather=self.weather,
                soil=self.soil,
                crops=[
                    {"crop_name": "wheat", "label": "wheat", "TSUM1": 800},
                    {"crop_name": "maize", "label": "maize", "TSUM1": 900},
                ],
                agromanagement=build_agromanagement(),
                site_parameters=self.site,
                name="best crop recommendation",
            )

        self.assertEqual(result["recommended_crop"]["label"], "maize")
        self.assertEqual(result["recommended_crop"]["expected_yield_estimate"], 6100.0)
        self.assertEqual(len(result["candidates"]), 2)

    def test_recommend_best_crop_requires_two_options(self):
        with self.assertRaises(CropSimulationError):
            self.service.recommend_best_crop(
                weather=self.weather,
                soil=self.soil,
                crops=[{"crop_name": "wheat", "TSUM1": 800}],
                agromanagement=build_agromanagement(),
                site_parameters=self.site,
            )

    def test_run_single_simulation_merges_irrigation_and_fertilization_recommendations(self):
        captured = {}

        def fake_run_simulation(**kwargs):
            captured.update(kwargs)
            return {
                "engine": "pcse",
                "model_name": "Wofost81_NWLP_CWB_CNB",
                "metrics": {
                    "yield_estimate": 5400.0,
                    "biomass": 9800.0,
                    "max_lai": 4.2,
                },
                "daily_output": [],
                "summary_output": [],
                "terminal_output": [],
            }

        with patch.object(self.service.manager, "run_simulation", side_effect=fake_run_simulation):
            self.service.run_single_simulation(
                weather=self.weather,
                soil=self.soil,
                crop_parameters=self.crop,
                agromanagement=build_agromanagement(),
                site_parameters=self.site,
                irrigation_recommendation={
                    "events": [
                        {
                            "date": "2026-04-25",
                            "amount": 2.5,
                            "efficiency": 0.8,
                        }
                    ]
                },
                fertilization_recommendation={
                    "events": [
                        {
                            "date": "2026-04-20",
                            "N_amount": 45,
                            "N_recovery": 0.7,
                        }
                    ]
                },
                name="managed run",
            )

        timed_events = captured["agromanagement"][0][date(2026, 4, 1)]["TimedEvents"]
        self.assertEqual(len(timed_events), 3)
        self.assertEqual(timed_events[1]["event_signal"], "irrigate")
        self.assertEqual(timed_events[1]["events_table"][0][date(2026, 4, 25)]["amount"], 2.5)
        self.assertEqual(timed_events[2]["event_signal"], "apply_n")
        self.assertEqual(
            timed_events[2]["events_table"][0][date(2026, 4, 20)]["N_amount"],
            45.0,
        )

    def test_raises_clear_error_when_pcse_is_unavailable(self):
        with patch("crop_simulation.services._load_pcse_bindings", return_value=None):
            with self.assertRaisesMessage(
                CropSimulationError,
                "PCSE is not installed or required PCSE classes could not be loaded.",
            ):
                self.service.run_single_simulation(
                    weather=self.weather,
                    soil=self.soil,
                    crop_parameters=self.crop,
                    agromanagement=build_agromanagement(),
                    site_parameters=self.site,
                    name="missing pcse",
                )


@skipUnless(
    importlib.util.find_spec("pcse") is not None,
    "pcse must be installed to run real WOFOST integration tests.",
)
class CropSimulationPcseIntegrationTests(TestCase):
    def setUp(self):
        os.environ["HOME"] = tempfile.mkdtemp(prefix="pcse-home-", dir="/tmp")
        from pcse import settings as pcse_settings
        from pcse.tests.db_input import (
            AgroManagementDataProvider,
            GridWeatherDataProvider,
            fetch_cropdata,
            fetch_sitedata,
            fetch_soildata,
        )

        def namedtuple_factory(cursor, row):
            fields = [column[0] for column in cursor.description]
            cls = namedtuple("Row", fields)
            return cls._make(row)

        self.grid = int(os.environ.get("PCSE_TEST_GRID", "31031"))
        self.crop_no = int(os.environ.get("PCSE_TEST_CROP_NO", "1"))
        self.year = int(os.environ.get("PCSE_TEST_YEAR", "2000"))

        self.connection = sqlite3.connect(
            os.path.join(pcse_settings.PCSE_USER_HOME, "pcse.db")
        )
        self.connection.row_factory = namedtuple_factory

        self.weather = GridWeatherDataProvider(
            self.connection,
            grid_no=self.grid,
        ).export()
        self.soil = fetch_soildata(self.connection, self.grid)
        self.site = fetch_sitedata(self.connection, self.grid, self.year)
        self.crop = fetch_cropdata(
            self.connection,
            self.grid,
            self.year,
            self.crop_no,
        )
        self.agromanagement = AgroManagementDataProvider(
            self.connection,
            self.grid,
            self.crop_no,
            self.year,
        )
        self.service = CropSimulationService(
            manager=PcseSimulationManager(model_name="Wofost72_WLP_CWB")
        )

    def tearDown(self):
        self.connection.close()

    def test_real_wofost_execute_full_service_path(self):
        result = self.service.run_single_simulation(
            weather=self.weather,
            soil=self.soil,
            crop_parameters=self.crop,
            agromanagement=self.agromanagement,
            site_parameters=self.site,
            name="pcse path",
        )

        scenario = SimulationScenario.objects.get()

        self.assertEqual(scenario.status, SimulationScenario.Status.SUCCESS)
        self.assertEqual(result["result"]["engine"], "pcse")
        self.assertIsNotNone(result["result"]["metrics"]["yield_estimate"])
        self.assertIsNotNone(result["result"]["metrics"]["biomass"])


class CropSimulationCanonicalSnapshotTests(TestCase):
    def setUp(self):
        self.location = SoilLocation.objects.create(latitude="35.700000", longitude="51.400000")
        self.weather = WeatherForecast.objects.create(
            location=self.location,
            forecast_date=date(2026, 4, 10),
            temperature_min=12.0,
            temperature_max=24.0,
            temperature_mean=18.0,
            humidity_mean=55.0,
            precipitation=1.0,
            et0=3.5,
        )
        self.plant = PlantCatalogSnapshot.objects.create(backend_plant_id=401, name="wheat")
        self.irrigation_method = IrrigationMethod.objects.create(name="drip")
        self.farm = SensorData.objects.create(
            farm_uuid="550e8400-e29b-41d4-a716-446655440000",
            center_location=self.location,
            weather_forecast=self.weather,
            irrigation_method=self.irrigation_method,
        )
        assign_farm_plants_from_backend_ids(self.farm, [self.plant.backend_plant_id])

    @patch("crop_simulation.services.build_ai_farm_snapshot")
    def test_build_simulation_payload_from_farm_uses_aggregated_metrics(self, mock_snapshot):
        from crop_simulation.services import build_simulation_payload_from_farm

        mock_snapshot.return_value = {
            "farm_uuid": str(self.farm.farm_uuid),
            "farm_metrics": {
                "resolved_metrics": {
                    "soil_moisture": 36.0,
                    "ndwi": 0.31,
                    "nitrogen": 21.0,
                    "phosphorus": 11.0,
                    "potassium": 17.0,
                    "soil_ph": 6.8,
                    "electrical_conductivity": 1.4,
                }
            },
            "source_metadata": {
                "farm_metrics": {
                    "canonical_source": "farmer_block_aggregated_snapshot",
                    "aggregation_strategy": "farmer_block_mean",
                }
            },
        }

        payload = build_simulation_payload_from_farm(farm_uuid=str(self.farm.farm_uuid), plant_name="wheat")

        self.assertEqual(payload["soil"]["soil_moisture"], 36.0)
        self.assertEqual(payload["site_parameters"]["NAVAILI"], 21.0)
        self.assertEqual(payload["soil"]["phosphorus"], 11.0)
        self.assertEqual(payload["source_metadata"]["farm_metrics"]["canonical_source"], "farmer_block_aggregated_snapshot")

    @patch("crop_simulation.services.build_ai_farm_snapshot")
    def test_build_simulation_payload_from_farm_handles_missing_block_metrics(self, mock_snapshot):
        from crop_simulation.services import build_simulation_payload_from_farm

        mock_snapshot.return_value = {
            "farm_uuid": str(self.farm.farm_uuid),
            "farm_metrics": {"resolved_metrics": {}},
            "source_metadata": {"farm_metrics": {"status": "missing"}},
        }

        payload = build_simulation_payload_from_farm(farm_uuid=str(self.farm.farm_uuid), plant_name="wheat")

        self.assertEqual(payload["site_parameters"]["WAV"], 40.0)
        self.assertEqual(payload["source_metadata"]["farm_metrics"]["status"], "missing")

    @patch("crop_simulation.services.build_ai_farm_snapshot")
    def test_run_single_simulation_stores_weather_provenance(self, mock_snapshot):
        mock_snapshot.return_value = {
            "farm_uuid": str(self.farm.farm_uuid),
            "farm_metrics": {"resolved_metrics": {"soil_moisture": 35.0, "ndwi": 0.3}},
            "source_metadata": {
                "farm_metrics": {"canonical_source": "farmer_block_aggregated_snapshot"},
                "weather": {"policy": "center_location_latest_forecast"},
            },
        }
        service = CropSimulationService()

        with patch.object(service.manager, "run_simulation", return_value={"engine": "pcse", "metrics": {}, "daily_output": [], "summary_output": [], "terminal_output": []}):
            result = service.run_single_simulation(
                farm_uuid=str(self.farm.farm_uuid),
                plant_name="wheat",
                agromanagement=build_agromanagement(),
            )

        self.assertEqual(result["result"]["source_metadata"]["weather"]["policy"], "center_location_latest_forecast")
        run = SimulationRun.objects.get()
        self.assertEqual(run.result_payload["source_metadata"]["farm_metrics"]["canonical_source"], "farmer_block_aggregated_snapshot")