"""
Run CSPO vs PPO-Lag on SafetyPointGoal1-v0 for comparison.
Reduced scale: 200K steps (paper uses 10M).
"""
import warnings
warnings.filterwarnings('ignore')
import os
os.environ['WANDB_MODE'] = 'disabled'

import omnisafe
import numpy as np

def run_experiment(algo, env_id, seed, total_steps, log_dir):
    agent = omnisafe.Agent(
        algo,
        env_id,
        custom_cfgs={
            'seed': seed,
            'train_cfgs': {
                'total_steps': total_steps,
                'vector_env_nums': 1,
                'torch_threads': 4,
                'device': 'cpu',
            },
            'algo_cfgs': {
                'steps_per_epoch': 20000,
                'update_iters': 10,
                'batch_size': 512,
            },
            'logger_cfgs': {
                'use_wandb': False,
                'use_tensorboard': False,
                'log_dir': log_dir,
                'window_lens': 10,
            },
        },
    )
    agent.learn()
    # Read final metrics from the logger
    final_ret = agent._logger.get_stats('Metrics/EpRet')
    final_cost = agent._logger.get_stats('Metrics/EpCost')
    return final_ret, final_cost

# Run CSPO on PointGoal
print("Running CSPO on SafetyPointGoal1-v0...")
cspo_ret, cspo_cost = run_experiment(
    'CSPO', 'SafetyPointGoal1-v0', seed=0,
    total_steps=200000, log_dir='/tmp/cspo_bench/cspo'
)
print(f"CSPO - Return: {cspo_ret:.2f}, Cost: {cspo_cost:.2f}")

# Run PPO-Lag on PointGoal
print("\nRunning PPO-Lag on SafetyPointGoal1-v0...")
ppolag_ret, ppolag_cost = run_experiment(
    'PPOLag', 'SafetyPointGoal1-v0', seed=0,
    total_steps=200000, log_dir='/tmp/cspo_bench/ppolag'
)
print(f"PPO-Lag - Return: {ppolag_ret:.2f}, Cost: {ppolag_cost:.2f}")

print("\nComparison:")
print(f"CSPO:    Return={cspo_ret:.2f}, Cost={cspo_cost:.2f}")
print(f"PPO-Lag: Return={ppolag_ret:.2f}, Cost={ppolag_cost:.2f}")