#!/usr/bin/env python
#  -*- coding: utf-8 -*-
__author__ = 'mayanqiong'

from datetime import datetime, time, timedelta
from typing import Optional, Union

import numpy as np
import pandas as pd
from pandas import DataFrame

from tqsdk.api import TqApi
from tqsdk import utils
from tqsdk.datetime import _get_trading_timestamp, _get_trade_timestamp, _get_trading_day_from_timestamp
from tqsdk.rangeset import _rangeset_slice, _rangeset_head
from tqsdk.tradeable import TqAccount, TqKq, TqSim


def twap_table(api: TqApi, symbol: str, target_pos: int, duration: int, min_volume_each_step: int, max_volume_each_step: int,
         account: Optional[Union[TqAccount, TqKq, TqSim]] = None):
    """
    返回基于 twap 策略的计划任务时间表。下单需要配合 TargetPosScheduler 使用。

    Args:
        api (TqApi): TqApi实例，该task依托于指定api下单/撤单

        symbol (str): 拟下单的合约 symbol, 格式为 交易所代码.合约代码,  例如 "SHFE.cu1801"

        target_pos (int): 目标持仓手数

        duration (int): 算法执行的时长，以秒为单位，时长可以跨非交易时间段，但是不可以跨交易日
        * 设置为 60*10, 可以是 10:10～10:15 + 10:30~10:35

        min_volume_each_step (int): 调整持仓手数最小值，每步调整的持仓手数默认在最小和最大值中产生

        max_volume_each_step (int): 调整持仓手数最大值，每步调整的持仓手数默认在最小和最大值中产生

        account (TqAccount/TqKq/TqSim): [可选]指定发送下单指令的账户实例, 多账户模式下，该参数必须指定

    Returns:
        pandas.DataFrame: 本函数返回一个 pandas.DataFrame 实例. 表示一份计划任务时间表。每一行表示一项目标持仓任务，包含以下列:

            + interval: 当前这项任务的持续时间长度，单位为秒
            + target_pos: 当前这项任务的目标持仓
            + price: 当前这项任务的下单价格模式，支持 PASSIVE（排队价），ACTIVE（对价），None（不下单，表示暂停一段时间）

    Example1::

        from tqsdk import TqApi, TargetPosScheduler
        from tqsdk.algorithm import twap_table

        api = TqApi(auth="信易账户,用户密码")
        quote = api.get_quote("CZCE.MA109")

        # 设置twap任务参数
        time_table = twap_table(api, "CZCE.MA109", -100, 600, 1, 5)  # 目标持仓 -100 手，600s 内完成
        print(time_table.to_string())

        target_pos_sch = TargetPosScheduler(api, "CZCE.MA109", time_table)
        # 启动循环
        while not target_pos_sch.is_finished():
            api.wait_update()
        api.close()


    Example2::

        from tqsdk import TqApi, TargetPosScheduler
        from tqsdk.algorithm import twap_table

        api = TqApi(auth="信易账户,用户密码")
        quote = api.get_quote("CZCE.MA109")

        # 设置 twap 任务参数，
        time_table = twap_table(api, "CZCE.MA109", -100, 600, 1, 5)  # 目标持仓 -100 手，600s 内完成

        # 定制化调整 time_table，例如希望第一项任务延迟 10s 再开始下单
        # 可以在 time_table 的头部加一行
        time_table = pandas.concat([
            DataFrame([[10, 10, None]], columns=['interval', 'target_pos', 'price']),
            time_table
        ], ignore_index=True)

        target_pos_sch = TargetPosScheduler(api, "CZCE.MA109", time_table)
        while not target_pos_sch.is_finished():
            api.wait_update()

        # 获取 target_pos_sch 实例所有的成交列表
        print(target_pos_sch.trades_df)

        # 利用成交列表，您可以计算出策略的各种表现指标，例如：
        average_trade_price = sum(scheduler.trades_df['price'] * scheduler.trades_df['volume']) / sum(scheduler.trades_df['volume'])
        print("成交均价:", average_trade_price)
        api.close()

    """
    account = api._account._check_valid(account)
    if account is None:
        raise Exception(f"多账户模式下, 需要指定账户实例 account")
    min_volume_each_step = int(min_volume_each_step)
    max_volume_each_step = int(max_volume_each_step)
    if max_volume_each_step <= 0 or min_volume_each_step <= 0:
        raise Exception("请调整参数, min_volume_each_step、max_volume_each_step 必须是大于 0 的整数。")
    if min_volume_each_step > max_volume_each_step:
        raise Exception("请调整参数, min_volume_each_step 必须小于 max_volume_each_step。")

    pos = api.get_position(symbol, account)
    target_pos = int(target_pos)
    delta_pos = target_pos - pos.pos
    volume = abs(delta_pos)  # 总的下单手数

    # 得到有效的手数序列和时间间隔序列
    if volume < max_volume_each_step:
        interval_list, volume_list = [duration], [volume]
    else:
        volume_list = _gen_random_list(sum_val=volume, min_val=min_volume_each_step, max_val=max_volume_each_step)

        interval = int(duration / len(volume_list))
        if interval < 3:
            raise Exception("请调整参数, 每次下单时间间隔不能小于3s, 将单次下单手数阈值调大或者增长下单时间。")
        min_interval = int(max(3, interval - 2))
        max_interval = int(interval * 2 - max(3, interval - 2)) + 1
        interval_list = _gen_random_list(sum_val=duration, min_val=min_interval, max_val=max_interval,
                                         length=len(volume_list))

    time_table = DataFrame(columns=['interval', 'volume', 'price'])
    for index, volume in enumerate(volume_list):
        assert interval_list[index] >= 3
        active_interval = 2
        append_time_table = pd.DataFrame([
            {"interval": interval_list[index] - active_interval, "volume": volume, "price": "PASSIVE"},
            {"interval": active_interval, "volume": 0, "price": "ACTIVE"}
        ])
        time_table = pd.concat([time_table, append_time_table], ignore_index=True)
    time_table['volume'] = time_table['volume'].mul(-1 if delta_pos < 0 else 1)
    time_table['target_pos'] = time_table['volume'].cumsum()
    time_table['target_pos'] = time_table['target_pos'].add(pos.pos)
    time_table.drop(columns=['volume'], inplace=True)
    time_table = time_table.astype({'target_pos': 'int64', 'interval': 'float64'})
    return time_table


