from __future__ import annotations
from typing import Dict, Any, List, Tuple
import math
import numpy as np
import pandas as pd

from .utils import heat_index_c


class RiskModel:
    """Hackathon-friendly, transparent risk model.

    Produces:
      - daily risk (0–100)
      - 12-month livability bands from climatology + warming offset
    """

    def __init__(self,
                 weights: Dict[str, float] | None = None,
                 heat_threshold_c: float = 38.0,
                 flood_daily_mm: float = 50.0,
                 wind_max_kph: float = 60.0):
        self.weights = weights or {
            "heat": 0.45,
            "humidity": 0.15,
            "precip": 0.20,
            "wind": 0.10,
            "air": 0.10,
        }
        self.heat_threshold = heat_threshold_c
        self.flood_daily_mm = flood_daily_mm
        self.wind_max_kph = wind_max_kph

    def daily_score(self, tmax_c: float, tmin_c: float, rh: float | None,
                    precip_mm: float, wind_kph: float, aqi: float | None) -> Tuple[float, Dict[str, float]]:
        # Heat component via heat index on Tmax with RH (fallback RH=40%)
        rh_eff = 40.0 if (rh is None or math.isnan(rh)) else rh
        hi = heat_index_c(tmax_c, rh_eff)
        heat_risk = np.clip((hi - self.heat_threshold) / 10.0, 0, 1)

        # Humidity discomfort
        hum_risk = np.clip((rh_eff - 70.0) / 30.0, 0, 1)

        # Precip flood proxy
        precip_risk = np.clip((precip_mm - self.flood_daily_mm) / 50.0, 0, 1)

        # Wind storm proxy (Beaufort 8+ ≈ >62 kph)
        wind_risk = np.clip((wind_kph - self.wind_max_kph) / 40.0, 0, 1)

        # Air quality (AQI 0–500). If None, neutral at 0.3
        if aqi is None or math.isnan(aqi):
            air_risk = 0.3
        else:
            air_risk = np.clip((aqi - 100.0) / 200.0, 0, 1)

        factors = {
            "heat": float(heat_risk),
            "humidity": float(hum_risk),
            "precip": float(precip_risk),
            "wind": float(wind_risk),
            "air": float(air_risk),
        }
        score01 = sum(self.weights[k] * factors[k] for k in self.weights)
        return float(round(score01 * 100, 1)), factors

    def band_from_score(self, score: float) -> str:
        if score < 25:
            return "SAFE"
        if score < 50:
            return "CAUTION"
        if score < 75:
            return "DANGER"
        return "EXTREME"

    def monthly_projection(self, monthly_t_mean: List[float], monthly_prcp: List[float], warming_offset_c: float = 0.4) -> List[Dict[str, Any]]:
        # Apply simple warming offset to means and map to bands via thresholds
        out = []
        for m in range(12):
            t = monthly_t_mean[m] + warming_offset_c
            p = monthly_prcp[m]
            # Map to synthetic risk
            heat_component = np.clip((t - 32.0) / 8.0, 0, 1)  # hotter than 32C pushes to danger
            rain_component = np.clip((p - 150.0) / 150.0, 0, 1)  # very wet months push flood risk
            score = (0.7 * heat_component + 0.3 * rain_component) * 100
            out.append({
                "month": m + 1,
                "temp_c": round(t, 1),
                "precip_mm": round(p, 1),
                "score": round(float(score), 1),
                "band": self.band_from_score(score),
            })
        return out