Ai/plant/gdd.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import date, timedelta
from typing import Any


DEFAULT_GROWTH_PROFILE = {
    "base_temperature": 10.0,
    "required_gdd_for_maturity": 1200.0,
    "stage_thresholds": {
        "flowering": 500.0,
        "fruiting": 850.0,
    },
    "current_cumulative_gdd": 0.0,
}


@dataclass
class HarvestPrediction:
    current_cumulative_gdd: float
    required_gdd_for_maturity: float
    remaining_gdd: float
    estimated_days_to_harvest: int
    predicted_harvest_date: str
    predicted_harvest_window: dict[str, str]
    daily_gdd_forecast: list[dict[str, float | str]]
    active_stage: str | None


def resolve_growth_profile(plant: Any | None) -> dict:
    profile = getattr(plant, "growth_profile", None) or {}
    stage_thresholds = {
        **DEFAULT_GROWTH_PROFILE["stage_thresholds"],
        **profile.get("stage_thresholds", {}),
    }
    return {
        **DEFAULT_GROWTH_PROFILE,
        **profile,
        "stage_thresholds": stage_thresholds,
    }


def calculate_daily_gdd(tmax: float, tmin: float, tbase: float) -> float:
    mean_temp = (tmax + tmin) / 2.0
    return round(max(mean_temp - tbase, 0.0), 3)


def determine_active_stage(current_cumulative_gdd: float, stage_thresholds: dict[str, float]) -> str | None:
    active_stage = None
    for stage, threshold in sorted(stage_thresholds.items(), key=lambda item: item[1]):
        if current_cumulative_gdd >= float(threshold):
            active_stage = stage
    return active_stage


def predict_harvest_from_forecasts(
    forecasts: list[Any],
    plant: Any | None,
) -> HarvestPrediction:
    profile = resolve_growth_profile(plant)
    base_temperature = float(profile.get("base_temperature", DEFAULT_GROWTH_PROFILE["base_temperature"]))
    required_gdd = float(profile.get("required_gdd_for_maturity", DEFAULT_GROWTH_PROFILE["required_gdd_for_maturity"]))
    current_cumulative_gdd = float(profile.get("current_cumulative_gdd", DEFAULT_GROWTH_PROFILE["current_cumulative_gdd"]))

    cumulative_gdd = current_cumulative_gdd
    daily_forecast: list[dict[str, float | str]] = []
    estimated_date = forecasts[-1].forecast_date if forecasts else date.today()

    for forecast in forecasts:
        tmax = float(getattr(forecast, "temperature_max", None) or getattr(forecast, "temperature_mean", 0.0))
        tmin = float(getattr(forecast, "temperature_min", None) or getattr(forecast, "temperature_mean", 0.0))
        daily_gdd = calculate_daily_gdd(tmax=tmax, tmin=tmin, tbase=base_temperature)
        cumulative_gdd += daily_gdd
        daily_forecast.append(
            {
                "date": forecast.forecast_date.isoformat(),
                "gdd": daily_gdd,
                "cumulative_gdd": round(cumulative_gdd, 3),
            }
        )
        if cumulative_gdd >= required_gdd:
            estimated_date = forecast.forecast_date
            break
    else:
        remaining_gdd_after_forecast = max(required_gdd - cumulative_gdd, 0.0)
        avg_gdd = sum(item["gdd"] for item in daily_forecast) / len(daily_forecast) if daily_forecast else 0.0
        extra_days = int(remaining_gdd_after_forecast / avg_gdd) + (1 if avg_gdd > 0 and remaining_gdd_after_forecast > 0 else 0)
        estimated_date = estimated_date + timedelta(days=max(extra_days, 0))

    remaining_gdd = max(required_gdd - current_cumulative_gdd, 0.0)
    estimated_days = max((estimated_date - date.today()).days, 0)
    active_stage = determine_active_stage(current_cumulative_gdd, profile.get("stage_thresholds", {}))

    return HarvestPrediction(
        current_cumulative_gdd=round(current_cumulative_gdd, 3),
        required_gdd_for_maturity=round(required_gdd, 3),
        remaining_gdd=round(remaining_gdd, 3),
        estimated_days_to_harvest=estimated_days,
        predicted_harvest_date=estimated_date.isoformat(),
        predicted_harvest_window={
            "start": (estimated_date - timedelta(days=3)).isoformat(),
            "end": (estimated_date + timedelta(days=3)).isoformat(),
        },
        daily_gdd_forecast=daily_forecast,
        active_stage=active_stage,
    )