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

from models.model_util import ModelUtil
from torch import nn


class TramNetV1Model(nn.Module):
    def __init__(self):
        super().__init__()
        channels = 64
        padding = 2
        self.conv1 = nn.Conv2d(1,6,5, padding=padding)
        self.pool = nn.MaxPool2d(2,2)
        self.conv2 = nn.Conv2d(6, channels//16, 5, padding=padding)
        self.conv3 = nn.Conv2d(channels//16, channels//8, 5, padding=padding)
        self.conv4 = nn.Conv2d(channels//8, channels//4, 5, padding=padding)
        self.conv5 = nn.Conv2d(channels//4, channels//2, 5, padding=padding)
        self.conv6 = nn.Conv2d(channels//2, channels, 5, padding=padding)
        self.conv7 = nn.Conv2d(channels, channels, 5, padding=padding)
        self.conv8 = nn.Conv2d(channels, channels, 5, padding=padding)
        self.conv9 = nn.Conv2d(channels, channels, 5, padding=padding)
        self.conv10 = nn.Conv2d(channels, channels, 5, padding=padding)
        self.pool1 = nn.MaxPool2d((3, 2), stride=(3,2))
        self.linear1 = nn.Conv2d(channels, channels//2, 1) 
        self.linear2 = nn.Conv2d(channels//2, channels//4, 1)
        self.linear3 = nn.Conv2d(channels//4, 1, 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.conv5(x))))
            x = self.pool((self.activation(self.conv6(x))))
            x = self.pool((self.activation(self.conv7(x))))
            x = self.pool1((self.activation(self.dropout(self.conv8(x)))))
            hidden = self.activation(self.linear1(x))
            hidden = self.activation(self.linear2(hidden))
            y_predicted = nn.Sigmoid()(self.linear3(hidden))
            if y_set is None:
                heatmaps.append(y_predicted)
            else:
                losses[i] = nn.BCELoss()(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.5):
        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]
                
                for y, row in enumerate(h[0]):
                    for x, cell in enumerate(row):
                        if cell > treshhold:
                            result["boxes"].append(
                                [
                                    int(width * x),
                                    int(height * y), 
                                    int((width) * (x+1)),
                                    int((height) * (y+1)),
                                ]
                            )
                            result["labels"].append(1)
                            result["scores"].append(cell)
                results.append(result)
            results_scaled.append(results)
        return results_scaled


class TramNetV1(ModelUtil):
    def __init__(self, device, filename="models/tram_net_v1.pth"):
        self.scales = ((352, 192), (480, 288), (608,384), (736, 480))
        self.device = device
        self.model = TramNetV1Model
        self.weights = None
        self.target_size = (64, 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)

