import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import pickle, numpy, math

# [(Trade_Time, PRICE, AMOUNT), ...]
# [(1620457034392, 56000, 0.5), (1620457034394, 56001, 0.05), ...]
# The trades usually have various time-interval. TODO: transform it to [(WeightedAvgPrice, Volume), ...] for every 1 minutes? 
# Time axis: [history, older, newer, ..., latest]
realtime_shortterm_dataset_aggtrade = []
realtime_shortterm_dataset_aggtrade_size = 1024*1024

# [(Snapshot_Time, [(buy_price1, amount1), ...] , [(sell_price1, amount1), ...]), ...]
# [(1620457034392, [(56000, 0.01), (55900, 1), (55700, 30), ...] , [(57000, 0.01), (57100, 1), ...] ), (1620457034394, [...]), ...]
# The snapshots should has almost identical time-interval. Good for LSTM. 
# Time axis: [history, older, newer, ..., latest]
realtime_shortterm_dataset_depth = []
realtime_shortterm_dataset_depth_size = 1024

# The trading thread would not start working, before finish analysing longterm dataset. 
# Time-interval for longterm dataset is 1 minute.
longterm_dataset = []


def load_realtime_dataset_on_start():
    global realtime_shortterm_dataset_aggtrade, realtime_shortterm_dataset_depth
    try:
        realtime_shortterm_dataset_aggtrade, realtime_shortterm_dataset_depth = pickle.load(open( "realtime_dataset_dump.pyobj", "rb"))
        print("Loaded {}+{} elements. ".format(len(realtime_shortterm_dataset_aggtrade), len(realtime_shortterm_dataset_depth)))
    except:
        print("No data to load. Skip data loading... ")


class LSTM_Shortterm_Predictor(nn.Module):
    def __init__(self, input_dim):
        super(LSTM_Shortterm_Predictor, self).__init__()
        self.lstm_idim = 10 # this is the length of depth_to_impulsive_score_vector()
        self.lstm_odim = 128

        # The input would be a tuple containing complex information. 
        # Firstly, serialize these information into a tuple. 
        self.serializer = nn.Linear(input_dim, self.lstm_idim)
        # TODO : remove this useless layer

        # The LSTM hidden states
        # with dimensionality hidden_dim.
        self.lstm = nn.LSTM(self.lstm_idim, self.lstm_odim)
        # TODO: this model is naive. solve the bug in dataset processing. 

        # The linear layer that maps from hidden state space to tag space
        self.out = nn.Linear(self.lstm_odim, 1)

    def forward(self, sample_seq):
        input_seq = sample_seq.view(len(sample_seq), 1, self.lstm_idim)
        lstm_in = self.serializer(input_seq)
        lstm_out, _ = self.lstm(lstm_in)
        predict_shortterm_trend = self.out(torch.tanh(lstm_out[-1:]))
        return predict_shortterm_trend

def depth_to_impulsive_score_vector(depth):
    _, buys, sells = depth
    def get_factors(array_of_pairs):
        values = [pair[0] for pair in array_of_pairs]
        weights = [pair[1] for pair in array_of_pairs]
        average = numpy.average(values, weights=weights)
        volume = numpy.sum(weights)
        variance = numpy.average((values-average)**2, weights=weights)
        leader_price, leader_weight = array_of_pairs[0]
        return (average, volume, math.sqrt(variance), leader_price, leader_weight)

    return get_factors(buys) + get_factors(sells)
    
model = LSTM_Shortterm_Predictor(10)
# optimizer = optim.SGD(model.parameters(), lr=0.1)
optimizer = optim.RMSprop(model.parameters(), lr=0.05, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0.75, centered=False)

def learn_once(depth_seq, trend_answer_score):
    # returns scalar loss
    input_seq = [depth_to_impulsive_score_vector(depth) for depth in depth_seq]

    model.zero_grad()
    i = torch.tensor(input_seq, dtype=torch.float)
    o = model(i)

    loss = abs(o - trend_answer_score)
    loss.backward(retain_graph=True)
    optimizer.step()

    return float(loss)

load_realtime_dataset_on_start()
losses = []
for i in range(2000):
    answer = numpy.average([realtime_shortterm_dataset_depth[i+di][1][0][0] for di in (256,512,768,1024,2000,5000)]) - realtime_shortterm_dataset_depth[i][1][0][0]
    print("answer=", answer)
    loss = learn_once(realtime_shortterm_dataset_depth[i:i+128], answer)
    print("Loss=", loss)
    losses += [loss]

#####################
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import numpy as np
x = numpy.arange(len(losses))
colors = cm.rainbow(np.linspace(0, 1, 10))
for index, y in enumerate(losses):
    plt.scatter(index, y, color=colors[1])
#for x, y in guess_xy:
#    plt.scatter(x, y, color=colors[6])
plt.show()




