"""Multi-location Open-Meteo client.

Fetches weather forecasts for all HK weather stations simultaneously,
extracting spatial gradients (cross-station deltas) that capture
HK's microclimate variability.

Uses the Open-Meteo multi-location API endpoint to batch-fetch data
for all stations in a single request.
"""

import numpy as np
import pandas as pd
from datetime import datetime
from typing import Dict, List, Optional, Tuple

from config import HK_COORDS


class SpatialWeatherClient:
    """
    Fetches weather data for all HK monitoring stations.

    Features extracted:
      - Per-station forecasts (temperature, wind, precipitation)
      - Cross-station gradients (urban-rural, coastal-inland, elevation)
      - Convergence/divergence indices (wind, pressure)
      - Spatial variability (station-to-station spread)
    """

    # Station groupings for spatial features
    STATION_GROUPS = {
        "urban": ["hko_headquarters"],
        "rural": ["sheung_shui", "ta_kwu_ling"] if "ta_kwu_ling" in HK_COORDS else ["sheung_shui"],
        "coastal": ["stanley", "cheung_chau"],
        "inland": ["chek_lap_kok"],
        "all": list(HK_COORDS.keys()),
    }

    # Orographic/geographic characteristics
    STATION_METADATA = {
        "hko_headquarters": {"elevation_m": 32, "dist_coast_km": 1.2, "urban": 1.0},
        "chek_lap_kok":     {"elevation_m": 5,  "dist_coast_km": 0.0, "urban": 0.3},
        "sheung_shui":      {"elevation_m": 10, "dist_coast_km": 15.0, "urban": 0.4},
        "stanley":          {"elevation_m": 30, "dist_coast_km": 0.3, "urban": 0.2},
        "cheung_chau":      {"elevation_m": 72, "dist_coast_km": 0.0, "urban": 0.1},
    }

    def __init__(self, cache_ttl: int = 3600):
        self.cache_ttl = cache_ttl
        self._cache: Dict[str, Tuple[datetime, pd.DataFrame]] = {}

    def fetch_all_stations(
        self,
        lead_days: int = 5,
        use_cache: bool = True,
    ) -> Dict[str, pd.DataFrame]:
        """
        Fetch daily forecasts for all HK stations.

        Uses Open-Meteo's multi-location endpoint to fetch all stations
        in a single HTTP request.

        Returns dict of {station_name: daily_forecast_df}
        """
        cache_key = f"all_{lead_days}"

        if use_cache and cache_key in self._cache:
            ts, data = self._cache[cache_key]
            if (datetime.now() - ts).total_seconds() < self.cache_ttl:
                return data

        try:
            import openmeteo_requests
            import requests_cache
            from retry_requests import retry

            cache = requests_cache.CachedSession('.openmeteo_spatial_cache', expire_after=self.cache_ttl)
            retry_session = retry(cache, retries=2, backoff_factor=0.2)
            client = openmeteo_requests.Client(session=retry_session)

            # Build multi-location params
            stations = [
                s for s in HK_COORDS.values()
                if isinstance(s, tuple) and len(s) == 2
            ]
            lats = [s[0] for s in stations]
            lons = [s[1] for s in stations]

            daily_vars = [
                "temperature_2m_max", "temperature_2m_min",
                "precipitation_sum", "precipitation_probability_max",
                "wind_speed_10m_max", "wind_gusts_10m_max",
                "weather_code",
            ]

            params = {
                "latitude": lats,
                "longitude": lons,
                "timezone": "Asia/Hong_Kong",
                "daily": daily_vars,
                "forecast_days": lead_days,
            }

            responses = client.weather_api(
                "https://api.open-meteo.com/v1/forecast", params=params
            )

            results = {}
            station_names = [k for k in HK_COORDS.keys() 
                           if isinstance(HK_COORDS[k], tuple) and len(HK_COORDS[k]) == 2]

            for i, (name, resp) in enumerate(zip(station_names, responses)):
                daily = resp.Daily()

                dates = pd.date_range(
                    start=pd.Timestamp(daily.Time(), unit="s", tz="UTC"),
                    end=pd.Timestamp(daily.TimeEnd(), unit="s", tz="UTC"),
                    freq=pd.Timedelta(seconds=daily.Interval()),
                    inclusive="left",
                )

                data = {"date": dates}
                for j in range(daily.VariablesLength()):
                    var = daily.Variables(j)
                    if j < len(daily_vars):
                        data[daily_vars[j]] = var.ValuesAsNumpy()

                df = pd.DataFrame(data).set_index("date")
                df.index = df.index.tz_convert("Asia/Hong_Kong")
                results[name] = df

            self._cache[cache_key] = (datetime.now(), results)
            return results

        except Exception as e:
            print(f"Spatial weather fetch error: {e}")
            return {}

    def extract_spatial_features(
        self,
        station_data: Dict[str, pd.DataFrame],
        day_index: int = 0,
    ) -> Dict[str, float]:
        """
        Extract spatial gradient features from multi-station forecasts.

        Returns dict of derived spatial features for a single forecast day.
        """
        if not station_data or len(station_data) < 2:
            return {}

        features = {}
        
        # Get values for the target day
        vals: Dict[str, Dict[str, float]] = {}
        for name, df in station_data.items():
            if day_index < len(df):
                row = df.iloc[day_index]
                vals[name] = {
                    "tmax": float(row.get("temperature_2m_max", np.nan)),
                    "tmin": float(row.get("temperature_2m_min", np.nan)),
                    "precip_prob": float(row.get("precipitation_probability_max", 0)),
                    "wind_max": float(row.get("wind_speed_10m_max", np.nan)),
                    "gust_max": float(row.get("wind_gusts_10m_max", np.nan)),
                }

        if len(vals) < 2:
            return features

        # === Temperature gradients ===
        tmax_vals = [v["tmax"] for v in vals.values() if not np.isnan(v["tmax"])]
        tmin_vals = [v["tmin"] for v in vals.values() if not np.isnan(v["tmin"])]

        if tmax_vals:
            features["tmax_range_hk"] = max(tmax_vals) - min(tmax_vals)
            features["tmax_std_hk"] = np.std(tmax_vals) if len(tmax_vals) > 1 else 0
        if tmin_vals:
            features["tmin_range_hk"] = max(tmin_vals) - min(tmin_vals)

        # === Urban Heat Island (urban - rural temperature delta) ===
        if "hko_headquarters" in vals and "sheung_shui" in vals:
            uhi_tmax = vals["hko_headquarters"]["tmax"] - vals["sheung_shui"]["tmax"]
            uhi_tmin = vals["hko_headquarters"]["tmin"] - vals["sheung_shui"]["tmin"]
            features["uhi_tmax_delta"] = uhi_tmax if not np.isnan(uhi_tmax) else 0
            features["uhi_tmin_delta"] = uhi_tmin if not np.isnan(uhi_tmin) else 0

        # === Coastal-inland temperature gradient ===
        coastal_tmax = [vals[s]["tmax"] for s in ["stanley", "cheung_chau"] 
                        if s in vals and not np.isnan(vals[s]["tmax"])]
        inland_tmax = [vals[s]["tmax"] for s in ["sheung_shui"] 
                       if s in vals and not np.isnan(vals[s]["tmax"])]
        if coastal_tmax and inland_tmax:
            features["coastal_inland_tmax_delta"] = np.mean(inland_tmax) - np.mean(coastal_tmax)

        # === Precipitation spatial heterogeneity ===
        precip_vals = [v["precip_prob"] for v in vals.values() if not np.isnan(v["precip_prob"])]
        if len(precip_vals) > 1:
            features["precip_prob_range"] = max(precip_vals) - min(precip_vals)
            features["precip_prob_std"] = np.std(precip_vals)

        # === Wind convergence index ===
        # High wind variability across stations → convergence
        wind_vals = [v["wind_max"] for v in vals.values() if not np.isnan(v["wind_max"])]
        gust_vals = [v["gust_max"] for v in vals.values() if not np.isnan(v["gust_max"])]
        if len(wind_vals) > 1:
            features["wind_range_hk"] = max(wind_vals) - min(wind_vals)
            features["wind_std_hk"] = np.std(wind_vals)
            # Convergence: high gust range relative to mean wind
            mean_wind = np.mean(wind_vals)
            gust_range = max(gust_vals) - min(gust_vals) if len(gust_vals) > 1 else 0
            features["convergence_index"] = gust_range / max(mean_wind, 0.1)

        # === Elevation-adjusted temperature ===
        # Higher elevations should be cooler (lapse rate ~0.65°C/100m)
        if "cheung_chau" in vals and "hko_headquarters" in vals:
            meta_cc = self.STATION_METADATA.get("cheung_chau", {})
            meta_hko = self.STATION_METADATA.get("hko_headquarters", {})
            elev_diff = meta_cc.get("elevation_m", 0) - meta_hko.get("elevation_m", 0)
            features["elevation_corrected_tmax"] = vals["cheung_chau"]["tmax"] + 0.0065 * elev_diff

        # === Distance-to-coast precipitation modifier ===
        # Inland stations often get less rain in summer convective events
        coastal_precip = [vals[s]["precip_prob"] for s in ["stanley", "cheung_chau"] 
                         if s in vals and not np.isnan(vals[s]["precip_prob"])]
        inland_precip = [vals[s]["precip_prob"] for s in ["sheung_shui"] 
                        if s in vals and not np.isnan(vals[s]["precip_prob"])]
        if coastal_precip and inland_precip:
            features["coastal_precip_excess"] = np.mean(coastal_precip) - np.mean(inland_precip)

        # === Composite spatial instability index (0-1) ===
        indicators = []
        for key, norm in [
            ("tmax_range_hk", 5.0),
            ("precip_prob_std", 30.0),
            ("wind_std_hk", 10.0),
            ("convergence_index", 5.0),
        ]:
            if key in features and not np.isnan(features[key]):
                indicators.append(np.clip(features[key] / norm, 0, 1))
        features["spatial_instability"] = float(np.mean(indicators)) if indicators else 0.0

        return features

    def convert_to_hourly(
        self,
        station_data: Dict[str, pd.DataFrame],
        station_name: str,
    ) -> Optional[pd.DataFrame]:
        """Estimate hourly data from daily for a specific station.

        Used to provide hourly-resolution features for the ML pipeline
        when the spatial fetch doesn't include hourly data.
        """
        if station_name not in station_data:
            return None

        daily = station_data[station_name]
        hours = []
        for day_idx in range(len(daily)):
            row = daily.iloc[day_idx]
            base_date = pd.Timestamp(row.name)
            for h in range(24):
                # Simple diurnal cycle model
                hour_frac = h / 24.0
                t_range = float(row.get("temperature_2m_max", 25) - row.get("temperature_2m_min", 20))
                t_hourly = float(row.get("temperature_2m_min", 20)) + t_range * np.sin(np.pi * (h - 6) / 12) ** 2
                hours.append({
                    "date": base_date + pd.Timedelta(hours=h),
                    "temperature_2m": t_hourly,
                    "relative_humidity_2m": 75 - 15 * np.sin(np.pi * (h - 6) / 12) ** 2,
                    "precipitation_probability": float(row.get("precipitation_probability_max", 0)) / 24,
                    "wind_speed_10m": float(row.get("wind_speed_10m_max", 10)) * 0.6,
                    "surface_pressure": 1013.0,
                    "cloud_cover": float(row.get("precipitation_probability_max", 0)) * 0.6,
                    "visibility": 15000.0,
                })
        df = pd.DataFrame(hours).set_index("date")
        df.index = pd.to_datetime(df.index)
        return df