import json
import matplotlib.pyplot as plt
import os
import pandas as pd

def generate_plots():
    if not os.path.exists("training_metrics.json"):
        print("training_metrics.json not found.")
        return

    with open("training_metrics.json", "r") as f:
        metrics = json.load(f)

    raw_history = metrics.get("log_history", [])
    data = [
        {"step": e["step"], "reward": e["reward"], "loss": e.get("loss")} 
        for e in raw_history 
        if "step" in e and "reward" in e
    ]

    if not data:
        print("No valid history data found.")
        return

    df = pd.DataFrame(data)
    os.makedirs("docs", exist_ok=True)

    # Plot Reward
    plt.figure(figsize=(10, 6))
    plt.plot(df["step"], df["reward"], color="#00f2fe", linewidth=2, label="Cumulative Reward")
    plt.fill_between(df["step"], df["reward"], color="#00f2fe", alpha=0.1)
    plt.title("EcoGrid RL Training: Reward Convergence")
    plt.xlabel("Step")
    plt.ylabel("Reward")
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.savefig("docs/reward_curve.png", dpi=150)
    plt.close()
    print("Saved docs/reward_curve.png")

    # Plot Loss (filtering out None)
    df_loss = df.dropna(subset=["loss"])
    if not df_loss.empty:
        plt.figure(figsize=(10, 6))
        plt.plot(df_loss["step"], df_loss["loss"], color="#ff4b4b", linewidth=2, label="Training Loss")
        plt.title("EcoGrid RL Training: Loss Curve")
        plt.xlabel("Step")
        plt.ylabel("Loss")
        plt.grid(True, alpha=0.3)
        plt.legend()
        plt.savefig("docs/loss_curve.png", dpi=150)
        plt.close()
        print("Saved docs/loss_curve.png")

if __name__ == "__main__":
    generate_plots()