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

from model_lstm import LSTM_SF, LSTM_MF
from model_gcn import WGCN_MF_2, WGCN_2

nhid    = 64
nhid_gcn = 128
dropout = 0.8

def get_model(model_name):
    if   model_name == 'LSTM_SF':
        return LSTM_SF(nhid=nhid,
                       nfeat=5,
                       dropout=0.8,
                       batch_first=True)
    elif model_name == 'LSTM_MF':
        ''' input: static_b, arith_b, seq_b '''
        return LSTM_MF(nhid=nhid,
                       nfeat_static=6,
                       nfeat_arith=4,
                       nfeat_seq=6,
                       dropout=0.8,
                       batch_first=True)
    elif model_name == 'GCN_SF':
        return WGCN_2(
            nfeat_seq=5, nhid_rnn=nhid_gcn,
            nhid_gcn=nhid_gcn, nout_gcn=32, nrelation=8,
            dropout=dropout
        )
    elif model_name == 'GCN_MF':
        return WGCN_MF_2(nfeat_static=5,
                         nfeat_arith=4,
                         nfeat_seq=5, nhid_rnn=nhid_gcn,
                         nhid_gcn=nhid_gcn, nout_gcn=32, nrelation=8,
                         dropout=dropout)
    else:
        return None

def fake_m_seq(model):
    seq_len, nfeat_static, nfeat_arith, nstock, nday = 5, 6, 4, 300, 20
    nfeat_seq = nfeat_static
    bs = nstock

    ndata = nfeat_seq * seq_len * nstock * nday
    x_seq = np.arange(ndata).reshape(-1, seq_len, nfeat_seq)
    x_arith = np.arange(nfeat_arith * nstock * nday).reshape(-1, nfeat_arith)
    # x_static = x_seq[:, -1, :]
    y = np.random.random(nstock * nday)
    x_seq, x_arith, y = \
        map(torch.Tensor, (x_seq, x_arith, y))
    print(x_seq.shape, x_arith.shape, y.shape)

    ''' Data '''
    dl_train = DataLoader(TensorDataset(x_arith, x_seq, y), batch_size=bs)
    arith_b, seq_b, y_b = next(iter(dl_train))
    static_b = seq_b[:, -1, :]
    print(static_b.shape, arith_b.shape, seq_b.shape, y_b.shape)

    ''' Model '''
    # model = get_model(nhid=64,
    #                   nfeat_static=nfeat_static,
    #                   nfeat_arith=nfeat_arith,
    #                   nfeat_seq=nfeat_seq,
    #                   batch_first=True)

    preds = model(static_b, arith_b, seq_b)
    print(preds.shape)

    ''' Loss '''
    loss_func = F.mse_loss

    loss = loss_func(preds, y_b)
    print(loss)

def fake_s_seq(model):
    seq_len, nfeat, nstock, nday = 5, 6, 300, 20
    bs = nstock

    ndata = nfeat * seq_len * nstock * nday
    x = np.arange(ndata).reshape(-1, seq_len, nfeat)
    y = np.random.random(nstock * nday)

    x = torch.Tensor(x)
    y = torch.LongTensor(y)
    print(x.shape, y.shape)

    ''' Data '''
    dl_train = DataLoader(TensorDataset(x, y), batch_size=bs)
    xb, yb = next(iter(dl_train))
    print(xb.shape, yb.shape)
    exit()

    ''' Model '''
    # model = get_model(nhid=64, nfeat=nfeat, batch_first=True)

    preds = model(xb)
    print(preds.shape)

    ''' Loss '''
    loss_func = F.mse_loss

    loss = loss_func(preds, yb)
    print(loss)

if __name__ == '__main__':


    # fake_m_seq(get_model('GRU_MF'))
    
    fake_s_seq(get_model('GRU_SF'))