# -*- coding: utf-8 -*-
"""multigrasp.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/11BvggyQ-xlzluou-FS_Y5hzl4QvmPBgm
"""

from __future__ import print_function
from __future__ import division
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
import glob
from torch.utils.data import Dataset, DataLoader
from PIL import Image
from skimage.draw import polygon
import torch.nn.functional as F
from scipy.ndimage.interpolation import shift, rotate
print("PyTorch Version: ",torch.__version__)
print("Torchvision Version: ",torchvision.__version__)

from google.colab import drive
drive.mount('/content/drive')

"""Hyperparameters"""

# Top level data directory. Here we assume the format of the directory conforms
#   to the ImageFolder structure
data_dir = "/content/drive/My Drive/cornell_grasps/"

# Regressor output size
num_classes = 343

# Models to choose from [resnet, alexnet, vgg, squeezenet, densenet, inception]
model_name = "alexnet"

# Batch size for training (change depending on how much memory you have)
batch_size = 128

# Number of epochs to train for
num_epochs = 20

# Flag for feature extracting. When False, we finetune the whole model,
#   when True we only update the reshaped layer params
feature_extract = True

# validation split
split = 0.9

"""Train helper functions"""

# convert rgb and depth images to three-channel rg-d image
def rgb_to_rgd(img_rgb, img_d, transforms):
    img_rgd_t = transforms['val'](img_rgb)
    img_d_t = transforms['depth'](img_d)
    img_d_t = np.clip((img_d_t - img_d_t.mean()), -1, 1)
    img_rgd_t[2] = img_d_t
    return img_rgd_t

# create 7x7x7 labels for multigrasp
def create_labels(rectangles):
    label = np.zeros((7,7,7))
    chosen_recs = []
    idxs = np.random.choice(len(rectangles), min(len(rectangles), 5), replace=False)
    chosen_recs = rectangles[idxs]
    centers = [box_center(i) for i in chosen_recs]
    for i, (x, y) in enumerate(centers):
        x_idx = int(x/640*7)
        y_idx = int(y/480*7)
        label[0, x_idx, y_idx] = 2
        label[1:, x_idx, y_idx] = box_to_grasp(chosen_recs[i])
    return label

# for a file with N grasp rectangles, returns a Nx4x2 array of the coordinates
def get_grasp_rectangles(file_name):
    with open(file_name, 'r') as f:
        coords = list(map(lambda coordinate: float(coordinate), f.read().strip().split()))
    rectangles = []
    i = 0
    while i < len(coords):
        rectangle = np.array([[coords[i], coords[i+1]],
                               [coords[i+2], coords[i+3]],
                               [coords[i+4], coords[i+5]],
                               [coords[i+6], coords[i+7]]])
        if not np.isnan(rectangle).any():
            rectangles.append(rectangle)
        i += 8
    return np.array(rectangles)

# get center of box
def box_center(box):
    x = (box[0,0] + (box[2,0] - box[0,0])/2)
    y = (box[0,1] + (box[2,1] - box[0,1])/2)
    return np.array([x, y])

# converts a grasp box to the 6-parameter grasp coordinates (scaled)
def box_to_grasp(box):
    x = (box[0,0] + (box[2,0] - box[0,0])/2)/100
    y = (box[0,1] + (box[2,1] - box[0,1])/2)/100
    theta = np.arctan2(box[1,1] - box[0,1], box[1,0] - box[0,0])
    h = ((box[3,0]-box[0,0])**2 + (box[3,1]-box[0,1])**2)**0.5/100
    w = ((box[1,0]-box[0,0])**2 + (box[1,1]-box[0,1])**2)**0.5/100
    return np.array([x, y, h, w, np.sin(2*theta), np.cos(2*theta)])

