# grid_env.py
import gymnasium as gym
import numpy as np
from typing import Dict, Tuple, List
from dataclasses import dataclass
import random


@dataclass
class GridConfig:
    """Smart grid configuration."""
    num_solar_farms: int = 3
    num_wind_farms: int = 2
    num_batteries: int = 2
    num_gas_peakers: int = 1
    num_dr_aggregators: int = 5
    total_demand_mw: float = 500
    renewable_capacity_mw: float = 600
    battery_capacity_mwh: float = 200
    gas_capacity_mw: float = 200
    carbon_grid_intensity: float = 350  # gCO2/kWh
    carbon_gas_intensity: float = 450  # gCO2/kWh
    episode_length: int = 288  # 24 hours @ 5-min steps
    dt_hours: float = 5 / 60


class SmartGridEnv(gym.Env):
    """Multi-agent smart grid environment."""

    def __init__(self, config: GridConfig = None):
        super().__init__()
        self.config = config or GridConfig()

        self.num_agents = (
            config.num_solar_farms + config.num_wind_farms
            + config.num_batteries + config.num_gas_peakers
            + config.num_dr_aggregators
        )

        obs_dim = (
            1  # demand
            + 1  # renewable forecast
            + 1  # battery soc
            + 1  # carbon intensity
            + 1  # price
            + 1  # time encoding
        )

        self.observation_space = gym.spaces.Box(
            low=-np.inf, high=np.inf, shape=(obs_dim,), dtype=np.float32
        )

        self.action_space = gym.spaces.Box(
            low=-1, high=1, shape=(1,), dtype=np.float32
        )

        self._init_state()

    def _init_state(self):
        """Initialize grid state."""
        self.demand = self.config.total_demand_mw
        self.renewable_output = 0.0
        self.battery_soc = 0.5  # 50% state of charge
        self.gas_output = 0.0
        self.carbon_intensity = self.config.carbon_grid_intensity
        self.price = 50.0  # $/MWh
        self.current_step = 0
        self.total_carbon = 0.0
        self.total_cost = 0.0
        self.renewable_used = 0.0
        self.renewable_curtailed = 0.0
        self.demand_fulfilled = 0.0

    def _get_demand_profile(self, step: int) -> float:
        """Get demand at current time step."""
        hour = (step * 5 / 60) % 24
        base = self.config.total_demand_mw
        daily_pattern = 1 + 0.3 * np.sin((hour - 6) * np.pi / 12)
        return base * daily_pattern

    def _get_renewable_forecast(self, step: int) -> float:
        """Get renewable generation forecast."""
        hour = (step * 5 / 60) % 24

        solar = max(0, np.sin((hour - 6) * np.pi / 12)) * 0.8
        wind = 0.3 + 0.2 * np.sin(hour * np.pi / 12)

        total = (solar * 0.6 + wind * 0.4) * self.config.renewable_capacity_mw
        return total

    def _get_obs(self) -> np.ndarray:
        """Get current observation."""
        hour = (self.current_step * 5 / 60) % 24
        time_sin = np.sin(2 * np.pi * hour / 24)
        time_cos = np.cos(2 * np.pi * hour / 24)

        return np.array([
            self.demand / self.config.total_demand_mw,
            self.renewable_output / self.config.renewable_capacity_mw,
            self.battery_soc,
            self.carbon_intensity / 500,
            self.price / 100,
            time_sin,
        ], dtype=np.float32)

    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict]:
        """Execute one time step."""
        action_value = float(action[0])

        demand = self._get_demand_profile(self.current_step)
        renewable_forecast = self._get_renewable_forecast(self.current_step)

        renewable_output = max(0, min(renewable_forecast, demand))
        renewable_curtailed = max(0, renewable_forecast - demand)

        battery_dispatch = 0.0
        if action_value > 0.3:
            battery_dispatch = min(
                action_value * self.config.battery_capacity_mwh / 4,
                self.battery_soc * self.config.battery_capacity_mwh,
            )
            self.battery_soc -= battery_dispatch / self.config.battery_capacity_mwh
        elif action_value < -0.3:
            charge_amount = abs(action_value) * self.config.battery_capacity_mwh / 4
            max_charge = (1 - self.battery_soc) * self.config.battery_capacity_mwh
            charge_amount = min(charge_amount, max_charge)
            self.battery_soc += charge_amount / self.config.battery_capacity_mwh

        remaining_demand = demand - renewable_output - battery_dispatch
        gas_output = max(0, min(remaining_demand, self.config.gas_capacity_mw))

        demand_fulfilled = min(demand, renewable_output + battery_dispatch + gas_output)

        gas_emissions = gas_output * self.config.dt_hours * self.config.carbon_gas_intensity / 1000
        grid_emissions = (demand - renewable_output) * self.config.dt_hours * self.carbon_intensity / 1000
        total_step_carbon = gas_emissions

        cost = (
            gas_output * self.config.dt_hours * self.price / 1000
            + renewable_curtailed * 0.5
            + max(0, demand - demand_fulfilled) * 10
        )

        self.total_carbon += total_step_carbon
        self.total_cost += cost
        self.renewable_used += renewable_output * self.config.dt_hours
        self.renewable_curtailed += renewable_curtailed * self.config.dt_hours
        self.demand_fulfilled += demand_fulfilled * self.config.dt_hours
        self.gas_output = gas_output
        self.renewable_output = renewable_output
        self.demand = demand

        reward = self._calculate_reward(
            demand_fulfilled, demand, renewable_output, renewable_forecast,
            total_step_carbon, cost,
        )

        self.current_step += 1
        terminated = self.current_step >= self.config.episode_length
        truncated = False

        info = {
            "demand_mw": demand,
            "renewable_mw": renewable_output,
            "battery_mw": battery_dispatch,
            "gas_mw": gas_output,
            "demand_fulfilled": demand_fulfilled,
            "carbon_kg": total_step_carbon,
            "cost_usd": cost,
            "renewable_utilization": renewable_output / max(renewable_forecast, 1),
            "reliability": demand_fulfilled / max(demand, 1),
        }

        return self._get_obs(), reward, terminated, truncated, info

    def _calculate_reward(
        self,
        demand_fulfilled: float,
        demand: float,
        renewable_output: float,
        renewable_forecast: float,
        carbon: float,
        cost: float,
    ) -> float:
        """Calculate multi-objective reward."""
        reliability_reward = 10.0 * (demand_fulfilled / max(demand, 1))

        renewable_reward = 5.0 * (renewable_output / max(renewable_forecast, 1))

        carbon_penalty = -carbon * 0.1

        cost_penalty = -cost * 0.01

        return reliability_reward + renewable_reward + carbon_penalty + cost_penalty

    def reset(self, seed=None, options=None):
        """Reset environment."""
        super().reset(seed=seed)
        self._init_state()
        return self._get_obs(), {}

