# CNN 
import torch.nn.functional as F
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.backends.cudnn as cudnn

# others
import numpy as np
import matplotlib.pyplot as plt
import time
import os
from PIL import Image
from tempfile import TemporaryDirectory
import time

# dataset 
import torchvision
from torchvision import datasets, models, transforms
from torchvision.datasets import Flowers102

# label 
from scipy.io import loadmat
import json

cudnn.benchmark = True
plt.ion()   # interactive mode

# Specify the data you want to download, the path and btach size, and the amount of training to do at once.
batch_size = 4
data_dir = '... /... /Data/flowers-102'

# Build the classes_name of the dataset
json_data = '{"21": "fire lily", "3": "canterbury bells", "45": "bolero deep blue", "1": "pink primrose", "34": "mexican aster", "27": "prince of wales feathers", "7": "moon orchid", "16": "globe-flower", "25": "grape hyacinth", "26": "corn poppy", "79": "toad lily", "39": "siam tulip", "24": "red ginger", "67": "spring crocus", "35": "alpine sea holly", "32": "garden phlox", "10": "globe thistle", "6": "tiger lily", "93": "ball moss", "33": "love in the mist", "9": "monkshood", "102": "blackberry lily", "14": "spear thistle", "19": "balloon flower", "100": "blanket flower", "13": "king protea", "49": "oxeye daisy", "15": "yellow iris", "61": "cautleya spicata", "31": "carnation", "64": "silverbush", "68": "bearded iris", "63": "black-eyed susan", "69": "windflower", "62": "japanese anemone", "20": "giant white arum lily", "38": "great masterwort", "4": "sweet pea", "86": "tree mallow", "101": "trumpet creeper", "42": "daffodil", "22": "pincushion flower", "2": "hard-leaved pocket orchid", "54": "sunflower", "66": "osteospermum", "70": "tree poppy", "85": "desert-rose", "99": "bromelia", "87": "magnolia", "5": "english marigold", "92": "bee balm", "28": "stemless gentian", "97": "mallow", "57": "gaura", "40": "lenten rose", "47": "marigold", "59": "orange dahlia", "48": "buttercup", "55": "pelargonium", "36": "ruby-lipped cattleya", "91": "hippeastrum", "29": "artichoke", "71": "gazania", "90": "canna lily", "18": "peruvian lily", "98": "mexican petunia", "8": "bird of paradise", "30": "sweet william", "17": "purple coneflower", "52": "wild pansy", "84": "columbine", "12": "colt\'s foot", "11": "snapdragon", "96": "camellia", "23": "fritillary", "50": "common dandelion", "44": "poinsettia", "53": "primula", "72": "azalea", "65": "californian poppy", "80": "anthurium", "76": "morning glory", "37": "cape flower", "56": "bishop of llandaff", "60": "pink-yellow dahlia", "82": "clematis", "58": "geranium", "75": "thorn apple", "41": "barbeton daisy", "95": "bougainvillea", "43": "sword lily", "83": "hibiscus", "78": "lotus lotus", "88": "cyclamen", "94": "foxglove", "81": "frangipani", "74": "rose", "89": "watercress", "73": "water lily", "46": "wallflower", "77": "passion flower", "51": "petunia"}'
# load data 
cat_to_name = json.loads(json_data)
# Turn the key into an int, because the label of the dataset starts from 0. But this json starts from 1, so we have to -1 
cat_to_name = {int(k)-1:v for k,v in cat_to_name.items()}
# sort and print 
class_names = dict(sorted(cat_to_name.items()))
print(class_names)

# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
    'train': transforms.Compose([
        # First, the image is randomly cropped and then resized. A random rectangular area is chosen and the image is cropped.
        # The cropped image is then resized to the specified size of 224x224 pixels.
        transforms.RandomResizedCrop(224),
        # Set the probability of image flipping, usually a number from 0 to 1, for example, 0.5, which means there's a 50% chance of flipping the image. Default value is 0.5
        transforms.RandomHorizontalFlip(),
        # Convert the image into a Tensor
        transforms.ToTensor(),
        # Normalize the image values using numerical normalization, where the first parameter is mean, and the second parameter is the standard deviation (std)
        # The reason for setting [0.485, 0.456, 0.406] can be referred from: https://www.geeksforgeeks.org/how-to-normalize-images-in-pytorch/
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        # This doesn't randomly select an area but directly resizes the entire image to fit the specified size.
        transforms.Resize(256),
        # Keep the central part of the image, then resize to meet the specified size.
        # Used for validation or test data to ensure that the test images have similar features, and don't have the randomness like RandomResizedCrop.
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}
# We download the training data into the data_dir/train folder, and use the data_transforms["train"] function for data transformation.
train_datasets = Flowers102(root=data_dir+"/train", split="train", download=True, transform=data_transforms["train"])
# We download the validation data into the data_dir/val folder, and use the data_transforms["val"] function for data transformation.
val_datasets = Flowers102(root=data_dir+"/val", split="val", download=True, transform=data_transforms["val"])

# Specify to download the flowers102 dataset, downloading both train and val datasets.
image_datasets = {x: Flowers102(root=data_dir, split=x, download=True, transform=data_transforms[x])
                    for x in ['train', 'val']}

# Convert to DataLoader format, and specify batch_size
dataloaders = {x: torch.utils.data.DataLoader(image_datasets[x], batch_size=batch_size,
                                             shuffle=True, num_workers=4)
              for x in ['train', 'val']}

device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

print("device: ",device)
print("image_datasets function call: ", dir(image_datasets["train"]))

def imshow(inp, title=None):
    """Display image for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    inp = np.clip(inp, 0, 1)
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # pause a bit so that plots are updated


# Get a batch of training data
inputs, classes = next(iter(dataloaders['train']))

# Make a grid from batch
out = torchvision.utils.make_grid(inputs)

# x.item() takes the value of the tensor, usually a number, and finds the English equivalent of that number in the class_names dic
imshow(out, title=[class_names[x.item()] for x in classes])

print(inputs.shape)
print("dataset_sizes: ",dataset_sizes)

import torch.nn as nn

class NewNet(nn.Module):
    def __init__(self):
        super(NewNet, self).__init__()
        # Layer 1: 3x3 kernel, depth = 32, 224-3+1=222 => 222x222 pixel
        self.conv1 = nn.Conv2d(3, 32, 3)
        self.bn1 = nn.BatchNorm2d(32)
        # self.dropout1 = nn.Dropout(0.5) # Apply dropout as needed

        # Layer 2: Max pooling with 2x2 kernel, 222/2=111 => 111x111 pixel
        self.pool = nn.MaxPool2d(2, 2)

        # Layer 3: 3x3 kernel, depth = 64, 111-3+1=109 => 109x109 pixel
        self.conv2 = nn.Conv2d(32, 64, 3)
        self.bn2 = nn.BatchNorm2d(64)
        # self.dropout2 = nn.Dropout(0.5) # Apply dropout as needed

        # Layer 4: Max pooling with 2x2 kernel, 109/2=54 => 54x54 pixel
        self.pool = nn.MaxPool2d(2, 2)

        # Layer 5: 3x3 kernel, depth = 128, 54-3+1=52 => 52x52 pixel
        self.conv3 = nn.Conv2d(64, 128, 3)
        self.bn3 = nn.BatchNorm2d(128)
        # self.dropout3 = nn.Dropout(0.5) # Apply dropout as needed

        # Layer 6: Max pooling with 2x2 kernel, 52/2=26 => 26x26 pixel
        self.pool = nn.MaxPool2d(2, 2)
        
        # Final input is 512, pixel is 26*26 => 128*26*26
        self.fc1 = nn.Linear(128 * 26 * 26, 2048) # <==== Adjust according to the hidden layer, 128 * 26 * 26
        self.fc2 = nn.Linear(2048, 1024)
        self.fc3 = nn.Linear(1024, 512)
        self.fc4 = nn.Linear(512, 102) # <==== 102 according to the number of classes in the dataset

    def forward(self, x):
        # We put the batch normalization before the activation function. 
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.pool(x)
        # x = self.dropout1(x) # Apply dropout as needed
        x = F.relu(self.bn2(self.conv2(x)))
        x = self.pool(x)
        # x = self.dropout2(x) # Apply dropout as needed
        x = F.relu(self.bn3(self.conv3(x)))
        x = self.pool(x)
        # x = self.dropout3(x) # Apply dropout as needed
        x = x.view(-1, 128 * 26 * 26) # <==== Adjust according to the hidden layer, 128 * 26 * 26
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        return x

net = NewNet()
net.to(device)

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

def train(epoch, start_time):
    net.train()
    cur_count = 0 
    running_loss = 0.0
    for batch_idx, data in enumerate(dataloaders["train"], 0):
        cur_count += len(data)
        inputs, labels = data[0].to(device), data[1].to(device)
        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs)
        outputs.to(device)
        
        loss = criterion(outputs, labels)
        loss.to(device)
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        if batch_idx % 100 == 99:
            print(f'[{epoch}, {batch_idx + 1:5d}] loss: {running_loss / 100:.3f} time elapsed: {round((time.time() - start_time))} sec.')
            running_loss = 0.0

def test(): 
    net.eval()  # set model to evaluation mode
    correct = 0
    total = 0
    class_correct = [0] * len(class_names)  
    class_total = [0] * len(class_names)  
    with torch.no_grad():
        for data in dataloaders["val"]:
            images, labels = data[0].to(device), data[1].to(device)
            outputs = net(images) 

            # select top 3 predictions
            _, predicted = torch.topk(outputs, 1, dim=1)
            
            # check if predicted labels are in true labels
            for i in range(len(labels)):
                total += 1
                class_total[labels[i]] += 1

                if labels[i] in predicted[i]:
                    correct += 1
                    class_correct[labels[i]] += 1

    class_accuracies = [class_correct[i] / class_total[i] for i in range(len(class_names))]
    accuracy = correct / total
    return accuracy, class_accuracies

num_epochs = 100 
start_time = time.time()

accuracy, class_accuracies = test()
print(f'Accuracy on test data (top-1): {100 * accuracy:.2f}%')
for epoch in range(0, num_epochs - 1):
    print(f"============ Epoch: {epoch} ==========")
    train(epoch, start_time)
    # every 5 epoch is completed, we will perform validation on the test set 
    if epoch % 5 == 0:
        accuracy, class_accuracies = test()
        # print accuracies
        print(f'Accuracy on test data (top-1): {100 * accuracy:.2f}%')

print(f'Finished Training. Total elapsed time: {round((time.time() - start_time) / 60, 1)} min')

# Set the start time for logging to monitor the training time for each epoch
since = time.time()

# Create a temporary folder to store the best model parameters
with TemporaryDirectory() as tempdir:
    best_model_params_path = os.path.join(tempdir, 'best_model_params.pt')
    # Haven't started training yet, but we first save the current model
    torch.save(model.state_dict(), best_model_params_path)
    best_acc = 0.0  # Set the current best accuracy to 0, it will be updated if a higher value is found to determine the best model

    for epoch in range(num_epochs):
        print(f'Epoch {epoch}/{num_epochs - 1}')
        print('-' * 10)

        # After training each epoch, proceed with validation
        for phase in ['train', 'val']:
            # Determine whether it's training or validation phase
            if phase == 'train':
                model.train()  # Set model to training mode
            else:
                model.eval()   # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                # Move to GPU
                inputs = inputs.to(device)
                labels = labels.to(device)

                # Zero the parameter gradients
                optimizer.zero_grad()

                # Forward propagation
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1) # Select the number with the highest prediction as the label
                    loss = criterion(outputs, labels) # Calculate the difference between the answer and prediction

                    # backward + optimize only if in training phase
                    # Perform backward propagation
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                # Since batch_size is 4, multiply loss by 4 to get the loss for a batch
                running_loss += loss.item() * inputs.size(0)
                # Calculate how many are correct in a batch
                running_corrects += torch.sum(preds == labels.data)

            # Adjust the learning rate only during training
            # scheduler is a learning rate (lr) adjuster used to modify the lr value during model training
            if phase == 'train':
                scheduler.step()

            # After an entire epoch of training, calculate the loss and accuracy for that epoch
            # Avg. loss = total loss / size of the entire dataset
            epoch_loss = running_loss / dataset_sizes[phase]
            # Avg. Acc = total number of correct answers / size of the entire dataset
            epoch_acc = running_corrects.float() / dataset_sizes[phase]

            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f} Time elapsed: {round((time.time() - since))} sec.')

            # If during validation, and accuracy is found to be better than the current best, save the model parameters
            if phase == 'val' and epoch_acc > best_acc:
                # Update the current best accuracy
                best_acc = epoch_acc
                # Deep copy the model
                torch.save(model.state_dict(), best_model_params_path)

        print()

    time_elapsed = time.time() - since
    print(f'Training complete in {time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s')
    print(f'Best val Acc: {best_acc:4f}')

    # Load the best model weights
    # Take out the best model so far to continue with the next epoch of training
    model.load_state_dict(torch.load(best_model_params_path))
return model

model_ft = models.resnet18(weights='IMAGENET1K_V1')

# num_ftrs is the number of input features for the last layer. 
# Retrieve the number of input features for the last layer
num_ftrs = model_ft.fc.in_features

# Here the size of each output sample is set to 102.
# model_ft.fc is the final layer of the model, used for classification.
# Creating the final layer ourselves, setting the input number as num_ftrs, and output number as 102 (since there are 102 classes in this case)
model_ft.fc = nn.Linear(num_ftrs, 102)
model_ft = model_ft.to(device)

# The loss function uses CrossEntropyLoss 
criterion = nn.CrossEntropyLoss()

# Observe that all parameters are being optimized
# Optimizer uses SGD with learning rate = 0.001, momentum = 0.9
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
# Multiply the learning rate by 0.1 every 7 epochs to decay the lr
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                       num_epochs=25)

model_conv = torchvision.models.resnet18(weights='IMAGENET1K_V1')

# !!! Add these two lines to set requires_grad to False, so that the parameters will not be updated 
for param in model_conv.parameters():
    param.requires_grad = False

num_ftrs = model_conv.fc.in_features
model_conv.fc = nn.Linear(num_ftrs, 102)
model_conv = model_conv.to(device)
criterion = nn.CrossEntropyLoss()
optimizer_conv = optim.SGD(model_conv.fc.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_conv, step_size=7, gamma=0.1)

# We start to train 
model_conv_SGD = train_model(model_conv, criterion, optimizer_conv,
                         exp_lr_scheduler, num_epochs=25)

from efficientnet_pytorch import EfficientNet

# Load the pre-trained EfficientNet-B5 model
model_ft = EfficientNet.from_pretrained('efficientnet-b5')

# We obtain the number of input features for the last layer
num_ftrs = model_ft._fc.in_features
# build a new layer, the input number is num_ftrs, and the output number is 102 (because there are 102 classes in this case)
model_ft.fc = nn.Linear(num_ftrs, 102)
# Put the model on the GPU
model_ft = model_ft.to(device)

# Loss function
criterion = nn.CrossEntropyLoss()

# Observe that all parameters are being optimized
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

# Decay LR by a factor of 0.1 every 7 epochs
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

# Start training 
model_ft_effb5 = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,
                       num_epochs=25)