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

from A_Model import get_model
from B_data_tool import get_static_data, get_seq_data
from B_date_tool import gen_date
from D_relation import gen_adj
from A_Global import path_stock, args


def get_data_rnn_sf(args={}):
    '''
    feat_label = {
        'feat_static': ['open', 'close', 'low', 'high', 'volume', 'money'],
        'feat_arith': ['r1d', 'r5d', 'r10d', 'r20d', 'r30d'],
        'feat_seq': ['open', 'close', 'low', 'high', 'volume', 'money'],
        'label': ['label']
    }
    '''
    '''1 Seq Data'''
    date_dict = gen_date(args)

    print('============ Prepare Data Dict - Seq    ============')
    seq_data_dict = get_seq_data(date_dict, args, 'feat_seq')

    return seq_data_dict


def get_data_rnn_mf(args={}):
    '''
    feat_label = {
        'feat_static': ['open', 'close', 'low', 'high', 'volume', 'money'],
        'feat_arith': ['r1d', 'r5d', 'r10d', 'r20d', 'r30d'],
        'feat_seq': ['open', 'close', 'low', 'high', 'volume', 'money'],
        'label': ['label']
    }
    '''
    '''1 Seq Data'''
    date_dict = gen_date(args)

    print('============ Prepare Data Dict - Arith ============')
    arith_data_dict    = get_static_data(date_dict, args, 'feat_arith')

    print('============ Prepare Data Dict - Seq    ============')
    seq_data_dict       = get_seq_data(date_dict, args, 'feat_seq')

    return arith_data_dict, seq_data_dict


def get_data_gcn_sf(args={}):

    seq_data_dict = get_data_rnn_sf(args)

    adj = gen_adj(args)

    return seq_data_dict, adj

def get_data_gcn_mf(args={}):

    arith_data_dict, seq_data_dict = get_data_rnn_mf(args)

    adj = gen_adj(args)

    return arith_data_dict, seq_data_dict, adj

def demo_mf():
    arith_data_dict, seq_data_dict = get_data_rnn_mf(args)

    dl_train = DataLoader(TensorDataset(arith_data_dict['x_train'], seq_data_dict['x_train'], seq_data_dict['y_train']),
                          batch_size=len(args['stocks']))
    print('===========================DEMO==============================')
    arith_b, seq_b, y_b = next(iter(dl_train))
    static_b = seq_b[:, -1, :]
    arith_b = arith_b.squeeze()
    print(static_b.shape, arith_b.shape, seq_b.shape, y_b.shape)

    ''' Model '''
    model = get_model('LSTM_MF')

    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)
    print('===========================DEMO==============================')

if __name__ == '__main__':
    from tool_data import get_stock_filted
    market              = 'NYSE'
    args['market']      = market
    stocks, stock2idx   = get_stock_filted(market)
    args['stocks']      = stocks

    arith_data_dict, seq_data_dict, adj = get_data_gcn_mf(args)

    dl_train = DataLoader(TensorDataset(arith_data_dict['x_train'], seq_data_dict['x_train'], seq_data_dict['y_train']),
                          batch_size=len(args['stocks']))
    print('===========================DEMO==============================')
    arith_b, seq_b, y_b = next(iter(dl_train))
    static_b = seq_b[:, -1, :]
    arith_b = arith_b.squeeze()
    print(static_b.shape, arith_b.shape, seq_b.shape, y_b.shape)

    ''' Model '''
    model = get_model('GCN_MF')

    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)
    print('===========================DEMO==============================')