import os
import numpy as np 
import pandas as pd 
from datetime import datetime
import time
#import random
from tqdm.autonotebook import tqdm


#Torch
import torch
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader

################# DETR FUCNTIONS FOR LOSS######################## 
import sys
sys.path.append('./detr/')

from detr.models.matcher import HungarianMatcher
from detr.models.detr import SetCriterion, DETR
#################################################################

import matplotlib.pyplot as plt

#Glob
from glob import glob

#added
import json
import itertools
from sklearn.model_selection import train_test_split
from PIL import Image, ImageOps
from numpy.random import rand, randint, ranf, randn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms, utils
import torchvision.transforms.functional as tf
from util.misc import (NestedTensor, nested_tensor_from_tensor_list,
                       accuracy, get_world_size, interpolate,
                       is_dist_avail_and_initialized)
%matplotlib inline

from  IPython.display import Image as ipImage
ipImage("../doc/KITTI_coordinate.png")

class KITTIData():
    def __init__(self, KITTIPath=None):
        if not KITTIPath==None:
            self.dataset_path = KITTIPath
        else:
            self.dataset_path = "../../../Dataset/KITTI/training/"
        self.img_folder_path = self.dataset_path + "image_2/"
        self.label_folder_path = self.dataset_path + "label_2/"
        self.calib_folder_path = self.dataset_path + "calib/"
        
        self.img_names = glob(self.img_folder_path + "*.png")
        for i in range(len(self.img_names)):
            self.img_names[i] = self.img_names[i].split('/')[-1][:-4]
        
        #print("Preparing", len(self.img_names), "images from KITTI.")
        #
        #self.annotations_list = pd.DataFrame()
        #names = ["type", "truncated", "occluded", "alpha", "left", "top", "right", "bottom",
        #         "height", "width", "length", "x", "y", "z", "rotation_y", ]
        #
        #for img_name in tqdm(self.img_names):
        #    label_path = self.label_folder_path + img_name + ".txt"
        #    annotations = pd.read_csv(label_path, names = names, sep='\s+')
        #    annotations["img_name"] = img_name
        #    self.annotations_list = self.annotations_list.append(annotations)
        #self.annotations_list = self.annotations_list[["img_name"]+names]
        #
        #print("Loading completed!")
        self.annotation_items = ["type", "truncated", "occluded", "alpha", "left", "top", "right", "bottom",
                                 "height", "width", "length", "x", "y", "z", "rotation_y", ]
        
    def __len__(self):
        return len(self.img_names)
        
    def line2P(self, l):
        P_elem = l.split()[1:]
        P = np.array(P_elem, dtype=np.float).reshape(3,-1)
        return P
    
    def readCalib(self, calib_name):
        calib_name = self.calib_folder_path + calib_name + ".txt"
        with open(calib_name) as f:
            P0 = self.line2P(f.readline())
            P1 = self.line2P(f.readline())
            P2 = self.line2P(f.readline())
            P3 = self.line2P(f.readline())
            R0_rect = self.line2P(f.readline())
            Tr_velo_to_cam = self.line2P(f.readline())
            Tr_imu_to_velo = self.line2P(f.readline())
        return P0, P1, P2, P3, R0_rect, Tr_velo_to_cam, Tr_imu_to_velo
    
    def getPILImg(self, img_name):
        img_path = self.img_folder_path + img_name + ".png"
        img = Image.open(img_path)
        return img
    
    def getAnnotation(self, img_name):
        #return self.annotations_list[self.annotations_list["img_name"]==img_name]
        label_path = self.label_folder_path + img_name + ".txt"
        annotation = pd.read_csv(label_path, names=self.annotation_items, sep='\s+')
        annotation["img_name"] = img_name
        return annotation
    
    def plot2Dbb(self, img_name):
        annotations = self.getAnnotation(img_name)
        img = self.getPILImg(img_name)
        
        plt.figure(figsize=(15, 4))
        plt.imshow(img)
        
        for _, anno in annotations.iterrows():
            _type,left,top,right,bottom = anno["type"],anno["left"],anno["top"],anno["right"],anno["bottom"]
            plt.plot([left,right,right,left,left],
                     [top,top,bottom,bottom,top],
                     color="r",linewidth=1
                )
            plt.text(left, top, _type, color="r")

    def calcSurcfaceAngle(self, pts3D):
        """
        return belows. Element of each list is [front, right, back, left]
        - list of cos(angle) for each surface
        """
        # 直方体の8点から、隠れた点を除いて出力する
        face_mask = np.array([[1,1,0,0,1,1,0,0], # front
                              [0,1,1,0,0,1,1,0], # right
                              [0,0,1,1,0,0,1,1], # back
                              [1,0,0,1,1,0,0,1]  # left
                             ])
        # 直方体の8点から、隠れた点を除いて出力する
        # face_idx = np.array([[0,1,5,4], # front face
        #                      [1,2,6,5], # right face
        #                      [2,3,7,6], # back face
        #                      [3,0,4,7]]) # left face
                             
        # 直方体の中心位置を計算
        c_rectangular = np.sum(pts3D, axis=1)/8
    
        # 各面の法線と面からカメラへの直線がなす角度が90度以下であればその面は見える
        cosForSfcs = []
        for i in range(4):
            p_in_plane = pts3D * face_mask[i] # 平面の4つの頂点
            c_plane = np.sum(p_in_plane, axis=1)/4 # 平面の中心
            normal = c_plane - c_rectangular# 平面の法線ベクトル
            c_to_O = -c_plane # 平面の中心からカメラ位置へのベクトル
            
            # cosを計算して保存
            cos = np.dot(c_to_O, normal)/np.linalg.norm(c_to_O)/np.linalg.norm(normal)
            cosForSfcs.append(cos)
            # 面の頂点を保存
            #ptsForSfcs.append(pts3D[:,face_idx[i]])
        return cosForSfcs
    
    def projectToImage(self, pts3D, P):
        P = np.array(P)
        
        ones = np.ones([1,pts3D.shape[1]])
        pts3D = np.append(pts3D, ones, axis=0)
        pts2D = np.dot(P, pts3D)
        pts2D[0] /= pts2D[2]
        pts2D[1] /= pts2D[2]
        pts2D = np.delete(pts2D, obj=2, axis=0)
        return pts2D
    
    def calc3Dcoord(self, obj, P):
        # 直方体の各頂点のxyz座標を計算
        face_idx = np.array([[0,1,5,4], # front face
                             [1,2,6,5], # right face
                             [2,3,7,6], # back face
                             [3,0,4,7]]) # left face
        ry = obj["rotation_y"]
        R = np.array([[ np.cos(ry), 0, np.sin(ry)],
                      [          0, 1,          0],
                      [-np.sin(ry), 0, np.cos(ry)]],
                     dtype=np.float)
        l = obj["length"]
        w = obj["width"]
        h = obj["height"]
        corners = np.array([[l/2, l/2, -l/2, -l/2, l/2, l/2, -l/2, -l/2],
                            [0,0,0,0,-h,-h,-h,-h],
                            [w/2, -w/2, -w/2, w/2, w/2, -w/2, -w/2, w/2]],
                           dtype=np.float).reshape(3,-1)
        pts3D = np.dot(R,corners)
        xyz = np.array([obj["x"], obj["y"], obj["z"]], dtype=np.float).reshape([3,1])
        pts3D += xyz
        return pts3D
    
    def devideToEachSurface(self, pts):
        # 直方体の各頂点のxyz座標を計算
        face_idx = np.array([[0,1,5,4], # front face
                             [1,2,6,5], # right face
                             [2,3,7,6], # back face
                             [3,0,4,7]]) # left face
        
        ptsForSfcs = []
        for i in range(4):
            ptsForSfcs.append(pts[:,face_idx[i]])

        return ptsForSfcs
    
    def calcPseudo3Dbb(self, obj, P):
        """
        return belows. if any surface is unvisible, each array stores np.nans
        - 4 corners of front or back surface as np.array shaped 2x4
        - 4 corners of left or right surface as np.array shaped 2x4
        - cos angle of front or back surface
        - cos angle of left or right surface
        """
        pts3D = self.calc3Dcoord(obj, P)
        pts2D = self.projectToImage(pts3D, P)
        pts2D_devided = self.devideToEachSurface(pts2D)
        cosForSfcs = self.calcSurcfaceAngle(pts3D)

        # 返り値の変数を用意、デフォルト値を設定
        pts2D_fr = np.zeros([2,4])
        pts2D_fr[:,:] = np.nan
        pts2D_lr = pts2D_fr.copy()
        cos_fr, cos_lr = np.nan, np.nan
        
        # 見えている面を選ぶ
        for i in range(0,3,2):
            if cosForSfcs[i]>0:
                pts2D_fr = pts2D_devided[i]
                cos_fr = cosForSfcs[i]
                break
        for i in range(1,4,2):
            if cosForSfcs[i]>0:
                pts2D_lr = pts2D_devided[i]
                cos_lr = cosForSfcs[i]
                break
        # pts2D_fr, lrをx座標でソートする。下の順番になる
        # left-bottom, left-top, right-bottom, right-top
        pts2D_fr = pts2D_fr[:,pts2D_fr[0].argsort()].T
        pts2D_lr = pts2D_lr[:,pts2D_lr[0].argsort()].T
        
        return pts2D_fr, pts2D_lr, cos_fr, cos_lr
    
    def array2Pseudo3Dbb(self, pts2D_fr, pts2D_lr, cos_fr, cos_lr):
        coord = np.append(pts2D_fr, pts2D_lr).reshape([8,2])
        pose = np.append(cos_fr, cos_lr)
        obj = {"coord":coord,
               "pose":pose}
        return obj
    
    @staticmethod
    def psuedo3Dbb2ccwh(pts2D_fr, pts2D_lr=np.zeros([0,2])):
        """
        convert array to 2Dbb(cx, cy, w, h)
        """
        array2D = np.append(pts2D_fr, pts2D_lr)
        array2D = array2D[np.logical_not(np.isnan(array2D))].reshape([-1, 2])
        x0 = np.min(array2D[:,0])
        x1 = np.max(array2D[:,0])
        y0 = np.min(array2D[:,1])
        y1 = np.max(array2D[:,1])

        c_x, c_y = (x0+x1)/2, (y0+y1)/2
        w, h = x1-x0, y1-y0
        
        return [c_x, c_y, w, h]
    
    def compute3Dbb(self, obj, P):
        face_idx = np.array([[1,2,6,5], # front face
                             [2,3,7,6], # left face
                             [3,4,8,7], # back face
                             [4,1,5,8]]) # right face
        
        corner3D = self.calc3Dcoord(obj, P)
        if any(corner3D[2,:]<0.1): # 対象物の前後距離<.1だったら出力しない
            corner2D = []
        else:
            corner2D = self.projectToImage(corner3D, P)
        
        return corner2D, face_idx
    
    def plot3Dbb(self, img_name):
        annotations = self.getAnnotation(img_name)
        img = self.getPILImg(img_name)
        
        plt.figure(figsize=(15, 4))
        plt.imshow(img)
        
        _, _, P2, _, _, _, _ = self.readCalib(img_name)
        
        for _, anno in annotations.iterrows():
            #if anno["type"] == "DontCare":
            if not anno["type"] in ['Car', 'Van', 'Truck']:
                continue
            corner2D, face_idx = self.compute3Dbb(anno, P2)
            if len(corner2D)==0:
                continue
            plt.plot(corner2D[0,[2,3,0,1,2,6,7,4,5,6,7,3,0,5,1,4,0]],
                     corner2D[1,[2,3,0,1,2,6,7,4,5,6,7,3,0,5,1,4,0]],
                     color="r",linewidth=1)
            # todo : 
            # plt.text(corner2D[0,0], corner2D[1,0]+10, "alpha : " + str(int(-anno["alpha"]/3.14*180)), color="green", size=15)
            
        plt.xlim([0,img.size[0]])
        plt.ylim([img.size[1],0])
    
    def plotPseudo3Dbb(self, img_name):
        annotations = self.getAnnotation(img_name)
        img = self.getPILImg(img_name)
        
        plt.figure(figsize=(15, 4))
        plt.imshow(img)
        
        _, _, P2, _, _, _, _ = self.readCalib(img_name)
        
        for _, anno in annotations.iterrows():
            #if anno["type"] == "DontCare":
            if not anno["type"] in ['Car', 'Van', 'Truck']:
                continue
            pts2D_fr, pts2D_lr, cos_fr, cos_lr = self.calcPseudo3Dbb(anno, P2)
            if cos_lr>0:
                plt.plot(pts2D_lr[[0,1,3,2,0],0],pts2D_lr[[0,1,3,2,0],1], "g", linewidth=2)
            if cos_fr>0:
                plt.plot(pts2D_fr[[0,1,3,2,0],0],pts2D_fr[[0,1,3,2,0],1], "r", linewidth=2)
            
        plt.xlim([0,img.size[0]])
        plt.ylim([img.size[1],0])

    #def __getitem__(self, idx):
    #    # Get data for pytorch
    #    # preprocess is done in self.getitem, and return except for gt
        # because batch process cannot treat not-fixed length
    #    img_name = self.img_names[idx]
    #    img = self.getPILImg(img_name)
    #    img = self.preprocessImage(img)
    #    
    #    # regularization of coordination
    #    targets = self.anno2listsObj(img_name)
    #    if len(targets["coords"])>0:
    #        targets["coords"] = torch.as_tensor(targets["coords"])
    #        targets["coords"] /= self.outSize
    #        targets["coords"] = targets["coords"].reshape([-1,16])
    #    else:
    #        targets["coords"] = torch.zeros([0,16]).float()
    #    targets["poses"] = torch.as_tensor(targets["poses"]).float()
    #    # todo : label設定する
    #    targets["labels"] = torch.zeros(targets["coords"].shape[-2]).long() # .shape[-1]はcoord数の16、.shape[-2]はlabel数
    #        
    #    # set bounding box for matching
    #    coords = targets["coords"].detach().numpy()
    #    bbox_normed = torch.as_tensor([self.psuedo3Dbb2ccwh(coord) for coord in coords], dtype=torch.float64).reshape([-1,4])
    #    targets["boxes"] = bbox_normed#*torch.as_tensor(np.append(self.outSize,self.outSize)) # アスペクト比を調整
    #    targets["boxes"] = targets["boxes"].float()
    #    targets["angles"] = torch.as_tensor(targets["angles"]).float()
    #    
    #    return img, targets

    
