%load_ext autoreload
%autoreload 2

'''step 1'''
# Load all necessary modules here, for clearness
import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
# from torchvision.datasets import MNIST
import torchvision
from torchvision import transforms
from torch.optim import lr_scheduler
# from tensorboardX import SummaryWriter
from collections import OrderedDict
from matplotlib import pyplot as plt

'''step 2'''
# Whether to put data in GPU according to GPU is available or not 
cuda = torch.cuda.is_available() 
device = torch.device("cuda:0") 
#  In case the default gpu does not have enough space, you can choose which device to use
torch.cuda.set_device(device) # device: id

# Since gpu in lab is not enough for your guys, we prefer to cpu computation
#device = torch.device("cuda:0") 

'''step 3'''
# However, in the same exp, seed for torch and numpy doesn't be the same.
SEED = 47

# Sets the seed for generating random numbers, including GPU and CPU
torch.manual_seed(SEED)

# Deterministic algorithm for convolutional ops
torch.backends.cudnn.deterministic = True
# Deterministic alogorithm for cudnn, otherwise, cuddn would choose the fastest algorithm for every 
# iteration ops, which cause variability and time consuming if input changes frequently.
torch.backends.cudnn.benchmark = False

# Seed the generator for Numpy
np.random.seed(SEED)

# Example to use kFold
from sklearn.model_selection import KFold
import numpy as np

train_transform = transforms.Compose([
    
])

dataset = torchvision.datasets.CIFAR10(root='./data',
                                      train=True, transform=train_transform,download=True)

data = dataset.train_data

# dataset.train_labels gets list object, we should transform to numpy for convinience
label = np.array(dataset.train_labels)

# set numpy random seed, we can get a determinate k-fold dataset
# np.random.seed(1)

kf = KFold(n_splits=5,shuffle=True)

for train_index, test_index in kf.split(data):
    print('train_index', train_index, 'test_index', test_index)
    train_data, train_label = data[train_index], label[train_index]
    test_data, test_label = data[test_index], label[test_index]

# here we use the last fold to be our trainset
dataset.train_data = train_data
dataset.train_labels = list(train_label)

def get_mean_std(dataset, ratio=0.01):
    """Get mean and std by sample ratio
    """
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=int(len(dataset)*ratio), 
                                             shuffle=True, num_workers=10)
    train = iter(dataloader).next()[0]
    mean = np.mean(train.numpy(), axis=(0,2,3))
    std = np.std(train.numpy(), axis=(0,2,3))
    return mean, std

# cifar10
train_dataset = torchvision.datasets.CIFAR10('./data', 
                                             train=True, download=False, 
                                             transform=transforms.ToTensor())
test_dataset = torchvision.datasets.CIFAR10('./data', 
                                           train=False, download=False, 
                                            transform=transforms.ToTensor())

train_mean, train_std = get_mean_std(train_dataset)

test_mean, test_std = get_mean_std(test_dataset)

print(train_mean, train_std)
print(test_mean,test_std)

'''step 4'''
# set hyper parameter 
batch_size = 32
n_epochs = 50
learning_rate = 1e-3

%matplotlib inline
from PIL import Image

# rotate 30°
transform_rotate = transforms.RandomRotation((30,30))

transform_horizontalflip = transforms.RandomHorizontalFlip(p=1)

transform_verticalflip = transforms.RandomVerticalFlip(p=1)

# the first image in cifar10 trainset
img = Image.open('./img/example.jpg')
plt.imshow(img)
plt.axis('off')
plt.show()

# a cat image
img1 = Image.open('./img/cat.jpeg')
plt.imshow(img1)
plt.axis('off')
plt.show()

# rotate the image
img2 = transform_rotate(img1)
plt.imshow(img2)
plt.axis('off')
plt.show()

# horizontal flip the image
img3 = transform_horizontalflip(img1)
plt.imshow(img3)
plt.axis('off')
plt.show()

# vertical flip the image
img4 = transform_verticalflip(img1)
plt.imshow(img4)
plt.axis('off')
plt.show()

'''step 5'''
'''
    the mean and variance below are from get_mean_std() function, every time you run the above function may get
    defferent value, because we use sampling
'''

