import os
from collections import OrderedDict, deque, namedtuple
from typing import List, Tuple, Iterator

import gym
import numpy as np
import torch
from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.utilities import DistributedType
from torch import Tensor, nn
from torch.optim import Adam, Optimizer
from torch.utils.data import DataLoader
from torch.utils.data.dataset import IterableDataset

class DQN(nn.Module):
    """Deep Q-Network. E' un semplice MLP che prende
    in input un vettore di osservazioni (lo stato) e restituisce in output
    i Q-value per le diverse azioni"""

    def __init__(self, 
                 obs_size, # dimensioni del vettore di osservazioni
                 n_actions, # numero di azioni
                 hidden_size = 128 # dimensione del livello nascosto
                ):
        super().__init__()
        # definiamo il MLP
        self.net = nn.Sequential(
            nn.Linear(obs_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, n_actions),
        )

    def forward(self, x):
        return self.net(x.float())

model = DQN(4, 2)
model(torch.zeros(1,4)).shape

# Una namedtuple per rappresentare una experience
Experience = namedtuple(
    "Experience",
    field_names=["state", "action", "reward", "done", "new_state"],
)

class ReplayBuffer:
    """Buffer che contine le esperienze già viste per effettuare il replay"""

    def __init__(self, capacity):
        # capacity indica la grandezza massima del buffer
        # deque è una coda accessibile da ambo i lati
        self.buffer = deque(maxlen=capacity)

    def __len__(self):
        return len(self.buffer)

    def append(self, experience):
        # permette di inserire una esperienza
        self.buffer.append(experience)

    def sample(self, batch_size):
        # campiona batch_size esperienze casualmente
        # ottiene batch_size indici senza sostituzione
        indices = np.random.choice(len(self.buffer), batch_size, replace=False)
        # preleva lista di stati, azioni, reward, booleani "done" e next_states
        states, actions, rewards, dones, next_states = zip(*(self.buffer[idx] for idx in indices))

        return ( #restituisci batch delle diverse quantità
            np.array(states),
            np.array(actions),
            np.array(rewards, dtype=np.float32),
            np.array(dones, dtype=np.bool),
            np.array(next_states),
        )

class RLDataset(IterableDataset):
    """Dataset di tipo Iterable. Si tratta di un dataset 'infinito'. 
    Verrà aggiornato in automatico quando inseriremo elementi dentro il buffer"""

    def __init__(self, 
                 buffer, # replay buffer
                 sample_size = 200 # numero di memories da campionare ad ogni step
                ):
        self.buffer = buffer 
        self.sample_size = sample_size

    def __iter__(self):
        # campiona le esperienze
        states, actions, rewards, dones, new_states = self.buffer.sample(self.sample_size)
        for i in range(len(dones)): #restituisci stati, azioni, reward, done e new_states, uno ad uno
            yield states[i], actions[i], rewards[i], dones[i], new_states[i]

class Agent:
    """Questa classe implementa l'agente di reinforcement learning
    e le sue interazioni con l'ambiente"""

    def __init__(self, 
                 env, #l'ambiente di gym
                 replay_buffer #il replay buffer
                ):
        self.env = env
        self.replay_buffer = replay_buffer
        self.episodes = -1
        self.reset()

    def reset(self):
        """Resetta l'ambiente e salva lo stato corrente"""
        self.state = self.env.reset()
        self.episodes+=1

    def get_action(self, 
                   net, # la Q-Net
                   epsilon, # probabilità di effettuare una azione casuale per bilanciare exploration/exploitation
                   device # il dispositivo sul quale effettuare la computazione
                  ):
        # con probabilità epsilon
        if np.random.random() < epsilon:
            # campiona una azione casuale
            action = self.env.action_space.sample()
        else:
            # trasforma l'ultimo stato in tensore
            state = torch.tensor([self.state])

            # metti lo stato sul device corretto
            if device not in ["cpu"]:
                state = state.cuda(device)

            # ottieni i q values
            q_values = net(state)
            # campiona l'azione con q-alue più alto
            _, action = torch.max(q_values, dim=1)
            # trasforma l'azione in intero
            action = int(action.item())

        return action

    @torch.no_grad() #disabilita i gradienti
    def play_step(
        self,
        net, # la q-net
        epsilon = 0.0, # probabilità di una azione random
        device = "cpu", # il device su cui lavorare
    ):
        # ottieni l'azione per lo stato corrente
        action = self.get_action(net, epsilon, device)

        # effettua uno step dell'environment di gym
        new_state, reward, done, _ = self.env.step(action)

        # costruiamo una nuova esperienza
        exp = Experience(self.state, action, reward, done, new_state)

        # aggiungiamola al buffer
        self.replay_buffer.append(exp)

        # conserviamo il nuovo stato come corrente
        self.state = new_state
        
        if done: # se l'episodio è finito
            self.reset() #resettiamo l'ambiente
        return reward, done
    
    @torch.no_grad()
    def play_episode(self, net, epsilon=0, device = "cpu", return_frames=True):
        """Questo metodo permette di eseguire un episodio completo. E' utile per visualizzare il comportamento dell'agente."""
        self.reset()
        total_reward = 0
        frames = []
        actions = []
        for t in count():
            # ottieni l'azione per lo stato corrente
            action = self.get_action(net, epsilon, device)
            
            # effettua uno step dell'environment di gym
            new_state, reward, done, _ = self.env.step(action)
            
            # conserviamo il nuovo stato come corrente
            self.state = new_state
            
            total_reward += reward
            if return_frames:
                frames.append(self.env.render(mode='rgb_array'))
            actions.append(action)
            
            if done:
                self.reset()
                break
        return frames, actions, total_reward

