# model_util.py
# Vojtěch Golda
from torchvision.models.detection import ssd300_vgg16, SSD300_VGG16_Weights 

import os
import json
import torch
import numpy as np
import cv2 as cv
from util.generators import dataset_generator
from torchvision.utils import draw_bounding_boxes
from util.running import running
from time import time


def preprocess(self, images_in):
    images = []
    for i in images_in:
        i = i.dot([0.33,0.34,0.33])
        images.append([cv.resize(i, self.target_size, cv.INTER_LANCZOS4)])

    return images
def get_data_bb(self, source):
    
    images = []
    ys = []
    box_count = 0

    for image, file in source:
            y = {
                "boxes": [],
                "scores": [],
                "labels": [],
            }

            
            for s in file["shapes"]:
                
                if s["label"] == "tram":
                        
                        #load coordinates
                        rectangle = s["points"]
                        x1 = int(rectangle[0][0] / file["imageWidth"] * self.target_size[1])
                        x2 = int(rectangle[1][0] / file["imageWidth"] * self.target_size[1])
                        if x1>x2:
                            x1, x2 = x2, x1
                        y1 = int(rectangle[0][1] / file["imageHeight"] * self.target_size[0])
                        y2 = int(rectangle[1][1] / file["imageHeight"] * self.target_size[0])
                        if y1>y2:
                            y1, y2 = y2, y1

                        y["boxes"].append([x1, y1, x2, y2])
                        y["scores"].append(1)
                        y["labels"].append(1)
                        box_count+=1
            if len(y["boxes"])>0: 
                images.append(image)
                ys.append(y)

    images = self.preprocess(images)
    return images, ys

def get_model(self, train=True):

    # Either load the model from a savefile or make a new one if that doesnt work
    try:
        if torch.cuda.is_available():
            model = torch.load(self.filename)
        else:
            model = torch.load(self.filename, map_location=torch.device('cpu'))
        model.to(self.device)
        model.train()
        # print(f"Successfully loaded model {self.filename}")
    except FileNotFoundError:
        if train:
            weights = self.weights.DEFAULT 
            model = self.model(weights=weights, box_score_tresh=0.9).to(self.device)
        else:
            raise FileNotFoundError(f"File {self.filename} not found.")
    
    return model




def batch2boxes(self, model, data):
    model.eval()
    upscale_ratios = []
    batch = []
    for frame, _ in data:
        batch.append(frame)
        upscale_ratio = (
            frame.shape[1] / self.target_size[0],
            frame.shape[0] / self.target_size[1]
        )
        upscale_ratios.append(upscale_ratio)

    with torch.no_grad():
        batch = torch.from_numpy(np.array(self.preprocess(batch)).astype(np.float32)).to(self.device)
        
        time1 = time()
        predictions = model(batch)
        time2 = time()
        running.time_total += time2-time1

        total_boxes = []

        for o, p in enumerate(predictions):
            boxes = []
            for i, b in enumerate(p["boxes"]): 
                if str(self.device) == "cpu":
                    b = b.detach().numpy()
                else:
                    b = b.cpu().numpy()
                b[0] *= upscale_ratios[o][0]
                b[1] *= upscale_ratios[o][1]
                b[2] *= upscale_ratios[o][0]
                b[3] *= upscale_ratios[o][1]
                if p["labels"][i] == 1: 
                    boxes.append(b.astype(int))
            total_boxes.append(boxes)
                    

        return total_boxes