# converts 6-parameter grasp coordinates to a grasp box
def grasp_to_box(grasp):
    x, y, h, w, sin, cos = grasp
    x, y, h, w = 100*x, 100*y, 100*h, 100*w
    theta = np.arctan2(sin, cos)/2
    edge1 = [x -w/2*np.cos(theta) +h/2*np.sin(theta), y -w/2*np.sin(theta) -h/2*np.cos(theta)]
    edge2 = [x +w/2*np.cos(theta) +h/2*np.sin(theta), y +w/2*np.sin(theta) -h/2*np.cos(theta)]
    edge3 = [x +w/2*np.cos(theta) -h/2*np.sin(theta), y +w/2*np.sin(theta) +h/2*np.cos(theta)]
    edge4 = [x -w/2*np.cos(theta) -h/2*np.sin(theta), y -w/2*np.sin(theta) +h/2*np.cos(theta)]
    return np.array([edge1, edge2, edge3, edge4])

# gets angle of a grasp box
def box_angle(box):
    return np.arctan2(box[1,0] - box[0,0], box[1,1] - box[0,1])

# rectangle metric between two rectangles
def rectangle_metric(rect1, rect2, angle_threshold = np.pi/6):
    ang1 = box_angle(rect1)
    ang2 = box_angle(rect2)
    if abs((ang1 - ang2 + np.pi/2) % np.pi - np.pi/2) > angle_threshold:
        return -1
    rr1, cc1 = polygon(rect1[:, 0], rect1[:, 1])
    rr2, cc2 = polygon(rect2[:, 0], rect2[:, 1])

    try:
        r_max = max(rr1.max(), rr2.max()) + 1
        c_max = max(cc1.max(), cc2.max()) + 1
    except:
        print("exception")
        # print("rect1", rect1)
        # print("rr1", rr1)
        # print("cc1", cc1)
        # print("rr2", cc2)
        # print("cc2", cc2)
        return -2

    canvas = np.zeros((r_max, c_max))
    canvas[rr1, cc1] += 1
    canvas[rr2, cc2] += 1
    union = np.sum(canvas > 0)
    if union == 0:
        print("union")
        return -3
    intersection = np.sum(canvas == 2)
    # print("good angle", abs((ang1 - ang2 + np.pi/2) % np.pi - np.pi/2))
    return intersection/union

# calculate maximum rectangle metric between pred_rec and rects
def max_rect_metric(pred_rec, rects):
    max_metric = 0
    # print('pred_rec:', pred_rec)
    test_max = -4
    for rec in rects:
        max_metric = max(max_metric, rectangle_metric(pred_rec, rec))
        test_max = max(test_max, rectangle_metric(pred_rec, rec))
    print('test_max', test_max)
    return max_metric

# convert model output to a predicted grasp
def get_pred_grasp(output):
    x, y = 0, 0
    currmax = -np.inf
    for i in range(7):
        for j in range(7):
            if output[0,i,j] > currmax:
                x, y = i, j
                currmax = output[0,i,j]
    return output[1:, x, y]

# multigrasp loss not being used in favor of masking
def custom_loss(preds, trues):
    heatmap_losses = 0
    grasp_losses = 0
    for i in range(len(trues)):
        # print('here', i, trues[0])
        pred = preds[i]
        true = trues[i]
        # heatmap_losses += F.mse_loss(pred[0], true[0])
        # print('true',i, true)
        for x in range(true[0].shape[0]):
            for y in range(true[0].shape[1]):
                if true[0,x,y] == 1:
                    grasp_losses += F.mse_loss(pred[1:,x,y], true[1:,x,y])
    return heatmap_losses + grasp_losses

# prevent loss from acting on grasps that are not in the same cells as the labeled grasps
def mask_outputs(outputs, labels):
    new_outputs = outputs.clone()
    for i in range(labels.shape[0]):
        for x in range(7):
            for y in range(7):
                if labels[i,0,x,y] == 0:
                    for j in range(1,7):
                        new_outputs[i,j,x,y] = 0
    return new_outputs