from itertools import count

class DQNLightning(LightningModule):
    def __init__(
        self,
        Agent = Agent,
        DQN = DQN,
        ReplayBuffer = ReplayBuffer,
        batch_size = 256, # il batch size con cui allenare la rete
        lr = 1e-2, # learning rate
        env = "CartPole-v0", # environment di gym
        gamma = 0.99, # iperprametro gamma
        sync_rate = 10, # ogni quanto aggiorare la target network
        replay_size = 1000, # dimensione massima del replay buffer
        warm_start_size = 1000, # quanti campioni usare per riempire il buffer all'inizio del training
        eps_last_frame = 1000, # dopo quanti frame iniziare a far decadere epsilon
        eps_start = 1.0, # epsilon iniziale
        eps_end = 0.01, # epsilon finale
        warm_start_steps = 1000, # massima reward nell'ambiente
        display_rate = 20 #ogni quanti episodi, mostrare un episodio a schermo
    ):
        super().__init__()
        self.save_hyperparameters() #salviamo gli iperparametri

        self.env = gym.make(self.hparams.env) #creiamo l'ambiente di gym
        obs_size = self.env.observation_space.shape[0] #numero di osservazioni
        n_actions = self.env.action_space.n #numero di azioni

        self.net = DQN(obs_size, n_actions) #Q-Net
        self.target_net = DQN(obs_size, n_actions) #rete target

        self.buffer = ReplayBuffer(self.hparams.replay_size) #buffer
        self.agent = Agent(self.env, self.buffer) #agente
        self.total_reward = 0 #reward totale
        self.episode_reward = 0 #reward dell'ultimo episodio
        self.populate(self.hparams.warm_start_steps) #riempie il memory buffer

    def populate(self, steps = 1000):
        for i in range(steps): # effettua 1000 step con policy random
            self.agent.play_step(self.net, epsilon=1.0)
        self.agent.episodes=0

    def forward(self, x):
        output = self.net(x)
        return output

    def dqn_mse_loss(self, batch):
        # calcola la loss MSE per allenare la rete
        states, actions, rewards, dones, next_states = batch

        # otteniamo i valori relativi alle azioni effettuate
        state_action_values = self.net(states).gather(1, actions.unsqueeze(-1)).squeeze(-1)

        with torch.no_grad():
            # stimiamo i valori degli stati futuri con la target net
            next_state_values = self.target_net(next_states).max(1)[0]
            # impostiamo a zero i valori degli ultimi step
            next_state_values[dones] = 0.0
            # scolleghiamo dal grafo computazionale
            next_state_values = next_state_values.detach()

        # calcoliamo i valori attesi degli stati
        expected_state_action_values = next_state_values * self.hparams.gamma + rewards

        # restituiamo la loss MSE tra gli stati stimati e quelli attesi
        return nn.MSELoss()(state_action_values, expected_state_action_values)

    def training_step(self, batch, batch_idx):
        device = self.get_device(batch) # otteniamo il device del batch
        epsilon = max( # calcoliamo epsilon usando il global_step corrente
            self.hparams.eps_end,
            self.hparams.eps_start - self.global_step / self.hparams.eps_last_frame,
        )

        # esegui un nuovo step in modo da continuare a riempire il buffer
        reward, done = self.agent.play_step(self.net, epsilon, device)
        self.episode_reward += reward # cumula la reward dell'episodio corrente

        # calcola la loss di training a partire dal batch
        loss = self.dqn_mse_loss(batch)

        # reshape necessario per alcuni tipi di training
        if self.trainer._distrib_type in {DistributedType.DP, DistributedType.DDP2}:
            loss = loss.unsqueeze(0)

        # se abbiamo finito
        if done:
            # imposta la reward totale
            self.total_reward = self.episode_reward
            self.episode_reward = 0 #reimposta la reward dell'episodio

        # aggiorna la target_net ogni sync_rate step
        if self.global_step % self.hparams.sync_rate == 0:
            self.target_net.load_state_dict(self.net.state_dict())

        log = {
            "total_reward": torch.tensor(self.total_reward).to(device),
            "reward": torch.tensor(reward).to(device),
            "train_loss": loss,
        }
        status = {
            "steps": torch.tensor(self.global_step).to(device),
            "total_reward": torch.tensor(self.total_reward).to(device),
        }
        
        self.log('train/loss', loss)
        self.log('train/total_reward', self.total_reward)
        self.log('train/episodes', self.agent.episodes)
        
        # ogni display_rate episodi, mostriamo un episodio completo a schermo
        if self.agent.episodes % self.hparams.display_rate == 0:
            frames, actions, total_reward = self.agent.play_episode(self.net, device=self.get_device(batch))
            vid_tensor = torch.from_numpy(np.stack(frames,0).transpose(0,3,1,2)/255).unsqueeze(0)
            self.logger.experiment.add_video('val/episode', vid_tensor, global_step=self.global_step, fps=30)
            
        return OrderedDict({"loss": loss, "log": log, "progress_bar": status})

    def configure_optimizers(self):
        # configura l'ottimizzatore
        optimizer = Adam(self.net.parameters(), lr=self.hparams.lr)
        return [optimizer]

    def __dataloader(self):
        # crea un RLDataset e lo lega al buffer
        dataset = RLDataset(self.buffer)
        dataloader = DataLoader( #costruisci il dataloader
            dataset=dataset,
            batch_size=self.hparams.batch_size,
        )
        return dataloader

    def train_dataloader(self):
        # restituisci il dataloader
        return self.__dataloader()

    def get_device(self, batch):
        # restituisce il device del batch
        return batch[0].device.index if self.on_gpu else "cpu"

