import argparse
import os
import gymnasium as gym
import numpy as np
import train_hdppo as m

def _warmup_encoder(agent, env_id, n_steps=200):
    env = gym.make(env_id)
    obs, _ = env.reset(seed=0)
    for i in range(n_steps):
        Hr, Hi = agent.encoder.encode(obs)
        action = agent.actor.greedy(Hr, Hi)
        obs, _, term, trunc, _ = env.step(action)
        if term or trunc:
            obs, _ = env.reset(seed=i + 1)
    env.close()

def load_agent_from_checkpoint(path, seed=42, warmup=True):
    data = np.load(path)
    cfg = dict(m.PENDULUM_CONFIG)
    cfg['D'] = int(data['D'])
    cfg['beta'] = float(data['beta_base'])
    cfg['fpe_phi_init'] = data['fpe_phi']
    if 'feat_lo' in data:
        cfg['feat_lo'] = data['feat_lo'].tolist()
        cfg['feat_hi'] = data['feat_hi'].tolist()
    agent = m.HDPPOAgentContinuous(cfg, seed=seed)
    agent.actor.W_re = data['W_actor_re'].copy()
    agent.actor.W_im = data['W_actor_im'].copy()
    if 'log_std' in data:
        agent.actor.log_std = data['log_std'].copy()
    if 'W_critic_re' in data:
        agent.critic.W_re = data['W_critic_re'].copy()
        agent.critic.W_im = data['W_critic_im'].copy()
        agent.critic.bias = float(data['critic_bias'])
    if warmup:
        _warmup_encoder(agent, 'Pendulum-v1')
    return agent

def resolve_weights(weights_arg):
    if os.path.isfile(weights_arg):
        return weights_arg
    try:
        from huggingface_hub import hf_hub_download
    except ImportError as exc:
        raise SystemExit('Install huggingface_hub to load remote checkpoints: pip install huggingface_hub') from exc
    return hf_hub_download(repo_id=weights_arg, filename='hdppo-Pendulum-v1/weights.npz')

def main():
    parser = argparse.ArgumentParser(description='Enjoy HD-PPO on Pendulum-v1')
    parser.add_argument('--weights', default='hdppo-Pendulum-v1/weights.npz')
    parser.add_argument('--episodes', type=int, default=5)
    parser.add_argument('--seed', type=int, default=10000)
    parser.add_argument('--render', action='store_true')
    args = parser.parse_args()
    weights = resolve_weights(args.weights)
    agent = load_agent_from_checkpoint(weights, seed=args.seed)
    if args.render:
        env = gym.make('Pendulum-v1', render_mode='human')
        for ep in range(args.episodes):
            obs, _ = env.reset(seed=args.seed + ep)
            done = False
            ep_r = 0.0
            while not done:
                Hr, Hi = agent.encoder.encode(obs)
                action = agent.actor.greedy(Hr, Hi)
                obs, reward, term, trunc, _ = env.step(action)
                ep_r += reward
                done = term or trunc
            print(f'episode {ep + 1}: return={ep_r:.1f}')
        env.close()
    else:
        stats = m.evaluate_agent(agent, n_episodes=args.episodes, seed_base=args.seed)
        print(f'mean={stats['mean_reward']:.2f}  ci95={stats['ci95_reward']:.2f}')
if __name__ == '__main__':
    main()