"""
Different neural network architectures for detecting the line
"""
from __future__ import print_function, division, absolute_import

import torch.nn as nn
import torch.nn.functional as F


class MlpNetwork(nn.Module):
    """
    MLP network for detecting the line
    :param input_dim: (int) 3 x H x W
    :param drop_p: (float) Dropout proba
    """

    def __init__(self, input_dim, num_output=6, drop_p=0.0):
        super(MlpNetwork, self).__init__()
        self.fc1 = nn.Linear(input_dim, 20)
        self.fc2 = nn.Linear(20, 4)
        self.fc3 = nn.Linear(4, num_output)
        self.drop_p = drop_p
        self.activation_fn = F.relu
        self._initializeWeights()

    def _initializeWeights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                pass
                # m.weight.data.normal_(0, 0.005)
                # m.weight.data.uniform_(-0.005, 0.005)
                # nn.init.xavier_uniform(m.weight.data, gain=nn.init.calculate_gain('relu'))
                # nn.init.kaiming_normal(m.weight.data)
                # nn.init.kaiming_uniform(m.weight.data)
                # m.bias.data.zero_()
                # m.bias.data.uniform_(-0.1, 0.1)

    def forward(self, x):
        # Flatten input
        x = x.view(x.size(0), -1)
        x = F.dropout(x, p=self.drop_p, training=self.training)
        x = self.activation_fn(self.fc1(x))
        x = self.activation_fn(self.fc2(x))
        x = self.activation_fn(self.fc3(x))
        return x


class ConvolutionalNetwork(nn.Module):
    def __init__(self, drop_p=0.0, num_output=6):
        super(ConvolutionalNetwork, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(3, 20, kernel_size=7, stride=2, padding=3),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, padding=1),

            nn.Conv2d(20, 20, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2)
        )

        # Factor = 1
        # self.fc1 = nn.Linear(20 * 20 * 39, 16)
        # Factor = 2
        # self.fc1 = nn.Linear(20 * 10 * 19, 16)
        # Factor = 4
        self.fc1 = nn.Linear(20 * 5 * 9, 16)
        self.fc2 = nn.Linear(16, num_output)
        self.drop_p = drop_p

    def forward(self, x):
        x = F.dropout(x, p=self.drop_p, training=self.training)
        x = self.conv_layers(x)
        # print(x.shape)
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x


class CustomNet(nn.Module):
    def __init__(self, num_output=6):
        super(CustomNet, self).__init__()

        self.model = nn.Sequential(
            nn.Conv2d(3, 20, kernel_size=7, stride=2, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),

            nn.Conv2d(20, 20, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        self.fc = nn.Sequential(
            nn.Linear(20 * 4 * 8, 32),
            # nn.Linear(20*9*18, 32),
            nn.ReLU(inplace=True),
            nn.Linear(32, 32),
            nn.ReLU(inplace=True),
            nn.Linear(32, num_output)
        )

    def forward(self, x):
        x = self.model(x)
        # print(x.shape)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x
