# -*- coding: utf-8 -*-
"""
Created on Wed Jan 15 15:01:34 2025

@author: pt
"""

import os
import random
import time
from dataclasses import dataclass
import argparse
import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions.normal import Normal
from torch.utils.tensorboard import SummaryWriter
import pandas as pd
import sys

# Add the path to the cleanrl_utils folder to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

@dataclass
class Args:
    exp_name: str = os.path.basename(__file__)[: -len(".py")]
    seed: int = 1
    torch_deterministic: bool = True
    cuda: bool = True
    track: bool = False
    wandb_project_name: str = "Kolmogorov-Arnold Networks for Online Reinforcement Learning"
    wandb_entity: str = None
    capture_video: bool = False
    save_model: bool = True
    upload_model: bool = False
    hf_entity: str = ""
    env_id: str = "HalfCheetah-v4"
    total_timesteps: int = 50000
    learning_rate: float = 3e-4
    num_envs: int = 1
    num_steps: int = 2048
    anneal_lr: bool = True
    gamma: float = 0.99
    gae_lambda: float = 0.95
    num_minibatches: int = 64
    update_epochs: int = 10
    norm_adv: bool = True
    clip_coef: float = 0.2
    clip_vloss: bool = True
    ent_coef: float = 0.0
    vf_coef: float = 0.5
    max_grad_norm: float = 0.5
    target_kl: float = None
    batch_size: int = 0
    minibatch_size: int = 0
    num_iterations: int = 0

def make_env(env_id, idx, capture_video, run_name, gamma):
    def thunk():
        if capture_video and idx == 0:
            env = gym.make(env_id, render_mode="rgb_array")
            env = gym.wrappers.RecordVideo(env, f"videos/{run_name}")
        else:
            env = gym.make(env_id)
        env = gym.wrappers.FlattenObservation(env)
        env = gym.wrappers.RecordEpisodeStatistics(env)
        env = gym.wrappers.ClipAction(env)
        env = gym.wrappers.NormalizeObservation(env)
        env = gym.wrappers.TransformObservation(env, lambda obs: np.clip(obs, -10, 10))
        env = gym.wrappers.NormalizeReward(env, gamma=gamma)
        env = gym.wrappers.TransformReward(env, lambda reward: np.clip(reward, -10, 10))
        return env

    return thunk



def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
    torch.nn.init.orthogonal_(layer.weight, std)
    torch.nn.init.constant_(layer.bias, bias_const)
    return layer