'''
    notice: we usually will not use the dataset above, because its transform function has been change the numpy
    data to tensor data, but those transformations such as filp, rotation, crop, pad should be done before the
    data transforms to tensor.
'''

#transform1
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
#     transforms.RandomRotation(10),
    transforms.ToTensor(), # Convert a PIL Image or numpy.ndarray to tensor.
    # Normalize a tensor image with mean 0.1307 and standard deviation 0.3081
    transforms.Normalize((0.4924044, 0.47831464, 0.44143882), (0.25063434, 0.2492162,  0.26660094))
])

# transform2
# train_transform = transforms.Compose([
#         transforms.RandomCrop(32, padding=4),
#         transforms.RandomHorizontalFlip(),
#         transforms.ToTensor(),
#         transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])


test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.49053448, 0.47128814, 0.43724576), (0.24659537, 0.24846372, 0.26557055))
])

train_dataset = torchvision.datasets.CIFAR10(root='./data', 
                            train=True, 
                            transform=train_transform,
                            download=True)

test_dataset = torchvision.datasets.CIFAR10(root='./data', 
                           train=False, 
                           transform=test_transform,
                           download=False)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, 
                                           batch_size=batch_size, 
                                           shuffle=False)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset, 
                                          batch_size=batch_size, 
                                          shuffle=False)

resnet18 = torchvision.models.resnet18()
print(resnet18)

'''step 6'''
class ResNet18(nn.Module):
    def __init__(self):
        super(ResNet18,self).__init__()
        original_model = torchvision.models.resnet18()
        original_model.conv1.stride = 1
        self.feature_extractor = nn.Sequential(
            *(list(original_model.children())[0:3]),
            *(list(original_model.children())[4:-2]),
            nn.AdaptiveAvgPool2d(1)
        )
        self.fc = nn.Linear(512,10)
        
    def forward(self, x):
        out1 = self.feature_extractor(x)
        out1 = out1.view(out1.size(0),-1)
        out2 = self.fc(out1)
        return out2

'''step 7'''
# create a model object
# model = torchvision.models.resnet18()
# model.avgpool = nn.AdaptiveAvgPool2d(1)
# model.fc = nn.Linear(512,10)
model = ResNet18()
model.to(device)
# Cross entropy
loss_fn = torch.nn.CrossEntropyLoss()
# l2_norm can be done in SGD
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9) 

'''step 8'''
def train(train_loader, model, loss_fn, optimizer,device):
    """train model using loss_fn and optimizer. When thid function is called, model trains for one epoch.
    Args:
        train_loader: train data
        model: prediction model
        loss_fn: loss function to judge the distance between target and outputs
        optimizer: optimize the loss function
    Returns:
        total_loss: loss
    """
    
    # set the module in training model, affecting module e.g., Dropout, BatchNorm, etc.
    model.train()
    
    total_loss = 0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad() # clear gradients of all optimized torch.Tensors'
        outputs = model(data) # make predictions 
        loss = loss_fn(outputs, target) # compute loss 
        total_loss += loss.item() # accumulate every batch loss in a epoch
        loss.backward() # compute gradient of loss over parameters 
            
        optimizer.step() # update parameters with gradient descent 
            
    average_loss = total_loss / batch_idx # average loss in this epoch
    
    return average_loss

'''step 9'''
def evaluate(loader, model, loss_fn, device):
    """test model's prediction performance on loader.  
    When thid function is called, model is evaluated.
    Args:
        loader: data for evaluation
        model: prediction model
        loss_fn: loss function to judge the distance between target and outputs
    Returns:
        total_loss
        accuracy
    """
    
    # context-manager that disabled gradient computation
    with torch.no_grad():
        
        # set the module in evaluation mode
        model.eval()
        
        correct = 0.0 # account correct amount of data
        total_loss = 0  # account loss
        
        for batch_idx, (data, target) in enumerate(loader):
            data, target = data.to(device), target.to(device)
            outputs = model(data) # make predictions 
            # return the maximum value of each row of the input tensor in the 
            # given dimension dim, the second return vale is the index location
            # of each maxium value found(argmax)
            _, predicted = torch.max(outputs, 1)
            # Detach: Returns a new Tensor, detached from the current graph.
            #The result will never require gradient.
            correct += (predicted == target).cpu().sum().detach().numpy()
            loss = loss_fn(outputs, target)  # compute loss 
            total_loss += loss.item() # accumulate every batch loss in a epoch
        
        accuracy = correct*100.0 / len(loader.dataset) # accuracy in a epoch
        average_loss = total_loss / len(loader)
        
    return average_loss, accuracy

