import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
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 Evaluation import evaluate_actor
from Panels import Panel_Env_Reward
from Algorithms import ZSPO

#%%
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 = ZSPO(expertise=0.01)
    agent.load_model("./Models/actor_initial.pth")

    panel = Panel_Env_Reward(link_function='BT', expertise = 0.01, size = 100)

    returns_curve = []
    stds_curve = []
    returns_queue = deque(maxlen=10)
    for policy_iter in range(max_iterations):
        perturb_vec = agent.perturb_actor()

        returns_0 = evaluate_actor(agent.actor, env, num_of_episodes=max_episodes_per_iteration, deterministic=1)
        returns_queue.extend(returns_0)
        
        # Evaluate +epsilon perturbation
        returns_plus = evaluate_actor(agent.actor_perturb, env, num_of_episodes=max_episodes_per_iteration, deterministic=1)

        # Now apply -epsilon to get antithetic version
        perturb_vec_for_minus = [v.clone() for v in perturb_vec]  # Safe copy
        with torch.no_grad():
            for module in agent.actor_perturb.modules():
                if isinstance(module, nn.Linear):
                    for param in module.parameters():
                        param.sub_(2 * perturb_vec_for_minus.pop(0) * agent.perturbation_dist)
            for name, param in agent.actor.named_parameters():
                if '.' not in name:
                    if not any(name.startswith(pn + '.') for pn, _ in agent.actor.named_modules()):
                        param.sub_(2 * perturb_vec_for_minus.pop(0) * agent.perturbation_dist)

        returns_minus = evaluate_actor(agent.actor_perturb, env, num_of_episodes=max_episodes_per_iteration, deterministic=1)

        # Restore back to original actor (not strictly needed if you re-copy each time)
        agent.actor_perturb.load_state_dict(agent.actor.state_dict())

        # Compute preference and train
        results, prob_truth = panel.batch_preference_from_reward(returns_minus, returns_plus)
        prob = np.mean(results)
        if prob > 0.5:
            results_plus, prob_plus = panel.batch_preference_from_reward(returns_0, returns_plus)
            if prob_plus > 0.5:
                agent.train(prob)
        elif prob < 0.5:
            results_minus, prob_minus = panel.batch_preference_from_reward(returns_0, returns_minus)
            if prob_minus > 0.5:
                agent.train(prob)

        print('Policy Iteration:', policy_iter,
              ',Returns:', round(np.mean(returns_0), 4), '+-', round(np.std(returns_0), 4),
              ',Prob:', round(np.mean(prob),4),
              ',Cumulative Returns:', round(np.mean(returns_queue), 4),
              ",Reward diff (mean1 - mean0):", round(np.mean(returns_plus) - np.mean(returns_minus),4))
        torch.save(agent.actor.state_dict(), './ZSPO/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('./ZSPO/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("./ZSPO/return_plot.png")  # Optional: save to file
    plt.show()

    env.close()
