from typing import List, Dict
from datetime import datetime, timedelta
import math
from .data_loader import NDVIPoint

class TemporalForecaster:
    def __init__(self, config_cls):
        self.cfg = config_cls

    def forecast_growth(
        self,
        ndvi_data: List[NDVIPoint],
        weather: Dict,
        soil: Dict,
        crop: str,
        start_date: datetime,
        seasons: int,
        scenario: str,
    ) -> Dict:
        if not ndvi_data:
            raise ValueError("No NDVI data available for forecasting")

        avg_ndvi = sum(p.ndvi for p in ndvi_data) / len(ndvi_data)
        max_ndvi = max(p.ndvi for p in ndvi_data)
        soil_score = soil.get("soil_score", 0.5)
        volatility = weather.get("volatility", 5.0)

        scenario_modifiers = {
            "baseline": 1.0,
            "wet_year": 1.1,
            "dry_year": 0.85,
            "heatwave": 0.8,
            "low_inputs": 0.9,
        }
        scen_mult = scenario_modifiers.get(scenario, 1.0)

        base_yield = (avg_ndvi * 0.6 + max_ndvi * 0.4) * 2.5
        yield_index = max(0.0, min(1.0, base_yield * soil_score * scen_mult))

        stress_risk = max(0.0, min(1.0, (volatility / 20.0) + (1 - soil_score) * 0.4))
        if scenario in ("dry_year", "heatwave"):
            stress_risk = min(1.0, stress_risk + 0.2)

        timeline = []
        day = start_date
        total_days = int(self.cfg.NDVI_LENGTH_DAYS * seasons / 2)
        for i in range(total_days):
            growth_phase = math.sin((i / total_days) * math.pi)
            growth_phase = max(0.0, growth_phase)
            growth_score = growth_phase * yield_index
            stage = self._stage_from_progress((i + 1) / total_days)
            timeline.append(
                {
                    "date": day.strftime("%Y-%m-%d"),
                    "growth_score": round(growth_score, 3),
                    "stage": stage,
                }
            )
            day += timedelta(days=1)

        suggestions = self._suggest_crops(crop, soil_score, scenario)

        return {
            "yield_index": round(yield_index, 3),
            "stress_risk": round(stress_risk, 3),
            "weather_volatility": weather.get("volatility", 0.0),
            "soil_score": soil_score,
            "growth_timeline": timeline,
            "crop_suggestions": suggestions,
        }

    def _stage_from_progress(self, frac: float) -> str:
        if frac < 0.15:
            return "Emergence"
        if frac < 0.35:
            return "Vegetative"
        if frac < 0.6:
            return "Flowering"
        if frac < 0.85:
            return "Grain filling"
        return "Maturity"

    def _suggest_crops(self, current_crop: str, soil_score: float, scenario: str):
        base_options = ["Maize", "Wheat", "Rice", "Soybean", "Cotton"]
        suggestions = []
        for c in base_options:
            if c == current_crop:
                continue
            score = soil_score
            if c in ("Rice",) and scenario == "dry_year":
                score -= 0.2
            if c in ("Cotton",) and scenario == "wet_year":
                score -= 0.1
            suggestions.append({"crop": c, "suitability": round(max(0.1, min(1.0, score)), 2)})
        suggestions.sort(key=lambda x: x["suitability"], reverse=True)
        return suggestions