"""Observation and action spaces for the SHIFT hospital microgrid: the sketch.

Week-1 work for the RL Environment seat. This fixes the *shape* of what the agent sees and
what it controls, built from ``configs/env.yaml``, so that the placeholder environment (next
week) and later the real twin from Simulations both expose exactly this interface.

Conventions
-----------
* One observation is one flat float32 vector. Every slot has a name, a unit and a range.
* One action is four floats: electrolyser, fuel cell, battery, shedding level.
* Names follow the Simulator I/O sheet where it has one (``P_PV``, ``SoC``, ``H2_level``,
  ``p_tank``, ``price``, ``CO2_int``, ``grid_on``); the Project Manual's interface list
  decides what is in and what is out.
* Ranges are bounds for scaling and sanity checks, not physical guarantees.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import yaml
from gymnasium import spaces

DEFAULT_CONFIG = Path(__file__).resolve().parents[2] / "configs" / "env.yaml"

# Shedding levels, as the manual lists them.
SHEDDING_NONE = 0        # nothing shed
SHEDDING_TIER3 = 1       # Tier 3 shed
SHEDDING_TIERS_2_3 = 2   # Tiers 2 and 3 shed; Tier 1 is never shed

ACTION_NAMES = ("u_ele", "u_fc", "u_batt", "shed")


@dataclass(frozen=True)
class Slot:
    """One named block of the observation vector."""

    name: str
    unit: str
    low: float
    high: float
    size: int = 1
    source: str = ""  # where the value comes from: twin state, data, calendar, forecaster


def load_config(path: str | Path | None = None) -> dict[str, Any]:
    """Read the environment settings (defaults to ``AICONTROL/configs/env.yaml``)."""
    with open(path or DEFAULT_CONFIG, encoding="utf-8") as f:
        return yaml.safe_load(f)


def observation_slots(cfg: dict[str, Any]) -> list[Slot]:
    """The observation layout, in order. Change the config, not this list, to resize things."""
    plant = cfg["plant"]
    obs = cfg["observation"]
    fc = cfg["forecast"]
    price_low, price_high = obs["price_range"]
    co2_low, co2_high = obs["co2_range"]

    slots = [
        Slot("time_of_day", "sin, cos", -1.0, 1.0, 2, "calendar"),
        Slot("time_of_year", "sin, cos", -1.0, 1.0, 2, "calendar"),
        Slot("P_PV", "kW", 0.0, plant["pv_kw"], 1, "twin: solar output now"),
        Slot("L_tier1", "kW", 0.0, plant["load_max_kw"], 1, "data: critical load now"),
        Slot("L_tier2", "kW", 0.0, plant["load_max_kw"], 1, "data: essential load now"),
        Slot("L_tier3", "kW", 0.0, plant["load_max_kw"], 1, "data: non-critical load now"),
        Slot("SoC", "fraction 0-1", 0.0, 1.0, 1, "twin: battery state of charge"),
        Slot("H2_level", "kg", 0.0, plant["h2_capacity_kg"], 1, "twin: hydrogen in the tank"),
        Slot("p_tank", "bar", 0.0, plant["tank_p_max_bar"], 1, "twin: tank pressure"),
        Slot("price", "currency/kWh", price_low, price_high, 1, "data: electricity price now"),
        Slot("CO2_int", "kg CO2/kWh", co2_low, co2_high, 1, "data: grid carbon intensity now"),
        Slot("grid_on", "0/1", 0.0, 1.0, 1, "data: grid available"),
        Slot("ele_on", "0/1", 0.0, 1.0, 1, "twin: electrolyser running"),
        Slot("fc_on", "0/1", 0.0, 1.0, 1, "twin: fuel cell running"),
    ]
    quantiles = ", ".join(fc["quantiles"])
    for quantity in fc["quantities"]:
        high = plant["pv_kw"] if quantity == "solar" else plant["load_max_kw"]
        for horizon in fc["horizons_min"]:
            slots.append(
                Slot(
                    f"fcst_{quantity}_{horizon}min",
                    "kW",
                    0.0,
                    high,
                    len(fc["quantiles"]),
                    f"forecaster: {quantiles} at +{horizon} min",
                )
            )
    return slots


def observation_size(cfg: dict[str, Any]) -> int:
    return sum(slot.size for slot in observation_slots(cfg))


def build_observation_space(cfg: dict[str, Any]) -> spaces.Box:
    slots = observation_slots(cfg)
    low = np.concatenate([np.full(s.size, s.low, dtype=np.float32) for s in slots])
    high = np.concatenate([np.full(s.size, s.high, dtype=np.float32) for s in slots])
    return spaces.Box(low=low, high=high, dtype=np.float32)


def build_action_space(cfg: dict[str, Any]) -> spaces.Box:
    """Four floats.

    ``u_ele``  in [0, 1]:   fraction of electrolyser rated power (0 = off)
    ``u_fc``   in [0, 1]:   fraction of fuel cell rated power (0 = off)
    ``u_batt`` in [-1, 1]:  fraction of battery max power; positive discharges, negative charges
    ``shed``   in [0, L-1]: rounded to a shedding level (0 nothing, 1 Tier 3, 2 Tiers 2 and 3)

    Stable-Baselines3 algorithms take either all-continuous or all-discrete actions, so the
    three-way shedding choice rides along as a continuous number and is rounded inside the
    environment. To confirm at the joint session with Simulations.
    """
    levels = cfg["action"]["shedding_levels"]
    low = np.array([0.0, 0.0, -1.0, 0.0], dtype=np.float32)
    high = np.array([1.0, 1.0, 1.0, float(levels - 1)], dtype=np.float32)
    return spaces.Box(low=low, high=high, dtype=np.float32)


def shedding_level(action: np.ndarray, cfg: dict[str, Any]) -> int:
    """Turn the fourth action number into a shedding level, clipped to the allowed ones."""
    levels = cfg["action"]["shedding_levels"]
    return int(np.clip(np.rint(action[3]), 0, levels - 1))


def describe(cfg: dict[str, Any]) -> str:
    """Markdown table of the observation layout, for the interface document."""
    lines = ["| # | name | size | unit | range | source |", "|---|---|---|---|---|---|"]
    index = 0
    for s in observation_slots(cfg):
        lines.append(
            f"| {index} | `{s.name}` | {s.size} | {s.unit} | {s.low:g} … {s.high:g} | {s.source} |"
        )
        index += s.size
    return "\n".join(lines)


if __name__ == "__main__":
    config = load_config()
    print(describe(config))
    print(f"\nobservation size: {observation_size(config)}")
    print(f"action space: {build_action_space(config)}")