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

import utils
from model import DuelingNetwork
from utils import ExperienceReplayMemory, AgentConfig, BranchingTensorEnv, init_seeds


class BranchingDQN(nn.Module):

    def __init__(self, obs, ac, config):
        super().__init__()

        self.q = DuelingNetwork(obs, ac)
        self.target = DuelingNetwork(obs, ac)

        self.target.load_state_dict(self.q.state_dict())

        self.target_net_update_freq = config.target_net_update_freq
        self.update_counter = 0
        self.config = config

    def get_action(self, x):
        with torch.no_grad():
            # a = self.q(x).max(1)[1]
            out = self.q(x)
            return torch.argmax(out).item()

    def update_policy(self, adam, memory, params):
        b_states, b_actions, b_rewards, b_next_states, b_masks = memory.sample(params.batch_size)

        states = torch.tensor(b_states).float()
        actions = torch.tensor(b_actions).long().reshape(states.shape[0], -1)
        rewards = torch.tensor(b_rewards).float().reshape(-1, 1)
        next_states = torch.tensor(b_next_states).float()
        masks = torch.tensor(b_masks).float().reshape(-1, 1)

        qvals = self.q(states)
        current_q_values = qvals.gather(1, actions).squeeze(-1)

        with torch.no_grad():
            argmax = torch.argmax(self.q(next_states), dim=1)

            max_next_q_vals = self.target(next_states).gather(1, argmax.unsqueeze(1))
            # max_next_q_vals = max_next_q_vals.mean(1, keepdim=True).expand(-1, max_next_q_vals.shape[1])   # Branching DQN

        expected_q_vals = rewards + max_next_q_vals * self.config.gamma * masks
        # print(expected_q_vals[:5])
        loss = F.mse_loss(expected_q_vals, current_q_values)

        # input(loss)

        # print('\n'*5)

        adam.zero_grad()
        loss.backward()

        for p in self.q.parameters():
            p.grad.data.clamp_(-1., 1.)
        adam.step()

        self.update_counter += 1
        if self.update_counter % self.target_net_update_freq == 0:
            self.update_counter = 0
            # self.target.load_state_dict(self.q.state_dict())
            self.target.soft_update(self.q, tau=self.config.target_net_update_tau)


if __name__ == "__main__":
    args = utils.arguments()

    init_seeds(500)

    env = BranchingTensorEnv(args.env)
    env.seed(1234)

    config = AgentConfig()
    memory = ExperienceReplayMemory(config.memory_size)
    agent = BranchingDQN(env.observation_space.shape[0], env.action_space.n, config)
    adam = optim.Adam(agent.q.parameters(), lr=config.lr)

    s = env.reset()
    ep_reward = 0.
    recap = []  # stores cumulative reward per episode (episode rewards)

    p_bar = tqdm(total=config.max_frames)
    for frame in range(config.max_frames):

        epsilon = config.epsilon_by_frame(frame)

        if np.random.random() > 0.5:
            action = agent.get_action(s)
        else:
            action = np.random.randint(0, env.action_space.n)

        ns, r, done, infos = env.step(action)
        ep_reward += r

        if done:
            ns = env.reset()
            recap.append(ep_reward)
            p_bar.set_description('Rew: {:.3f}'.format(ep_reward))
            ep_reward = 0.

        memory.push((s.reshape(-1).numpy().tolist(), action, r, ns.reshape(-1).numpy().tolist(), 0. if done else 1.))
        s = ns

        p_bar.update(1)

        # Update online and target model only from specified number of frames
        if frame > config.learning_starts:
            agent.update_policy(adam, memory, config)

        if frame % 1000 == 0:
            utils.save(agent, recap, args)

    p_bar.close()
