import os
import sys
import numpy as np

# Add root directory to sys.path for absolute imports
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from openenv.tasks import Task, random_policy, heuristic_policy
from openenv.models import Action
from agent.baselines import ea_only_policy, fr_only_policy, ps_only_policy
from infer import policy as rl_policy


# =========================
# REAL KPI EVALUATOR
# =========================
def evaluate_real_kpis(policy_func):

    task = Task("medium")
    env = task.env
    obs = env.reset().to_array()

    # Economic assumptions
    peak_tariff = 20.0
    fr_price = 0.1

    demand_history = []
    net_history = []

    real_ea = 0
    real_fr = 0
    total_deg = 0
    voltage_violations = 0
    max_line_loading = 0.0

    done = False

    while not done:

        idx = env.t

        action_val = policy_func(obs)

        # step
        next_obs, _, done, _ = env.step(Action(power=action_val))

        # ✅ actual power used (important for realism)
        power = action_val  # (correct since no ramp/clip beyond bounds)

        # -------------------------
        # GRID CONSTRAINTS TRACKING
        # -------------------------
        if env.v_pu < 0.95 or env.v_pu > 1.05:
            voltage_violations += 1
        if env.line_loading > max_line_loading:
            max_line_loading = env.line_loading

        # -------------------------
        # DEGRADATION
        # -------------------------
        deg_coeff = 0.02
        total_deg += deg_coeff * (abs(power) ** 1.3)

        # -------------------------
        # DEMAND TRACKING
        # -------------------------
        d = env.demand_series[idx]

        demand_history.append(d)
        net_history.append(d + power)

        obs = next_obs.to_array()

    # -------------------------
    # PEAK SHAVING
    # -------------------------
    orig_peak = np.max(demand_history)
    new_peak = np.max(net_history)

    ps_savings = max(0, (orig_peak - new_peak) * peak_tariff)

    peak_reduction = (orig_peak - new_peak) / (orig_peak + 1e-6)

    # -------------------------
    # FINAL ECONOMICS & GRID
    # -------------------------
    real_ea = env.total_EA
    real_fr = env.total_FR
    
    net_profit = (real_ea + real_fr + ps_savings) - total_deg

    return net_profit, peak_reduction, voltage_violations, max_line_loading


# =========================
# RUN
# =========================
if __name__ == "__main__":

    policies = {
        "Random": random_policy,
        "Heuristic": heuristic_policy,
        "EA-Only": ea_only_policy,
        "FR-Only": fr_only_policy,
        "PS-Only": ps_only_policy,
        "RL Agent (Ours)": rl_policy
    }

    print(f"{'Policy':<20} | {'Avg Net Profit ($)':<18} | {'Peak Reduc %':<12}")
    print("-" * 60)

    NUM_EPISODES = 5

    for name, pol in policies.items():

        profits = []
        peaks = []
        v_viols = []
        l_loads = []

        for _ in range(NUM_EPISODES):
            p, pr, v_v, l_l = evaluate_real_kpis(pol)
            profits.append(p)
            peaks.append(pr)
            v_viols.append(v_v)
            l_loads.append(l_l)

        # Add grid constraint information to the print output
        print(f"{name:<20} | {np.mean(profits):>18.2f} | {np.mean(peaks)*100:>11.1f}% | V-Viols: {np.mean(v_viols):.1f} | MaxLoad: {np.mean(l_loads):.1f}%")