import gymnasium as gym
import torch.optim as optim
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import json
import random
from torch.distributions import Normal
from collections import deque
import matplotlib.pyplot as plt

"""
Global constants
"""
SEED = 31
MAX_STEPS = 500

STATE_DIM = 4
ACTION_DIM = 2

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

"""
Import Self-Defined Module
"""

from Algorithms import ZPG
from Evaluation import evaluate_actor
from Panels import Panel_Env_Reward

#%%
if __name__ == "__main__":
    max_iterations = 50
    max_episodes_per_iteration = 100

    env = gym.make('CartPole-v1', max_episode_steps=MAX_STEPS)
    env.reset(seed=SEED)

    agent = ZPG(expertise=0.001)
    agent.load_model("./Models/actor_pretrain.pth")

    panel = Panel_Env_Reward(link_function='L', expertise=0.001)

    returns_curve = []
    stds_curve = []
    returns_queue = deque(maxlen=10)
    for policy_iter in range(max_iterations):
        returns_0 = evaluate_actor(agent.actor, env, num_of_episodes=max_episodes_per_iteration, deterministic=1)
        returns_queue.extend(returns_0)

        perturb_vec = agent.perturb_actor()
        returns_1 = evaluate_actor(agent.actor_perturb, env, num_of_episodes=max_episodes_per_iteration, deterministic=1)

        results, prob = panel.batch_preference_from_reward(returns_0, returns_1)
        agent.train(np.mean(results))

        print('Policy Iteration:', policy_iter,
              ',Returns:', round(np.mean(returns_0), 4), '+-', round(np.std(returns_0) / np.sqrt(max_episodes_per_iteration), 4),
              ',Prob:', round(np.mean(prob),4),
              ',Cumulative Returns:', round(np.mean(returns_queue), 4),
              ",Reward diff (mean1 - mean0):", round(np.mean(returns_1) - np.mean(returns_0),4))
        torch.save(agent.actor.state_dict(), './ZPG/actor_' + str(policy_iter) + '.pth')

        returns_curve.append(float(np.mean(returns_0)))
        stds_curve.append(float(np.std(returns_0)))

    data = {
        'returns': returns_curve,
        'stds': stds_curve
    }
    with open('./ZPG/data.json', 'w') as f:
        json.dump(data, f, indent=4)

    plt.figure(figsize=(10, 5))
    plt.plot(returns_curve)
    plt.xlabel('Episode')
    plt.ylabel('Return')
    plt.title('Return Over Training Episodes')
    plt.grid()
    plt.tight_layout()
    plt.savefig("./ZPG/return_plot.png")  # Optional: save to file
    plt.show()

    env.close()