class GraspDataset(Dataset):
    """Grasp dataset."""

    def __init__(self, file_path, transforms, start=0.0, end=1.0, ds_rotate=0, random_shift=False, random_rotate=False):
        """
        :param file_path: Cornell Dataset directory.
        :param start: If splitting the dataset, start at this fraction [0,1]
        :param end: If splitting the dataset, finish at this fraction
        :param ds_rotate: If splitting the dataset, rotate the list of items by this fraction first
        :param kwargs: kwargs for GraspDatasetBase
        """
        graspf = glob.glob(os.path.join(file_path, '*', 'pcd*cpos.txt'))
        graspf.sort()
        l = len(graspf)
        if l == 0:
            raise FileNotFoundError('No dataset files found. Check path: {}'.format(file_path))

        if ds_rotate:
            graspf = graspf[int(l*ds_rotate):] + graspf[:int(l*ds_rotate)]

        depthf = [f.replace('cpos.txt', 'd.tiff') for f in graspf]
        rgbf = [f.replace('d.tiff', 'r.png') for f in depthf]

        self.transforms = transforms
        self.grasp_files = graspf[int(l*start):int(l*end)]
        self.depth_files = depthf[int(l*start):int(l*end)]
        self.rgb_files = rgbf[int(l*start):int(l*end)]
        self.random_shift = random_shift
        self.random_rotate = random_rotate

    def __len__(self):
        return len(self.grasp_files)

    def __getitem__(self, idx):
        img_rgb = Image.open(self.rgb_files[idx])
        img_d = Image.open(self.depth_files[idx])
        x_shift = 0
        y_shift = 0
        if self.random_shift:
            x_shift = np.random.randint(-50, 51)
            y_shift = np.random.randint(-50, 51)
        img_rgb = shift(img_rgb, [y_shift, x_shift, 0])
        img_d = shift(img_d, [y_shift, x_shift])
        theta_rotate = 0
        if self.random_rotate:
            theta_rotate = np.random.randint(0, 360)
        img_rgb = rotate(img_rgb, theta_rotate, reshape=False)
        img_d = rotate(img_d, theta_rotate, reshape=False)
        img_rgd = rgb_to_rgd(Image.fromarray(img_rgb), Image.fromarray(img_d), self.transforms)

        rectangles = get_grasp_rectangles(self.grasp_files[idx])
        if self.random_shift:
            rectangles[:, :, 0] += x_shift
            rectangles[:, :, 1] += y_shift

        if self.random_rotate:
            for i in range(rectangles.shape[0]):
                for j in range(rectangles.shape[1]):
                    new_row = rectangles[i,j,0]-640/2, 480/2-rectangles[i,j,1]
                    arm = (new_row[0]**2 + new_row[1]**2)**0.5
                    theta = np.arctan2(new_row[1], new_row[0]) + np.pi/180*theta_rotate
                    new_row = 640/2 + arm*np.cos(theta), 480/2 - arm*np.sin(theta)
                    rectangles[i,j] = new_row

        return idx, img_rgd, torch.from_numpy(create_labels(rectangles)).float(), (x_shift, y_shift, theta_rotate), img_rgb
    def get_rectangles(self, idx, shifts):
        x_shift, y_shift, theta_rotate = shifts
        rectangles = get_grasp_rectangles(self.grasp_files[idx])
        if self.random_shift:
            rectangles[:, :, 0] += x_shift
            rectangles[:, :, 1] += y_shift

        if self.random_rotate:
            for i in range(rectangles.shape[0]):
                for j in range(rectangles.shape[1]):
                    new_row = rectangles[i,j,0]-640/2, 480/2-rectangles[i,j,1]
                    arm = (new_row[0]**2 + new_row[1]**2)**0.5
                    theta = np.arctan2(new_row[1], new_row[0]) + np.pi/180*theta_rotate
                    new_row = 640/2 + arm*np.cos(theta), 480/2 - arm*np.sin(theta)
                    rectangles[i,j] = new_row
        return rectangles

"""Train function"""

