# model_util2.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



def get_data_hm(self, source):
    crop_tolerance = 0.3
    padding_tolerance = 0.6
    images_scaled = [[] for _ in self.scales]
    boxes_scaled = [[] for _ in self.scales]
        
    images = []
    ys = []
    box_count = 0

    for image, file in source:
        if len(file["shapes"])>1:
            continue
        for s_index, scale in enumerate(self.scales):
            confidence_matrix = np.zeros((1,scale[1]//self.target_size[1], scale[0]//self.target_size[0]))
            
            
            for s in file["shapes"]:
                
                if s["label"] == "tram":
                        
                        #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]
                        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]
                        if y1>y2:
                            y1, y2 = y2, y1
                         
                        x = x1 + (x2-x1)/2
                        y = y1 + (y2-y1)/2
                        
                        y_cell = int(y)
                        x_cell = int(x) 

                        # validate cropped portion
                        if (
                            x_cell - x1 > crop_tolerance or
                            x2 - x_cell - 1 > crop_tolerance or
                            y_cell - y1 > crop_tolerance or
                            y2 - y_cell - 1 > crop_tolerance
                        ): continue
                        # validate padding
                        if (
                            x1 - x_cell > padding_tolerance or
                            x_cell - x2 + 1 > padding_tolerance or
                            y1 - y_cell > padding_tolerance or
                            y_cell - y2 + 1 > padding_tolerance
                        ): continue

                        confidence_matrix[0][y_cell][x_cell] = 1
                        box_count+=1
                        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():
        predictions = model(batch_set)
        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
