import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
%matplotlib inline

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

# Hyper parametri 
input_size = 32*32    #  =1024 sono le dimensioni delle immagini
num_classes = 10      #  devo classificare immagini di numeri
num_epochs = 6        #  quanti giri completi vengono fatti
batch_size = 4        #  questo no so come sia stato scelto
learning_rate = 0.001 #  piccolo
hidden_size1  = 220   #  scelto da me
hidden_size2  = 184   #  scelto da me 

transform =  transforms.Compose( [transforms.ToTensor(), transforms.Normalize( (0.5,0.5,0.5), (0.5,0.5,0.5)  )  ] )


train_dataset = torchvision.datasets.CIFAR10(root= './data/', train= True, 
                                           transform =transform,            # uso le trasformazioni indicate sopra
                                          download = True)

test_dataset = torchvision.datasets.CIFAR10(root= './data/', train= False, 
                                           transform =transform,
                                          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)

# Pytorch lavora con numeri, qui assegno i nomi corrispondenti
#             0       1       2       3      4      5       6       7       8        9 
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

def imshow(img):
    img = img /2 + 0.5 # "denormalizza" forse per avere dei colori migliori
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1,2,0)))


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

class ConvNet(nn.Module):

    
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1  = nn.Conv2d(3, 6, 5)     # input channel size, output channel size,  Kernel size 5 (5x5)
        self.pool   = nn.MaxPool2d(2,2)      # kernel size, stride  (e' un quadrato 4x4 diviso in 4 parti 2x2)
        self.conv2  = nn.Conv2d(6, 16, 5)    # input channel size = out di prima = 6, scelgo l'out= 16 e kernel size 5 (5x5)  
        self.fc1    = nn.Linear(16*5*5, hidden_size1) # fully connected (spiegato dopo) ingresso, e 120 in uscita
        self.fc2    = nn.Linear(hidden_size1, hidden_size2)     # in ingresso prende l'uscita del precedente e in output un tensor 1D con 84 ingressi
        self.fc3    = nn.Linear(hidden_size2, 10)      # in uscita ho solo le 10 classi di oggetti (riga 38)            
        
    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))  #  MaxPool(ReLU(convoluzione(immagine))
        x = self.pool(F.relu(self.conv2(x)))  #  MaxPool(ReLU(convoluzione( ))
        x = x.view(-1, 16*5*5)                #  metto l'immagine processata in un vettore 1D
        x = F.relu(self.fc1(x))               #  ReLU(fc())
        x = F.relu(self.fc2(x))               #  ReLU(fc())
        x = self.fc3(x)                       #  qui non faccio la ReLU, perche' poi nella cross entropy c'e' softmax
        return x



model = ConvNet().to(device)                  # metto sul device    
    
criterion = nn.CrossEntropyLoss()             # serve per fare la LOSS, include la softmax per l'ultimo strato
optimizer = torch.optim.SGD(model.parameters(), lr = learning_rate)
    
n_total_steps  = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # forma [4, 3, 32, 32] -> [4, 3, 1024]       4 immagini, 3 canali di colore, 32 x 32 
        # input layer: 3 canali di input, 6 canali di uscita, 5 kernel size
        images = images.to(device)
        labels = labels.to(device)
        
        # forward
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # backward e ottimizzazione dei pesi
        optimizer.zero_grad()
        loss.backward()              # chain rule
        optimizer.step()             # step di ottimizzazione dato il learning rate
        
        if (i+1) % 1000 ==0:
            print(f'epoch {epoch+1} / {num_epochs}, step {i+1}/{n_total_steps}, loss= {loss.item():.4f}')     
print("Training Terminato")

# qui faccio la fase di testing:

with torch.no_grad():
    n_correct = 0
    n_samples = 0
    n_class_correct = [0 for i in range(10)]  # e' una lista uso una list-comprehension
    n_class_samples = [0 for i in range(10)]  # non chiaro cosa sia questo
    
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        
        _, predicted = torch.max(outputs, 1)
        n_samples += labels.size(0)      # non serve una quadra?
        n_correct += (predicted == labels).sum().item()       # questa da capire bene 
        
        for i in range (batch_size):
            label = labels[i]        # gira nel batch, label della singola immagine
            pred = predicted[i]      # guarda la riga 90
            
            if (label == pred):
                n_class_correct[label] += 1  # fa un istogramma
            n_class_samples [label] +=1      # ne ho trovsata una un piu' della classe label
    
    acc = 100.0 * n_correct / n_samples      # percentuale di corrette sul totale dei sample
    print(f'Accuracy  = {acc}:.4f')
        
    for i in range(10):
        acc = 100.0 *n_class_correct[i]/ n_class_samples[i]  # accuracy della singola classe
        print(f'Accuracy della classe {classes[i]}: {acc} %'  )
    
        