from pytorch_lightning.loggers import TensorBoardLogger

model = DQNLightning()

logger = TensorBoardLogger("tb_logs_rl", name="cartpole")

trainer = Trainer(
    gpus=1,
    max_epochs=10000, #facciamo training per 10000 step
    logger=logger,
    enable_progress_bar=False
)

trainer.fit(model)

avg_reward = 0
for _ in range(100):
    _, _, total_reward = model.agent.play_episode(model.net, return_frames=False)
    avg_reward+=total_reward
avg_reward/=100
print(avg_reward)

avg_reward = 0
for _ in range(100):
    _, _, total_reward = model.agent.play_episode(model.net, return_frames=False, epsilon=1)
    avg_reward+=total_reward
avg_reward/=100
print(avg_reward)

from pytorch_lightning.loggers import TensorBoardLogger

model = DQNLightning(env='MountainCar-v0')

logger = TensorBoardLogger("tb_logs_rl", name="mountain_car")

trainer = Trainer(
    gpus=1,
    max_epochs=10000, #facciamo training per 10000 step
    logger=logger,
    enable_progress_bar=False
)

trainer.fit(model)

from pytorch_lightning.loggers import TensorBoardLogger

model = DQNLightning(env='Acrobot-v1')

logger = TensorBoardLogger("tb_logs_rl", name="acrobot")

trainer = Trainer(
    gpus=1,
    max_epochs=10000, #facciamo training per 10000 step
    logger=logger,
    enable_progress_bar=False
)

trainer.fit(model)