# marl_agent.py
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Tuple
import numpy as np
from stable_baselines3 import PPO
from stable_baselines3.common.policies import ActorCriticPolicy
from stable_baselines3.common.callbacks import BaseCallback


class CarbonRewardShaping(BaseCallback):
    """Callback for carbon-aware reward shaping."""

    def __init__(self, carbon_budget: float = 1000, verbose: int = 0):
        super().__init__(verbose)
        self.carbon_budget = carbon_budget
        self.total_carbon = 0

    def _on_step(self) -> bool:
        info = self.locals.get("info", {})
        if "carbon_kg" in info:
            self.total_carbon += info["carbon_kg"]

        if self.total_carbon > self.carbon_budget:
            penalty = -100
            self.locals["rewards"] = self.locals["rewards"] + penalty

        return True

    def _on_rollout_end(self):
        utilization = 0
        if "renewable_utilization" in self.locals.get("info", {}):
            utilization = self.locals["info"]["renewable_utilization"]

        self.model.logger.record("carbon/total_kg", self.total_carbon)
        self.model.logger.record("carbon/budget_remaining",
                                  max(0, self.carbon_budget - self.total_carbon))


class MultiAgentGridController:
    """Controls multiple agents in the smart grid."""

    def __init__(self, env, num_agents: int = 13):
        self.env = env
        self.num_agents = num_agents

        self.agent_configs = {
            "solar": {"policy": "MlpPolicy", "lr": 3e-4, "ent_coef": 0.01},
            "wind": {"policy": "MlpPolicy", "lr": 3e-4, "ent_coef": 0.01},
            "battery": {"policy": "MlpPolicy", "lr": 1e-3, "ent_coef": 0.05},
            "gas": {"policy": "MlpPolicy", "lr": 3e-4, "ent_coef": 0.02},
            "dr": {"policy": "MlpPolicy", "lr": 5e-4, "ent_coef": 0.03},
        }

        self.agents = {}
        self._create_agents()

    def _create_agents(self):
        """Create individual agent controllers."""
        for agent_id in range(self.num_agents):
            agent_type = self._get_agent_type(agent_id)
            config = self.agent_configs[agent_type]

            agent = PPO(
                config["policy"],
                self.env,
                learning_rate=config["lr"],
                n_steps=2048,
                batch_size=64,
                n_epochs=10,
                gamma=0.99,
                gae_lambda=0.95,
                clip_range=0.2,
                ent_coef=config["ent_coef"],
                verbose=0,
            )
            self.agents[agent_id] = {"model": agent, "type": agent_type}

    def _get_agent_type(self, agent_id: int) -> str:
        """Map agent ID to type."""
        config = self.env.config
        if agent_id < config.num_solar_farms:
            return "solar"
        elif agent_id < config.num_solar_farms + config.num_wind_farms:
            return "wind"
        elif agent_id < config.num_solar_farms + config.num_wind_farms + config.num_batteries:
            return "battery"
        elif agent_id < (config.num_solar_farms + config.num_wind_farms +
                         config.num_batteries + config.num_gas_peakers):
            return "gas"
        else:
            return "dr"

    def train(self, total_timesteps: int = 500000):
        """Train all agents."""
        carbon_callback = CarbonRewardShaping(carbon_budget=500)

        for agent_id, agent_info in self.agents.items():
            print(f"Training agent {agent_id} ({agent_info['type']})")
            agent_info["model"].learn(
                total_timesteps=total_timesteps // self.num_agents,
                callback=[carbon_callback],
            )

    def get_collective_action(self, observations: Dict) -> Dict:
        """Get coordinated actions from all agents."""
        actions = {}
        for agent_id, obs in observations.items():
            if agent_id in self.agents:
                model = self.agents[agent_id]["model"]
                action, _ = model.predict(obs, deterministic=True)
                actions[agent_id] = action
        return actions

    def evaluate(self, n_episodes: int = 100) -> Dict:
        """Evaluate collective performance."""
        results = {
            "total_carbon": [],
            "total_cost": [],
            "reliability": [],
            "renewable_utilization": [],
        }

        for _ in range(n_episodes):
            obs, _ = self.env.reset()
            episode_results = {
                "carbon": 0, "cost": 0, "reliability": 0, "renewable": 0,
            }

            for step in range(self.env.config.episode_length):
                actions = {0: self.agents[0]["model"].predict(obs, deterministic=True)[0]}
                obs, reward, terminated, truncated, info = self.env.step(actions[0])

                episode_results["carbon"] += info.get("carbon_kg", 0)
                episode_results["cost"] += info.get("cost_usd", 0)
                episode_results["reliability"] += info.get("reliability", 0)
                episode_results["renewable"] += info.get("renewable_utilization", 0)

            results["total_carbon"].append(episode_results["carbon"])
            results["total_cost"].append(episode_results["cost"])
            results["reliability"].append(episode_results["reliability"] / self.env.config.episode_length)
            results["renewable_utilization"].append(episode_results["renewable"] / self.env.config.episode_length)

        return {
            "avg_carbon_kg": np.mean(results["total_carbon"]),
            "avg_cost_usd": np.mean(results["total_cost"]),
            "avg_reliability": np.mean(results["reliability"]),
            "avg_renewable_utilization": np.mean(results["renewable_utilization"]),
        }