# Package Imports
import os
import numpy as np
import random
import torch
import torchvision
from torch.utils.data.dataset import Dataset
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
import matplotlib.pyplot as plt
import numpy as np
from torch import nn
import torchvision.models as models
from torch import optim
from PIL import Image
import json
import os
from torch.nn import NLLLoss
from torch import optim
import time
import pandas as pd


use_gpu = True

%matplotlib inline
l = plt.imread("data/the_little_mermaid/valid/1/scene60481.png")
r = plt.imread("data/the_little_mermaid/train/0/scene81721.png")
fig, ax = plt.subplots(1,2,figsize=(20,10))

ax[0].imshow(l)
ax[0].axis("off")
ax[1].imshow(r)
ax[1].axis("off")

film = "the_little_mermaid"
film_dir = os.path.join("data", film)

def train_test_split(data_dir):
    train_dir = data_dir + '/train'
    valid_dir = data_dir + '/valid'
    test_dir = data_dir + '/test'
    
    files = np.asarray(os.listdir(data_dir))
    if not os.path.isdir(train_dir):
        os.mkdir(train_dir)

    if not os.path.isdir(valid_dir):
        os.mkdir(valid_dir)

    if not os.path.isdir(test_dir):
        os.mkdir(test_dir)

    for file in files:    
        if file.endswith(".png"):
            d_ten = random.randint(1, 4)
            if d_ten == 1:
                # put it in the training folder
                os.rename(os.path.join(data_dir, file), os.path.join(train_dir, file))

            if d_ten == 2:
                # put it in the validation folder
                os.rename(os.path.join(data_dir, file), os.path.join(valid_dir, file))

            if d_ten > 2:
                os.rename(os.path.join(data_dir, file), os.path.join(test_dir, file))

train_test_split(film_dir)

def load_data(film_dir):
    train_dir = film_dir + '/train'
    valid_dir = film_dir +'/valid'

    xforms = {
        'train': transforms.Compose([
            transforms.RandomRotation(30),
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])
        ]),
        'validate': transforms.Compose([
            transforms.Resize(255),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406],
                                 [0.229, 0.224, 0.225])
        ])
    }

    # Data loading

    image_datasets = {
        'train': ImageFolder(train_dir, transform=xforms['train']),
        'validate': ImageFolder(valid_dir, transform=xforms['validate']),
    }


    # Data batching
    dataloaders = {
        'train': DataLoader(image_datasets['train'], batch_size=64, shuffle=True),
        'validate': DataLoader(image_datasets['validate'], batch_size=64, shuffle=True),
    }
    
    return dataloaders, image_datasets

dataloaders, image_datasets = load_data(film_dir)

def map_labels(data_dir):

    with open(os.path.join(film_dir, 'class_dict.json'), 'r') as f:
        class_dict = json.load(f)

    return class_dict

class_dict = map_labels(film_dir)

num_of_classes = len(next(os.walk(film_dir + '/train'))[1])

num_of_classes

model = models.vgg19(pretrained=True)

# this freezes the model's features so we are only training the classifier
for param in model.parameters():
    param.requires_grad = False

input_features = 25088
hidden_units = 4096
out_features = 2

model.classifier = nn.Sequential(
    nn.Linear(input_features, hidden_units),
    nn.ReLU(inplace=True),
    nn.Dropout(p=0.5, inplace=False),
    nn.Linear(hidden_units, hidden_units),
    nn.ReLU(inplace=True),
    nn.Dropout(p=0.5, inplace=False),
    nn.Linear(hidden_units, out_features),
    nn.LogSoftmax(dim=1)
)

# learning rate
alpha = 0.0001

optimizer = optim.Adam(model.classifier.parameters(), lr=alpha)

criterion = NLLLoss()

train_losses = []
validate_losses = []

