from .model import Darknet19
import torch
import torch.nn as nn


class YoloV2(nn.Module):
    """Yolo version 2; It is an extented version of 
    yolo v1 capable of detecting 9000 object"""
    def __init__(self, modelUrl):
        super(YoloV2, self).__init__()
        self.modelUrl = modelUrl

        self.darknet19 = Darknet19(modelUrl)
        self.darknet19.loadWeights()
        self.weights = self.darknet19.layers
        arch = self.darknet19.arch

        self.path1 = self.makeSequence(arch[0])
        self.parallel1 = self.makeSequence(arch[1])
        self.parallel2 = self.makeSequence(arch[2])
        self.path2 = self.makeSequence(arch[3])

    def makeSequence(self, arch):
        layers = []
        for id, name in enumerate(arch):
            layers.append(self.weights[name])
        return nn.ModuleList(layers)

    def forward(self, input):
        out = input
        for layer in self.path1:
            out = layer(out)
        out1 = out.clone()
        for layer in self.parallel1:
            out1 = layer(out1)
        out2 = out.clone()
        for layer in self.parallel2:
            out2 = layer(out2)
        final = torch.cat([out2, out1], dim=1)
        for layer in self.path2:
            final = layer(final)
        return final


from keras.models import load_model
import torch
import tensorflow as tf
import json
import torch.nn as nn
import torch.autograd as autograd
import numpy as np

class scale_to_depth(nn.Module):
    def __init__(self, block_size=1):
        super(scale_to_depth, self).__init__()
        self.block_size = block_size        
    def forward(self, input):
        batch_size, in_channels, in_height, in_width = input.size()
        channels = in_channels * (self.block_size ** 2)

        out_height = int(in_height / self.block_size)
        out_width = int(in_width / self.block_size)

        input_view = input.contiguous().view(
            batch_size, in_channels, self.block_size, self.block_size,
            out_height, out_width)

        shuffle_out = input_view.permute(0, 1, 4, 2, 5, 3).contiguous()
        return shuffle_out.view(batch_size, channels, out_height, out_width)

class Darknet19:
    """This is the model to create the pretrained darknet19"""
    def __init__(self, modelUrl):
        super(Darknet19, self).__init__()
        self.modelUrl = modelUrl
        self.layers = {}
        self.lid = {}
        self.lod = {}
        self.lin = {}
        self.arch = self.makeArch()

    def makeArch(self):
        path1 = ["conv2d_1", 
                "batch_normalization_1",
                "leaky_re_lu_1",
                "max_pooling2d_1",
                "conv2d_2", 
                "batch_normalization_2",
                "leaky_re_lu_2",
                "max_pooling2d_2",
                "conv2d_3", 
                "batch_normalization_3",
                "leaky_re_lu_3",
                "conv2d_4", 
                "batch_normalization_4",
                "leaky_re_lu_4",
                "conv2d_5", 
                "batch_normalization_5",
                "leaky_re_lu_5",
                "max_pooling2d_3",
                "conv2d_6", 
                "batch_normalization_6",
                "leaky_re_lu_6",
                "conv2d_7", 
                "batch_normalization_7",
                "leaky_re_lu_7",
                "conv2d_8", 
                "batch_normalization_8",
                "leaky_re_lu_8",
                "max_pooling2d_4",
                "conv2d_9", 
                "batch_normalization_9",
                "leaky_re_lu_9",
                "conv2d_10", 
                "batch_normalization_10",
                "leaky_re_lu_10",
                "conv2d_11", 
                "batch_normalization_11",
                "leaky_re_lu_11",
                "conv2d_12", 
                "batch_normalization_12",
                "leaky_re_lu_12",
                "conv2d_13", 
                "batch_normalization_13",
                "leaky_re_lu_13"
            ]
        paralle1 = ["max_pooling2d_5",
                "conv2d_14", 
                "batch_normalization_14",
                "leaky_re_lu_14",
                "conv2d_15", 
                "batch_normalization_15",
                "leaky_re_lu_15",
                "conv2d_16", 
                "batch_normalization_16",
                "leaky_re_lu_16",
                "conv2d_17", 
                "batch_normalization_17",
                "leaky_re_lu_17",
                "conv2d_18", 
                "batch_normalization_18",
                "leaky_re_lu_18",
                "conv2d_19", 
                "batch_normalization_19",
                "leaky_re_lu_19",
                "conv2d_20", 
                "batch_normalization_20",
                "leaky_re_lu_20"
            ]
        paralle2 = ["conv2d_21", 
                "batch_normalization_21",
                "leaky_re_lu_21",
                "space_to_depth_x2"
            ]
        path2 = [
                "conv2d_22", 
                "batch_normalization_22",
                "leaky_re_lu_22",
                "conv2d_23"
            ]
        return path1, paralle1, paralle2, path2


    def loadWeights(self):
        model = load_model(self.modelUrl)        
        j = json.loads(model.to_json())


        for i, layer in enumerate(j['config']['layers']):
            ln = layer['name']
            l = model.get_layer(name=layer['name'])
            if layer['class_name'] != 'Concatenate':
                self.lid[ln] = l.input_shape[3]
            else:
                self.lid[ln] = l.input_shape[0][3]
            self.lod[ln] = l.output_shape[3]
            w = l.get_weights()
            if layer['class_name'] == 'Conv2D':
                filter_size = layer['config']['kernel_size'][0]
                if filter_size == 3:
                    self.layers[ln] = nn.Conv2d(self.lid[ln],self.lod[ln], 
                        filter_size,padding=1,stride=1,bias=False)
                elif filter_size==1:
                    self.layers[ln] = nn.Conv2d(self.lid[ln],self.lod[ln], 
                        filter_size,padding=0,stride=1,bias=False)
                self.layers[ln].weight.data = torch.from_numpy(w[0].transpose((3,2,0,1)))

            elif layer['class_name'] == 'BatchNormalization':
                self.layers[ln] = nn.BatchNorm2d(self.lid[ln])
                self.layers[ln].weight.data = torch.from_numpy(w[0])
                self.layers[ln].bias.data = torch.from_numpy(w[1])
                self.layers[ln].running_mean.data = torch.from_numpy(w[2])
                self.layers[ln].running_var.data = torch.from_numpy(w[3])

            elif layer['class_name'] == 'LeakyReLU':
                self.layers[ln] = nn.LeakyReLU(.1)
            elif layer['class_name'] == 'MaxPooling2D':
                self.layers[ln] = nn.MaxPool2d(2, 2)
            elif layer['class_name'] == 'Lambda':
                self.layers[ln] = scale_to_depth(2)