#!/usr/bin/env python
#
# file: $ISIP_EXP/tuh_dpath/exp_0074/scripts/model.py
#
# revision history:
#  20190925 (TE): first version
#
# usage:
#
# This script hold the model architecture
#------------------------------------------------------------------------------

# import pytorch modules
#
import torch
import torch.nn as nn

# import modules
#
import os
import random

# for reproducibility, we seed the rng
#
SEED1 = 1337
DEF_NUM_FEATS = 2
NUM_NODES = 32
NUM_CLASSES = 3

#-----------------------------------------------------------------------------
#
# helper functions are listed here
#
#-----------------------------------------------------------------------------

# function: set_seed
#
# arguments: seed - the seed for all the rng
#
# returns: none
#
# this method seeds all the random number generators and makes
# the results deterministic
#
def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
#
# end of method

#------------------------------------------------------------------------------
#
# the model is defined here
#
#------------------------------------------------------------------------------

# define the PyTorch MLP model
#
class Model(nn.Module):

    # function: init
    #
    # arguments: input_size - int representing size of input
    #            hidden_size - number of nodes in the hidden layer
    #            num_classes - number of classes to classify
    #
    # return: none
    #
    # This method is the main function.
    #
    def __init__(self, input_size, hidden_size, num_classes, dropout = 0.3):

        # inherit the superclass properties/methods
        #
        super(Model, self).__init__()

        # define the model
        #
        self.linear1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(hidden_size, int(hidden_size / 2))
        self.linear3 = nn.Linear(int(hidden_size / 2), num_classes)
        self.dropout = nn.Dropout(dropout)
    #
    # end of function


    # function: forward
    #
    # arguments: data - the input to the model
    #
    # return: out - the output of the model
    #
    # This method feeds the data through the network
    #
    def forward(self, data):

        # pass data to the layers
        #
        x = self.linear1(data)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.linear2(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.linear3(x)
        
        # return the output
        #
        return x
    #
    # end of method
#
# end of class

#
# end of file