import json
import numpy as np
import gymnasium as gym
from stable_baselines3 import PPO
from rl_env import KidneyExchangeEnv

def evaluate_policy(env, model, num_episodes=10, seed=0):
    returns = []
    for ep in range(num_episodes):
        obs, info = env.reset(seed=seed + ep * 100)
        done = False
        ep_return = 0.0
        while not done:
            if hasattr(env, 'action_masks'):
                masks = env.action_masks()
                action, _ = model.predict(obs, deterministic=True, action_masks=masks)
            else:
                action, _ = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            ep_return += reward
            done = terminated or truncated
        returns.append(ep_return)
    return np.mean(returns)

def evaluate_random(env, num_episodes=10, seed=0):
    returns = []
    for ep in range(num_episodes):
        obs, info = env.reset(seed=seed + ep * 100)
        done = False
        ep_return = 0.0
        while not done:
            if hasattr(env, 'action_masks'):
                masks = env.action_masks()
                valid_actions = [a for a, m in enumerate(masks) if m]
                if not valid_actions:
                    action = env.action_space.n - 1
                else:
                    action = np.random.choice(valid_actions)
            else:
                action = env.action_space.sample()
            obs, reward, terminated, truncated, info = env.step(action)
            ep_return += reward
            done = terminated or truncated
        returns.append(ep_return)
    return np.mean(returns)

def evaluate_greedy(env, num_episodes=10, seed=0):
    returns = []
    for ep in range(num_episodes):
        obs, info = env.reset(seed=seed + ep * 100)
        done = False
        ep_return = 0.0
        while not done:
            if hasattr(env, 'action_masks'):
                masks = env.action_masks()
                valid_actions = [a for a, m in enumerate(masks) if m]
                # Greedy: pick the first valid action that is not STOP, else STOP
                valid_non_stop = [a for a in valid_actions if a < env.action_space.n - 1]
                if valid_non_stop:
                    action = valid_non_stop[0]
                else:
                    action = env.action_space.n - 1
            else:
                action = 0 # Fallback
            obs, reward, terminated, truncated, info = env.step(action)
            ep_return += reward
            done = terminated or truncated
        returns.append(ep_return)
    return np.mean(returns)

def main():
    env = KidneyExchangeEnv()
    try:
        from sb3_contrib import MaskablePPO
        model = MaskablePPO.load("policy.zip")
    except:
        model = PPO.load("policy.zip")

    seeds = [42, 100, 2023, 555, 999]
    episodes_per_seed = 10

    agent_means = []
    random_means = []
    greedy_means = []

    for s in seeds:
        agent_mean = evaluate_policy(env, model, num_episodes=episodes_per_seed, seed=s)
        rnd_mean = evaluate_random(env, num_episodes=episodes_per_seed, seed=s)
        grd_mean = evaluate_greedy(env, num_episodes=episodes_per_seed, seed=s)
        
        agent_means.append(agent_mean)
        random_means.append(rnd_mean)
        greedy_means.append(grd_mean)
        
        print(f"[eval] seed={s} agent={agent_mean:.3f} random={rnd_mean:.3f} greedy={grd_mean:.3f}", flush=True)

    agent_return_mean = float(np.mean(agent_means))
    agent_return_std = float(np.std(agent_means))
    baseline_random = float(np.mean(random_means))
    baseline_greedy = float(np.mean(greedy_means))
    
    beats_baselines = bool(agent_return_mean > baseline_random and agent_return_mean > baseline_greedy)

    results = {
        "agent_return_mean": agent_return_mean,
        "agent_return_std": agent_return_std,
        "baseline_random": baseline_random,
        "baseline_greedy": baseline_greedy,
        "seeds": seeds,
        "episodes": episodes_per_seed,
        "beats_baselines": beats_baselines
    }

    with open("rl_achievability.json", "w") as f:
        json.dump(results, f, indent=2)

    verdict = {
        "algo": "PPO",
        "library": "sb3",
        "agent_return_mean": agent_return_mean,
        "agent_return_std": agent_return_std,
        "baseline_random": baseline_random,
        "baseline_greedy": baseline_greedy,
        "beats_baselines": beats_baselines,
        "seeds": seeds
    }
    print("VERDICT", json.dumps(verdict))

if __name__ == "__main__":
    main()