import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import ast

rewards = []
step_idx = 0
stage_boundaries = []
current_stage = None

with open("content/train_v2.log", "r") as f:
    for line in f:
        line = line.strip()
        if "Starting Stage 1" in line:
            current_stage = ("Stage 1\n(Easy)", step_idx)
        elif "Starting Stage 2" in line:
            current_stage = ("Stage 2\n(Medium)", step_idx)
            stage_boundaries.append(("Stage 2\n(Medium)", step_idx))
        elif "Starting Stage 3" in line:
            stage_boundaries.append(("Stage 3\n(Hard+Chaos)", step_idx))
        if line.startswith("{'loss':") and "'reward':" in line:
            try:
                data = ast.literal_eval(line)
                if 'reward' in data:
                    rewards.append(float(data['reward']))
                    step_idx += 1
            except:
                pass

fig, ax = plt.subplots(figsize=(12, 6))
ax.plot(range(len(rewards)), rewards, color='#4C72B0', linewidth=1.5, alpha=0.8, label='Reward per step')

# Smoothed trend line
if len(rewards) > 10:
    import numpy as np
    window = max(10, len(rewards) // 15)
    smoothed = np.convolve(rewards, np.ones(window)/window, mode='valid')
    ax.plot(range(window-1, len(rewards)), smoothed, color='#DD8452', linewidth=2.5, label=f'Trend (window={window})')

# Stage markers
colors = ['#2ca02c', '#ff7f0e', '#d62728']
for i, (label, boundary) in enumerate(stage_boundaries):
    ax.axvline(x=boundary, color=colors[i], linestyle='--', linewidth=1.8, alpha=0.8)
    ax.text(boundary + 2, ax.get_ylim()[0] if rewards else -30, label,
            color=colors[i], fontsize=9, verticalalignment='bottom')

ax.axhline(0, color='black', linestyle=':', linewidth=1.2, alpha=0.5, label='Zero reward')
ax.set_title('V2 Curriculum Training — Reward Progression\n(Easy → Medium → Hard+Chaos)', fontsize=14, fontweight='bold')
ax.set_xlabel('Training Steps', fontsize=12)
ax.set_ylabel('Reward Score', fontsize=12)
ax.legend(fontsize=10)
ax.grid(True, linestyle='--', alpha=0.4)
plt.tight_layout()
plt.savefig("content/reward_graph_v2.png", dpi=150)
print(f"Done! {len(rewards)} data points plotted. Saved to content/reward_graph_v2.png")