import gymnasium as gym
from gymnasium import spaces
import numpy as np
from battery import Battery

class CooptimEnv(gym.Env):
    def __init__(self, input_data, config):
        super().__init__()
        self.input_data = input_data.copy()
        self.config = config
        self.current_step = 0
        
        # Initialize Battery from Config
        bat_cfg = config["battery"]
        self.battery = Battery(
            e_max_mwh=bat_cfg["e_max_mwh"],
            p_ch_max=bat_cfg["p_ch_max_mw"],
            p_dis_max=bat_cfg["p_dis_max_mw"],
            eta_ch=bat_cfg["eta_ch"],
            eta_dis=bat_cfg["eta_dis"]
        )

        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(input_data.shape[1] + 1,), dtype=np.float32
        )

        # Action: [Market Choice, Power Fraction]
        self.action_space = spaces.Box(
            low=np.array([0.0, -1.0]), high=np.array([1.0, 1.0]), dtype=np.float32
        )

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self.current_step = 0
        soc = self.battery.reset()
        obs = np.append(self.input_data.iloc[self.current_step].values, soc)
        return obs.astype(np.float32), {}

    def step(self, action):
        market_choice, p_fraction = action
        inp = self.input_data.iloc[self.current_step]
        
        c_energy = self.config["columns"]["energy"]
        c_fcr = self.config["columns"]["fcr"]
        
        reward = 0.0
        dt_hours = 0.25 # 15-min intervals

        if market_choice < 0.6: 
            # --- ARBITRAGE MODE ---
            soc = self.battery.step(p_fraction, dt_hours)
            # Revenue = flow * price (Positive for sell, Negative for buy)
            revenue = p_fraction * inp[c_energy] * dt_hours
            # Throughput penalty to prevent excessive micro-trading
            t_penalty = self.config["throughput_penalty"]["c_eur_per_mwh"]
            costs = abs(p_fraction * self.battery.p_dis_max) * dt_hours * t_penalty
            reward = revenue - costs
        else:
            # --- ANCILLARY MODE ---
            soc = self.battery.soc 
            # Derate Ancillary during training to encourage exploration of Arb peaks
            reward = (1.0 * inp[c_fcr] * dt_hours) * 0.5

        self.current_step += 1
        terminated = self.current_step >= len(self.input_data)
        truncated = False

        # --- TERMINAL SOC PENALTY ---
        # If the battery is empty at end-of-day, apply a massive penalty
        if terminated:
            target_soc = self.config["end_of_day"]["min_soc_mwh"]
            soc_error = abs(self.battery.soc - target_soc)
            reward -= (soc_error * 100.0) 

        if not terminated:
            next_obs = np.append(self.input_data.iloc[self.current_step].values, soc)
        else:
            next_obs = np.zeros(self.observation_space.shape)

        return next_obs.astype(np.float32), float(reward), terminated, truncated, {}