# -*- coding: utf-8 -*-
from __future__ import division
import os
import numpy as np
import torch
from torch import optim
from torch.nn.utils import clip_grad_norm_
from model import DQN, LCR_Model
from tqdm import tqdm



class Agent():
    def __init__(self, args, env):
        self.action_space = env.action_space()
        self.atoms = args.atoms
        self.Vmin = args.V_min
        self.Vmax = args.V_max
        self.support = torch.linspace(args.V_min, args.V_max, self.atoms).to(device=args.device)  # Support (range) of z
        self.delta_z = (args.V_max - args.V_min) / (self.atoms - 1)
        self.batch_size = args.batch_size
        self.n = args.multi_step
        self.discount = args.discount
        self.norm_clip = args.norm_clip
        self.use_lcr = args.use_lcr
        if self.use_lcr:
            self.K = args.K
            self.gradient_steps = args.gradient_steps
            self.lcr_batch_size = args.lcr_batch_size
            self.lcr_net = LCR_Model(self.K).to(args.device)
            self.mse_loss = torch.nn.MSELoss()

        self.online_net = DQN(args, self.action_space).to(device=args.device)
        if args.model:  # Load pretrained model if provided
            if os.path.isfile(args.model):
                state_dict = torch.load(args.model, map_location='cpu')  # Always load tensors onto CPU by default, will shift to GPU if necessary
                if 'conv1.weight' in state_dict.keys():
                    for old_key, new_key in (('conv1.weight', 'convs.0.weight'), ('conv1.bias', 'convs.0.bias'), ('conv2.weight', 'convs.2.weight'), ('conv2.bias', 'convs.2.bias'), ('conv3.weight', 'convs.4.weight'), ('conv3.bias', 'convs.4.bias')):
                        state_dict[new_key] = state_dict[old_key]  # Re-map state dict for old pretrained models
                        del state_dict[old_key]  # Delete old keys for strict load_state_dict
                self.online_net.load_state_dict(state_dict)
                print("Loading pretrained model: " + args.model)
            else:  # Raise error if incorrect model path provided
                raise FileNotFoundError(args.model)

        self.online_net.train()

        self.target_net = DQN(args, self.action_space).to(device=args.device)
        self.update_target_net()
        self.target_net.train()
        for param in self.target_net.parameters():
            param.requires_grad = False

        self.optimiser = optim.Adam(self.online_net.parameters(), lr=args.learning_rate, eps=args.adam_eps)

    # Resets noisy weights in all linear layers (of online net only)
    def reset_noise(self):
        self.online_net.reset_noise()

    # Acts based on single state (no batch)
    def act(self, state):
        with torch.no_grad():
            return (self.online_net(state.unsqueeze(0)) * self.support).sum(2).argmax(1).item()

    # Acts with an ε-greedy policy (used for evaluation only)
    def act_e_greedy(self, state, epsilon=0.001):  # High ε can reduce evaluation scores drastically
        return np.random.randint(0, self.action_space) if np.random.random() < epsilon else self.act(state)

    def learn(self, mem):
        # Sample transitions
        idxs, states, actions, returns, next_states, nonterminals, weights = mem.sample(self.batch_size)

        # Calculate current state probabilities (online network noise already sampled)
        log_ps = self.online_net(states, log=True)  # Log probabilities log p(s_t, ·; θonline)
        log_ps_a = log_ps[range(self.batch_size), actions]  # log p(s_t, a_t; θonline)

        with torch.no_grad():
            # Calculate nth next state probabilities
            pns = self.online_net(next_states)  # Probabilities p(s_t+n, ·; θonline)
            dns = self.support.expand_as(pns) * pns  # Distribution d_t+n = (z, p(s_t+n, ·; θonline))
            argmax_indices_ns = dns.sum(2).argmax(1)  # Perform argmax action selection using online network: argmax_a[(z, p(s_t+n, a; θonline))]
            self.target_net.reset_noise()  # Sample new target net noise
            pns = self.target_net(next_states)  # Probabilities p(s_t+n, ·; θtarget)
            pns_a = pns[range(self.batch_size), argmax_indices_ns]  # Double-Q probabilities p(s_t+n, argmax_a[(z, p(s_t+n, a; θonline))]; θtarget)

            # Compute Tz (Bellman operator T applied to z)
            Tz = returns.unsqueeze(1) + nonterminals * (self.discount ** self.n) * self.support.unsqueeze(0)  # Tz = R^n + (γ^n)z (accounting for terminal states)
            Tz = Tz.clamp(min=self.Vmin, max=self.Vmax)  # Clamp between supported values
            # Compute L2 projection of Tz onto fixed support z
            b = (Tz - self.Vmin) / self.delta_z  # b = (Tz - Vmin) / Δz
            l, u = b.floor().to(torch.int64), b.ceil().to(torch.int64)
            # Fix disappearing probability mass when l = b = u (b is int)
            l[(u > 0) * (l == u)] -= 1
            u[(l < (self.atoms - 1)) * (l == u)] += 1

            # Distribute probability of Tz
            m = states.new_zeros(self.batch_size, self.atoms)
            offset = torch.linspace(0, ((self.batch_size - 1) * self.atoms), self.batch_size).unsqueeze(1).expand(self.batch_size, self.atoms).to(actions)
            m.view(-1).index_add_(0, (l + offset).view(-1), (pns_a * (u.float() - b)).view(-1))  # m_l = m_l + p(s_t+n, a*)(u - b)
            m.view(-1).index_add_(0, (u + offset).view(-1), (pns_a * (b - l.float())).view(-1))  # m_u = m_u + p(s_t+n, a*)(b - l)

        loss = -torch.sum(m * log_ps_a, 1)  # Cross-entropy loss (minimises DKL(m||p(s_t, a_t)))
        self.online_net.zero_grad()
        (weights * loss).mean().backward()  # Backpropagate importance-weighted minibatch loss
        clip_grad_norm_(self.online_net.parameters(), self.norm_clip)  # Clip gradients by L2 norm
        self.optimiser.step()

        detach_loss = loss.detach().cpu().numpy()
        mem.update_priorities(idxs, detach_loss)  # Update priorities of sampled transitions
        return detach_loss.mean()

    def lcr_learn(self, args, lcr_mem):
        # Sample transitions
        states = lcr_mem.get(self.lcr_batch_size)
        # states = lcr_mem.sample(self.lcr_batch_size)
        self.lcr_net.reset_weights()
        X_dim = states[0].shape
        x_actual = torch.zeros(size=(self.lcr_batch_size,) + X_dim, device=args.device)
        x_nearest = torch.zeros(size=(self.lcr_batch_size,) + (self.K, ) + X_dim, device=args.device)
        idxs = np.arange(self.lcr_batch_size)
        for id in idxs:
            x_all = torch.zeros(size=(self.K + 1,) + X_dim, device=args.device)
            if id >= self.K // 2 and self.lcr_batch_size - id > self.K // 2 + (self.K % 2):
                for count, j in enumerate(range(id - self.K // 2, id + self.K // 2 + (self.K % 2) + 1)):
                    x_all[count] = states[j]
                x_actual[id] = x_all[self.K // 2]
                # Calculate nearest X's
                x_nearest[id] = torch.concat((x_all[0:self.K//2], x_all[self.K//2+1:]))
            elif id < self.K // 2:
                for j in range(0, self.K + 1):
                    x_all[j] = states[j]
                x_actual[id] = x_all[0]
                # Calculate nearest X's
                x_nearest[id] = x_all[1:]
            else:
                for count, j in enumerate(range(self.lcr_batch_size - self.K - 1, self.lcr_batch_size)):
                    x_all[count] = states[j]
                x_actual[id] = x_all[-1]
                # Calculate nearest X's
                x_nearest[id] = x_all[:-1]
        # Calculate R
        with torch.no_grad():
            r_actual = self.online_net.conv_output(x_actual.to(args.device))
        lcr_losses = 0
        for _ in tqdm(range(self.gradient_steps)):
            r_nearest = torch.empty(size=(self.lcr_batch_size, self.online_net.conv_output_size, self.K),
                                    device=args.device)
            for k in range(self.K):
                r_nearest[:, :, k] = self.online_net.conv_output(x_nearest[:, k, :].to(args.device))
            r_predicted = torch.squeeze(self.lcr_net(r_nearest)).to(args.device)
            lcr_loss = self.mse_loss(r_predicted, r_actual).to(args.device)
            lcr_optimizer = optim.Adam(list(self.online_net.parameters()) + list(self.lcr_net.parameters()),
                                       lr=args.lcr_lr, eps=args.adam_eps)
            self.online_net.zero_grad()
            self.lcr_net.zero_grad()
            lcr_loss.mean().backward()  # Backpropagate loss
            lcr_losses += lcr_loss.detach().cpu().numpy()
            lcr_optimizer.step()
        return lcr_losses/self.gradient_steps

    def update_target_net(self):
        self.target_net.load_state_dict(self.online_net.state_dict())

    # Save model parameters on current device (don't move model between devices)
    def save(self, path, name='model.pth'):
        torch.save(self.online_net.state_dict(), os.path.join(path, name))

    # Evaluates Q-value based on single state (no batch)
    def evaluate_q(self, state):
        with torch.no_grad():
            return (self.online_net(state.unsqueeze(0)) * self.support).sum(2).max(1)[0].item()

    def train(self):
        self.online_net.train()

    def eval(self):
        self.online_net.eval()
