# scripts/helpers/engines/physics_engine.py
from __future__ import annotations
import numpy as np
import pandas as pd
from datetime import datetime, timedelta
from typing import Any
from .base_engine import BaseEngine
from ..utils import set_global_seed
from .. import features


class PhysicsEngine(BaseEngine):
    """
    Engine for physics-style data generation (e.g. weather simulation).

    Handles:
    - Iteration over multiple regions / cities
    - Daily temporal generation (loop over days)
    - Physical dependency chain (temperature → humidity → rain → pressure)
    - Optional world randomness and feature noise
    """

    def __init__(self, config: dict[str, Any]):
        super().__init__(config)

        # Physical / world parameters
        self.days = config.get("days", 365)
        self.start_date = config.get("start_date", "2024-01-01")
        self.world_randomness = config.get("world_randomness", 1.0)
        self.noise = config.get("noise", None) or {}
        self.cities = config.get("cities", {}) or config.get("regions", {})
        self.output_path = config.get("output_path", None)
        self.as_df = config.get("as_df", True)
        self.filename = config.get("filename", "weather.csv")

        set_global_seed(self.seed)

    # ---------------------------------------------------------
    # Core generator
    # ---------------------------------------------------------
    def _generate_city_weather(self, city: str, params: dict[str, float]) -> list[dict[str, Any]]:
        """
        Generate one city's daily weather series.
        """
        records = []
        start = datetime.fromisoformat(self.start_date)

        for offset in range(self.days):
            date = start + timedelta(days=offset)
            doy = date.timetuple().tm_yday

            temp = features.temperature(
                doy,
                params.get("temp_shift", 0),
                noise=self.noise.get("temperature"),
            )
            hum = features.humidity(temp, noise=self.noise.get("humidity"))
            rain = features.rain_mm(
                temp,
                hum,
                params.get("rain_factor", 1.0),
                self.world_randomness,
                noise=self.noise.get("rain_mm"),
            )
            pres = features.pressure(rain, noise=self.noise.get("pressure"))

            # Optional derived feature
            dew = None
            if "dew_point" in self.noise or "dew_point" in params:
                dew = features.dew_point(temp, hum, noise=self.noise.get("dew_point"))

            record = {
                "date": date.strftime("%Y-%m-%d"),
                "city": city,
                "temp_c": round(temp, 1),
                "humidity": round(hum, 1),
                "press_air": round(pres, 1),
                "rain_mm": round(rain, 2),
            }
            if dew is not None:
                record["dew_point"] = round(dew, 1)

            records.append(record)

        return records

    # ---------------------------------------------------------
    # Main entry
    # ---------------------------------------------------------
    def generate(self) -> pd.DataFrame:
        """
        Generate a full physical weather dataset for all configured cities.
        """
        all_records: list[dict[str, Any]] = []

        if not self.cities:
            raise ValueError("No cities or regions defined in config.")

        for city, params in self.cities.items():
            city_records = self._generate_city_weather(city, params)
            all_records.extend(city_records)

        df = pd.DataFrame(all_records)

        # ✅ Optional derived features (if defined)
        df = self.apply_derived_features(df, self.config)

        # ✅ Optional saving
        if self.output_path:
            import os
            os.makedirs(self.output_path, exist_ok=True)
            filename = self.filename or f"physics_{self.scenario}.csv"
            file_path = f"{self.output_path}/{filename}"
            df.to_csv(file_path, index=False)
            print(f"✅ Physics dataset saved → {file_path} ({len(df)} rows)")
            return df if self.as_df else None

        return df