KITTIDataset = KITTIData()
i = 0
KITTIDataset.plot2Dbb(KITTIDataset.img_names[i])
KITTIDataset.plot3Dbb(KITTIDataset.img_names[i])
KITTIDataset.plotPseudo3Dbb(KITTIDataset.img_names[i])

class ImagePreprocessor(KITTIData):
    def __init__(self, training=False, fliplr=False, shear_x=False, shear_y=False,
                 scale=False, crop=False, offset=False):
        super().__init__()
        
        self.training = training
        self.augFlags = {
            "fliplr":fliplr,
            "shear_x":shear_x,
            "shear_y":shear_y,
            "scale":scale,
            "crop":crop,
            "offset":offset
            }
        self.augParams = {
            "fliplr":False,
            "shear_x":0,
            "shear_y":0,
            "scale_x":1,
            "scale_y":1,
            "offset":[0,0]
            }
        self.outSize = np.array([450,135]) # x,y
        self.suppression_captureRate = .15
    
    def setAugParams(self, img):
        self.inSize = np.array(img.size) # x,y
        
        if self.training:
            if self.augFlags["fliplr"]:
                self.augParams["fliplr"] = rand()>.5
            if self.augFlags["shear_x"]:
                self.augParams["shear_x"] = (rand()-.5)*2*self.augFlags["shear_x"]
            if self.augFlags["shear_y"]:
                self.augParams["shear_y"] = (rand()-.5)*2*self.augFlags["shear_y"]
            if self.augFlags["scale"]:
                self.augParams["crop_x"] = randint((2*self.augFlags["scale"]-1)*self.inSize[0],self.inSize[0])
                self.augParams["crop_y"] = randint((2*self.augFlags["scale"]-1)*self.inSize[1],self.inSize[1])
                self.augParams["offset_x"] = randint(0,self.inSize[0]-self.augParams["crop_x"])
                self.augParams["offset_y"] = randint(0,self.inSize[1]-self.augParams["crop_y"])
    
    def preprocessImage(self, PILimg):
        img = PILimg
        self.setAugParams(img)
        
        if self.training:
            if self.augFlags["fliplr"] and self.augParams["fliplr"]:
                img = ImageOps.mirror(img)
            affined = tf.affine(img, angle=0,
                                    translate=(0,0),
                                    scale=self.augFlags["scale"],
                                    shear=[self.augParams["shear_x"],self.augParams["shear_y"]]
                                   )
            cropped = tf.crop(affined, self.augParams["offset_y"], self.augParams["offset_x"],
                              self.augParams["crop_y"], self.augParams["crop_x"])
        else:
            cropped = img
        resized = transforms.Resize(self.outSize[::-1])(cropped)
        resized = transforms.ToTensor()(resized)
        
        return resized
    
    def captured_rate(self, bb_ccwh, img_size):
        # determine the (x, y)-coordinates of the intersection rectangle
        c_x, c_y, w, h = bb_ccwh
        img_w, img_h = img_size[0], img_size[1]
        
        area_bb = w*h
        # intersection area
        x0 = max(0, c_x-w/2)
        x1 = min(img_w, c_x+w/2)
        y0 = max(0, c_y-h/2)
        y1 = min(img_h, c_y+h/2)
        area_intersection = (x1-x0)*(y1-y0)
        
        if x1<x0 or y1<y0:
            area_intersection = 0
        
        return area_intersection/(area_bb+.00001)
    
    def preprocessAnno(self, pts2D_fr, pts2D_lr, angle):
        array = np.append(pts2D_fr, pts2D_lr).reshape([8,2])
        # fliplr, scale, shear_x, shear_y, offset, crop
        if self.training:
            if self.augFlags["fliplr"] and self.augParams["fliplr"]:
                array[:,0] = self.inSize[0] - array[:,0]
                array[0:2,:],array[2:4,:],array[4:6,:],array[6:8,:] =\
                     array[2:4,:].copy(),array[0:2,:].copy(),array[6:8,:].copy(),array[4:6,:].copy()
                angle = np.pi - angle
            if self.augFlags["scale"]:
                array = array*self.augFlags["scale"] + self.inSize*(1-self.augFlags["scale"])/2
            if self.augFlags["shear_x"]:
                off_y = np.array(array[:,1]-self.inSize[1]/2)
                array[:,0] -= off_y*np.tan(self.augParams["shear_x"]/180*np.pi)
            if self.augFlags["shear_y"]:
                off_x = np.array(array[:,0]-self.inSize[0]/2)
                array[:,1] -= off_x*np.tan(self.augParams["shear_y"]/180*np.pi)
            if self.augFlags["scale"]:
                array -= np.array([self.augParams["offset_x"], self.augParams["offset_y"]])
                array *= self.outSize/np.array([self.augParams["crop_x"], self.augParams["crop_y"]])
        else:
            array *= self.outSize/self.inSize
        pts2D_fr = array[:4,:]
        pts2D_lr = array[4:,:]
        return pts2D_fr, pts2D_lr, angle
    
    def anno2objsList(self, img_name):
        # return list which contains objects
        annotations = self.getAnnotation(img_name)
        _, _, P2, _, _, _, _ = self.readCalib(img_name)
        objects = []
        for _, anno in annotations.iterrows():
            #if anno["type"] == "DontCare":
            if not anno["type"] in ['Car', 'Van', 'Truck']:
                continue
            angle = -anno["alpha"]
            pts2D_fr, pts2D_lr, cos_fr, cos_lr = self.calcPseudo3Dbb(anno, P2)
            pts2D_fr_aug, pts2D_lr_aug, angle = self.preprocessAnno(pts2D_fr, pts2D_lr, angle)
            angle = (angle)%(2*np.pi)
            
            # 元画像に写っている面積が小さければ出力しない
            # 変換後の画像に写っている面積が小さければ出力しない
            bb2D_aug = self.psuedo3Dbb2ccwh(pts2D_fr_aug, pts2D_lr_aug)
            bb2D = self.psuedo3Dbb2ccwh(pts2D_fr, pts2D_lr)
            if self.captured_rate(bb2D, self.inSize)<self.suppression_captureRate:
                continue
            elif self.captured_rate(bb2D_aug, self.outSize)<self.suppression_captureRate:
                continue
            obj = self.array2Pseudo3Dbb(pts2D_fr_aug, pts2D_lr_aug, cos_fr, cos_lr)
            objects.append(obj)
        return objects
    
    def anno2listsObj(self, img_name):
        # return dictionary which contains lists of object informatation
        annotations = self.getAnnotation(img_name)
        _, _, P2, _, _, _, _ = self.readCalib(img_name)
        listObject = {"coords":[], "poses":[], "angles":[]}
        for _, anno in annotations.iterrows():
            #if anno["type"] == "DontCare":
            if not anno["type"] in ['Car', 'Van', 'Truck']:
                continue
            angle = -anno["alpha"]
            pts2D_fr, pts2D_lr, cos_fr, cos_lr = self.calcPseudo3Dbb(anno, P2)
            pts2D_fr_aug, pts2D_lr_aug, angle = self.preprocessAnno(pts2D_fr, pts2D_lr, angle)
            angle = (angle)%(2*np.pi)
            # 元画像に写っている面積が小さければ出力しない
            # 変換後の画像に写っている面積が小さければ出力しない
            bb2D_aug = self.psuedo3Dbb2ccwh(pts2D_fr_aug, pts2D_lr_aug)
            bb2D = self.psuedo3Dbb2ccwh(pts2D_fr, pts2D_lr)
            if self.captured_rate(bb2D, self.inSize)<self.suppression_captureRate:
                continue
            elif self.captured_rate(bb2D_aug, self.outSize)<self.suppression_captureRate:
                continue
            
            listObject["coords"].append(np.append(pts2D_fr_aug, pts2D_lr_aug).reshape([8,2]))
            listObject["poses"].append(np.append(cos_fr, cos_lr))
            listObject["angles"].append(angle)
        return listObject
    
    def plotPseudo3Dbb(self, img_name):
        img = self.getPILImg(img_name)
        img = self.preprocessImage(img)
        img = np.rollaxis(np.array(img),0,3)
        
        plt.figure(figsize=(15, 4))
        plt.imshow(img)
        
        objs = self.anno2objsList(img_name)
        for obj in objs:
            if not np.isnan(obj["pose"][1]):
                plt.plot(obj["coord"][[4,5,7,6,4],0],obj["coord"][[4,5,7,6,4],1], "g", linewidth=1)
            if not np.isnan(obj["pose"][0]):
                plt.plot(obj["coord"][[0,1,3,2,0],0],obj["coord"][[0,1,3,2,0],1], "r", linewidth=2)
        
        plt.xlim([0,img.shape[1]])
        plt.ylim([img.shape[0],0])

    def __getitem__(self, idx):
        # Get data for pytorch
        # preprocess is done in self.getitem, and return except for gt
        # because batch process cannot treat not-fixed length
        img_name = self.img_names[idx]
        img = self.getPILImg(img_name)
        img = self.preprocessImage(img)
        
        # regularization of coordination
        targets = self.anno2listsObj(img_name)
        if len(targets["coords"])>0:
            targets["coords"] = torch.as_tensor(targets["coords"])
            targets["coords"] /= self.outSize
            targets["coords"] = targets["coords"].reshape([-1,16]).float()
        else:
            targets["coords"] = torch.zeros([0,16]).float()
        targets["poses"] = torch.as_tensor(targets["poses"]).reshape([-1,2]).float()
        # todo : label設定する
        targets["labels"] = torch.zeros(targets["coords"].shape[-2]).long() # .shape[-1]はcoord数の16、.shape[-2]はlabel数
        targets["angles"] = torch.as_tensor(targets["angles"])
        targets["angles_c"] = torch.cos(targets["angles"]).reshape([-1,1]).float()
        targets["angles_s"] = torch.sin(targets["angles"]).reshape([-1,1]).float()
        del targets["angles"]
            
        # set bounding box for matching
        coords = targets["coords"].detach().numpy()
        bbox_normed = torch.as_tensor([self.psuedo3Dbb2ccwh(coord) for coord in coords]).reshape([-1,4]).float()
        targets["boxes"] = bbox_normed*torch.as_tensor([1,135/450]*2) # アスペクト比を調整
        return img, targets

