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


from Model.gcn_layers import GraphConvolution

# def get_model_WGCN_MF(nhid, nfeat, nrelation, batch_first=True):
#     return WGCN_MF(nrelation=nrelation, nfeat=nfeat, nhid=nhid, nclass=10)

def get_model_WGCN_MF(nfeat_static, nfeat_arith, nfeat_seq, nhid_rnn, nhid_gcn, nout_gcn, nrelation, dropout=0):
    return WGCN_MF_2(nfeat_static=nfeat_static,
                 nfeat_arith=nfeat_arith,
                 nfeat_seq=nfeat_seq, nhid_rnn=nhid_rnn,
                 nhid_gcn=nhid_gcn, nout_gcn=nout_gcn, nrelation=nrelation,
                 dropout=dropout)

# def get_model_WGCN(nhid, nfeat, nrelation, batch_first=True):
#     return WGCN(nrelation=nrelation, nfeat=nfeat, nhid_gcn=nhid, nout_gcn=12)

def get_model_WGCN(nfeat_seq, nhid_rnn, nhid_gcn, nout_gcn, nrelation, dropout, batch_first=True):
    return WGCN_2(
        nfeat_seq=nfeat_seq, nhid_rnn=nhid_rnn,
        nhid_gcn=nhid_gcn, nout_gcn=nout_gcn, nrelation=nrelation,
        dropout=dropout
    )

def get_model_GCN(nfeat_seq, nhid_rnn, nhid_gcn, nout_gcn, dropout, batch_first=True):

    pass


class WGCN_MF_2(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,
                 nfeat_static,
                 nfeat_arith,
                 nfeat_seq, nhid_rnn,
                 nhid_gcn, nout_gcn, nrelation,
                 dropout=0, batch_first=True):
        '''

        :param nfeat_static:
        :param nfeat_arith:
        :param nfeat_seq:
        :param nhid_rnn:    rnn 隐藏层输出的维度，也是 gcn 输入的维度
        :param nhid_gcn:
        :param nout_gcn:
        :param nrelation:
        :param dropout:
        :param batch_first:
        '''
        super(WGCN_MF_2, self).__init__()

        self.nhid_rnn = nhid_rnn

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

        self.gc1 = GraphConvolution(nrelation, nhid_rnn, nhid_gcn)
        self.gc2 = GraphConvolution(nrelation, nhid_gcn, nout_gcn)

        self.l_gcn      = nn.Linear(nout_gcn, 1)
        self.l_rnn      = nn.Linear(nhid_rnn, 1)
        self.l_arith    = nn.Linear(nfeat_arith, 1)
        self.l_static   = nn.Linear(nfeat_static, 1)
        self.l_0        = nn.Linear(4, 1)

        self.dropout = nn.Dropout(dropout)

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

    def weight_init(self):
        self.l_gcn.reset_parameters()
        self.l_rnn.reset_parameters()
        self.l_arith.reset_parameters()
        self.l_static.reset_parameters()
        self.l_0.reset_parameters()
        for name, param in self.rnn.named_parameters():
            if name.startswith("weight"):
                nn.init.uniform_(param)
            else:
                nn.init.zeros_(param)

    def forward(self, static_b, arith_b, seq_b, adj):
        # print('seq_b - ', seq_b)
        output, (h_n, c_n) = self.rnn(seq_b, None)
        h_n = h_n.view(-1, self.nhid_rnn)
        h_n = self.dropout(h_n)
        # print('h_n -', h_n)

        x_gcn    = F.relu(self.gc1(h_n, adj))
        x_gcn    = F.relu(self.gc2(x_gcn, adj))

        # print('x_gcn -', x_gcn)
        gcn_    = F.relu(self.l_gcn(x_gcn))
        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_, gcn_), 1)
        res = F.relu(self.l_0(fc_input))
        # print(res)
        # exit()
        res = res.squeeze()
        # print('res', res)
        return res

class WGCN_2(nn.Module):
    def __init__(self, nfeat_seq, nhid_rnn, nhid_gcn, nout_gcn, nrelation, dropout=0.1):
        super(WGCN_2, self).__init__()
        self.nhid_rnn = nhid_rnn
        
        self.gc1 = GraphConvolution(nrelation, nhid_rnn, nhid_gcn)
        self.gc2 = GraphConvolution(nrelation, nhid_gcn, nout_gcn)

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

        self.fc = nn.Linear(nout_gcn, 1)

        self.dropout = nn.Dropout(dropout)

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

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

    def forward(self, seq_b, adj):

        output, (h_n, c_n) = self.rnn(seq_b, None)
        h_n = h_n.view(-1, self.nhid_rnn)
        h_n = self.dropout(h_n)

        x_gcn = F.relu(self.gc1(h_n, adj))
        x_gcn = F.relu(self.gc2(x_gcn, adj))

        output = self.fc(x_gcn)
        # print(fc_input.shape, output.shape)
        return output.squeeze()

def test_WGCN():

    '''
    在一个周期内，adj是固定的
    '''
    nstock, nday, seq_len, nfeat_seq = 300, 20, 5, 6
    bs = nstock
    nrelation = 10

    ndata = nstock * nday * seq_len * nfeat_seq
    x_seq = np.arange(ndata).reshape(-1, seq_len, nfeat_seq) # (nstock * nday, seq_len, nfeat)
    adj = np.ones((nstock, nstock, nrelation))
    # 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, y, adj = \
        map(torch.Tensor, (x_seq, y, adj))
    print(x_seq.shape, y.shape, adj.shape)

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

    ''' Model '''
    model = get_model_WGCN(
        nfeat_seq=nfeat_seq, nhid_rnn=64,
        nhid_gcn=64, nout_gcn=32, nrelation=nrelation,
        dropout=0
    )

    preds = model(seq_b, adj)
    print(preds.shape)

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

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

def test_WGCN_MF():

    nfeat_arith = 7
    nstock, nday, seq_len, nfeat_seq = 300, 20, 5, 6
    bs = nstock
    nrelation = 10

    adj = np.ones((nstock, nstock, nrelation))
    ndata = nstock * nday * seq_len * nfeat_seq
    x_seq = np.arange(ndata).reshape(-1, seq_len, nfeat_seq)  # (nstock * nday, seq_len, nfeat)
    x_arith = np.arange(nfeat_arith * nstock * nday).reshape(-1, nfeat_arith)
    y = np.random.random(nstock * nday)
    x_arith, x_seq, y, adj = \
        map(torch.Tensor, (x_arith, x_seq, y, adj))
    print(x_seq.shape, y.shape, adj.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_WGCN_MF(
        nfeat_static=nfeat_seq,
        nfeat_arith=nfeat_arith,
        nfeat_seq=nfeat_seq, nhid_rnn=64,
        nhid_gcn=64, nout_gcn=32, nrelation=nrelation,
        )
    preds = model(static_b, arith_b, seq_b, adj)
    print(preds.shape)

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

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

if __name__ == '__main__':

    # test_WGCN()
    test_WGCN_MF()
    pass


