import sys
import os
import cv2
import torch
import torch.nn as nn

def op_name(op_name, m):
    m.op_name = op_name
    return m

def conv_bn_relu(name, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, momentum = 0.1, track_running_stats=True):
    return nn.Sequential(
        op_name(name, nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, False)),
        op_name(name + '/bn', nn.BatchNorm2d(out_channels, momentum = momentum, track_running_stats = track_running_stats)),
        op_name(name + '/relu', nn.ReLU(inplace=True)),
    )


# Define a Basic resnet block
class BasicResnetBlock(nn.Module):
    def __init__(self, name, in_channels, out_channels, kernel_size=3, stride=2, padding=1, direct_plus=False):
        super(BasicResnetBlock, self).__init__()

        self.op_name = name

        self.conv_block = nn.Sequential(
            conv_bn_relu(name=name + '/conv1', in_channels=in_channels, out_channels=out_channels,
                         kernel_size=kernel_size, stride=stride, padding=padding),

            conv_bn_relu(name=name + '/conv2', in_channels=out_channels, out_channels=out_channels,
                         kernel_size=3, stride=1, padding=1)
        )

        if direct_plus and (in_channels == out_channels) and (stride == 1):
            self.sc_block = None
        else:
            self.sc_block = conv_bn_relu(name=name + '/sc_conv', in_channels=in_channels, out_channels=out_channels,
                                         kernel_size=1, stride=stride, padding=0)

    def forward(self, x):
        if self.sc_block is None:
            out = x + self.conv_block(x)
        else:
            out = self.conv_block(x) + self.sc_block(x)
        return out

class Flatten(nn.Module):
    def __init__(self, axis):
        super(Flatten, self).__init__()
        self.axis = axis

    def forward(self, x):
        assert self.axis == 1
        x = x.reshape(x.shape[0], -1)
        # x = x.view(-1, 1152)
        return x

def flatten(name, axis):
    return op_name(name, Flatten(axis))

def linear_bn_relu(name, in_features, out_features, momentum = 0.1, track_running_stats=True):
    return nn.Sequential(
        op_name(name + '/FC', nn.Linear(in_features, out_features)),
        op_name(name + '/bn', nn.BatchNorm1d(out_features, momentum = momentum, track_running_stats = track_running_stats)),
        op_name(name + '/relu', nn.ReLU(inplace=True)),
    )

class LeftEye(nn.Module):
    def __init__(self, name, in_channels, out_channels):
        super(LeftEye, self).__init__()
        self.op_name = name

        op_list = []

        op_list += [conv_bn_relu(name + '/first_conv', in_channels, 24, kernel_size = 5, stride = 2, padding = 2) ]

        ch_num = [24, 32, 64, 96, 128]

        op_list += [BasicResnetBlock(name + '/stage%d'%(i + 1), ch_num[i], ch_num[i+1]) for i in range(len(ch_num) - 1) ]

        op_list += [flatten(name + '/flatten', 1),
                    linear_bn_relu(name + '/FC1', 1152, 256),
                    op_name(name + '/FC2', nn.Linear(256, out_channels))]

        self.conv_block = nn.Sequential(*op_list)

    def forward(self, x):
        return self.conv_block(x)

def get_left_eye_symbol(symbol_name='BigResNet', input_nc = 3, output_nc = 17 * 2):
    return LeftEye(symbol_name, input_nc, output_nc)

if __name__=='__main__':
    pass