#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import print_function

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


# deterministic training
random.seed(0)
np.random.seed(0)
torch.manual_seed(0)


def weight_init(m):
    if isinstance(m, nn.Linear):
        nn.init.kaiming_normal_(m.weight)


#  class for linear layers (see class below)
class Linear(nn.Module):
    def __init__(self, linear_size, p_dropout=0.5):
        super(Linear, self).__init__()
        self.l_size = linear_size

        self.relu = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p_dropout)

        self.w1 = nn.Linear(self.l_size, self.l_size)
        self.batch_norm1 = nn.BatchNorm1d(self.l_size)

        self.w2 = nn.Linear(self.l_size, self.l_size)
        self.batch_norm2 = nn.BatchNorm1d(self.l_size)

    def forward(self, x):
        y = self.w1(x)
        y = self.batch_norm1(y)
        y = self.relu(y)
        y = self.dropout(y)

        y = self.w2(y)
        y = self.batch_norm2(y)
        y = self.relu(y)
        y = self.dropout(y)

        out = x + y

        return out


class LinearModel(nn.Module):

    # define the components of the network
    def __init__(
        self,
        linear_size=1024,
        num_stage=2,
        p_dropout=0.5,
        input_size=24,
        output_size=12,
        drop_inp=0,
    ):
        super(LinearModel, self).__init__()

        self.linear_size = linear_size
        self.p_dropout = p_dropout
        self.num_stage = num_stage
        self.input_size = input_size
        self.output_size = output_size

        # process input to linear size
        self.w1 = nn.Linear(self.input_size, self.linear_size)
        self.batch_norm1 = nn.BatchNorm1d(self.linear_size)

        # linear layers
        self.linear_stages = []
        for l in range(num_stage):
            self.linear_stages.append(Linear(self.linear_size, self.p_dropout))
        self.linear_stages = nn.ModuleList(self.linear_stages)

        # post-processing
        self.w2 = nn.Linear(self.linear_size, self.output_size)

        self.relu = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(self.p_dropout)

        self.drop_inp = drop_inp
        # self.drop_inp = nn.Dropout(p=drop_inp)

    # this function assembles the network
    def forward(self, x):
        # pre-processing
        if self.training:
            # drop when A is True
            A_array = torch.rand(x.size()).cuda() < self.drop_inp
            no_drop = torch.rand(x.size(0)) < 0.5
            A_array[no_drop] = False
            x = x.masked_fill(A_array, 0)
        # iu#x = self.drop_inp(x)  # HACK HACK HACK

        y = self.w1(x)
        y = self.batch_norm1(y)
        y = self.relu(y)
        y = self.dropout(y)

        # linear layers
        for i in range(self.num_stage):
            y = self.linear_stages[i](y)

        # post-processing
        y = self.w2(y)

        return y