def train(epochs):
    
    model.train()
    # activate the gpu
    if use_gpu:
        model.cuda()

    steps = 0
    t0 = time.time()

    for e in range (epochs):

        running_loss = 0
        
        for X, y in dataloaders['train']:
            steps += 1
            
            # move features to gpu
            if use_gpu:
                X, y = X.to('cuda'), y.to('cuda')

            optimizer.zero_grad()

            output = model.forward(X)
            loss = criterion(output, y)
            loss.backward()
            optimizer.step()
            
            # a timer
            running_loss += loss.item()
            t2 = time.time()
            percent_complete = steps/ (len(dataloaders['train']) * epochs)
            time_elapsed = t2-t0
            estimated_time_left = time_elapsed/percent_complete - time_elapsed

            print("{}% complete. {}m{}s elapsed. {}m{}s left.".format(round(percent_complete * 100), 
                                                              *divmod(int(time_elapsed), 60), 
                                                              *divmod(int(estimated_time_left), 60)), end='\r')
        else:
            # Validation Loss and Accuracy
            validate_loss = 0
            
            global accuracy
            accuracy = 0

            with torch.no_grad():
                model.eval()
                for X, y in dataloaders['validate']:

                    if use_gpu:
                        X, y = X.to('cuda'), y.to('cuda')

                    log_probabilities = model.forward(X)
                    validate_loss += criterion(log_probabilities, y).item()

                    # we can convert the log probabilities to probabilities using torch.exp
                    # torch.exp is e^x
                    probabilities = torch.exp(log_probabilities)

                    # returns true if the label is the highest probablility class predicted by the model 
                    equality = (y.data == probabilities.max(dim=1)[1])

                    # .type?
                    accuracy += equality.type(torch.FloatTensor).mean()

            model.train()

            train_losses.append(running_loss/len(dataloaders['train']))
            validate_losses.append(validate_loss/len(dataloaders['validate']))

            print("Epoch: {}/{}..".format(e+1, epochs),
                  "Training Loss: {:.3f}.. ".format(running_loss/len(dataloaders['train'])),
                  "Validate Loss: {:.3f}.. ".format(validate_loss/len(dataloaders['validate'])),
                  "Validate Accuracy: {:.3f}".format(accuracy/len(dataloaders['validate'])))
        

train(10)

%matplotlib inline
fig, ax = plt.subplots(facecolor=('1'))
ax.set_facecolor('1')
ax.set_title('Loss vs time', color='0')
ax.set_xlabel('epochs', color='1')
ax.set_ylabel('average loss', color='1')
ax.plot(train_losses, '#009d8c', label="training losses")
ax.plot(validate_losses, '#d94f45', label="validate losses")
ax.legend(frameon=True)
ax.tick_params(labelcolor='0')

def save(path):
    model.class_to_idx = image_datasets['train'].class_to_idx

    checkpoint = {'model_state_dict': model.state_dict(),
                  'optimizer_state_dict': optimizer.state_dict(),
                  'classifier': model.classifier,
                  'class_to_idx': model.class_to_idx,
                  'alpha': alpha
                 }

    torch.save(checkpoint, path)

# MAKE SURE THIS IS COMMENTED OUT IF YOU AREN'T TRAINING ARIEL OR YOU WILL WIPE HER CHECKPOINT

save(os.path.join(film_dir,'checkpoint.pth'))

# Image Processing

def process_image(image):

    image = Image.open(image)
    
    # Process a PIL image for use in a PyTorch model
    xform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    return xform(image)

def imshow(image, ax=None, title=None):

    if ax is None:
        fig, ax = plt.subplots()
    
    # PyTorch tensors assume the color channel is the first dimension
    # but matplotlib assumes is the third dimension
    image = image.numpy().transpose((1, 2, 0))
    
    # Undo preprocessing
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = std * image + mean
    
    # Image needs to be clipped between 0 and 1 or it looks like noise when displayed
    image = np.clip(image, 0, 1)
    ax.imshow(image)
    ax.axis('off')
    return ax

imshow(process_image("data/the_little_mermaid/valid/1/scene47521.png"))

