import os
from typing import List, Dict
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from .data_loader import NDVIPoint

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

    def generate_growth_visualization(
        self,
        ndvi_series: List[NDVIPoint],
        forecast: Dict,
        rotation_plan: List[Dict],
        output_path: str,
        scenario: str,
    ):
        if not ndvi_series:
            raise ValueError("NDVI series is empty")

        dates = [p.date for p in ndvi_series]
        ndvi_values = [p.ndvi for p in ndvi_series]
        evi_values = [p.evi for p in ndvi_series]

        timeline = forecast.get("growth_timeline", [])

        fig, axes = plt.subplots(2, 2, figsize=(10, 6))
        ax1, ax2, ax3, ax4 = axes.ravel()

        ax1.plot(dates, ndvi_values, label="NDVI")
        ax1.plot(dates, evi_values, linestyle="--", label="EVI")
        ax1.set_title("Historical Vegetation Indices")
        ax1.set_xlabel("Date")
        ax1.set_ylabel("Index")
        ax1.legend()
        ax1.grid(True, alpha=0.3)

        if timeline:
            f_dates = [t["date"] for t in timeline]
            growth_scores = [t["growth_score"] for t in timeline]
            ax2.plot(f_dates, growth_scores)
            ax2.set_title("Forecast Growth Trajectory")
            ax2.set_xlabel("Date")
            ax2.set_ylabel("Growth score")
            ax2.tick_params(axis="x", rotation=45)
            ax2.grid(True, alpha=0.3)

        crops = [r["crop"] for r in rotation_plan]
        health = [r.get("rotation_health", 0.5) for r in rotation_plan]
        ax3.bar(range(len(crops)), health)
        ax3.set_xticks(range(len(crops)))
        ax3.set_xticklabels(crops, rotation=30, ha="right")
        ax3.set_ylim(0, 1)
        ax3.set_title("Rotation Health by Season")
        ax3.set_ylabel("Health index")

        yield_index = forecast.get("yield_index", 0.0)
        stress_risk = forecast.get("stress_risk", 0.0)
        soil_score = forecast.get("soil_score", 0.0)
        crop_suggestions = forecast.get("crop_suggestions", [])[:3]

        ax4.axis("off")
        lines = [
            f"Scenario: {scenario}",
            f"Yield index: {yield_index:.2f}",
            f"Stress risk: {stress_risk:.2f}",
            f"Soil score: {soil_score:.2f}",
            "",
            "Top crop suggestions:",
        ]
        for s in crop_suggestions:
            lines.append(f"- {s['crop']}: {s['suitability']:.2f}")
        ax4.text(0.02, 0.98, "\n".join(lines), va="top", ha="left")

        plt.tight_layout()
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        fig.savefig(output_path, dpi=120, bbox_inches="tight")
        plt.close(fig)