import os.path
from datetime import datetime, timedelta

import numpy as np
import pandas as pd
from hftbacktest import EXCH_EVENT, LOCAL_EVENT
from numba import njit

date_from = 20240501
date_to = 20240531

# Target symbol used to generate order latency based on its feed data latency.
# To obtain realistic backtesting results, it is best to use the actual historical order latency for each pair.
# For example purposes, we use the feed latency of one pair as the order latency.
symbol = 'SOLUSDT'

# Order latency can differ significantly from feed latency. To artificially generate order latency, multiply the given
# multiplier by the order entry latency and order response latency.
mul_entry = 4
mul_resp = 3

# The path for the converted npz files for the Rust version.
npz_path = '.'

# The path for the generated latency files for the Rust version.
latency_path = '.'


@njit
def generate_order_latency_nb(data, order_latency, mul_entry, offset_entry, mul_resp, offset_resp):
    for i in range(len(data)):
        exch_ts = data[i].exch_ts
        local_ts = data[i].local_ts
        feed_latency = local_ts - exch_ts
        order_entry_latency = mul_entry * feed_latency + offset_entry
        order_resp_latency = mul_resp * feed_latency + offset_resp

        req_ts = local_ts
        order_exch_ts = req_ts + order_entry_latency
        resp_ts = order_exch_ts + order_resp_latency

        order_latency[i].req_timestamp = req_ts
        order_latency[i].exch_timestamp = order_exch_ts
        order_latency[i].resp_timestamp = resp_ts


def generate_order_latency(feed_file, output_file=None, mul_entry=1, offset_entry=0, mul_resp=1, offset_resp=0):
    data = np.load(feed_file)['data']
    df = pd.DataFrame(data)
    df = df[(df['ev'] & EXCH_EVENT == EXCH_EVENT) | (df['ev'] & LOCAL_EVENT == LOCAL_EVENT)]
    s = (df['local_ts'] / 1_000_000_000).astype(int)
    df = df.groupby(s).last()
    data = df.to_records(index=False)

    order_latency = np.zeros(
        len(data),
        dtype=[
            ('req_ts', '<i8'),
            ('exch_ts', '<i8'),
            ('resp_ts', '<i8'),
            ('_padding', '<i8')
        ]
    )
    generate_order_latency_nb(data, order_latency, mul_entry, offset_entry, mul_resp, offset_resp)

    if output_file is not None:
        np.savez_compressed(output_file, data=order_latency)

    return order_latency


date = datetime.strptime(str(date_from), '%Y%m%d')
date_to = datetime.strptime(str(date_to), '%Y%m%d')
while date <= date_to:
    yyyymmdd = date.strftime('%Y%m%d')
    print(f'Generating order latency for {yyyymmdd} from the feed latency')
    try:
        generate_order_latency(
            os.path.join(npz_path, f'{symbol}_{yyyymmdd}.npz'),
            output_file=os.path.join(latency_path, f'latency_{yyyymmdd}.npz'),
            mul_entry=mul_entry,
            mul_resp=mul_resp
        )
    except Exception as e:
        print(e, yyyymmdd)
    date += timedelta(days=1)