def load_from_checkpoint(checkpoint_path):

    if use_gpu:
        device = torch.device("cuda")
        checkpoint = torch.load(checkpoint_path, map_location="cuda:0")
    else:
        device = torch.device("cpu")
        checkpoint = torch.load(checkpoint_, map_location=device)
    
    # reconstruct model and custom classifier from checkpoint
    model = models.vgg19(pretrained=True)

    # freeze vgg features
    # https://stackoverflow.com/questions/51748138/pytorch-how-to-set-requires-grad-false
    for param in model.parameters():
        param.requires_grad = False

    model.classifier = checkpoint['classifier']

    model.load_state_dict(checkpoint['model_state_dict'])

    model.to(device)

    # reconstruct optimizer from checkpoint
    optimizer = optim.Adam(model.classifier.parameters(), lr=checkpoint['alpha'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    
    # load {class : index} dictionary
    model.class_to_idx = checkpoint['class_to_idx']
    
    # invert dictionary for {index : class}
    # https://therenegadecoder.com/code/how-to-invert-a-dictionary-in-python/
    return model, optimizer

def predict(image_path, model, topk=1):
    idx_to_class = dict(zip(model.class_to_idx.values(), model.class_to_idx.keys()))
   
    
    image = process_image(image_path)
    image = image.unsqueeze_(0)
    
    # put model into eval mode for predictions
    model.eval()
    
    # Implement the code to predict the class from an image file
    with torch.no_grad():
        if use_gpu:
            log_probabilities = model.forward(image.cuda())
        else:
            log_probabilities = model.forward(image)
        
    probabilities = torch.exp(log_probabilities)
    topks = probabilities.topk(topk)
    idxs = topks.indices.tolist()[0]
    predicted_classes = [idx_to_class[i] for i in idxs]
        
    return topks.values, predicted_classes

cp = os.path.join(film_dir, "checkpoint.pth")
model, optimizer = load_from_checkpoint(cp)
P, classes = predict("data/the_little_mermaid/valid/1/scene47521.png", model, topk=1)

P

classes

def test_sort(film_dir, cutoff=0.5):
    
    checkpoint_path = os.path.join(film_dir, "checkpoint.pth")
    
    test_dir = film_dir + '/test'
    if not os.path.isdir(test_dir):
        os.mkdir(test_dir)
    files = np.asarray(os.listdir(test_dir))

    if not os.path.isdir(test_dir + '/1'):
        os.mkdir(os.path.join(test_dir + '/1'))

    if not os.path.isdir(test_dir + '/0'):
        os.mkdir(test_dir + '/0')

    model, optimizer = load_from_checkpoint(checkpoint_path)

    for file in files:
        if file.endswith(".png"):
            P, classes = predict(os.path.join(test_dir, file), model, topk=1)

            if classes[0] == '1' and P[0] > cutoff:
                os.rename(os.path.join(test_dir, file), os.path.join(test_dir + '/1', file))
            else:
                os.rename(os.path.join(test_dir, file), os.path.join(test_dir + '/0', file))

films = ["snow_white", "cinderella", "sleeping_beauty", "the_little_mermaid", "beauty_and_the_beast", "aladdin", "pocahontas", "mulan", "the_princess_and_the_frog", "tangled", "brave", "frozen"]

def split_all_film_data():
    for film in films:
        path = os.path.join('data', film)
        train_test_split(path)

#split_all_film_data()

d = {'train_pos': [31, 40, 28, 68, 49, 36, 61, 66, 70, 96, 76, 46], 
     'train_neg': [164, 176, 147, 155, 156, 196, 143, 169, 183, 201, 171, 219],
     'valid_pos': [43, 40, 37, 61, 49, 36, 61, 66, 70, 97, 87, 60],
     'valid_neg': [156, 186, 187, 135, 184, 230, 151, 145, 210, 168, 177, 261],
     'color1': ["#0058a2", "#f8aadc", "#f3ae47", "#d94f45", "#78492a", "#f9d938", "#00bce1", "#28346e", "#8fd18e", "#be84c6", "#f75620", "#fff6db"],
     'color2': ["#fff28c", "#7eafd0", "#bf4574", "#009d8c", "#ffda4d", "#45b5a9", "#d88745", "#89b74b", "#dceaa7", "#ecec96", "#025467", "#a9e4f9"]
    }

df = pd.DataFrame(data=d, index=films)

df

# Sanity Check

def sanity_check(path, top_k):

    P, classes = predict(path, model, topk=top_k)
    image = process_image(path)

    if use_gpu:
        probs = P.cpu().numpy()[0]
    else:
        probs = P.numpy()[0]

    names = [class_dict[i] for i in classes]
    
    colors = [df['color2'][int(i)] for i in classes]

    fig, axs = plt.subplots(2, figsize=(4, 8), facecolor=('1'))
    axs[0].set_title(names[0], color='0')
    axs[1].set_facecolor('1')
    imshow(process_image(path),axs[0])
    y_pos = np.arange(len(names))
    plt.yticks(y_pos, names)
    bar_graph = axs[1].barh(y_pos, probs, color=colors)
    plt.gca().invert_yaxis()
    axs[1].tick_params(labelcolor='0')

sanity_check("data/the_little_mermaid/valid/1/scene47521.png", 2)

def sample_dist_plot(pos, neg, title):
    ind = np.arange(len(pos))    # the x locations for the groups
    width = 0.6    # the width of the bars: can also be len(x) sequence

    p1 = plt.barh(ind, pos, width, color=df['color1'])
    p2 = plt.barh(ind, neg, width,
                 left=pos, color=df['color2'])
    plt.xlabel('Positive and negative samples by film')
    plt.title(title)
    plt.yticks(ind, films)
    plt.xticks(np.arange(0, 350, 25))

    plt.show()

sample_dist_plot(df['train_pos'], df['train_neg'], "Training set sample distribution")

sample_dist_plot(df['valid_pos'], df['valid_neg'], "Validation set sample distribution")

for film in films:

    model.classifier = nn.Sequential(
        nn.Linear(input_features, hidden_units),
        nn.ReLU(inplace=True),
        nn.Dropout(p=0.5, inplace=False),
        nn.Linear(hidden_units, hidden_units),
        nn.ReLU(inplace=True),
        nn.Dropout(p=0.5, inplace=False),
        nn.Linear(hidden_units, out_features),
        nn.LogSoftmax(dim=1)
    )

    optimizer = optim.Adam(model.classifier.parameters(), lr=alpha)

    train_losses = []
    validate_losses = []

    # set path to film
    film_dir = os.path.join('data', film)

    # load the data
    dataloaders, image_datasets = load_data(film_dir)

    # map labels
    class_dict = map_labels(film_dir)

    # train the model
    train(10)
    
    # save the checkpoint
    save(os.path.join(film_dir,'checkpoint.pth'))

    #set checkpoint_path

    test_sort(film_dir, cutoff=0.6)

out_features = 12

model.classifier = nn.Sequential(
        nn.Linear(input_features, hidden_units),
        nn.ReLU(inplace=True),
        nn.Dropout(p=0.5, inplace=False),
        nn.Linear(hidden_units, hidden_units),
        nn.ReLU(inplace=True),
        nn.Dropout(p=0.5, inplace=False),
        nn.Linear(hidden_units, out_features),
        nn.LogSoftmax(dim=1)
)

optimizer = optim.Adam(model.classifier.parameters(), lr=alpha)

train_losses = []
validate_losses = []

# set path to film
film_dir = os.path.join('data', 'multi_class')

# load the data
dataloaders, image_datasets = load_data(film_dir)

# map labels
class_dict = map_labels(film_dir)

# train the model
train(20)
    
# save the checkpoint
save(os.path.join(film_dir,'checkpoint.pth'))

sanity_check("data/multi_class/valid/0/scene14641.png", 5)

sanity_check("data/multi_class/valid/11/scene75121.png", 5)

sanity_check("data/multi_class/valid/2/scene41641.png", 5)

sanity_check("data/multi_class/valid/5/scene78121.png", 5)

sanity_check("data/multi_class/valid/6/scene83401.png", 5)

sanity_check("data/multi_class/valid/8/scene81481.png", 5)

sanity_check("data/multi_class/valid/8/scene38041.png", 5)