Ai/crop_simulation/water_stress.py

from __future__ import annotations

from statistics import mean
from typing import Any

from farm_data.models import SensorData

from .growth_simulation import GrowthSimulationError, _run_simulation, _safe_float, build_growth_context


def _clamp(value: float, lower: float, upper: float) -> float:
    return max(lower, min(upper, value))


def _level_for_index(water_stress: int) -> str:
    if water_stress <= 20:
        return "پایین"
    if water_stress <= 45:
        return "متوسط"
    return "بالا"


def _stage_sensitivity(dvs: float) -> tuple[str, float]:
    if dvs < 0.2:
        return "establishment", 0.9
    if dvs < 1.0:
        return "vegetative", 1.0
    if dvs < 1.3:
        return "flowering", 1.2
    if dvs < 2.0:
        return "reproductive", 1.1
    return "maturity", 0.85


def _compute_water_stress_index(
    *,
    daily_output: list[dict[str, Any]],
    soil_parameters: dict[str, Any],
) -> tuple[int, dict[str, Any]]:
    latest = daily_output[-1] if daily_output else {}
    recent_window = daily_output[-3:] if daily_output else []

    smfcf = _safe_float(soil_parameters.get("SMFCF"), 0.34)
    smw = _safe_float(soil_parameters.get("SMW"), 0.14)
    rdmsol = max(_safe_float(soil_parameters.get("RDMSOL"), 120.0), 1.0)

    latest_sm = _safe_float(latest.get("SM"), 0.0)
    available_water_ratio = _clamp((latest_sm - smw) / max(smfcf - smw, 0.01), 0.0, 1.0)
    moisture_deficit = (1.0 - available_water_ratio) * 65.0

    recent_et0 = mean(_safe_float(item.get("ET0"), 0.0) for item in recent_window) if recent_window else 0.0
    et0_pressure = _clamp((recent_et0 / 0.45) * 18.0, 0.0, 18.0)

    recent_rain = sum(_safe_float(item.get("RAIN"), 0.0) for item in recent_window)
    rainfall_relief = _clamp(recent_rain * 2.5, 0.0, 15.0)

    moisture_trend = 0.0
    if len(recent_window) >= 2:
        moisture_trend = max(
            (_safe_float(recent_window[0].get("SM"), latest_sm) - latest_sm) * 100.0,
            0.0,
        )
    trend_pressure = _clamp(moisture_trend * 1.6, 0.0, 12.0)

    stage_code, stage_multiplier = _stage_sensitivity(_safe_float(latest.get("DVS"), 0.0))
    root_depth_relief = _clamp(((rdmsol - 60.0) / 60.0) * 6.0, 0.0, 6.0)

    raw_score = ((moisture_deficit + et0_pressure + trend_pressure - rainfall_relief - root_depth_relief) *
        stage_multiplier)
    water_stress = int(round(_clamp(raw_score, 0.0, 100.0)))

    return water_stress, {
        "soilMoisturePercent": round(latest_sm * 100.0, 2),
        "availableWaterRatio": round(available_water_ratio, 4),
        "fieldCapacity": round(smfcf, 4),
        "wiltingPoint": round(smw, 4),
        "rootDepthCm": round(rdmsol, 2),
        "recentEt0": round(recent_et0, 4),
        "recentRain": round(recent_rain, 2),
        "soilMoistureDrop": round(moisture_trend, 2),
        "developmentStage": round(_safe_float(latest.get("DVS"), 0.0), 4),
        "stageCode": stage_code,
        "stageSensitivity": round(stage_multiplier, 2),
        "engine": "crop_simulation",
        "formula": (
            "stress = clamp(((moisture_deficit + et0_pressure + trend_pressure - "
            "rainfall_relief - root_depth_relief) * stage_sensitivity), 0, 100)"
        ),
    }


class WaterStressSimulationService:
    def _resolve_plant_name(self, *, farm_uuid: str, plant_name: str | None) -> str:
        if plant_name:
            return plant_name

        sensor = SensorData.objects.prefetch_related("plants").filter(farm_uuid=farm_uuid).first()
        if sensor is None:
            raise GrowthSimulationError("Farm not found.")

        plant = sensor.plants.first()
        if plant is None:
            raise GrowthSimulationError("Plant not found for the selected farm.")
        return plant.name

    def get_water_stress(self, *, farm_uuid: str, plant_name: str | None = None) -> dict[str, Any]:
        resolved_plant_name = self._resolve_plant_name(farm_uuid=farm_uuid, plant_name=plant_name)
        context = build_growth_context(
            {
                "farm_uuid": farm_uuid,
                "plant_name": resolved_plant_name,
            }
        )
        simulation_result, _scenario_id, simulation_warning = _run_simulation(context)
        daily_output = simulation_result.get("daily_output") or []
        if not daily_output:
            raise GrowthSimulationError("Water stress simulation produced no daily output.")

        water_stress, source_metric = _compute_water_stress_index(
            daily_output=daily_output,
            soil_parameters=context.soil_parameters,
        )
        if simulation_warning:
            source_metric["simulationWarning"] = simulation_warning

        return {
            "farm_uuid": str(farm_uuid),
            "plant_name": context.plant_name,
            "waterStressIndex": water_stress,
            "level": _level_for_index(water_stress),
            "sourceMetric": source_metric,
        }