vehicleAug = ImagePreprocessor(training=True, fliplr=True, shear_x=False, shear_y=10,
                 scale=0.7, crop=False, offset=False)
i = 0
vehicleAug.plotPseudo3Dbb(KITTIDataset.img_names[i])

img, targets = vehicleAug[0]

plt.figure(figsize=(12,4))
plt.imshow(np.rollaxis(np.array(img),0,3))
FrontOrRear = targets["coords"][:,:8].reshape([-1,4,2])*torch.tensor([450,135])
LeftOrRight = targets["coords"][:,8:].reshape([-1,4,2])*torch.tensor([450,135])
pose = targets["poses"]
angle = torch.atan2(targets["angles_s"], targets["angles_c"])*180/3.1415
x0, x1 = targets["boxes"][:,0]-.5*targets["boxes"][:,2], targets["boxes"][:,0]+.5*targets["boxes"][:,2]
y0, y1 = targets["boxes"][:,1]-.5*targets["boxes"][:,3], targets["boxes"][:,1]+.5*targets["boxes"][:,3]
x0, x1, y0, y1 = x0*450, x1*450, y0*450, y1*450
for i in range(FrontOrRear.shape[0]):
    plt.plot([x0[i], x0[i], x1[i], x1[i], x0[i]], [y0[i], y1[i], y1[i], y0[i], y0[i]], "b")
    tmp = False
    if not torch.isnan(pose[i,1]):
        plt.plot(LeftOrRight[i,:,0], LeftOrRight[i,:,1], "g")
        plt.plot(LeftOrRight[i,0,0], LeftOrRight[i,0,1], "go")
        plt.text(LeftOrRight[i,1,0], LeftOrRight[i,1,1], '{:.2f}'.format(pose[i,1]), size=15, color="g")
    if not torch.isnan(pose[i,0]):
        plt.plot(FrontOrRear[i,:,0], FrontOrRear[i,:,1], "r")
        plt.plot(FrontOrRear[i,0,0], FrontOrRear[i,0,1], "ro")
        plt.text(FrontOrRear[i,1,0], FrontOrRear[i,1,1], '{:.2f}'.format(pose[i,0]), size=15, color="r")
        tmp = True
    if tmp:
        x, y = FrontOrRear[i,1,0], FrontOrRear[i,1,1]
    else:
        x, y, = LeftOrRight[i,1,0], LeftOrRight[i,1,1]
    plt.text(x, y-10, '{:.2f}'.format(angle[i,0]), size=15, color="lime")