def train_model(model, dataloaders, criterion, optimizer, num_epochs=25):
    since = time.time()

    val_acc_history = []

    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0

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

        # Each epoch has a training and validation phase
        for phase in ['train', 'val']:
            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 idx, inputs, labels, shifts, _ in dataloaders[phase]:
                # inputs = inputs.to(device)
                # labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()


                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == 'train'):
                    # Get model outputs and calculate loss
                    outputs = model(inputs)
                    outputs = outputs.view(-1, 7, 7, 7)
                    masked_outputs = mask_outputs(outputs, labels)
                    loss = criterion(masked_outputs, labels)
                    print("loss", loss)


                    # _, preds = torch.max(outputs, 1)

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

                # get rectangles
                if phase == 'val':
                    rectangles = dataloaders[phase].dataset.get_rectangles(idx, shifts)
                    pred_grasp = get_pred_grasp(outputs[0])
                    pred_boxes = grasp_to_box(pred_grasp)

                    print('best metric', max_rect_metric(pred_boxes, rectangles))
                    running_corrects += np.sum(max_rect_metric(pred_boxes, rectangles) > 0.25)

                # statistics
                running_loss += loss.item() * inputs.size(0)
                # running_corrects += torch.sum(preds == labels.data)
                

            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = 0
            if phase == 'val':
                epoch_acc = running_corrects / len(dataloaders[phase].dataset)


            print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
            if phase == 'val':
                val_acc_history.append(epoch_acc)

        print()

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

    # load best model weights
    model.load_state_dict(best_model_wts)
    return model, val_acc_history

# Helper function for modifying AlexNet
def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False

"""Create model"""

def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):
    # Initialize these variables which will be set in this if statement. Each of these
    #   variables is model specific.
    model_ft = None
    input_size = 0

    if model_name == "resnet":
        """ Resnet18
        """
        model_ft = models.resnet18(pretrained=use_pretrained)
        set_parameter_requires_grad(model_ft, feature_extract)
        num_ftrs = model_ft.fc.in_features
        model_ft.fc = nn.Linear(num_ftrs, num_classes)
        input_size = 224

    elif model_name == "alexnet":
        """ Alexnet
        """
        model_ft = models.alexnet(pretrained=use_pretrained)
        set_parameter_requires_grad(model_ft, feature_extract)
        num_ftrs = model_ft.classifier[6].in_features
        model_ft.classifier[6] = nn.Linear(num_ftrs,num_classes)
        input_size = 224


    else:
        print("Invalid model name, exiting...")
        exit()

    return model_ft, input_size

# Initialize the model for this run
model_ft, input_size = initialize_model(model_name, num_classes, feature_extract, use_pretrained=True)

# Print the model we just instantiated
print(model_ft)

