import json
import matplotlib.pyplot as plt
from pathlib import Path

def plot_training_results(model_dir="./rulebreaker_model"):
    """
    Reads the trainer_state.json saved by Hugging Face TRL
    and plots the reward curve.
    """
    state_path = Path(model_dir) / "trainer_state.json"
    
    if not state_path.exists():
        print(f"Could not find {state_path}.")
        print("Make sure the training has finished and saved the model!")
        return

    # Load the training logs
    with open(state_path, "r") as f:
        state = json.load(f)

    log_history = state.get("log_history", [])
    
    steps = []
    rewards = []

    # Extract step and reward from the logs
    for log in log_history:
        if "reward" in log and "step" in log:
            steps.append(log["step"])
            rewards.append(log["reward"])

    if not steps:
        print("No reward data found in the logs.")
        return

    # Try to load baselines from training_logs/baselines.json
    random_baseline = None
    greedy_baseline = None
    baselines_path = Path("./training_logs/baselines.json")
    if baselines_path.exists():
        try:
            with open(baselines_path, "r") as f:
                baselines = json.load(f)
            random_baseline = baselines.get("random", {}).get("avg_lawyer_reward")
            greedy_baseline = baselines.get("greedy", {}).get("avg_lawyer_reward")
            print(f"Loaded baselines from {baselines_path}")
        except (json.JSONDecodeError, KeyError) as e:
            print(f"Warning: Could not load baselines: {e}")

    # Create a beautiful plot
    plt.figure(figsize=(10, 6))
    plt.plot(steps, rewards, marker='o', linestyle='-', color='#1f77b4', linewidth=2, markersize=6)
    
    # Add titles and labels
    plt.title('RL Training Progress: Lawyer Agent Reward', fontsize=16, fontweight='bold', pad=20)
    plt.xlabel('Training Steps', fontsize=12)
    plt.ylabel('Average Reward', fontsize=12)
    
    # Add baseline references (from file or hardcoded fallback)
    rand_val = random_baseline if random_baseline is not None else 0.037
    greed_val = greedy_baseline if greedy_baseline is not None else 0.996
    plt.axhline(y=rand_val, color='red', linestyle='--', alpha=0.5, label=f'Random Baseline (~{rand_val:.2f})')
    plt.axhline(y=greed_val, color='green', linestyle='--', alpha=0.5, label=f'Greedy Baseline (~{greed_val:.2f})')
    
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.legend(loc='lower right', fontsize=10)
    
    # Save the plot
    output_file = "reward_curve.png"
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"✅ Successfully created graph: {output_file}")
    
    try:
        plt.show() # Will display inline if run in Colab
    except Exception:
        pass

if __name__ == "__main__":
    plot_training_results()