class ComplexLinear(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(ComplexLinear, self).__init__()
        # 初始化参数
        self.real_weights = nn.Parameter(torch.randn(input_dim, output_dim) * 0.01 + 0.5)
        self.real_bias = nn.Parameter(torch.randn(output_dim) * 0.01 + 0.5)
        self.imag_bias = nn.Parameter(torch.randn(output_dim) * 0.01 + 0.5)
        self.lambda_k1 = nn.Parameter(torch.randn(output_dim) * 0.001 + 0.01)
        self.lambda_k2 = nn.Parameter(torch.randn(output_dim) * 0.001 + 0.01)

    def forward(self, x):
        # 检查输入是否是复数
        if torch.is_complex(x):
            x_real = x.real
            x_imag = x.imag
        else:
            x_real = x
            x_imag = torch.zeros_like(x)

        real_1 = torch.matmul(x_real, self.real_weights) + self.real_bias
        real_1 = torch.pow(real_1, 2)

        complex_output = real_1 + self.imag_bias**2
        complex_output_1 = real_1 / complex_output
        output_1 = complex_output_1 * self.lambda_k1

        complex_output_2 = self.imag_bias / complex_output
        output_2 = complex_output_2 * self.lambda_k2
        return output_1 + output_2



class Agent(nn.Module):
    def __init__(self, envs, actor_mean_layers=1, critic_layers=2):
        super().__init__()
        
        input_dim = np.array(envs.single_observation_space.shape).prod()
        action_dim = np.prod(envs.single_action_space.shape)

        if critic_layers == 1:
            self.critic = nn.Sequential(
                ComplexLinear(input_dim, 64),
                # layer_init(nn.Linear(64, 1), std=1.0),
                nn.Linear(64, 1)
            )
        else:
            self.critic = nn.Sequential(
                ComplexLinear(input_dim, 64),
                ComplexLinear(64, 64),
                layer_init(nn.Linear(64, 1), std=1.0),
            )


        if actor_mean_layers == 1:
            self.actor_mean = nn.Sequential(
                ComplexLinear(input_dim, 64),
                # nn.Tanh(),
                # layer_init(nn.Linear(64, np.prod(envs.single_action_space.shape)), std=0.01),
                nn.Linear(64, np.prod(envs.single_action_space.shape))
            )
        else:
            self.actor_mean = nn.Sequential(
                ComplexLinear(input_dim, 64),
                # nn.Tanh(),
                ComplexLinear(64, 64),
                # nn.Tanh(),
                layer_init(nn.Linear(64, np.prod(envs.single_action_space.shape)), std=0.01),
            )
        
        self.actor_logstd = nn.Parameter(torch.zeros(1, np.prod(envs.single_action_space.shape)))

    def get_value(self, x):
        return self.critic(x)

    def get_action_and_value(self, x, action=None):
        action_mean = self.actor_mean(x)
        action_logstd = self.actor_logstd.expand_as(action_mean)
        action_std = torch.exp(action_logstd)
        probs = Normal(action_mean, action_std)
        if action is None:
            action = probs.sample()
        return action, probs.log_prob(action).sum(1), probs.entropy().sum(1), self.critic(x)







if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    # parser.add_argument("-a", "--algorithm", type=int, required=True, help="Algorithm version to train (1 or 2)")
    parser.add_argument("-a", "--algorithm", type=int, default=1, help="Algorithm version to train (1 or 2)")
    args = parser.parse_args()

    seeds = range(5)
    configs = [
        {"actor_mean_layers": 1, "critic_layers": 2},
        {"actor_mean_layers": 1, "critic_layers": 1}
    ]
    # env_ids = ["HalfCheetah-v4", "Walker2d-v4", "Hopper-v4", "InvertedPendulum-v4", "Swimmer-v4", "Pusher-v4"]
    env_ids = ["Swimmer-v4"]
    for env_id in env_ids:
        config = configs[args.algorithm - 1]
        for seed in seeds:
            train_args = Args()
            train_args.env_id = env_id
            train_args.seed = seed
            train_args.batch_size = int(train_args.num_envs * train_args.num_steps)
            train_args.minibatch_size = int(train_args.batch_size // train_args.num_minibatches)
            train_args.num_iterations = train_args.total_timesteps // train_args.batch_size
            run_name = f"{train_args.env_id}__{train_args.exp_name}__{train_args.seed}__{int(time.time())}__am{config['actor_mean_layers']}__cr{config['critic_layers']}"
            writer = SummaryWriter(f"runs/{run_name}")

            # TRY NOT TO MODIFY: seeding
            random.seed(train_args.seed)
            np.random.seed(train_args.seed)
            torch.manual_seed(train_args.seed)
            torch.backends.cudnn.deterministic = train_args.torch_deterministic

            device = torch.device("cuda" if torch.cuda.is_available() and train_args.cuda else "cpu")

            # env setup
            envs = gym.vector.SyncVectorEnv(
                [make_env(train_args.env_id, i, train_args.capture_video, run_name, train_args.gamma) for i in range(train_args.num_envs)]
            )
            assert isinstance(envs.single_action_space, gym.spaces.Box), "only continuous action space is supported"

            agent = Agent(envs, actor_mean_layers=config["actor_mean_layers"], critic_layers=config["critic_layers"]).to(device)
            optimizer = optim.Adam(agent.parameters(), lr=train_args.learning_rate, eps=1e-5)

            obs = torch.zeros((train_args.num_steps, train_args.num_envs) + envs.single_observation_space.shape).to(device)
            actions = torch.zeros((train_args.num_steps, train_args.num_envs) + envs.single_action_space.shape).to(device)
            logprobs = torch.zeros((train_args.num_steps, train_args.num_envs)).to(device)
            rewards = torch.zeros((train_args.num_steps, train_args.num_envs)).to(device)
            dones = torch.zeros((train_args.num_steps, train_args.num_envs)).to(device)
            values = torch.zeros((train_args.num_steps, train_args.num_envs)).to(device)

            global_step = 0
            start_time = time.time()
            next_obs, _ = envs.reset(seed=train_args.seed)
            next_obs = torch.Tensor(next_obs).to(device)
            next_done = torch.zeros(train_args.num_envs).to(device)

            for iteration in range(1, train_args.num_iterations + 1):
                if train_args.anneal_lr:
                    frac = 1.0 - (iteration - 1.0) / train_args.num_iterations
                    lrnow = frac * train_args.learning_rate
                    optimizer.param_groups[0]["lr"] = lrnow

                for step in range(0, train_args.num_steps):
                    global_step += train_args.num_envs
                    obs[step] = next_obs
                    dones[step] = next_done

                    with torch.no_grad():
                        action, logprob, _, value = agent.get_action_and_value(next_obs)
                        values[step] = value.flatten()
                    actions[step] = action
                    logprobs[step] = logprob

                    next_obs, reward, terminations, truncations, infos = envs.step(action.cpu().numpy())
                    next_done = np.logical_or(terminations, truncations)
                    rewards[step] = torch.tensor(reward).to(device).view(-1)
                    next_obs, next_done = torch.Tensor(next_obs).to(device), torch.Tensor(next_done).to(device)

                    if "final_info" in infos:
                        reward_mean = []
                        for info in infos["final_info"]:
                            if info and "episode" in info:
                                print(f"global_step={global_step}, episodic_return={info['episode']['r']}")
                                reward_mean.append(info["episode"]["r"])
                                writer.add_scalar("charts/episodic_length", info["episode"]["l"], global_step)

                                # Save to CSV
                                csv_data = {
                                    "env_id": train_args.env_id,
                                    "global_step": global_step,
                                    "episodic_length": info["episode"]["l"],
                                    "episodic_return": info["episode"]["r"],
                                    "seed": train_args.seed,
                                }
                                csv_filename = f"runs/cauchy/{train_args.env_id}_cauchy_results.csv"
                                pd.DataFrame([csv_data]).to_csv(csv_filename, mode="a", header=not os.path.exists(csv_filename), index=False)
                        writer.add_scalar("charts/episodic_return", np.median(reward_mean), global_step)

                with torch.no_grad():
                    next_value = agent.get_value(next_obs).reshape(1, -1)
                    advantages = torch.zeros_like(rewards).to(device)
                    lastgaelam = 0
                    for t in reversed(range(train_args.num_steps)):
                        if t == train_args.num_steps - 1:
                            nextnonterminal = 1.0 - next_done
                            nextvalues = next_value
                        else:
                            nextnonterminal = 1.0 - dones[t + 1]
                            nextvalues = values[t + 1]
                        delta = rewards[t] + train_args.gamma * nextvalues * nextnonterminal - values[t]
                        advantages[t] = lastgaelam = delta + train_args.gamma * train_args.gae_lambda * nextnonterminal * lastgaelam
                    returns = advantages + values

                # flatten the batch
                b_obs = obs.reshape((-1,) + envs.single_observation_space.shape)
                b_logprobs = logprobs.reshape(-1)
                b_actions = actions.reshape((-1,) + envs.single_action_space.shape)
                b_advantages = advantages.reshape(-1)
                b_returns = returns.reshape(-1)
                b_values = values.reshape(-1)

                # Optimizing the policy and value network
                b_inds = np.arange(train_args.batch_size)
                clipfracs = []
                for epoch in range(train_args.update_epochs):
                    np.random.shuffle(b_inds)
                    for start in range(0, train_args.batch_size, train_args.minibatch_size):
                        end = start + train_args.minibatch_size
                        mb_inds = b_inds[start:end]

                        _, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions[mb_inds])
                        logratio = newlogprob - b_logprobs[mb_inds]
                        ratio = logratio.exp()

                        with torch.no_grad():
                            # calculate approx_kl http://joschu.net/blog/kl-approx.html
                            old_approx_kl = (-logratio).mean()
                            approx_kl = ((ratio - 1) - logratio).mean()
                            clipfracs += [((ratio - 1.0).abs() > train_args.clip_coef).float().mean().item()]

                        mb_advantages = b_advantages[mb_inds]
                        if train_args.norm_adv:
                            mb_advantages = (mb_advantages - mb_advantages.mean()) / (mb_advantages.std() + 1e-8)

                        # Policy loss
                        pg_loss1 = -mb_advantages * ratio
                        pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - train_args.clip_coef, 1 + train_args.clip_coef)
                        pg_loss = torch.max(pg_loss1, pg_loss2).mean()

                        # Value loss
                        newvalue = newvalue.view(-1)
                        if train_args.clip_vloss:
                            v_loss_unclipped = (newvalue - b_returns[mb_inds]) ** 2
                            v_clipped = b_values[mb_inds] + torch.clamp(
                                newvalue - b_values[mb_inds],
                                -train_args.clip_coef,
                                train_args.clip_coef,
                            )
                            v_loss_clipped = (v_clipped - b_returns[mb_inds]) ** 2
                            v_loss_max = torch.max(v_loss_unclipped, v_loss_clipped)
                            v_loss = 0.5 * v_loss_max.mean()
                        else:
                            v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()

                        entropy_loss = entropy.mean()
                        loss = pg_loss - train_args.ent_coef * entropy_loss + v_loss * train_args.vf_coef

                        optimizer.zero_grad()
                        loss.backward()
                        nn.utils.clip_grad_norm_(agent.parameters(), train_args.max_grad_norm)
                        optimizer.step()

                    if train_args.target_kl is not None and approx_kl > train_args.target_kl:
                        break

                y_pred, y_true = b_values.cpu().numpy(), b_returns.cpu().numpy()
                var_y = np.var(y_true)
                explained_var = np.nan if var_y == 0 else 1 - np.var(y_true - y_pred) / var_y

                # TRY NOT TO MODIFY: record rewards for plotting purposes
                writer.add_scalar("charts/learning_rate", optimizer.param_groups[0]["lr"], global_step)
                writer.add_scalar("losses/value_loss", v_loss.item(), global_step)
                writer.add_scalar("losses/policy_loss", pg_loss.item(), global_step)
                writer.add_scalar("losses/entropy", entropy_loss.item(), global_step)
                writer.add_scalar("losses/old_approx_kl", old_approx_kl.item(), global_step)
                writer.add_scalar("losses/approx_kl", approx_kl.item(), global_step)
                writer.add_scalar("losses/clipfrac", np.mean(clipfracs), global_step)
                writer.add_scalar("losses/explained_variance", explained_var, global_step)
                print("SPS:", int(global_step / (time.time() - start_time)))
                writer.add_scalar("charts/SPS", int(global_step / (time.time() - start_time)), global_step)

            if train_args.save_model:
                model_path = f"runs/{run_name}/{train_args.exp_name}.cleanrl_model"
                torch.save(agent.state_dict(), model_path)
                print(f"model saved to {model_path}")
            envs.close()
            writer.close()
