import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import time

import torch.nn.functional as F 

import pytorch_lightning as pl # PL
from pytorch_lightning import Trainer 

%matplotlib inline

# device config
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# Hyper parametri 
input_size = 28*28      #  =784 sono le dimensioni delle immagini
hidden_size  = 500      #  scelto da me
num_classes = 10        #  devo classificare immagini di numeri
num_epochs = 2          #  quanti giri completi vengono fatti
batch_size = 100        #  questo non so come sia stato scelto
learning_rate = 0.001   #  piccolo

#train_dataset = torchvision.datasets.MNIST(root= './data/', train= True, transform =transforms.ToTensor(), download = True)
#test_dataset = torchvision.datasets.MNIST(root= './data/', train= False, transform =transforms.ToTensor(), download = False) # ho gia' scaricato tutto con il train_datast

#train_loader = torch.utils.data.DataLoader(dataset= train_dataset, batch_size = batch_size, shuffle=True)
#test_loader = torch.utils.data.DataLoader (dataset=  test_dataset, batch_size = batch_size, shuffle=False)


############  MODELLO ####################


#class NeuralNet(nn.Module):             #vecchi
class LitNeuralNet(pl.LightningModule):  # nome scelto da noi, pl.LightningModule e' la versione super di nn.Module    
    def __init__(self, input_size, hidden_size, num_classes):
        super(LitNeuralNet, self).__init__()
        self.input_size = input_size
        self.l1 = nn.Linear(input_size, hidden_size)  
        self.relu = nn.ReLU()
        self.l2 = nn.Linear(hidden_size, num_classes)
       
    def forward (self, x):    # x e' l'input.
        out = self.l1(x)      # 
        out = self.relu(out)  # 
        out = self.l2(out)
        return out            # 

    
    def training_step(self, batch, batch_idx):  # PL non uso piu' i loop!!!! 
        #x,y = batch           # unpack
        images, labels  =batch # PL unpack
        images = images.reshape(-1,28*28)   # non e' piu' necessairo to(device)
        
        #y_hats = self(x)     # ????
        outputs  = self(images)  # PL questo e' il forward pass! usiamo self perche' usiamo direttamente questo modulo
        
        #loss = F.cross_entropy(y_hat, y)
        loss =F.cross_entropy(outputs, labels)
        
        tensorboard_logs = {'train_loss':loss}    # devo metterlo anche nel validatio_epoch_end
        return {'loss':loss, 'log':tensorboard_logs}   # PL i nomi delle chiavi sono FISSI: e' log non Log
        #return {'loss': loss}
    
    
    
    def configure_optimizers(self):  # PL ma il nome e' fissato da PL? penso di si'
        #optimizer = torch.optim.Adam(self.parameters(), lr = learning_rate)  # non ho capito se serve o no...   
        return torch.optim.Adam(self.parameters(), lr= 0.001)  #  PL self e' l'istanza del modello
    
    def train_dataloader(self):
        #dataset = MNIST(os.getcwd(), train=True, download=True, transform =transforms.ToTensor())
        #loader = DataLoader (dataset, batch_size= 32, num_workers=1, shuffle=True)   #lui ha messo num_workers = 4
        train_dataset = torchvision.datasets.MNIST(root= './data/', train= True, transform =transforms.ToTensor(), download = True)
        train_loader = torch.utils.data.DataLoader(dataset= train_dataset, batch_size = batch_size, num_workers=4, shuffle=True)
        return train_loader
    
    # il nome deve essere questo!
    def val_dataloader(self):
        val_dataset = torchvision.datasets.MNIST(root= './data/', train= False, transform =transforms.ToTensor(), download = False) # ho gia' scaricato tutto con il train_datast
        val_loader = torch.utils.data.DataLoader (dataset= val_dataset, batch_size = batch_size, num_workers=4, shuffle=False)
        return val_loader    
    
    
    # questo viene eseguito dopo ogni epoch di validazione
    def validation_epoch_end(self, outputs):
        avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
     
        tensorboard_logs  = {'tavg_val_loss': avg_loss} # PL per ora non lo guardiamo
        return {'tavg_val_loss': avg_loss, 'log':tensorboard_logs}
        #return {'val_loss': avg_loss}    
    

start = time.time()
#trainer = Trainer(max_epochs = num_epochs, fast_dev_run = False)  # PL Trainer e' importato da PL
trainer = Trainer(gpus=1,max_epochs = num_epochs, fast_dev_run = False)  # PL Trainer e' importato da PL

model = LitNeuralNet(input_size, hidden_size, num_classes)        # PL istanzio il modello  
trainer.fit(model)                                                # faccio il training sul modello
print("Tempo training+test= ", time.time()-start)