# Create transformation functions for data
data_transforms = {
    # 'train': transforms.Compose([
    #     transforms.RandomResizedCrop(input_size),
    #     transforms.RandomHorizontalFlip(),
    #     transforms.ToTensor(),
    #     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    # ]),
    'train': transforms.Compose([
        # transforms.Resize(input_size),
        transforms.CenterCrop(input_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        # transforms.Resize(input_size),
        transforms.CenterCrop(input_size),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'depth': transforms.Compose([
        # transforms.Resize(input_size),
        transforms.CenterCrop(input_size),
        transforms.ToTensor(),
    ]),
}

"""Load Datasets"""

train_dataset = GraspDataset(data_dir, data_transforms, start=0.0, end=split, random_shift=True, random_rotate=True)
train_data = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=4
)
val_dataset = GraspDataset(data_dir, data_transforms, start=split, end=1.0)
val_data = DataLoader(
    val_dataset,
    batch_size=1,
    shuffle=False,
    num_workers=4
)
dataloaders_dict = {'train': train_data, 'val': val_data}

"""Train all or only some weights"""

# Gather the parameters to be optimized/updated in this run. If we are
#  finetuning we will be updating all parameters. However, if we are
#  doing feature extract method, we will only update the parameters
#  that we have just initialized, i.e. the parameters with requires_grad
#  is True.
params_to_update = model_ft.parameters()
print("Params to learn:")
feature_extract = True
if feature_extract:
    params_to_update = []
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            params_to_update.append(param)
            print("\t",name)
else:
    for param in model_ft.parameters():
        param.requires_grad = True
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            print("\t",name)

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

# Setup the loss fxn
criterion = nn.MSELoss()
# criterion = custom_loss

# Train and evaluate
model_ft, hist1 = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, num_epochs=40)

# Train all weights in the network to fine-tune results
params_to_update = model_ft.parameters()
print("Params to learn:")
feature_extract = False
if feature_extract:
    params_to_update = []
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            params_to_update.append(param)
            print("\t",name)
else:
    for param in model_ft.parameters():
        param.requires_grad = True
    for name,param in model_ft.named_parameters():
        if param.requires_grad == True:
            print("\t",name)

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

# Setup the loss fxn
criterion = nn.MSELoss()

# Train and evaluate
model_ft, hist2 = train_model(model_ft, dataloaders_dict, criterion, optimizer_ft, num_epochs=num_epochs)

"""Create Accuracy Plots"""

plt.plot(hist1+hist2)
# plt.plot(hist2)
plt.title("Validation Accuracy per Epoch")
plt.ylabel("accuracy")
plt.xlabel("epoch")
# plt.show()
plt.savefig("model_3_accuracy.png")

"""Test the trained model"""

img_idx = 45
x1, x2, x3, shifts, img_rgb = train_dataset.__getitem__(img_idx)
plt.imshow(x2.squeeze().permute(1,2,0))
print(x2.shape)

with torch.no_grad():
    output = model_ft(x2.unsqueeze(0)).view(-1,7,7,7)
    pred_grasp = get_pred_grasp(output.squeeze())
    pred_box = grasp_to_box(pred_grasp)

given_recs = train_dataset.get_rectangles(img_idx, shifts)

# new_point = np.array([640, 320])
# new_row = new_point[0]-640/2, new_point[1]-480/2
# arm = (new_row[0]**2 + new_row[1]**2)**0.5
# # print(np.arctan2(new_row[1], new_row[0]))
# theta = np.arctan2(new_row[1], new_row[0]) - np.pi/180*90
# new_row = 640/2 + arm*np.cos(theta), 480/2 + arm*np.sin(theta)
# new_point = new_row

img = Image.open(train_dataset.rgb_files[img_idx])
# print(img)
# print(type(img))
# print(Image.fromarray(rotate(img, 30, reshape=False)))
# plt.imshow(shift(img, [40, 20, 0]))
# print(np.max(rotate(img, 30, reshape=False)))
# plt.imshow(rotate(img, 30, reshape=False))
# plt.imshow(img)

plt.imshow(img_rgb)

# x = pred_box[:,0]
# y = pred_box[:,1]

for rec in given_recs:
    xs = rec[:,0]
    ys = rec[:,1]
    plt.plot(xs[:2],ys[:2], 'b', linewidth=1)
    plt.plot(xs[2:],ys[2:], 'b', linewidth=1)
    plt.plot(xs[1:3],ys[1:3], 'r', linewidth=1)
    plt.plot([xs[0],xs[3]],[ys[0],ys[3]], 'r', linewidth=1)
# plt.plot(x[:2],y[:2], 'g', linewidth=3)
# plt.plot(x[2:],y[2:], 'g', linewidth=3)
# plt.plot(x[1:3],y[1:3], 'y', linewidth=3)
# plt.plot([x[0],x[3]],[y[0],y[3]], 'y', linewidth=3)

# plt.scatter(new_point[0], new_point[1])
# plt.scatter(grasp[:,0], grasp[:,1])
# plt.scatter(pred_box[:,0], pred_box[:,1])
# plt.scatter(given_rec[:2,0], given_rec[:2,1])
# print(pred_box[1])
# plt.savefig('grasp6_labels.png')