"""Heuristic (clinical protocol) baseline agent for NephroSim environments.

Implements standard clinical protocols as baseline agents:
- UreaClearing: Fixed Qb=350, Qd=500 (standard prescription)
- UltrafiltrationControl: Linear UF rate with BP-triggered reduction
- PhosphateMgmt: Fixed binder dose based on pre-dialysis phosphate
- FullDialysisSession: Combined standard protocol
"""

import gymnasium as gym
import numpy as np

import nephrosim  # noqa: F401


def _heuristic_action_urea(obs: np.ndarray) -> np.ndarray:
    """Standard prescription: maintain default flow rates."""
    # Action [0,0] means no adjustment from current rates
    return np.array([0.0, 0.0], dtype=np.float32)


def _heuristic_action_uf(obs: np.ndarray) -> np.ndarray:
    """Clinical protocol: reduce UF rate if BP drops."""
    sbp = obs[0]  # systolic BP (already in mmHg range after observation)
    total_removed = obs[4]
    target = obs[5]
    elapsed = obs[6]

    # If BP is dropping toward hypotension threshold, reduce UF
    if sbp < 100:
        return np.array([-0.5], dtype=np.float32)  # reduce UF rate
    elif sbp < 110:
        return np.array([-0.2], dtype=np.float32)  # slight reduction
    elif total_removed < target * elapsed * 1.1:
        return np.array([0.1], dtype=np.float32)  # increase slightly if behind
    return np.array([0.0], dtype=np.float32)  # maintain


def _heuristic_action_phosphate(obs: np.ndarray) -> np.ndarray:
    """Fixed binder dose: 1000 mg if phosphate > 4.5, else 500 mg."""
    serum_p = obs[0]
    if serum_p > 5.5:
        return np.array([0.75], dtype=np.float32)  # 1500 mg
    elif serum_p > 4.5:
        return np.array([0.5], dtype=np.float32)  # 1000 mg
    elif serum_p > 3.5:
        return np.array([0.25], dtype=np.float32)  # 500 mg
    return np.array([0.1], dtype=np.float32)  # 200 mg (minimal)


def _heuristic_action_full(obs: np.ndarray) -> np.ndarray:
    """Combined clinical protocol for full session."""
    sbp = obs[3]  # systolic BP

    qb_adj = 0.0
    qd_adj = 0.0
    uf_adj = 0.0

    if sbp < 100:
        uf_adj = -0.5
    elif sbp < 110:
        uf_adj = -0.2

    return np.array([qb_adj, qd_adj, uf_adj], dtype=np.float32)


HEURISTIC_MAP = {
    "nephrosim/UreaClearing-v0": _heuristic_action_urea,
    "nephrosim/UltrafiltrationControl-v0": _heuristic_action_uf,
    "nephrosim/PhosphateMgmt-v0": _heuristic_action_phosphate,
    "nephrosim/FullDialysisSession-v0": _heuristic_action_full,
}


def evaluate_heuristic(
    env_id: str,
    n_episodes: int = 50,
    seed: int = 42,
) -> dict:
    """Evaluate the heuristic (clinical protocol) agent.

    Args:
        env_id: Gymnasium environment ID.
        n_episodes: Number of evaluation episodes.
        seed: Random seed.

    Returns:
        Dict with mean_reward, std_reward, episode_rewards.
    """
    heuristic_fn = HEURISTIC_MAP.get(env_id)
    if heuristic_fn is None:
        raise ValueError(f"No heuristic defined for {env_id}")

    rng = np.random.default_rng(seed)
    episode_rewards = []

    for ep in range(n_episodes):
        env = gym.make(env_id)
        obs, _ = env.reset(seed=int(rng.integers(0, 2**31)))
        total_reward = 0.0
        done = False

        while not done:
            action = heuristic_fn(obs)
            obs, reward, terminated, truncated, info = env.step(action)
            total_reward += float(reward)
            done = terminated or truncated

        episode_rewards.append(total_reward)
        env.close()

    return {
        "mean_reward": float(np.mean(episode_rewards)),
        "std_reward": float(np.std(episode_rewards)),
        "episode_rewards": episode_rewards,
    }
