import gymnasium as gym
import numpy as np
import torch
import json
import random
from collections import deque
import matplotlib.pyplot as plt

"""
Global constants
"""
SEED = 42
MAX_STEPS = 1000

STATE_DIM = 17
ACTION_DIM = 6
ACTION_HIGH = torch.FloatTensor(np.ones(ACTION_DIM))
ACTION_LOW = - torch.FloatTensor(np.ones(ACTION_DIM))

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

"""
Import Self-Defined Module
"""
from Networks import Reward
from Evaluation import evaluate_actor
from Algorithms import PPO
from Panels import Panel_Env_Reward

#%%
if __name__ == "__main__":
    max_iterations = 200
    max_episodes_per_iteration = 5

    env = gym.make('HalfCheetah-v5', max_episode_steps=MAX_STEPS)
    env.reset(seed=SEED)

    agent = PPO()
    agent.load_model("./Models/actor_initial.pth")

    panel = Panel_Env_Reward(link_function='L', expertise = 0.01)
    reward_net = Reward(state_dim=STATE_DIM, action_dim=ACTION_DIM)
    reward_net.load_state_dict(torch.load('./Models/reward_net.pth', weights_only=True))

    returns_curve = []
    stds_curve = []
    returns_queue = deque(maxlen=max_episodes_per_iteration)
    policy_iter = 0
    memory = {'states': [], 'actions': [], 'rewards': [], 'log_probs': [], 'dones': [], 'values': []}
    for episode in range(max_iterations * max_episodes_per_iteration):
        state, _ = env.reset()
        total_reward = 0

        while 1:
            action, log_prob = agent.select_action(state)
            value = agent.critic(torch.tensor(state, dtype=torch.float32).unsqueeze(0)).item()

            next_state, reward, terminated, truncated, _ = env.step(action)
            reward_prox = reward_net(torch.FloatTensor(state).unsqueeze(0), torch.FloatTensor(action).unsqueeze(0)).item()
            done = terminated or truncated

            memory['states'].append(state)
            memory['actions'].append(action)
            memory['rewards'].append(reward_prox)
            memory['log_probs'].append(log_prob)
            memory['dones'].append(done)
            memory['values'].append(value)

            total_reward += reward
            state = next_state

            if done:
                state, _ = env.reset()
                break

        next_value = agent.critic(torch.tensor(state, dtype=torch.float32).unsqueeze(0)).item()
        advantages, returns = agent.compute_gae(memory['rewards'], memory['dones'], memory['values'], next_value)
        memory['advantages'] = advantages
        memory['returns'] = returns

        returns_queue.append(total_reward)

        if (episode + 1) % max_episodes_per_iteration == 0:
            returns = evaluate_actor(agent.actor, env, num_of_episodes=max_episodes_per_iteration, deterministic=1)
            returns_curve.append(float(np.mean(returns)))
            stds_curve.append(float(np.std(returns)))
            agent.train(memory)
            memory = {'states': [], 'actions': [], 'rewards': [], 'log_probs': [], 'dones': [], 'values': []}
            print('Policy Iteration:', policy_iter,
                  ',Returns:', round(np.mean(returns), 4), '+-', round(np.std(returns), 4)
                  )
            torch.save(agent.actor.state_dict(), './RMPPO/actor_' + str(policy_iter) + '.pth')
            policy_iter += 1

    data = {
        'returns': returns_curve,
        'stds': stds_curve
    }
    with open('./RMPPO/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("./RMPPO/return_plot.png")  # Optional: save to file
    plt.show()

    env.close()