# prendo e modifico il codice di prima secondo quanto scritto nella lezione 14
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import lr_scheduler  # nuovo non ancora  caricato prima
import torchvision
#import torchvision.transforms as transforms
from torchvision import datasets, models, transforms  # 
import matplotlib.pyplot as plt
import time                      # per controllare l'orario
import os                        # per interagire con l'os
import copy                      # per fare una deep copy del modello che e' un oggetto complesso (non vogliamo shallow copy)
import sys                       # per bloccare (c'e' una funzione tipo  stop del fortran)
%matplotlib inline

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

mean = np.array([0.485, 0.456, 0.406])   # queste non so come sono state scelte
std = np.array([0.229, 0.224, 0.225])    # idem


# Hyper parametri 
input_size = 32*32    #  =1024 sono le dimensioni delle immagini
num_classes = 10      #  devo classificare immagini di numeri
num_epochs = 6        #  quanti giri completi vengono fatti
batch_size = 4        #  questo no so come sia stato scelto
learning_rate = 0.001 #  piccolo
hidden_size1  = 220   #  scelto da me
hidden_size2  = 184   #  scelto da me 


# qui faccio un dizionario che contiene le trasformazioni da applicare
data_transforms = {
    'train': transforms.Compose([transforms.RandomResizedCrop(224), 
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean,std)])
    ,
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean,std)
    ])
}

# import data
data_dir = 'data/hymenoptera_data'         # dove ci sono le api? Imenotteri
sets  = ['train', 'val']                   # lista che contiene mi sa che poi non lo usa


# Questo sotto e' da pensare un momento
# e' una list comprehension (in un dictionary)
# dove in realta' ci sono 2 directory, poteva mettere il path... e forse scriveva meno.
# in ogni caso la parte che non mi e' chiarissima sono i : dopo la x
# no e' banale: e' un dictionary e la x sono le chiavi e le datasets.Image... sono i valori
# tutti ottenuti con una list comprehension finale
image_datasets= {x:datasets.ImageFolder( os.path.join(data_dir, x), data_transforms[x] )
                 for x in ['train', 'val'] }

dataloaders  = {x: torch.utils.data.DataLoader(image_datasets[x],
                    batch_size = batch_size,
                    shuffle= True,
                    num_workers= 0) for x in  ['train', 'val']}


dataset_sizes  = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names  = image_datasets['train'].classes      # .classes e' un attributo 

print(class_names)
#sys.exit()                      # per fermare qui


######### funzione che contiene il training loop  ##############

def train_model(model, criterion, optimizer, scheduler, num_epochs=25): 
    since = time.time()     # penso chen questo sia il momento d'inizio

    best_model_wts = copy.deepcopy(model.state_dict())  # fa una deep copy del modello
    best_acc = 0.0                                      # 
    
    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs-1}')
        print('-'* 10)                             # disegna una riga orizzontale
        
        # In ogni epoca si ha una fase di training e validation:
        for phase in ['train', 'val']:
            if phase== 'train':
                model.train()              # set model to train mode  #######################
            else:
                model.eval()               # set model to evaluation mode
            
            running_loss = 0.0             # in qualche modo voglio vedere live se il training va bene
            running_corrects =0
            

            for inputs, labels in dataloaders[phase]:       # distinguo tra le fasi train/val
                inputs = inputs.to(device)
                labels = labels.to(device)
            
                ################  forward pass
                with torch.set_grad_enabled(phase == 'train'):
                    outputs =  model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs,labels)
            
                ################ backward pass
                    if phase == 'train':
                        optimizer.zero_grad()      # non mi seve quando faccio l'ottimizzazione
                        loss.backward()            # chain rule
                        optimizer.step()           # mi muovo in funzione del gradiente
                 
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds==labels.data)  
            
            if phase =='train':
                scheduler.step()    # =========================  questo aggiorna il learning_rate
            
            epoch_loss =running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            
            print(f'{phase} Loss: {epoch_loss:.4f}  Acc: {epoch_acc:.4f} ')

            # deep copy model
            
            if phase =='val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())   # questo copia il dizionario associato al miglior modello
        
        print()
    time_elapsed = time.time() -since
    print(f'Training completato in {time_elapsed//60:.0f}  minuti {time_elapsed%60:.0f}s ')
           
    model.load_state_dict(best_model_wts)    
    return model    
        
######### importiamo il MODELLO ############### 

model = models.resnet18(pretrained=True )       # qui il modello e' importato da torchvision e lo mette nella cache

# se voglio bloccare tutti i parametri tranne quelli dell'ultimo layer basta che 
# li blocco:!

for param in model.parameters():
    param.requires_grad = False     # non fa il gradiente!

num_ftrs = model.fc.in_features                 # ho preso il layer fully connected (ultimo) e queste sono le input features?

# ora creo un nuovo layer e lo assegno come ultimo layer. 
# come faccio a sapere che l'ultimo layer si chiama `fc`?

model.fc = nn.Linear(num_ftrs, 2)  # ho solo 2 classi in uscita.  DI DEFAULT ha requires_grad = True
model.to(device)                    

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001 )   # occhio che lui aveva importato in modo diverso e chiama solo optim.SGD

############ SCHEDULER che fa una update del lr #######

step_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1) # ogni 7 epoch lr =lr*gamma (diventa 1/10)