def vwap_table(api: TqApi, symbol: str, target_pos: int, duration: float,
               account: Optional[Union[TqAccount, TqKq, TqSim]] = None):
    """
    返回基于 vwap 策略的计划任务时间表。下单需要配合 TargetPosScheduler 使用。

    调用 vwap_table 函数，根据以下逻辑生成 time_table：

    1. 根据 target_pos - 当前合约的净持仓，得到总的需要调整手数
    2. 请求 symbol 合约的 ``1min`` K 线
    3. 采样取用最近 10 日内，以合约当前行情时间的下一分钟为起点，每日 duration / 60 根 K 线, \
    例如当前合约时间为 14:35:35，那么采样是会使用 14:36:00 开始的分钟线 K 线
    4. 按日期分组，分别计算交易日内，每根 K 线成交量占总成交量的比例
    5. 计算最近 10 日内相同分钟内的成交量占比的算术平均数，将第 1 步得到的总调整手数按照得到的比例分配
    6. 每一分钟，前 58s 以追加价格下单，后 2s 以对价价格下单

    Args:
        api (TqApi): TqApi实例，该task依托于指定api下单/撤单

        symbol (str): 拟下单的合约 symbol, 格式为 交易所代码.合约代码,  例如 "SHFE.cu2201"

        target_pos (int): 目标持仓手数

        duration (int): 算法执行的时长，以秒为单位，必须是 60 的整数倍，时长可以跨非交易时间段，但是不可以跨交易日
        * 设置为 60*10, 可以是 10:10～10:15 + 10:30~10:35

        account (TqAccount/TqKq/TqSim): [可选]指定发送下单指令的账户实例, 多账户模式下，该参数必须指定

    Returns:
        pandas.DataFrame: 本函数返回一个 pandas.DataFrame 实例. 表示一份计划任务时间表。每一行表示一项目标持仓任务，包含以下列:

            + interval: 当前这项任务的持续时间长度，单位为秒
            + target_pos: 当前这项任务的目标持仓
            + price: 当前这项任务的下单价格模式，支持 PASSIVE（排队价），ACTIVE（对价），None（不下单，表示暂停一段时间）

    Example1::

        from tqsdk import TqApi, TargetPosScheduler
        from tqsdk.algorithm import vwap_table

        api = TqApi(auth="信易账户,用户密码")
        quote = api.get_quote("CZCE.MA109")

        # 设置 vwap 任务参数
        time_table = vwap_table(api, "CZCE.MA109", -100, 600)  # 目标持仓 -100 手，600s 内完成
        print(time_table.to_string())

        target_pos_sch = TargetPosScheduler(api, "CZCE.MA109", time_table)
        # 启动循环
        while not target_pos_sch.is_finished():
            api.wait_update()
        api.close()


    """
    account = api._account._check_valid(account)
    if account is None:
        raise Exception(f"多账户模式下, 需要指定账户实例 account")

    TIME_CELL = 60  # 等时长下单的时间单元, 单位: 秒
    HISTORY_DAY_LENGTH = 10  # 使用多少天的历史数据用来计算每个时间单元的下单手数

    if duration % TIME_CELL or duration < 60:
        raise Exception(f"duration {duration} 参数应该为 {TIME_CELL} 的整数倍")

    pos = account.get_position(symbol)
    target_pos = int(target_pos)
    delta_pos = target_pos - pos.pos
    target_volume = abs(delta_pos)  # 总的下单手数
    if target_volume == 0:
        return DataFrame(columns=['interval', 'target_pos', 'price'])

    # 获取 Kline
    klines = api.get_kline_serial(symbol, TIME_CELL, data_length=int(10 * 60 * 60 / TIME_CELL * HISTORY_DAY_LENGTH))
    klines["time"] = klines.datetime.apply(lambda x: datetime.fromtimestamp(x // 1000000000).time())  # k线时间
    klines["date"] = klines.datetime.apply(lambda x: datetime.fromtimestamp(_get_trading_day_from_timestamp(x) // 1000000000).date())  # k线交易日

    quote = api.get_quote(symbol)
    # 当前交易日完整的交易时间段
    trading_timestamp = _get_trading_timestamp(quote, quote.datetime)
    trading_timestamp_nano_range = trading_timestamp['night'] + trading_timestamp['day']  # 当前交易日完整的交易时间段
    # 当前时间 行情时间
    current_timestamp_nano = _get_trade_timestamp(quote.datetime, float('nan'))
    if not trading_timestamp_nano_range[0][0] <= current_timestamp_nano < trading_timestamp_nano_range[-1][1]:
        raise Exception("当前时间不在指定的交易时间段内")

    current_datetime = datetime.fromtimestamp(current_timestamp_nano//1000000000)
    # 下一分钟的开始时间
    next_datetime = current_datetime.replace(second=0) + timedelta(minutes=1)
    start_datetime_nano = int(next_datetime.timestamp()) * 1000000000
    r = _rangeset_head(_rangeset_slice(trading_timestamp_nano_range, start_datetime_nano), int(duration * 1e9))
    if not (r and trading_timestamp_nano_range[0][0] <= r[-1][-1] < trading_timestamp_nano_range[-1][1]):
        raise Exception("指定时间段超出当前交易日")

    start_datetime = datetime.fromtimestamp(start_datetime_nano // 1000000000)
    end_datetime = datetime.fromtimestamp((r[-1][-1] - 1) // 1000000000)
    time_slot_start = time(start_datetime.hour, start_datetime.minute)  # 计划交易时段起始时间点
    time_slot_end = time(end_datetime.hour, end_datetime.minute)  # 计划交易时段终点时间点
    if time_slot_end > time_slot_start:  # 判断是否类似 23:00:00 开始， 01:00:00 结束这样跨天的情况
        klines = klines[(klines["time"] >= time_slot_start) & (klines["time"] <= time_slot_end)]
    else:
        klines = klines[(klines["time"] >= time_slot_start) | (klines["time"] <= time_slot_end)]

    # 获取在预设交易时间段内的所有K线, 即时间位于 time_slot_start 到 time_slot_end 之间的数据
    need_date = klines['date'].drop_duplicates()[-HISTORY_DAY_LENGTH:]
    klines = klines[klines['date'].isin(need_date)]

    grouped_datetime = klines.groupby(['date', 'time'])['volume'].sum()
    # 计算每个交易日内的预设交易时间段内的成交量总和(level=0: 表示按第一级索引"data"来分组)后,将每根k线的成交量除以所在交易日内的总成交量,计算其所占比例
    volume_percent = grouped_datetime / grouped_datetime.groupby(level=0).sum()
    predicted_percent = volume_percent.groupby(level=1).mean()  # 将历史上相同时间单元的成交量占比使用算数平均计算出预测值

    # 计算每个时间单元的成交量预测值
    time_table = DataFrame(columns=['interval', 'volume', 'price'])
    volume_left = target_volume  # 剩余手数
    percent_left = 1  # 剩余百分比
    for index, value in predicted_percent.items():
        volume = round(target_volume * (value / percent_left))
        volume_left -= volume
        percent_left -= value
        append_time_table = pd.DataFrame([
            {"interval": TIME_CELL - 2, "volume": volume, "price": "PASSIVE"},
            {"interval": 2, "volume": 0, "price": "ACTIVE"}
        ])
        time_table = pd.concat([time_table, append_time_table], ignore_index=True)

    time_table['volume'] = time_table['volume'].mul(np.sign(delta_pos))
    time_table['target_pos'] = time_table['volume'].cumsum()
    time_table['target_pos'] = time_table['target_pos'].add(pos.pos)
    time_table.drop(columns=['volume'], inplace=True)
    time_table = time_table.astype({'target_pos': 'int64', 'interval': 'float64'})
    return time_table


def _gen_random_list(sum_val: int, min_val: int, max_val: int, length: int = None):
    """
    生成随机列表，参数应该满足：min_val * length <= sum_val <= max_val * length
    :param int sum_val: 列表元素之和
    :param int min_val: 列表元素最小值
    :param int max_val: 列表元素最大值
    :param int length: 列表长度，如果没有指定，则返回的列表长度没有指定
    :return: 整型列表，满足 sum(list) = sum_val, len(list) == length, min_val < any_item(list) < max_val
    """
    if length is None:
        length = sum_val * 2 // (min_val + max_val)  # 先确定 ist 长度，interval 大小，再生成 volume_list 随机列表
        # 例如：sum = 16 min_val = 11 max_val = 15，不满足 min_val * length <= sum_val <= max_val * length
        assert min_val * length <= sum_val <= max_val * length + min_val
    else:
        assert min_val * length <= sum_val <= max_val * length

    result_list = [min_val for _ in range(length)]
    if sum(result_list) == sum_val:
        return result_list  # 全部最小值刚好满足，可以提前退出
    result_rest_value = sum_val - min_val * length  # 剩余可以填充的个数
    result_rest_position = (max_val - min_val) * length  # 剩余需要填充的位置个数
    if sum_val > max_val * length:
        result_list.append(0)
        result_rest_position += min_val
    result_rest_list = _gen_shuffle_list(result_rest_value, result_rest_position)
    for i in range(len(result_list)):
        start = (max_val - min_val) * i
        end = (max_val - min_val) * (i + 1) if start < (max_val - min_val) * length else result_rest_position
        result_list[i] += sum(result_rest_list[start:end])
    assert len(result_list) == length or len(result_list) == length + 1
    assert sum(result_list) == sum_val
    return result_list


def _gen_shuffle_list(x: int, n: int):
    """从 n 个位置中随机选中 x 个"""
    assert x <= n
    result_list = [1 for _ in range(x)] + [0 for _ in range(n - x)]
    utils.RD.shuffle(result_list)
    return result_list
