from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from gibil.classes.models import ForecastKind, PowerForecastPoint, PowerForecastRun, WeatherForecastPoint
from gibil.classes.oracle.config import EnergyForecastConfig
from gibil.classes.predictors.solar_baseline import BaselineSolarProductionOracle
from gibil.classes.predictors.math_utils import dot, fit_ridge_regression, quantile
from gibil.classes.sigen.store import SigenStore
from gibil.classes.weather.store import WeatherStore


class RollingSolarRegressionOracle:
    """Forecasts solar production with a rolling ridge regression."""

    model_version = "rolling_solar_regression_v1"

    def __init__(
        self,
        weather_store: WeatherStore,
        sigen_store: SigenStore,
        config: EnergyForecastConfig,
    ) -> None:
        self.weather_store = weather_store
        self.sigen_store = sigen_store
        self.config = config

    def forecast(self, issued_at: datetime | None = None) -> PowerForecastRun:
        if issued_at is None:
            issued_at = datetime.now(timezone.utc)

        weather_points = self.weather_store.load_latest_forecast_points(
            start_at=issued_at,
            end_at=issued_at + timedelta(hours=self.config.horizon_hours),
        )
        training_samples = self.sigen_store.load_solar_training_samples(
            lookback=timedelta(days=self.config.solar_training_days)
        )
        model = self._fit_model(training_samples)
        if model is None:
            return BaselineSolarProductionOracle(
                weather_store=self.weather_store,
                sigen_store=self.sigen_store,
                config=self.config,
            ).forecast(issued_at=issued_at)

        points = [
            self._forecast_point(
                weather_point=point,
                issued_at=issued_at,
                model=model,
                training_sample_count=len(training_samples),
            )
            for point in weather_points
        ]

        return PowerForecastRun(
            issued_at=issued_at,
            kind=ForecastKind.SOLAR,
            source="rolling_solar_regression_oracle",
            model_version=self.model_version,
            points=points,
        )

    def _fit_model(
        self,
        samples: list[dict[str, float | int | object]],
    ) -> "_SolarRegressionModel | None":
        if len(samples) < self.config.solar_min_training_samples:
            return None

        features = [
            self._features(
                radiation=float(sample["shortwave_radiation_w_m2"]),
                cloud_cover=float(sample["cloud_cover_pct"]),
            )
            for sample in samples
        ]
        targets = [float(sample["solar_power_w"]) for sample in samples]

        coefficients = fit_ridge_regression(
            features,
            targets,
            ridge_lambda=self.config.solar_ridge_lambda,
        )
        if coefficients is None:
            return None

        residuals = [
            target - dot(coefficients, feature)
            for feature, target in zip(features, targets)
        ]
        return _SolarRegressionModel(
            coefficients=coefficients,
            residual_p10=quantile(residuals, 0.10),
            residual_p90=quantile(residuals, 0.90),
            peak_w=self._solar_peak_w(),
        )

    def _forecast_point(
        self,
        weather_point: WeatherForecastPoint,
        issued_at: datetime,
        model: "_SolarRegressionModel",
        training_sample_count: int,
    ) -> PowerForecastPoint:
        radiation = max(weather_point.shortwave_radiation_w_m2 or 0.0, 0.0)
        cloud_cover = self._cloud_cover(weather_point.cloud_cover_pct)
        expected = model.predict(self._features(radiation, cloud_cover))
        expected *= self.config.solar_scale
        p10 = max(0.0, expected + model.residual_p10)
        p90 = min(model.peak_w, expected + model.residual_p90)
        if p90 < expected:
            p90 = expected
        if p10 > expected:
            p10 = expected

        return PowerForecastPoint(
            target_at=weather_point.target_at,
            horizon_minutes=self._horizon_minutes(issued_at, weather_point.target_at),
            expected_power_w=expected,
            p10_power_w=p10,
            p50_power_w=expected,
            p90_power_w=p90,
            confidence=0.45,
            source="rolling_solar_regression",
            model_version=self.model_version,
            metadata={
                "shortwave_radiation_w_m2": weather_point.shortwave_radiation_w_m2,
                "cloud_cover_pct": weather_point.cloud_cover_pct,
                "temperature_c": weather_point.temperature_c,
                "solar_peak_w": model.peak_w,
                "training_sample_count": training_sample_count,
                "residual_p10_w": model.residual_p10,
                "residual_p90_w": model.residual_p90,
            },
        )

    def _features(self, radiation: float, cloud_cover: float) -> list[float]:
        radiation_kw = radiation / 1000.0
        cloud = cloud_cover / 100.0
        clear = 1.0 - cloud
        return [
            1.0,
            radiation_kw,
            radiation_kw * clear,
            radiation_kw * cloud,
            cloud,
        ]

    def _cloud_cover(self, value: float | None) -> float:
        if value is None:
            return 0.0
        return min(max(value, 0.0), 100.0)

    def _solar_peak_w(self) -> float:
        recent_peak = self.sigen_store.load_recent_solar_peak_w()
        if recent_peak is None or recent_peak <= 0:
            return self.config.fallback_solar_peak_w
        return recent_peak * max(self.config.solar_peak_headroom, 1.0)

    def _horizon_minutes(self, issued_at: datetime, target_at: datetime) -> int:
        return max(0, round((target_at - issued_at).total_seconds() / 60))


@dataclass(frozen=True)
class _SolarRegressionModel:
    coefficients: list[float]
    residual_p10: float
    residual_p90: float
    peak_w: float

    def predict(self, features: list[float]) -> float:
        return min(max(dot(self.coefficients, features), 0.0), self.peak_w)