print("targets : ", targets)

len(vehicleAug)

img_list = vehicleAug.img_names
train_list, val_list = train_test_split(img_list, test_size=0.05, shuffle=True, random_state=0)
print("train dataset has ", len(train_list), "images.")
print("val dataset has ", len(val_list), "images.")

# Gets the GPU if there is one, otherwise the cpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

model = DETRModel(num_classes=num_classes,num_queries=num_queries)
model = model.to(device)
criterion = SetVehicleCriterion(num_classes-1, matcher, weight_dict, eos_coef = null_class_coef, losses=losses)
criterion = criterion.to(device)

for param in model.model.backbone.parameters():
    param.required_grad = False
best_map = np.inf
history = pd.DataFrame()

for epoch in range(EPOCHS):
    train_loss = train_fn(train_loader, model,criterion, optimizer,device,scheduler=scheduler,epoch=epoch,history=history)
    eval_loss = eval_fn(val_loader, model,criterion, device, epoch, history)
    
    #valid_map_score = calculate_final_score(predictions,0.5,form='coco')

    print('|EPOCH {}| LEARNING RATE {:.0e}| TRAIN_LOSS {:.5f}| EVAL_LOSS {:.5f}|'.format(
        epoch+1,scheduler.get_lr()[0], train_loss.avg,eval_loss.avg))

    if eval_loss.avg < best_map:
        best_map = eval_loss.avg
        print('Best model found in Epoch {}........Saving Model'.format(epoch+1))
        torch.save(model.state_dict(), '../model/0214_epoch{}_evalLoss{:.5f}.pth'.format(epoch+1,eval_loss.avg))