import os
import torch as T
import torch.nn.functional as F
import numpy as np
from rl_models.buffer import ReplayBuffer
from rl_models.networks import ActorNetwork, CriticNetwork, ValueNetwork

if T.cuda.is_available():
    print("Using GPU")
else:
    print("Using CPU")


class Agent():
    def __init__(self, config=None, alpha=0.0003, beta=0.0003, input_dims=[8],
                 env=None, gamma=0.99, n_actions=2, max_size=1000000, tau=0.005,
                 update_interval=1, layer1_size=256, layer2_size=256, batch_size=256, reward_scale=2,
                 chkpt_dir='tmp/sac'):
        if config is not None:
            # SAC params
            self.batch_size = config['SAC']['batch_size']
            self.layer1_size = config['SAC']['layer1_size']
            self.layer2_size = config['SAC']['layer2_size']
            self.gamma = config['SAC']['gamma']
            self.tau = config['SAC']['tau']
            self.alpha = config['SAC']['alpha']
            self.beta = config['SAC']['beta']
            # Experiment
            self.buffer_max_size = config['Experiment']['buffer_memory_size']
            self.update_interval = config['Experiment']['learn_every_n_steps']
            self.scale = config['Experiment']['reward_scale']

        else:
            self.gamma = gamma
            self.tau = tau
            self.buffer_max_size = max_size
            self.batch_size = batch_size
            self.update_interval = update_interval
            self.alpha = alpha
            self.beta = beta

            self.layer1_size = layer1_size
            self.layer2_size = layer2_size
            self.scale = reward_scale
        self.env = env
        self.input_dims = input_dims
        self.n_actions = n_actions
        self.chkpt_dir = chkpt_dir

        if config is not None and 'chkpt_dir' in config["Experiment"].values():
            self.chkpt_dir = config['chkpt_dir']

        self.memory = ReplayBuffer(self.buffer_max_size, self.input_dims, self.n_actions)
        self.actor = ActorNetwork(self.alpha, self.input_dims, n_actions=self.n_actions,
                                  name='actor', max_action=self.env.action_space.high+1,
                                  fc1_dims=self.layer1_size, fc2_dims=self.layer2_size,
                                  chkpt_dir=self.chkpt_dir)
        self.critic_1 = CriticNetwork(self.beta, self.input_dims, n_actions=self.n_actions,
                                      name='critic_1', fc1_dims=self.layer1_size, fc2_dims=self.layer2_size,
                                      chkpt_dir=self.chkpt_dir)
        self.critic_2 = CriticNetwork(self.beta, self.input_dims, n_actions=self.n_actions,
                                      name='critic_2', fc1_dims=self.layer1_size, fc2_dims=self.layer2_size,
                                      chkpt_dir=self.chkpt_dir)
        self.value = ValueNetwork(self.beta, self.input_dims, name='value', fc1_dims=self.layer1_size,
                                  fc2_dims=self.layer2_size,
                                  chkpt_dir=self.chkpt_dir)
        self.target_value = ValueNetwork(self.beta, self.input_dims, name='target_value', fc1_dims=self.layer1_size,
                                         fc2_dims=self.layer2_size, chkpt_dir=self.chkpt_dir)
        self.update_network_parameters(tau=self.tau)

    def choose_action(self, observation):
        state = T.Tensor([observation]).to(self.actor.device)
        actions, _ = self.actor.sample_normal(state, reparameterize=False)

        return actions.cpu().detach().numpy()[0]

    def remember(self, state, action, reward, new_state, done):
        self.memory.store_transition(state, action, reward, new_state, done)

    def update_network_parameters(self, tau=None):
        if tau is None:
            tau = self.tau

        target_value_params = self.target_value.named_parameters()
        value_params = self.value.named_parameters()

        target_value_state_dict = dict(target_value_params)
        value_state_dict = dict(value_params)

        for name in value_state_dict:
            value_state_dict[name] = tau * value_state_dict[name].clone() + \
                                     (1 - tau) * target_value_state_dict[name].clone()

        self.target_value.load_state_dict(value_state_dict)

    def save_models(self):
        print('.... saving models ....')
        self.actor.save_checkpoint()
        self.value.save_checkpoint()
        self.target_value.save_checkpoint()
        self.critic_1.save_checkpoint()
        self.critic_2.save_checkpoint()

    def load_models(self):
        print('.... loading models ....')
        self.actor.load_checkpoint()
        self.value.load_checkpoint()
        self.target_value.load_checkpoint()
        self.critic_1.load_checkpoint()
        self.critic_2.load_checkpoint()

    def learn(self, episode=None):
        if episode is not None:  # online training
            state, action, reward, new_state, done = episode
            state = np.asarray([state])
            action = np.asarray([action])
            new_state = np.asarray([new_state])

        else:  # offline training
            if self.memory.mem_cntr < self.batch_size:
                return
            state, action, reward, new_state, done = \
                self.memory.sample_buffer(self.batch_size)

        reward = T.tensor(reward, dtype=T.float).to(self.actor.device)
        done = T.tensor(done).to(self.actor.device)
        state_ = T.tensor(new_state, dtype=T.float).to(self.actor.device)
        state = T.tensor(state, dtype=T.float).to(self.actor.device)
        action = T.tensor(action, dtype=T.float).to(self.actor.device)

        value = self.value(state).view(-1)
        value_ = self.target_value(state_).view(-1)
        value_[done] = 0.0

        actions, log_probs = self.actor.sample_normal(state, reparameterize=False)
        log_probs = log_probs.view(-1)
        q1_new_policy = self.critic_1.forward(state, actions)
        q2_new_policy = self.critic_2.forward(state, actions)
        critic_value = T.min(q1_new_policy, q2_new_policy)
        critic_value = critic_value.view(-1)

        self.value.optimizer.zero_grad()
        value_target = critic_value - log_probs
        value_loss = 0.5 * F.mse_loss(value, value_target)
        value_loss.backward(retain_graph=True)
        self.value.optimizer.step()

        actions, log_probs = self.actor.sample_normal(state, reparameterize=True)
        log_probs = log_probs.view(-1)
        q1_new_policy = self.critic_1.forward(state, actions)
        q2_new_policy = self.critic_2.forward(state, actions)
        critic_value = T.min(q1_new_policy, q2_new_policy)
        critic_value = critic_value.view(-1)

        actor_loss = log_probs - critic_value
        actor_loss = T.mean(actor_loss)
        self.actor.optimizer.zero_grad()
        actor_loss.backward(retain_graph=True)
        self.actor.optimizer.step()

        self.critic_1.optimizer.zero_grad()
        self.critic_2.optimizer.zero_grad()
        q_hat = self.scale * reward + self.gamma * value_
        q1_old_policy = self.critic_1.forward(state, action).view(-1)
        q2_old_policy = self.critic_2.forward(state, action).view(-1)
        critic_1_loss = 0.5 * F.mse_loss(q1_old_policy, q_hat)
        critic_2_loss = 0.5 * F.mse_loss(q2_old_policy, q_hat)

        critic_loss = critic_1_loss + critic_2_loss
        critic_loss.backward()
        self.critic_1.optimizer.step()
        self.critic_2.optimizer.step()

        self.update_network_parameters()
