import os, sys, time, datetime, random
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.autograd import Variable
import torch

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from PIL import Image
import cv2
from sort import *
import torch.optim as optim
from models import YoloBlock
from math import floor
from utils import *

S = [13, 26, 52]
anchors = variables.ANCHORS
scaled_anchors = torch.tensor(anchors).to(variables.DEVICE) / (
        1 / torch.tensor(S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
).to(variables.DEVICE)
model = YoloBlock(no_classes=variables.NO_CLASSES).to(variables.DEVICE)
optimizer = optim.Adam(
    model.parameters(), lr=variables.LR, weight_decay=variables.WEIGHT_DECAY
)
load_model('trained_model/118_epoch_cars_kitti.pth.tar', model, optimizer, variables.LR)

model.cuda()
model.eval()

def detect_image(img):
    img_arr = np.array(img)
    transforms = variables.train_transforms(image=img_arr, bboxes=[])
    img_arr = transforms["image"]
    img_arr = img_arr.reshape(1, 3, variables.IMG_SHAPE, variables.IMG_SHAPE).to(variables.DEVICE)


    with torch.no_grad():
        out = model(img_arr)
        frames = [[] for _ in range(img_arr.shape[0])]
        for i in range(3):
            batch_size, A, S, _, _ = out[i].shape
            anchor = scaled_anchors[i]
            boxes_scale_i = CTF(
                out[i], anchor, SPLIT_SIZE=S, is_preds=True
            )
            for idx, (box) in enumerate(boxes_scale_i):
                frames[idx] += box

        model.train()
        out = NMS(frames[0], 0.5, 0.5, frame_format="1_point")
    if len(out) == 0:
        return
    else:
        print(torch.as_tensor(out))
        return torch.as_tensor(out)


def play(video_path):
    cv2.namedWindow('Stream', cv2.WINDOW_NORMAL)
    cv2.resizeWindow('Stream', (800, 600))
    video = cv2.VideoCapture(video_path)
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    ret, frame = video.read()
    vw, vh = frame.shape[1], frame.shape[0]
    outvideo = cv2.VideoWriter(video_path.replace(".mp4", "-det.mp4"), fourcc, 20.0, (vw, vh))
    motion_tracker = Sort()
    colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (128, 0, 0), (0, 128, 0), (0, 0, 128),
              (128, 0, 128), (128, 128, 0)]
    for ii in range(400):
        ret, frame = video.read()
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        pilimg = Image.fromarray(frame)
        detections = detect_image(pilimg)
        img = np.array(pilimg)

        pad_x = max(img.shape[0] - img.shape[1], 0) * (variables.IMG_SHAPE / max(img.shape))
        pad_y = max(img.shape[1] - img.shape[0], 0) * (variables.IMG_SHAPE / max(img.shape))
        unpad_h = variables.IMG_SHAPE - pad_y
        unpad_w = variables.IMG_SHAPE - pad_x
        if detections is not None:
            tracked_objects = motion_tracker.update(detections.cpu())

            for x1, y1, x2, y2, obj_id, cls_pred in tracked_objects:
                try:
                    box_h = int(((y2 - y1) / unpad_h) * img.shape[0])
                    box_w = int(((x2 - x1) / unpad_w) * img.shape[1])
                    y1 = int(((y1 - pad_y // 2) / unpad_h) * img.shape[0])
                    x1 = int(((x1 - pad_x // 2) / unpad_w) * img.shape[1])
                    color = colors[int(obj_id) % len(colors)]
                    cls = variables.CLASSES[int(cls_pred)]
                    cv2.rectangle(frame, (x1, y1), (x1+box_w, y1+box_h), color, 4)
                    cv2.rectangle(frame, (x1, y1-35), (x1+len(cls)*19+80, y1), color, -1)
                    cv2.putText(frame, cls + "-" + str(int(obj_id)), (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (255,255,255), 3)
                except:
                    pass
        cv2.imshow('Stream', frame)
        outvideo.write(frame)
        ch = 0xFF & cv2.waitKey(1)
        if ch == 2700:
            break


if __name__ == '__main__':
    # video_path = 'data/dashcam.mp4'
    # play(video_path)
    # load image and get detections
    img_path = r"D:\deeplearning\data\kitti_single\testing\image_2\000003.png"
    img_size = 416
    prev_time = time.time()
    img = Image.open(img_path)
    detections = detect_image(img)
    inference_time = datetime.timedelta(seconds=time.time() - prev_time)
    print('Inference Time: %s' % (inference_time))
    # Get bounding-box colors
    cmap = plt.get_cmap('tab20b')
    colors = [cmap(i) for i in np.linspace(0, 1, 20)]
    img = np.array(img)
    plt.figure()
    fig, ax = plt.subplots(1, figsize=(12, 9))
    ax.imshow(img)
    pad_x = max(img.shape[0] - img.shape[1], 0) * (img_size / max(img.shape))
    pad_y = max(img.shape[1] - img.shape[0], 0) * (img_size / max(img.shape))
    unpad_h = img_size - pad_y
    unpad_w = img_size - pad_x
    if detections is not None:
        unique_labels = detections[:, -1].cpu().unique()
        n_cls_preds = len(unique_labels)
        bbox_colors = random.sample(colors, n_cls_preds)
        # browse detections and draw bounding boxes
        for cls_pred, confidence, x, y, box_w, box_h in detections:
            box_h = box_h * img.shape[0]
            box_w = box_w * img.shape[1]
            y1 = y * img.shape[0]
            x1 = x * img.shape[1]
            color = "red"
            bbox = patches.Rectangle((x1, y1), box_w, box_h,
                                     linewidth=2, edgecolor=color, facecolor='none')
            ax.add_patch(bbox)
            plt.text(x1, y1, s= f"{variables.CLASSES[int(cls_pred)]}, {floor(confidence.item()*100)}%",
                     color='white', verticalalignment='top',
                     bbox={'color': color, 'pad': 0})
    plt.axis('off')
    # save image
    # plt.savefig(img_path.replace(".jpg", "-det.jpg"),
    #             bbox_inches='tight', pad_inches=0.0)
    plt.show()