import numpy as np
import torch
import torch.optim as optim
import torch.nn.functional as F

from collections import namedtuple

from assets.RL.AtariNet import DQN, ExtendedDQN
from assets.memory.ReplayBuffer import ReplayBuffer
from assets.helperFunctions.timestamps import print_timestamp as print_ts


class DQN_module():
    """
    A wrapper class that augments the AtariNet.DQN class by optmizing methods.
    """
    def __init__(self, batch_size, gamma, history_length, size_replaybuffer,
                 optim_learning_rate, dim_actions):
        """
        update_cnt: Target network update counter.
        net:        The q-value network.
        target_net: The target network.
        optimizer:  The optimizer used to update the q-value network.
        memory:     The memory class for the replay buffer
        """
        self.device = self._setup_torch()
        self.batch_size = batch_size
        self.gamma = gamma
        self.history_length = history_length
        self.dim_actions = dim_actions
        self.size_replaybuffer = size_replaybuffer
        self.optim_learning_rate = optim_learning_rate
        self.loss = 0
        self.update_cnt = 0  # Target network update_counter

        self.state_q_values = 0

        self.net, self.target_net, self.optimizer = self._build_model()
        self.print_architecture()
        self.memory = ReplayBuffer(self.size_replaybuffer)

    def _build_model(self):
        """
        Initializing 2 networks and an Adam optimizer.
        """
        net = ExtendedDQN(self.history_length, self.dim_actions).to(self.device)
        target_net = ExtendedDQN(self.history_length, self.dim_actions).to(self.device)
        optimizer = optim.Adam(net.parameters(), lr=self.optim_learning_rate)
        return net, target_net, optimizer

    def print_architecture(self):
        print_ts("Network: \n{}".format(self.net))
        print_ts("Optimizer: \n{}".format(self.optimizer))
        print_ts("Target Network: \n{}".format(self.target_net))

    def _setup_torch(self):
        """
        Setting GPU if available. Else, use the CPU.
        """
        # Initalizing
        device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
        torch.set_printoptions(linewidth=750, profile="full")
        print_ts("Performing calculations on {}".format(device))
        return device

    def predict_q_values(self, state):
        with torch.no_grad():
            state_tensor = torch.tensor(state,
                                        device=self.device,
                                        dtype=torch.float,
                                        requires_grad=False)
            return self.net(state_tensor)

    def resize(self):
        """
        From original DQN pendulum tutorial with PyTorch
        """
        return T.Compose([T.ToPILImage(),
                            T.Resize(40, interpolation=Image.CUBIC),
                            T.ToTensor()])

    # ##########################################################################
    # Optimizing the network
    # ##########################################################################

    def optimize(self):
        """
        optimizes the model. currently only trains the actions
        # TODO: extend Q-Update function to the x and y coordinates
        """

        # Sample and unzip batch from experience replay buffer
        batch = self.sample_batch()
        self.unzip_batch(batch)

        # Calculate Q-values
        self.calculate_q_values()

        # calculate td targets of the actions, x&y coordinates
        self.td_target = self.calculate_td_target(self.next_state_q_max)

        # Compute the loss
        self.compute_loss()

        # optimize model
        self.optimize_model()

        # print("Optimization finished")
        return self.loss

    def sample_batch(self):
        """
        Sample from batch.
        """
        Transition = namedtuple('Transition',
                                ('state', 'action', 'reward', 'next_state'))
        transitions = self.memory.sample(self.batch_size)
        return Transition(*zip(*transitions))

    def unzip_batch(self, batch):
        """
        Get the batches from the transition tuple
        np.concatenate concatenate all the batch data into a single ndarray
        """

        self.non_final_mask = torch.tensor(tuple(map(lambda s: s is not None,
            batch.next_state)), device=self.device, dtype=torch.uint8)

        tensor_next_state = [torch.as_tensor(s, device=self.device, dtype=torch.float) for s in batch.next_state if s is not None]
        self.non_final_next_states = torch.cat(tensor_next_state)

        self.state_batch = torch.as_tensor(np.concatenate(batch.state),
                                           device=self.device,
                                           dtype=torch.float)
        self.action_batch = torch.tensor(batch.action,
                                         device=self.device,
                                         dtype=torch.long)
        self.reward_batch = torch.as_tensor(batch.reward,
                                            device=self.device,
                                            dtype=torch.float)


    def calculate_q_values(self):
        """
        """
        # forward pass
        self.state_q_values = self.net(self.state_batch)
        print(self.state_q_values)

        # gather action values with respect to the chosen action
        self.state_q_values = self.state_q_values.gather(1, self.action_batch)


        # compute action values of the next state over all actions and take max
        self.next_state_q_max = torch.zeros(self.batch_size, device=self.device)
        self.next_state_q_max[self.non_final_mask] = self.target_net(self.non_final_next_states).max(1)[0].detach()

        # self.next_state_q_values = self.target_net(self.next_state_batch)
        # self.next_state_q_max = self.next_state_q_values.max(1)[0].detach()

    def calculate_td_target(self, q_values_best_next):
        """
        Calculating the TD-target for given q-values.
        """
        return (q_values_best_next * self.gamma) + self.reward_batch

    def q_value_analysis_on_batch(self, q_value_tensor):
        """
        """
        q_max = q_value_tensor.detach().numpy().max(axis=1).mean()
        q_min = q_value_tensor.detach().numpy().min(axis=1).mean()
        q_mean = q_value_tensor.detach().numpy().mean(axis=1).mean()
        q_var = q_value_tensor.detach().numpy().var(axis=1).mean()
        q_argmax = q_value_tensor.detach().numpy().argmax(axis=1)
        q_span = q_max - q_min
        return q_max, q_min, q_mean, q_var, q_span, q_argmax

    def q_value_analysis(self, q_value_tensor):
        """
        """
        q_max = q_value_tensor.detach().numpy().max()
        q_min = q_value_tensor.detach().numpy().min()
        q_mean = q_value_tensor.detach().numpy().mean()
        q_var = q_value_tensor.detach().numpy().var()
        q_argmax = q_value_tensor.detach().numpy().argmax()
        q_span = q_max - q_min
        return q_max, q_min, q_mean, q_var, q_span, q_argmax

    def compute_loss(self):
        """
        Calculating the loss between TD-target Q-values.
        """
        self.loss = F.mse_loss(self.state_q_values,
                               self.td_target.unsqueeze(1))

    def optimize_model(self):
        """
        This step is used to update the parameter weights:
        loss.backwards propagates the backward gradients for all net weights.
        optimizer.step updates all net weights with the assigned gradients.
        optimizer.zero_grad sets the gradients to 0 for the next update cycle.
        """
        self.optimizer.zero_grad()
        self.loss.backward()
        self.optimizer.step()

    def update_target_net(self):
        """
        updates weights of the target network, i.e. copies model weights to it
        """
        self.transfer_weights(self.net, self.target_net)
        self.update_cnt += 1

    # ##########################################################################
    # Weight methods
    # ##########################################################################

    def save(self, save_path):
        torch.save(self.net.state_dict(), save_path)

    def load(self, load_path):
        print("Model loaded from {}".format(load_path))
        self.net.load_state_dict(torch.load(load_path, map_location=self.device))
        self.net.eval()


    def transfer_weights(self, src_net, target_net):
        """
        updates weights of the target network, i.e. copies model weights to it
        """
        target_net.load_state_dict(src_net.state_dict())
        self.update_cnt += 1

    def get_net_weights(self):
        """
        Get the weights of the neural network
        """
        return self.net.state_dict()

    def set_net_weights(self, weights):
        """
        Set the weights of the neural network
        """
        self.net.load_state_dict

    def get_target_net_weights(self):
        return self.target_net.state_dict()

    def set_target_net_weights(self, weights):
        self.target_net.load_state_dict