'''step 10'''
def fit(train_loader, val_loader, model, loss_fn, optimizer, n_epochs, device):
    """train and val model here, we use train_epoch to train model and 
    val_epoch to val model prediction performance
    Args: 
        train_loader: train data
        val_loader: validation data
        model: prediction model
        loss_fn: loss function to judge the distance between target and outputs
        optimizer: optimize the loss function
        n_epochs: training epochs
    Returns:
        train_accs: accuracy of train n_epochs, a list
        train_losses: loss of n_epochs, a list
    """
    
    
    train_accs = [] # save train accuracy every epoch
    train_losses = [] # save train loss every epoch
    
    test_accs = []
    test_losses = []
    
    scheduler = lr_scheduler.StepLR(optimizer,step_size=6,gamma=0.1)
    
    for epoch in range(n_epochs): # train for n_epochs 
        # train model on training datasets, optimize loss function and update model parameters
        
#         # change the learning rate at any epoch you want 
#         if n_epochs % 6 == 0 and n_epochs != 0:
#             lr = lr * 0.1
#             for param_group in optimizer.param_groups:
#                 param_groups['lr'] = lr
        
        train_loss= train(train_loader, model, loss_fn, optimizer, device=device)
        
        # evaluate model performance on train dataset
        _, train_accuracy = evaluate(train_loader, model, loss_fn, device=device)
        
        # change the learning rate by scheduler
        scheduler.step()
        
        
        message = 'Epoch: {}/{}. Train set: Average loss: {:.4f}, Accuracy: {:.4f}'.format(epoch+1, \
                                                                n_epochs, train_loss, train_accuracy)
        print(message)
    
        # save loss, accuracy
        train_accs.append(train_accuracy)
        train_losses.append(train_loss)
        show_curve(train_accs,'tranin_accs')
        show_curve(train_losses,'train_losses')
    
        # evaluate model performance on val dataset
        val_loss, val_accuracy = evaluate(val_loader, model, loss_fn, device=device)
        
        
        test_accs.append(val_accuracy)
        test_losses.append(val_loss)
        show_curve(test_accs,'test_accs')
        show_curve(test_losses,'test_losses')
        
        
        message = 'Epoch: {}/{}. Validation set: Average loss: {:.4f}, Accuracy: {:.4f}'.format(epoch+1, \
                                                                n_epochs, val_loss, val_accuracy)
        print(message)
            
    return train_accs, train_losses

'''step 10'''
def show_curve(ys, title):
    """plot curlve for Loss and Accuacy
    
    !!YOU CAN READ THIS LATER, if you are interested
    
    Args:
        ys: loss or acc list
        title: Loss or Accuracy
    """
    x = np.array(range(len(ys)))
    y = np.array(ys)
    plt.plot(x, y, c='b')
    plt.axis()
    plt.title('{} Curve:'.format(title))
    plt.xlabel('Epoch')
    plt.ylabel('{} Value'.format(title))
    plt.show()

'''step 12: without tuning lr'''
train_accs, train_losses = fit(train_loader, test_loader, model, loss_fn, optimizer, n_epochs, device=device)

#         # change the learning rate at any epoch you want 
#         if n_epochs % 6 == 0 and n_epochs != 0:
#             lr = lr * 0.1
#             for param_group in optimizer.param_groups:
#                 param_groups['lr'] = lr

# to save the parameters of a model
torch.save(model.state_dict(), './params/resnet18_params.pt')

# to save the model
torch.save(model, './params/resnet18.pt')

model1 = torchvision.models.resnet18()
print(model1.state_dict())

# generate a parameters file of model1
torch.save(model1.state_dict(),'./params/resnet18_params.pt')