import torch
import torchvision
import cv2
import numpy as np
import os
import glob as glob
import csv
import pickle
import argparse
import matplotlib.pyplot as plt
from PIL import Image
from torchvision.transforms import functional
import time
from classifier import Net

classes = [
    "speed limit 20 (prohibitory)",
    "speed limit 30 (prohibitory)",
    "speed limit 50 (prohibitory)",
    "speed limit 60 (prohibitory)",
    "speed limit 70 (prohibitory)",
    "speed limit 80 (prohibitory)",
    "restriction ends 80 (other)",
    "speed limit 100 (prohibitory)",
    "speed limit 120 (prohibitory)",
    "no overtaking (prohibitory)",
    "no overtaking (trucks) (prohibitory)",
    "priority at next intersection (danger)",
    "priority road (other)",
    "give way (other)",
    "stop (other)",
    "no traffic both ways (prohibitory)",
    "no trucks (prohibitory)",
    "no entry (other)",
    "danger (danger)",
    "bend left (danger)",
    "bend right (danger)",
    "bend (danger)",
    "uneven road (danger)",
    "slippery road (danger)",
    "road narrows (danger)",
    "construction (danger)",
    "traffic signal (danger)",
    "pedestrian crossing (danger)",
    "school crossing (danger)",
    "cycles crossing (danger)",
    "snow (danger)",
    "animals (danger)",
    "restriction ends (other)",
    "go right (mandatory)",
    "go left (mandatory)",
    "go straight (mandatory)",
    "go right or straight (mandatory)",
    "go left or straight (mandatory)",
    "keep right (mandatory)",
    "keep left (mandatory)",
    "roundabout (mandatory)",
    "restriction ends (overtaking) (other)",
    "restriction ends (overtaking (trucks)) (other)"
]

def classify(model, patch):
    patch = cv2.resize(patch, (112, 112))
    patch = functional.to_tensor(patch)
    patch = patch[None, :]
    model.eval()
    with torch.no_grad():
        prediction = model(patch.to(device))

    return int(prediction[0].argmax(axis=1)[0])

def predict(model, classify_model, img, classes):
    model.eval()
    with torch.no_grad():
        prediction = model([img.to(device)])
    
    boxes = prediction[0]['boxes']
    scores = prediction[0]['scores']
    
    keep = torchvision.ops.nms(boxes, scores, 0.1) #Apply non maximum supression
    
    img = img.permute(1,2,0)  # C,H,W_H,W,C, for drawing
    img = (img * 255).byte().data.cpu()
    img = np.array(img)  # tensor to ndarray
    #Convert np array img to right format.
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    for k in range(len(keep)):
        xmin = round(prediction[0]['boxes'][k][0].item())
        ymin = round(prediction[0]['boxes'][k][1].item())
        xmax = round(prediction[0]['boxes'][k][2].item())
        ymax = round(prediction[0]['boxes'][k][3].item())

        patch = img[ymin:ymax, xmin:xmax].copy()
        label = classify(classify_model, patch)

        score = prediction[0]['scores'][k].item()
        print("\n===============\n")
        print("Label is: {} ({})\n===\n(Xmin, Ymin, Xmax, Ymax) = ({}, {}, {}, {}) \n===".format(label, classes[int(label)], xmin, ymin, xmax, ymax))

        #color = list(np.random.random(size=3)*256)
        color = (0, 255, 0)

        cv2.rectangle(img, (xmin, ymin), (xmax, ymax), color, thickness=2)
        
        if classes:
            cv2.putText(img, classes[label]+"-"+str(round(score,2)), (xmin, ymin), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, thickness=2)
        else:
            cv2.putText(img, str(label)+"-"+str(round(score,2)), (xmin, ymin), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, thickness=2)
            
        print("Class: "+ str(label))
        print("Score: "+ str(score))
        
    return img



ap = argparse.ArgumentParser()
ap.add_argument('-i', '--input', required=True,
                help = 'Path to input video')
ap.add_argument('-o', '--output', required=True,
                help = 'Path to output video')
ap.add_argument('-m', '--model', required=True,
                help = 'path to pytorch detector net')
ap.add_argument('-c', '--classifier', required=True,
                help = 'path to pytorch classifier net')
args = ap.parse_args()




win_name = "Stream"
cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)

vs = cv2.VideoCapture(args.input)
writer = None
(W, H) = (None, None)


# Load model 

num_classes = 43
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False, progress=True, num_classes=num_classes, pretrained_backbone=True)

model = torch.load(args.model)

# Load classifier
classifier = Net(43)
classifier = torch.load(args.classifier)


#device = torch.device('cuda')
device = torch.device('cpu')

model.to(device)
classifier.to(device)

print("Successful load of models")

while True:
    (grabbed, frame) = vs.read()

    if not grabbed:
        break

    if W is None or H is None:
        (H, W) = frame.shape[:2]

    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    tensor = functional.to_tensor(frame)
    
    start = time.time()
    out = predict(model, classifier, tensor, classes)
    end = time.time()
    print("Inference time: "+str(end-start))
    print("\n===============\n")
    """
    cv2.imshow(win_name, out)
    if cv2.waitKey(1) == ord('q'):
        break
    
    """
    if writer is None:
        fourcc = cv2.VideoWriter_fourcc(*"MJPG")
        writer = cv2.VideoWriter(args.output, fourcc, 30, (frame.shape[1], frame.shape[0]), True)
    writer.write(out)

writer.release()
vs.release()
    


