import os
import numpy as np
import pandas as pd

# ==========================================
# 1.(Internal Structure)
# ==========================================
TIME_COLS = ["step", "month", "day", "hour", "minute"]
ENV_COLS = ["out_temp", "out_rh"]
GLOBAL_REWARD_COLS = ["power_kw"] 
GLOBAL_COLS = TIME_COLS + ENV_COLS + GLOBAL_REWARD_COLS

ZONE_LIST = ["core", "p1", "p2", "p3", "p4"]
ZONE_STATE_TEMPLATE = ["temp", "occ", "rh"]
ZONE_ACTION_TEMPLATE = ["htg", "clg"]

def get_full_schema():
    states = []
    actions = []
    for name in ZONE_LIST:
        states += [f"{prefix}_{name}" for prefix in ZONE_STATE_TEMPLATE]
        actions += [f"{prefix}_{name}" for prefix in ZONE_ACTION_TEMPLATE]
    return GLOBAL_COLS + states + actions

# ==========================================
# 2. PHYSICS UTILITIES
# ==========================================
def batch_calculate_rh(temp_array: np.ndarray, dewpoint_array: np.ndarray) -> np.ndarray:
    """August-Roche-Magnus approximation for Relative Humidity."""
    A, B = 17.625, 243.04
    rh = 100 * np.exp((A * dewpoint_array / (B + dewpoint_array)) - 
                      (A * temp_array / (B + temp_array)))
    return np.clip(rh, 0.0, 100.0)

# ==========================================
# 3. MAIN GENERATOR FUNCTION
# ==========================================
def save_dt_training_data(df_raw: pd.DataFrame, out_dir: str, location: str):
    dt_df = pd.DataFrame()

    dt_df['step'] = df_raw.get('step', range(len(df_raw)))
    dt_df['month'] = df_raw.get('month', 1)
    dt_df['day'] = df_raw.get('day_of_month', 1) 
    dt_df['hour'] = df_raw.get('hour', 0)
    dt_df['minute'] = (dt_df['step'] % 4) * 15
    dt_df['out_temp'] = df_raw['outdoor_temp']
    dt_df['out_rh'] = batch_calculate_rh(
        df_raw['outdoor_temp'].values, 
        df_raw['outdoor_dewpoint'].values
    )

    dt_df['power_kw'] = df_raw['elec_power'] / 1000.0


    for zone in ZONE_LIST:
        s_name = "core" if zone == "core" else f"perim{zone[-1]}"
        
        # States
        dt_df[f"temp_{zone}"] = df_raw[f"{s_name}_temp"]
        dt_df[f"occ_{zone}"]  = df_raw[f"{s_name}_occ_count"]
        dt_df[f"rh_{zone}"]   = df_raw[f"{s_name}_rh"]
        dt_df[f"htg_{zone}"] = df_raw.get("setpoint_htg", 21.0)
        dt_df[f"clg_{zone}"] = df_raw.get("setpoint_clg", 24.0)


    ALL_COLUMNS = get_full_schema()
    dt_df = dt_df[ALL_COLUMNS]


    os.makedirs(out_dir, exist_ok=True)
    filename = f"{location}_ComfortDT_Training.csv"
    save_path = os.path.join(out_dir, filename)
    dt_df.to_csv(save_path, index=False)

    print(f" DT Data Saved: {filename} | Shape: {dt_df.shape}")
    return save_path