# unihvac/rewards.py
from __future__ import annotations

from dataclasses import dataclass, asdict
from typing import Dict, Any, Tuple, Optional
import numpy as np
import pandas as pd


@dataclass(frozen=True)
class RewardConfig:
    version: str = "v_ashrae"
    prefer_step_kwh_cols: Tuple[str, ...] = (
        "HVAC_elec_kWh_step",
        "hvac_kWh_step",
        "elec_kWh_step",
    )

    elec_power_col: str = "elec_power"
    
    comfort_col: str = "ppd_weighted"
    w_energy: float = 1.0
    w_comfort: float = 0.1


def config_to_meta(cfg: RewardConfig) -> Dict[str, Any]:
    return asdict(cfg)

def compute_reward_components(df: pd.DataFrame, timestep_hours: float, cfg: RewardConfig) -> Tuple[np.ndarray, np.ndarray]:
    if df is None or len(df) == 0:
        return np.zeros((0,), dtype=np.float32), np.zeros((0,), dtype=np.float32)
    energy_kwh = np.zeros(len(df), dtype=np.float32)
    found_energy = False
    
    for col in cfg.prefer_step_kwh_cols:
        if col in df.columns:
            energy_kwh = df[col].fillna(0.0).astype(np.float32).values
            found_energy = True
            break
    
    if not found_energy and cfg.elec_power_col in df.columns:
        power_w = df[cfg.elec_power_col].fillna(0.0).astype(np.float32).values
        energy_kwh = (power_w / 1000.0) * timestep_hours
    comfort_val = np.zeros(len(df), dtype=np.float32)
    if cfg.comfort_col in df.columns:
        comfort_val = df[cfg.comfort_col].fillna(0.0).astype(np.float32).values
    r_energy = -1.0 * energy_kwh
    r_comfort = -1.0 * comfort_val 

    return r_energy.astype(np.float32), r_comfort.astype(np.float32)

def compute_terminals(df: pd.DataFrame) -> np.ndarray:
    T = 0 if df is None else len(df)
    terminals = np.zeros((T,), dtype=np.int8)
    if T > 0:
        terminals[-1] = 1
    return terminals