# unihvac/tables.py
from __future__ import annotations
import pandas as pd
import numpy as np

def print_monthly_tables_split(df: pd.DataFrame, location: str, time_step_hours: float):
    """
    Table 1: Monthly HVAC electricity + temps
    Table 2: Monthly occupancy
    """
    drop_cols = [c for c in df.columns if c.startswith("has_") or c == "feature_registry"]
    df_clean = df.drop(columns=drop_cols, errors="ignore").copy()

    if "month" not in df_clean.columns:
        return

    df_clean["month"] = df_clean["month"].round().astype(int)

    # 1. Electricity / Energy Calculation
    energy_cols = []
    peak_cols = []
    if "elec_power" in df_clean.columns:
        if "elec_power_kw" not in df_clean.columns:
            df_clean["elec_power_kw"] = df_clean["elec_power"] / 1000.0
        if "elec_energy_kwh" not in df_clean.columns:
            df_clean["elec_energy_kwh"] = df_clean["elec_power_kw"] * time_step_hours
        energy_cols.append("elec_energy_kwh")
        peak_cols.append("elec_power_kw")

    # 2. Temperature Aggregation
    temp_cols = [c for c in ["outdoor_temp", "core_temp", "perim1_temp", "perim2_temp", "perim3_temp", "perim4_temp"] 
                 if c in df_clean.columns]
    
    agg1 = {c: "sum" for c in energy_cols}
    agg1.update({c: "max" for c in peak_cols})
    agg1.update({c: "mean" for c in temp_cols})
    tbl1 = df_clean.groupby("month").agg(agg1).sort_index()

    # 3. Occupancy Aggregation
    occ_cols = [c for c in df_clean.columns if c.endswith("_occ_count")]
    tbl2 = df_clean.groupby("month")[occ_cols].mean().sort_index() if occ_cols else pd.DataFrame()
    if not tbl2.empty:
        tbl2["occ_mean_total"] = tbl2.sum(axis=1)

    print("\n" + "=" * 110)
    print(f"MONTHLY ELECTRICITY + TEMPERATURE — {location}")
    print("=" * 110)
    print(tbl1.round(2).to_string())
    print("\n" + "=" * 110)
    print(f"MONTHLY OCCUPANCY — {location}")
    print("=" * 110)
    print(tbl2.round(3).to_string())
    print("=" * 110 + "\n")


def print_monthly_tables_extra(df: pd.DataFrame, location: str) -> None:
 
    d = df.copy()
    if "month" not in d.columns:
        return
    d["month"] = d["month"].round().astype(int)

    violation_cols = [c for c in ["comfort_violation_degCh", "comfort_violation_fixed_degCh"] if c in d.columns]
    tbl_sums = d.groupby("month")[violation_cols].sum()


    occ_cols = [c for c in d.columns if c.endswith("_occ_count")]
    total_occ = d[occ_cols].sum(axis=1)
    is_occupied = total_occ > 1e-6
    d_occ = d[is_occupied].copy()


    def person_weighted_ppd(group):
        occ = group[occ_cols].sum(axis=1)
        raw_ppd = group["ppd_weighted"]
        return (raw_ppd * occ).sum() / occ.sum() if occ.sum() > 0 else np.nan

    if not d_occ.empty and "ppd_weighted" in d_occ.columns:
        ppd_monthly = d_occ.groupby("month", group_keys=False).apply(person_weighted_ppd)
        ppd_monthly = ppd_monthly.clip(lower=5.0)
        
        pmv_monthly = d_occ.groupby("month")["pmv_weighted"].mean()
        rh_monthly = d_occ.groupby("month")["rh_weighted"].mean()

        tbl_means = pd.DataFrame({
            "ppd_weighted": ppd_monthly,
            "pmv_weighted": pmv_monthly,
            "rh_weighted_%": rh_monthly
        })
        tbl3a = pd.concat([tbl_sums, tbl_means], axis=1).sort_index()
    else:
        tbl3a = tbl_sums


    outdoor_vars = [c for c in ["outdoor_temp", "outdoor_dewpoint", "outdoor_wetbulb"] if c in d.columns]
    tbl3b = d.groupby("month")[outdoor_vars].mean().sort_index() if outdoor_vars else None


    print("\n" + "=" * 110)
    print(f"MONTHLY COMFORT OUTCOMES (Occupancy Weighted) — {location}")
    print("=" * 110)
    print(tbl3a.round(3).to_string())
    print("=" * 110)


    if tbl3b is not None:
        print("\n" + "=" * 110)
        print(f"MONTHLY OUTDOOR CONDITIONS — {location}")
        print("=" * 110)
        print(tbl3b.round(3).to_string())
        print("=" * 110)