from ctpbee import CtpbeeApi
from ctpbee.constant import *
from define import *
import numpy as np
import helper


class CTA(CtpbeeApi):
    def __init__(self, name,model,order_symbol):
        super().__init__(name)
        self.k_list = []
        self.model = model
        self.order_symbol = order_symbol
        self.max_order = MAX_ORDER
        # cell and hidden state of the LSTM
        self.states = None
        # Episode start signals are used to reset the lstm states
        self.episode_starts = np.ones((1,), dtype=bool)
     
    
    def on_init(self, init: bool) -> None:  # 初始化完成回调 
        self.info("init successful")

    def on_tick(self, tick: TickData) -> None:
        self.info("on_tick")
        '''
        if len(self.k_list)<20:
            print(tick.datetime, tick.last_price)  # 打印tick时间戳以及最新价格 
            return
        tick_data = dict()
        tick_data["price"] = tick.last_price
        tick_data["buy_price_1"] = tick.bid_price_1
        tick_data["buy_price_2"] = tick.bid_price_2
        tick_data["buy_price_3"] = tick.bid_price_3
        tick_data["buy_price_4"] = tick.bid_price_4
        tick_data["buy_price_5"] = tick.bid_price_5
        tick_data["sell_price_1"] = tick.ask_price_1
        tick_data["sell_price_2"] = tick.ask_price_2
        tick_data["sell_price_3"] = tick.ask_price_3
        tick_data["sell_price_4"] = tick.ask_price_4
        tick_data["sell_price_5"] = tick.ask_price_5
        tick_data["buy_volume_1"] = tick.bid_volume_1
        tick_data["buy_volume_2"] = tick.bid_volume_2
        tick_data["buy_volume_3"] = tick.bid_volume_3
        tick_data["buy_volume_4"] = tick.bid_volume_4
        tick_data["buy_volume_5"] = tick.bid_volume_5
        tick_data["sell_volume_1"] = tick.ask_volume_1
        tick_data["sell_volume_2"] = tick.ask_volume_2
        tick_data["sell_volume_3"] = tick.ask_volume_3
        tick_data["sell_volume_4"] = tick.ask_volume_4
        tick_data["sell_volume_5"] = tick.ask_volume_5
        position = self.center.get_position(self.order_symbol)
        holder_data = {"buy_order":position.long_volume,"sell_order":position.short_volume,"buy_price":position.long_price,"sell_price":position.short_price}
        obs = helper.get_obs(tick_data,self.k_list,holder_data)
        action,self.states=self.model.predict(obs,state=self.states, episode_start=self.episode_starts,deterministic=True)
        self.episode_starts = False ##dones
        active_value = action[0]
        
        print(f'action : {active_value} {holder_data["sell_order"]} {holder_data["buy_order"]}')
        if active_value > 1 and active_value <= 2:
            price = tick.last_price
            if tick.ask_price_1 != tick.last_price:
                price = tick.ask_price_1
            if position.short_volume>0:
                print(f'sell_close : {price}')
                self.action.sell_close(price, 1, tick)
            else:
                if position.long_volume < self.max_order:
                    print(f'buy_open : {price}')
                    self.action.buy_open(price, 1, tick)
            
        elif active_value > 2 and active_value <= 3:
            price = tick.last_price
            if tick.bid_price_1 != tick.last_price:
                price = tick.bid_price_1
            if position.long_volume>0:
                print(f'buy_close : {price}')
                self.action.buy_close(price, 1, tick)
            else:
                if position.short_volume < self.max_order:
                    print(f'sell_open : {price}')
                    self.action.sell_open(price, 1, tick)
         '''   

    def on_contract(self, contract: ContractData) -> None:
        if contract.local_symbol == self.order_symbol:
            self.action.subscribe(contract.local_symbol)  # 订阅行情 
            print("合约乘数: ", contract.size)
    
    def on_bar(self, bar: BarData) -> None:
        print("on_bar bar.interval",bar.interval,len(self.k_list))
        if bar.interval == 1:
            
            if len(self.k_list)>=K_LIST_LEN:
                self.k_list.pop(0)
            self.k_list.append({"open":bar.open_price,"close":bar.close_price,"high":bar.high_price,"low":bar.low_price})