# tram_net_v2.py
# Vojtěch Golda
from torchvision.models.detection import ssd300_vgg16, SSD300_VGG16_Weights 
from util.model_util3 import get_data_hm, get_model, batch2boxes

from models.model_util import ModelUtil
from torch import nn
import numpy as np

class TramNetV2Model(nn.Module):
    def __init__(self):
        super().__init__()
        channels = 256
        padding = 1
        self.stride_window_size = (96, 96)

        self.conv1 = nn.Conv2d(1,6,3, padding=padding)
        self.pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6, channels//16, 3, padding=padding)
        self.conv3 = nn.Conv2d(channels//16, channels//8, 3, padding=padding)
        self.conv4 = nn.Conv2d(channels//8, channels//4, 3, padding=padding)
        self.conv5 = nn.Conv2d(channels//4, channels//2, 3, padding=padding)
        self.conv6 = nn.Conv2d(channels//2, channels, 3, padding=padding)
        self.conv7 = nn.Conv2d(channels, channels, 3, padding=padding)
        self.conv8 = nn.Conv2d(channels, channels, 3, padding=padding)
        self.conv9 = nn.Conv2d(channels, channels, 3, padding=padding)
        self.conv10 = nn.Conv2d(channels, channels, 3, padding=padding)
        self.pool1 = nn.MaxPool2d((3, 3), stride=(1,1), padding=(1,1))
        self.linear1 = nn.Conv2d(channels, channels//2, 1) 
        self.linear2 = nn.Conv2d(channels//2, channels//4, 1)
        self.linear3 = nn.Conv2d(channels//4, 4, 1)
        self.activation = nn.ReLU()
        self.dropout = lambda x:x

        

    def forward(self, x_set, y_set=None):
        heatmaps = []
        losses = {}
        for i, x in enumerate(x_set):
            if len(x.shape) < 3: # empty scale
                heatmaps.append([])
                losses[i] = 0
                continue
            x = self.pool((self.activation(self.dropout(self.conv2(self.activation(self.dropout(self.conv1(x))))))))
            x = self.pool((self.activation(self.dropout(self.conv4(self.activation(self.dropout(self.conv3(x))))))))
            x = self.pool((self.activation(self.dropout(self.conv6(self.activation(self.dropout(self.conv5(x))))))))
            x = self.pool((self.activation(self.dropout(self.conv8(self.activation(self.dropout(self.conv7(x))))))))
            x = self.pool((self.activation(self.conv9(x))))
            x = self.pool1((self.activation(self.dropout(self.conv10(x)))))
            hidden = self.activation(self.linear1(x))
            hidden = self.activation(self.linear2(hidden))
            y_predicted = self.activation(self.linear3(hidden))
            if y_set is None:
                heatmaps.append(y_predicted)
            else:
                losses[i] = nn.MSELoss()(y_predicted, y_set[i]) 
            
        if y_set is None:
            return self.heatmaps_to_boxes(x_set, heatmaps)
        else:
            return losses
            

    def heatmaps_to_boxes(self, images_scaled, heatmaps_scaled, treshhold=0.2):
        results_scaled = []
        for images, heatmaps in zip(images_scaled, heatmaps_scaled): 
            results = []
            for i, h in zip(images, heatmaps):
                result = {
                    "boxes": [], "labels": [], "scores": []
                }
                height = i.shape[1]/h.shape[1]
                width = i.shape[2]/h.shape[2]
                h = h.cpu() 
                horizontal_compenzation = (self.stride_window_size[0] - width) / width / 2 
                vertical_compenzation = (self.stride_window_size[1] - height) / height / 2 
                for category in range(len(h)):
                    while True:
                        max_coords = np.unravel_index(np.argmax(h[category], axis=None), h[category].shape)
                        max_value = h[category][max_coords[0]][max_coords[1]]
                        h[category][max_coords[0]][max_coords[1]] = 0
                        x = max_coords[1]
                        y = max_coords[0]
                        if max_value > treshhold:
                                result["boxes"].append(
                                    [
                                        int(width * (x - horizontal_compenzation)),
                                        int(height * (y - vertical_compenzation)), 
                                        int((width) * (x + 1 + horizontal_compenzation)),
                                        int((height) * (y + 1 + vertical_compenzation)),
                                    ]
                                )
                                result["labels"].append(category)
                                result["scores"].append(max_value)

                        else:
                            break
                results.append(result)
            results_scaled.append(results)
        return results_scaled


class TramNetV2(ModelUtil):
    def __init__(self, device, filename="models/tram_net_v2.pth"):
        #self.scales = ((64,32), (160, 80), (288, 144), (416, 208), (544, 272), (672, 336), (800, 400), (928, 464), (1056, 528) )
        self.scales = ((32,32), (64, 64), (96, 96), (128, 128), (160, 160), (192, 192), (224, 224), (256, 256))
        self.device = device
        self.model = TramNetV2Model
        self.weights = None
        self.target_size = (32,32)
        self.stride_window_size = (96,96)
        self.filename = filename


    def get_data(*args, **kwargs):
        return get_data_hm(*args, **kwargs)

    def get_model(*args, **kwargs):
        return get_model(*args, **kwargs)

    def batch2boxes(*args, **kwargs):
        return batch2boxes(*args, **kwargs)

