#!/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 numpy as np
import os
import random

# for reproducibility, we seed the rng
#
SEED1 = 1337
def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
set_seed(SEED1)

#------------------------------------------------------------------------------
#
# 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):

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

        # define the model
        #
        self.neural_net = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, num_classes))
    #
    # 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):

        # return the output
        #
        return self.neural_net(data)
    #
    # end of method
#
# end of class