#!/usr/bin/python3
# https://binance-docs.github.io/apidocs/spot/en/#test-new-order-trade

# we have one thread to maintain the dataset, another thread would run the pytorch model every 5 seconds, to see if it's time to place order. 

import json, time
from websocket import create_connection
import sys, pickle
from signal import *

# [(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 = -1#1024*1024

# [(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 = -1 #1024

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

def save_realtime_dataset_on_exit():
    global realtime_shortterm_dataset_aggtrade, realtime_shortterm_dataset_depth
    print("Saving {}+{} elements...".format(len(realtime_shortterm_dataset_aggtrade), len(realtime_shortterm_dataset_depth)))
    pickle.dump((realtime_shortterm_dataset_aggtrade, realtime_shortterm_dataset_depth), open( "realtime_dataset_dump.pyobj", "wb"))
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... ")

def dataset_maintain_thread_main():
    global realtime_shortterm_dataset_aggtrade, realtime_shortterm_dataset_aggtrade_size, realtime_shortterm_dataset_depth, realtime_shortterm_dataset_depth_size

    ws = create_connection("wss://stream.binance.com:9443/ws/btcusdt")
    ws.send(json.dumps({
        "method": "SUBSCRIBE",
        "params": ["btcusdt@aggTrade", "btcusdt@depth"],
        "id": 1,
    }))

    subs_response = ws.recv()
    if '"result":null' not in subs_response:
        raise RuntimeError("Failed to subscribe: server says: " + subs_response)
    
    _debug_tmp = 0
    while True:
        message = ws.recv()
        # print("DEBUG: message=", message)
        parsed = json.loads(message)
        if 'p' in parsed.keys():
            # Is a trade message
            trade_time, price, amount = parsed.get('T'), parsed.get('p'), parsed.get('q')
            realtime_shortterm_dataset_aggtrade.append((trade_time, price, amount))
            if realtime_shortterm_dataset_aggtrade_size != -1 and len(realtime_shortterm_dataset_aggtrade) > realtime_shortterm_dataset_aggtrade_size:
                del realtime_shortterm_dataset_aggtrade[0]
        elif 'b' in parsed.keys():
            # Is a depth snapshot update
            snapshot_time, buy_makers, sell_makers = parsed.get('E'), parsed.get('b'), parsed.get('a')
            buy_makes_parsed = [(float(ele[0]), float(ele[1])) for ele in buy_makers]
            sell_makes_parsed = [(float(ele[0]), float(ele[1])) for ele in sell_makers]
            realtime_shortterm_dataset_depth.append((snapshot_time, buy_makes_parsed, sell_makes_parsed))
            if realtime_shortterm_dataset_depth_size != -1 and len(realtime_shortterm_dataset_depth) > realtime_shortterm_dataset_depth_size:
                del realtime_shortterm_dataset_depth[0]
        else:
            print("ERROR: unexpected message (maybe ping message): ", message)

        ###### DEBUG ONLY
        if int(time.time()) % 2000 == 0 and _debug_tmp != int(time.time()): # 2000 sec
            _debug_tmp = int(time.time())
            print("DEBUG: len(depth) = ", len(realtime_shortterm_dataset_depth))
        #     print("DEBUG: dumping realtime_shortterm_dataset_aggtrade================")
        #     print(realtime_shortterm_dataset_aggtrade)
        #     print("DEBUG: dumping realtime_shortterm_dataset_depth================")
        #     print(realtime_shortterm_dataset_depth)
        #     print("DEBUG: dumping END ==|||||||||||||||||||||||||||||================")
    
    ws.close()

def sigint_handler(sig, frame):
    print('Got SIGINT! Saving realtime dataset... ')
    save_realtime_dataset_on_exit()
    sys.exit(0)

for sig in (SIGABRT, SIGILL, SIGINT, SIGSEGV, SIGTERM):
    signal(sig, sigint_handler)

load_realtime_dataset_on_start()
try:
    dataset_maintain_thread_main()
except:
    save_realtime_dataset_on_exit()
    raise


