import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
import numpy as np

class LSTM_SF(nn.Module):
    ''' Pure LSTM for stock price forecasting
    :parameter
        nhid, nfeat, batch_first

    :input  - (batch_size, seq_len, nfeat)
    :output - (batch_size)
    '''
    def __init__(self, nhid, nfeat, dropout=0, batch_first=True):
        super(LSTM_SF, self).__init__()

        self.hidden_size = nhid

        self.rnn = nn.LSTM(
            num_layers=1,
            input_size=nfeat,
            hidden_size=nhid,
            batch_first=batch_first)
        self.l_0 = nn.Linear(nhid, 1)

        self.dropout = nn.Dropout(dropout)

    def forward(self, xb):

        output, (h_n, c_n) = self.rnn(xb, None)
        # print(h_n.shape)
        h_n = h_n.view(-1, self.hidden_size)
        # print(h_n.shape)
        ''' print((h_n == output[:,-1,:]).all()) - tensor(True) '''
        h_n = self.dropout(h_n)
        res = F.relu(self.l_0(h_n))
        # print(res.shape)
        res = res.squeeze()
        return res

class LSTM_MF(nn.Module):
    ''' Pure LSTM for stock price forecasting
    :parameter
        nhid, nfeat, batch_first

    :input  - (batch_size, seq_len, nfeat)
    :output - (batch_size)
    '''
    def __init__(self, nhid, nfeat_static, nfeat_arith, nfeat_seq, dropout=0, batch_first=True):
        super(LSTM_MF, self).__init__()

        self.hidden_size = nhid

        self.rnn = nn.LSTM(
            num_layers=1,
            input_size=nfeat_seq,
            hidden_size=nhid,
            batch_first=batch_first)

        for name, param in self.rnn.named_parameters():
            if name.startswith("weight"):
                nn.init.uniform_(param)
            else:
                nn.init.zeros_(param)

        self.l_rnn      = nn.Linear(nhid, 1)
        self.l_arith    = nn.Linear(nfeat_arith, 1)
        self.l_static   = nn.Linear(nfeat_static, 1)
        self.l_0        = nn.Linear(3, 1)

        self.dropout = nn.Dropout(dropout)

    def forward(self, static_b, arith_b, seq_b):
        '''
        :param xb: (nstock, seq_len, nfeat)
        :param fb: (nstock, nfeat)
        :return:
        '''
        output, (h_n, c_n) = self.rnn(seq_b, None)
        h_n = h_n.view(-1, self.hidden_size)
        h_n = self.dropout(h_n)

        rnn_    = F.relu(self.l_rnn(h_n))
        arith_  = F.relu(self.l_arith(arith_b))
        static_ = F.relu(self.l_static(static_b))

        fc_input = torch.cat((rnn_, arith_, static_), 1)
        res = F.relu(self.l_0(fc_input))
        res = res.squeeze()
        return res

if __name__ == '__main__':

    pass