# model_util3.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 get_data_hm(self, source):
    keywords = {"train": 0, "bus": 1, "car": 2, "truck": 3} # hashmap to int 
    crop_tolerance = (0.6, 0.9)
    padding_tolerance = (0.6, 0.9)
    images_scaled = [[] for _ in self.scales]
    boxes_scaled = [[] for _ in self.scales]
    tolerance_check = False



    box_count = 0
    image_count = 0
    
    for image, file in source:
        image_count+=1
        if file["imageHeight"]>file["imageWidth"]:
            continue

        for s_index, scale in enumerate(self.scales):
            confidence_matrix = np.zeros((
                len(keywords),
                scale[1]//self.target_size[1], 
                scale[0]//self.target_size[0]))
            has_at_least_one = False
            for s in file["shapes"]:
                
                if s["label"] in keywords:
                        
                         
                        #load coordinates
                        rectangle = s["points"]
                        x1 = rectangle[0][0] / file["imageWidth"] * confidence_matrix.shape[2]
                        x2 = rectangle[1][0] / file["imageWidth"] * confidence_matrix.shape[2] + x1
                        if x1>x2:
                            x1, x2 = x2, x1
                        y1 = rectangle[0][1] / file["imageHeight"] * confidence_matrix.shape[1]
                        y2 = rectangle[1][1] / file["imageHeight"] * confidence_matrix.shape[1] + y1
                        if y1>y2:
                            y1, y2 = y2, y1
                         
                        x = x1 + (x2-x1)/2
                        y = y1 + (y2-y1)/2
                        
                        if not tolerance_check:
                            y_cell = int(y)
                            x_cell = int(x) 
                            has_at_least_one = True
                            confidence_matrix[keywords[s["label"]]][y_cell][x_cell] = 20
                            box_count+=1

                        if tolerance_check:
                            v_multi = (self.stride_window_size[1] - self.target_size[1]) / self.target_size[1] / 2
                            h_multi = (self.stride_window_size[0] - self.target_size[0]) / self.target_size[0] / 2
                            for y_cell in range(confidence_matrix.shape[1]):
                                for x_cell in range(confidence_matrix.shape[2]):
                                    # validate cropped portion
                                    if (
                                        x_cell - x1 - h_multi > crop_tolerance[0] * h_multi or
                                        x2 - x_cell - 1 - h_multi > crop_tolerance[0] * h_multi  or
                                        y_cell - y1 - v_multi > crop_tolerance[1] * v_multi or 
                                        y2 - y_cell - 1 - v_multi > crop_tolerance[1] * v_multi
                                    ): continue
                                    # validate padding
                                    if (
                                        x1 - x_cell + h_multi > padding_tolerance[0] * h_multi or
                                        x_cell - x2 + 1 + h_multi > padding_tolerance[0] * h_multi or
                                        y1 - y_cell + v_multi > padding_tolerance[1] * v_multi or
                                        y_cell - y2 + 1 + v_multi > padding_tolerance[1] * v_multi
                                    ): continue
                                    has_at_least_one = True
                                    confidence_matrix[keywords[s["label"]]][y_cell][x_cell] = 10
                                    box_count+=1

            if has_at_least_one:
                i = image.dot([0.33,0.34,0.33])
                images_scaled[s_index].append([cv.resize(
                                i,
                                (scale[0], scale[1]),
                                cv.INTER_LANCZOS4)
                ])
                boxes_scaled[s_index].append(confidence_matrix)
    images_scaled = [np.array(images) for images in images_scaled]
    return images_scaled, boxes_scaled

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:
            model = self.model().to(self.device)
        else:
            raise FileNotFoundError(f"File {self.filename} not found.")
    
    return model




def batch2boxes(self, model, data):
    model.eval()
    upscale_ratios = [[] for _ in self.scales]
    batch_set = [[] for _ in self.scales]
    for frame, _ in data:
        frame = frame.dot([0.33,0.34,0.33])
        for scale, batch, upscale_ratio in zip(self.scales, batch_set, upscale_ratios):
            batch.append([cv.resize(
                frame,
                (scale[0], scale[1]),
                cv.INTER_LANCZOS4)
            ])
            upscale_ratio_spec = (
                frame.shape[1] / scale[0],
                frame.shape[0] / scale[1]
            )
            upscale_ratio.append(upscale_ratio_spec)

    for i in range(len(batch_set)):
        batch_set[i] = torch.from_numpy(np.array((batch_set[i])).astype(np.float32)).to(self.device)

    with torch.no_grad():
        time1 = time()
        predictions = model(batch_set)
        time2 = time()
        running.time_total += time2-time1

        total_boxes = []
        
        # for each image
        for o, p_set in enumerate(zip(*predictions)):
            boxes = []
            # for each scale
            for u, p in enumerate(p_set):
                # for each bbox 
                for i, b in enumerate(p["boxes"]): 
                    b[0] *= upscale_ratios[u][o][0]
                    b[1] *= upscale_ratios[u][o][1]
                    b[2] *= upscale_ratios[u][o][0]
                    b[3] *= upscale_ratios[u][o][1]
                    if p["labels"][i] == 1: 
                        boxes.append(np.array(b).astype(int))
            total_boxes.append(boxes)
                    

        return total_boxes
