import numpy as np
import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from twin.market import MarketTwin
from twin.load import LoadTwin
from twin.solar import SolarTwin
from twin.grid import GridTwin
from twin.battery import BatteryTwin



class EnergyEnv:
    """
    Reinforcement Learning Environment for:
    Market Arbitrage + Grid Stability + Battery Optimization

    Core Design:
    - RL handles arbitrage decisions
    - Grid intelligence handles ancillary + safety
    - Battery executes constrained actions
    """

    def __init__(
        self,
        price_data_path,
        episode_length=96  # 96 = 24 hours at 15-min resolution
    ):

        # Digital twins
        self.market = MarketTwin(price_data_path)
        self.load = LoadTwin()
        self.solar = SolarTwin()
        self.grid = GridTwin()
        self.battery = BatteryTwin()

        self.episode_length = episode_length
        self.timestep = 0

        self.total_profit = 0.0

    # -------------------------------------------------------

    def reset(self):
        """Reset environment for new RL episode"""

        self.market.reset()
        self.load.reset()
        self.solar.reset()
        self.grid.reset()
        self.battery.reset()

        self.timestep = 0
        self.total_profit = 0.0

        return self._get_state()

    # -------------------------------------------------------

    def step(self, action):
        """
        Execute one environment step.

        Parameters
        ----------
        action : float
            RL agent output ∈ [-1,1]
            -1 → discharge
             0 → idle
            +1 → charge

        Returns
        -------
        tuple
            next_state, reward, done, info
        """

        hour = (self.timestep // 4) % 24

        # --- Market ---
        price = self.market.step()
        market_features = self.market.get_price_state()

        # --- Load & Solar ---
        load = self.load.step(hour)
        solar = self.solar.step(hour)

        # --- Grid Intelligence ---
        grid_info = self.grid.grid_state(
            price=price,
            volatility=market_features["volatility"],
            load=load,
            solar=solar
        )

        reserve_ratio = grid_info["reserve_ratio"]

        # --- Battery Execution ---
        soc, power, deg_cost = self.battery.step(
            action=action,
            reserve_ratio=reserve_ratio
        )

        # --- Revenue Computation ---
        # prev_soc = self.battery.prev_soc
        # curr_soc = self.battery.soc

        # Net asset change = energy gained * future selling value
        # price_ma = self.market.price_ma
        # arbitrage_profit = -power * (price - price_ma) * 0.25

        arbitrage_profit = -power * price * 0.25

        ancillary_reward = reserve_ratio * self.battery.max_power * 200 * 0.25

        self.total_profit += arbitrage_profit + ancillary_reward

        # --- Risk Penalty ---
        grid_risk_penalty = 50 * grid_info["failure_probability"]

        # --- Final Reward ---
        reward = (
            arbitrage_profit +
            ancillary_reward -
            grid_risk_penalty -
            deg_cost
        )

        # --- Step Increment ---
        self.timestep += 1
        done = self.timestep >= self.episode_length

        next_state = self._get_state()

        info = {
            "price": price,
            "load": load,
            "solar": solar,
            "soc": soc,
            "grid_stress": grid_info["grid_stress"],
            "predicted_stress": grid_info["predicted_stress"],
            "reserve_ratio": reserve_ratio,
            "failure_probability": grid_info["failure_probability"],
            "arbitrage_profit": arbitrage_profit,
            "ancillary_reward": ancillary_reward,
            "total_profit": self.total_profit,
            "power": power
        }

        return next_state, reward, done, info

    # -------------------------------------------------------

    def _get_state(self):
        """
        Construct RL state vector.

        Returns
        -------
        numpy.ndarray
            State vector for RL agent
        """

        market_state = self.market.get_price_state()
        battery_state = self.battery.state()

        state = np.array([
            market_state["price"] / 6000,      # normalized price
            market_state["volatility"],        # volatility
            market_state["trend"],             # price trend
            battery_state["soc_fraction"],     # battery SoC
            self.timestep / self.episode_length  # time of day
        ], dtype=np.float32)

        return state