server/battery_environment.py

import json
import os
import pandas as pd
import numpy as np
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

try:
    from ..models import BatteryAction, BatteryObservation
    from ..rl_env import CooptimEnv
except (ModuleNotFoundError, ImportError):
    import sys
    import os
    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from models import BatteryAction, BatteryObservation
    from rl_env import CooptimEnv
    from models import BatteryAction, BatteryObservation
    from rl_env import CooptimEnv

class BatteryEnvironment(Environment):
    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self, mode="medium"):
        self._state = State(episode_id=str(uuid4()), step_count=0)
        self.mode = mode
        
        # Load config
        config_path = os.path.join(os.path.dirname(__file__), "..", "config.json")
        with open(config_path, "r") as f:
            self.config = json.load(f)
            
        # Modify config based on mode (difficulty)
        if mode == "hard":
            self.config["battery"]["eta_ch"] = 0.8  # Degraded efficiency
            self.config["battery"]["eta_dis"] = 0.8
            self.config["end_of_day"]["min_soc_mwh"] = 15.0 # Strict constraint
        elif mode == "easy":
            self.config["throughput_penalty"]["c_eur_per_mwh"] = 0.0 # No penalty
            
        # Example dummy data for openenv. In real life we'd load correct CSV
        date_rng = pd.date_range(start='2025-01-01', end='2025-01-02', freq='15min')
        df = pd.DataFrame(index=date_rng)
        
        # Adjust market conditions based on mode
        if mode == "easy":
            df[self.config["columns"]["energy"]] = np.random.uniform(50, 60, size=len(date_rng))
            df[self.config["columns"]["fcr"]] = np.random.uniform(200, 300, size=len(date_rng)) # High ancillary
        elif mode == "medium":
            df[self.config["columns"]["energy"]] = np.random.uniform(-100, 500, size=len(date_rng)) # Volatile
            df[self.config["columns"]["fcr"]] = np.random.uniform(10, 50, size=len(date_rng))
        else: # hard
            df[self.config["columns"]["energy"]] = np.random.uniform(-50, 200, size=len(date_rng))
            df[self.config["columns"]["fcr"]] = np.random.uniform(10, 50, size=len(date_rng))
            
        self.gym_env = CooptimEnv(df, self.config)

    def reset(self) -> BatteryObservation:
        self._state = State(episode_id=str(uuid4()), step_count=0)
        obs_array, _ = self.gym_env.reset()
        return self._get_obs(obs_array, 0.0, False)

    def step(self, action: BatteryAction) -> BatteryObservation:
        self._state.step_count += 1
        
        act = np.array([action.market_choice, action.p_fraction])
        obs_array, reward, terminated, truncated, _ = self.gym_env.step(act)
        
        done = terminated or truncated
        return self._get_obs(obs_array, reward, done)
        
    def _get_obs(self, obs_array, reward, done) -> BatteryObservation:
        idx = min(self.gym_env.current_step, len(self.gym_env.input_data) - 1)
        prices = self.gym_env.input_data.iloc[idx]
        e_col = self.gym_env.config["columns"]["energy"]
        fcr_col = self.gym_env.config["columns"]["fcr"]
        
        return BatteryObservation(
            energy_price=float(prices.get(e_col, 0.0)),
            fcr_price=float(prices.get(fcr_col, 0.0)),
            soc=float(self.gym_env.battery.soc),
            reward=reward,
            done=done,
            metadata={"step": self._state.step_count}
        )

    @property
    def state(self) -> State:
        return self._state