#for epoch in range(100):
#    train()   # optimizer.step()
#    evaluate()
#    scheduler.Step()
    
model = train_model(model, criterion, optimizer, step_lr_scheduler, num_epochs=20)

# qui alla fine ho trovato il modello ottimale, e potrei salvarlo volendo.

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

from torch.utils.tensorboard import SummaryWriter   ##############  TENSORBOARD
import sys 
import torch.nn.functional as F

##### costruisco un writer ####################
writer  =SummaryWriter('runs/mnist') # come argomento serve la dir dove vanno salvati i file
#####    fine writer   ########################


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

input_size = 28*28    
hidden_size  = 100    
num_classes = 10      
num_epochs = 2        
batch_size = 100      
learning_rate = 0.001 

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) # 

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)

examples = iter(test_loader)
examples_data, examples_targets = examples.next()
#for i in range(6):
#    plt.subplot(2,3,i+1)
#    plt.imshow(example_data[i][0], cmap='gray')

    
img_grid  = torchvision.utils.make_grid(examples_data)
writer.add_image('Immagini di Mnist' , img_grid)    
writer.close()     # questo assicura che tutti gli output sono flushati
#sys.exit()    # per non dover fare tutto il training    

############  MODELLO ####################
class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.l1 = nn.Linear(input_size, hidden_size)  
        self.relu = nn.ReLU()
        self.l2 = nn.Linear(hidden_size, num_classes)
       
    def forward (self, x):
        out = self.l1(x)       
        out = self.relu(out)   
        out = self.l2(out)
        return out             

############### ISTANZIO MODELLO, Back e Forw ########
model = NeuralNet(input_size, hidden_size, num_classes).to(device)
criterion  = nn.CrossEntropyLoss()   # non ha bisogno di parametri
optimizer = torch.optim.Adam(model.parameters(), lr = learning_rate)        


######  Aggiungo un altro grafo alla dashboard di Tensorboard,
###### ora uso il metodo     add_graph   (prima ho usato add_image)
writer.add_graph(model, examples_data.reshape(-1,28*28).to(device))
writer.close()
#sys.exit()

############### TRAINING LOOP ############
n_total_steps = len(train_loader)


running_loss = 0.0      # questo e' il valore aggiornato in tempo reale
running_correct = 0     # idem

for epoch in range(num_epochs):                           
   for i, (images,labels) in enumerate (train_loader):         
        images = images.reshape(-1, 28*28).to(device)     
        labels = labels.to(device)
              
        outputs = model (images)                           
        loss = criterion(outputs, labels)
                 
        optimizer.zero_grad()     
        loss.backward()            
        optimizer.step()          
        
################ da mandare a TENSORBOARD ####################        
        running_loss += loss.item()  # aggiorno il totale
        _, predicted  =torch.max(outputs.data, 1)
        running_correct += (predicted == labels).sum().item()
################ ##################### ########################        
        
        
        if (i+1) % 100 ==0:
            #print(f'epoch {epoch+1} / {num_epochs}, step {i+1}/{n_total_steps}, loss= {loss.item():.4f}')        

############## qui aggiungo alla dashboard un nuovo oggetto TENSORBOARD ######
            writer.add_scalar('Training loss', running_loss/100, epoch* n_total_steps + i)
            writer.add_scalar('Accuracy', running_correct/100, epoch* n_total_steps + i) 
            running_loss = 0.0     # riazzero
            running_correct = 0    # riazzero
            writer.close()
######################################################            
            
            

        
        
        
################ per Tensorboard ########################
labels2 = []   #occhio che aveva gia' definito labels sotto, ho messo un 2 Tensorboard
preds = []
#########################################################
        
############ TEST LOOP #################
with torch.no_grad():    
    n_correct = 0         
    n_samples = 0           
    for images, labels in test_loader:
        images = images.reshape(-1, 28*28).to(device)
        labels = labels.to(device)
        
        outputs = model(images)    
        
        _, predictions = torch.max(outputs,1)  
        n_samples += labels.shape[0]           
        n_correct += (predictions == labels).sum().item()
        
        class_predictions = [F.softmax(output, dim=0) for output in outputs]  # ho bisogno di probabilita', quindi serve softmax
        preds.append(class_predictions)
        labels2.append(predicted)   # per TENSORBOARD
        
    preds = torch.cat ([torch.stack(batch) for batch in preds ]) # tensorboard 2D oggetto 1000, 1
    labels2 = torch.cat(labels2)   # Tensorboard concateno le liste in un oggetto 1D        1000
    
    classes = range(10)   # tutte le possibili classi  0-9  Tensorboard
    for i in classes:
        labels_i = labels2 == i    # Tensorboard    non chiaro cosa fa
        preds_i = preds[:,i]      # Tensorboard 
        writer.add_pr_curve(str(i), labels_i, preds_i, global_step =0) # TEnsorboard
        writer.close()       # chiudo il writer
    
    
    acc = 100.0  * n_correct / n_samples       
    print(f'accuracy ={acc}')