#!/usr/bin/env python
# -*- coding: utf-8 -*-
__date__ = '2018/4/3 21:24'
__author__ = 'ooo'

import numpy as np
import torch
import torch.nn as nn
import math
import torch.nn.functional as F


class FPNet(nn.Module):
    def __init__(self, indepth, outdepth, stages):
        super(FPNet, self).__init__()
        self.stages = stages
        self.P6 = nn.MaxPool2d(kernel_size=1, stride=2)
        self.P5_conv1 = nn.Conv2d(indepth[3], outdepth, kernel_size=1, stride=1)
        self.P5_conv2 = nn.Sequential(
            SamePad2d(kernel_size=3, stride=1),
            nn.Conv2d(outdepth, outdepth, kernel_size=3, stride=1),
        )
        self.P4_conv1 = nn.Conv2d(indepth[2], outdepth, kernel_size=1, stride=1)
        self.P4_conv2 = nn.Sequential(
            SamePad2d(kernel_size=3, stride=1),
            nn.Conv2d(outdepth, outdepth, kernel_size=3, stride=1),
        )
        self.P3_conv1 = nn.Conv2d(indepth[1], outdepth, kernel_size=1, stride=1)
        self.P3_conv2 = nn.Sequential(
            SamePad2d(kernel_size=3, stride=1),
            nn.Conv2d(outdepth, outdepth, kernel_size=3, stride=1),
        )
        self.P2_conv1 = nn.Conv2d(indepth[0], outdepth, kernel_size=1, stride=1)
        self.P2_conv2 = nn.Sequential(
            SamePad2d(kernel_size=3, stride=1),
            nn.Conv2d(outdepth, outdepth, kernel_size=3, stride=1),
        )

    def forward(self, feature_maps):
        _, C2, C3, C4, C5, _ = feature_maps
        p5_out = self.P5_conv1(C5)
        p4_out = self.P4_conv1(C4) + F.upsample(p5_out, scale_factor=2)
        p3_out = self.P3_conv1(C3) + F.upsample(p4_out, scale_factor=2)
        p2_out = self.P2_conv1(C2) + F.upsample(p3_out, scale_factor=2)

        p5_out = self.P5_conv2(p5_out)
        p4_out = self.P4_conv2(p4_out)
        p3_out = self.P3_conv2(p3_out)
        p2_out = self.P2_conv2(p2_out)

        # P6 is used for the 5th anchor scale in RPN. Generated by
        # subsampling from P5 with stride of 2.
        p6_out = self.P6(p5_out)

        return [p2_out, p3_out, p4_out, p5_out, p6_out]


class SamePad2d(nn.Module):
    """Mimics tensorflow's 'SAME' padding.
    """

    def __init__(self, kernel_size, stride):
        super(SamePad2d, self).__init__()
        self.kernel_size = torch.nn.modules.utils._pair(kernel_size)
        self.stride = torch.nn.modules.utils._pair(stride)

    def forward(self, input):
        in_width = input.size()[2]
        in_height = input.size()[3]
        out_width = math.ceil(float(in_width) / float(self.stride[0]))
        out_height = math.ceil(float(in_height) / float(self.stride[1]))
        pad_along_width = ((out_width - 1) * self.stride[0] +
                           self.kernel_size[0] - in_width)
        pad_along_height = ((out_height - 1) * self.stride[1] +
                            self.kernel_size[1] - in_height)
        pad_left = math.floor(pad_along_width / 2)
        pad_top = math.floor(pad_along_height / 2)
        pad_right = pad_along_width - pad_left
        pad_bottom = pad_along_height - pad_top
        return F.pad(input, (pad_left, pad_right, pad_top, pad_bottom), 'constant', 0)

    def __repr__(self):
        return self.__class__.__name__