from kungfu.wingchun.constants import *
import kungfu.yijinjing.time as kft

# 鏌滃彴
source = "sim"
# 璐︽埛
account = "simTest"

# 绛栫暐鍙傛暟
MAX_POSITION = 2000
ORDER_VOLUME = 100

# 鑾峰彇鎸佷粨
def get_position(instrument_id, book):
    long_position = None
    for pos in book.long_positions.values():
        if instrument_id == pos.instrument_id:
            long_position = pos
    return long_position
    
# 鍚姩鍓嶅洖璋冿紝娣诲姞浜ゆ槗璐︽埛锛岃闃呰鎯咃紝绛栫暐鍒濆鍖栬绠楃瓑
def pre_start(context):
    context.log.info("preparing strategy")

    context.subscribe(source, ["600000", "600004", "600006", "600007"], Exchange.SSE)
    context.subscribe(source, ["000001", "000004", "000005", "000006"], Exchange.SZE)
    context.add_account(source, account)

# 蹇収鏁版嵁鍥炶皟
def on_quote(context, quote, location, dest):
    # context.log.info(f"{quote=}")
    # 璺宠繃闆嗗悎绔炰环鏃堕棿
    # if "T" not in quote.trading_phase_code:
    #     return

    # 绛栫暐閫昏緫
    ask_volume_sum = sum(x for x in quote.ask_volume)
    bid_volume_sum = sum(x for x in quote.bid_volume)
    
    if bid_volume_sum > 2 * ask_volume_sum and ask_volume_sum > 0:
        position = get_position(quote.instrument_id, context.book) 
        if not position or position.volume < MAX_POSITION:
            order_id = context.insert_order(quote.instrument_id, quote.exchange_id, source, account, quote.ask_price[0], ORDER_VOLUME,
                                            PriceType.Limit, Side.Buy, Offset.Open)
            if order_id:
                # 閫氳繃娣诲姞鏃堕棿鍥炶皟锛屽湪浜旂浠ュ悗鎾ゅ崟, 濡傛灉5m绉掑唴鏈兘鎴愪氦锛岃鍗曚細琚挙
                context.add_timer(context.now() + 5 * kft.NANO_PER_SECOND, lambda context, event: context.cancel_order(order_id))

    if ask_volume_sum > 2 * bid_volume_sum and bid_volume_sum > 0:
        position = get_position(quote.instrument_id, context.book) 
        if position and position.volume > ORDER_VOLUME:
            order_id = context.insert_order(quote.instrument_id, quote.exchange_id, source, account, quote.bid_price[0], ORDER_VOLUME,
                                            PriceType.Limit, Side.Sell, Offset.Close)
            if order_id:
                context.add_timer(context.now() + 5 * kft.NANO_PER_SECOND, lambda context, event: context.cancel_order(order_id))


# 璁㈠崟鍥炴姤
def on_order(context, order, location, dest):
    pass

# 鎴愪氦鍥炴姤
def on_trade(context, trade, location, dest):
    context.log.info(f"{trade=}")    
    pass

# 绛栫暐杩涚▼閫€鍑哄墠鏂规硶
def post_stop(context):
    book = context.book
    # context.log.info("trades: \n", list(book.trades.items()))
    context.log.info(f"{book.long_positions=}")

    pass