from vnpy_ctastrategy import (
    CtaTemplate,
    TargetPosTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,
)
import pandas as pd

class vip09(TargetPosTemplate):
    """"""
    author = "松鼠Quant"
    
    Length = 30
    N=5
    X=10
    TS = 45
    lots=1

    current_bar = 0

    Highup=0
    Lowdown=0

    HigherAfterEntry = float('inf')
    LowerAfterEntry = -float('inf')
    liQKA = 0
    DliqPoint = 0
    KliqPoint = 0
    OBVValue=[0]
    MAOBV=[0]
    M2=[0]
    
    cond1=[0]*2
    cond2=[0]*2
    cond3=[0]*2
    cond4=[0]*2
    cond5=[0]*2
    
    kong_cond1=[0]*2
    kong_cond2=[0]*2
    kong_cond3=[0]*2
    kong_cond4=[0]*2
    kong_cond5=[0]*2
    
    MA1=0
    MA2=0
    MA3=0
    
    

    parameters = ["Length",'N','X',"TS",'lots']
    variables = ["HigherAfterEntry","LowerAfterEntry","liQKA","DliqPoint","KliqPoint"]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        """"""
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)

        #df: pd.DataFrame = pd.read_csv(r"index_contract.csv")
        #self.contracts_sizes = {row.vt_symbol.split('.')[0]: row.contract_size for _, row in df.iterrows()}

        self.bg = BarGenerator(self.on_bar)
        self.am = ArrayManager(200)
        self.pos=0

    def on_init(self):
        """
        Callback when strategy is inited.
        """
        self.write_log("策略初始化")
        self.load_bar(1)

    def on_start(self):
        """
        Callback when strategy is started.
        """
        self.write_log("策略启动")

    def on_stop(self):
        """
        Callback when strategy is stopped.
        """
        self.write_log("策略停止")

    def on_tick(self, tick: TickData):
        """
        Callback of new tick data update.
        """
        self.bg.update_tick(tick)
    
    def on_bar(self, bar: BarData):
        """
        Callback of new bar data update.
        """
        
        am = self.am
        am.update_bar(bar)
        if not am.inited:
            return
        # 如果历史数据还没有初始化完毕，则直接返回
        # 注意：am 是交易系统中的一个数据管理器，可能用于管理历史数据
        # bar 是当前的K线数据，其中包含了开盘价、最高价、最低价、收盘价和成交量等信息

        if self.current_bar == 0:
            self.OBVValue = [0]
        else:
            if am.close_array[-1] > am.close_array[-2]:
                self.OBVValue += [self.OBVValue[-1] + am.volume_array[-1]]
            elif am.close_array[-1] < am.close_array[-2]:
                self.OBVValue += [self.OBVValue[-1] - am.volume_array[-1]]

        # 根据收盘价的涨跌更新 OBV 值
        # 如果当前收盘价大于昨日收盘价，则 OBV 值加上当前成交量
        # 如果当前收盘价小于昨日收盘价，则 OBV 值减去当前成交量
        # 注意：OBV 是一种基于成交量变化来反映买卖力量的技术指标，用于判断市场的买卖压力

        self.current_bar=self.current_bar+1 #Bar线计数
        #print(bar.datetime,"current_bar:",self.current_bar)

        # 记录当前Bar的计数器值，并打印当前Bar的时间和计数器值

        if(self.current_bar<self.Length*3):
            return

        # 如果当前Bar的计数器值小于三倍的Length，则直接返回
        # 这可能是因为需要一定数量的历史数据才能计算出所需的指标值

        # print(f'{self.OBVValue},OBVValue')
        self.MA1=am.ema(self.Length,array=True)
        self.MA2=am.ema(self.Length*2,array=True)
        self.MA3=am.ema(self.Length*3,array=True)

        # 计算三条指定长度的指数移动平均线（EMA）
        # 分别记为MA1、MA2和MA3

        self.MAOBV += [(sum(self.OBVValue[-self.Length:]) / self.Length)]
        self.M2 += [(sum(self.OBVValue[-self.X:]) / self.X)]

        # 计算 OBV 和 M2 的移动平均值

        self.cond1=self.MA1 if am.close_array[-1] > self.MA1[-2] else self.cond1
        self.cond2=self.MA1 if self.M2[-1] > self.M2[-2] else self.cond2
        self.cond3=self.MA2 if am.close_array[-1] > self.MA2[-1] and self.MAOBV[-1] > self.MAOBV[-2] else self.cond3
        self.cond4=self.MA2 if am.close_array[-1] > self.MA2[-1] else self.cond4
        self.cond5=self.MA3 if am.close_array[-1] > self.MA3[-1] and self.MAOBV[-1] > self.MAOBV[-2] else self.cond5

        # 根据条件更新多头条件指标cond1到cond5
        # 根据条件更新空头条件指标kong_cond1到kong_cond5

        self.kong_cond1=self.MA1 if am.close_array[-1] < self.MA1[-2] else self.kong_cond1
        self.kong_cond2=self.MA1 if self.M2[-1] < self.M2[-2] else self.kong_cond2
        self.kong_cond3=self.MA2 if am.close_array[-1] < self.MA2[-1] and self.MAOBV[-1] < self.MAOBV[-2] else self.kong_cond3
        self.kong_cond4=self.MA2 if am.close_array[-1] < self.MA2[-1] else self.kong_cond4
        self.kong_cond5=self.MA3 if am.close_array[-1] < self.MA3[-1] and self.MAOBV[-1] < self.MAOBV[-2] else self.kong_cond5

        # 更新空头条件指标kong_cond1到kong_cond5

        self.Highup=max(am.high_array[-self.N:-1]) 
        self.Lowdown=min(am.low_array[-self.N:-1])

        # 计算最近N个Bar的最高价和最低价，分别记为Highup和Lowdown

        Long_line=am.close_array[-1]>self.MA3[-1] and self.MAOBV[-1]>self.MAOBV[-2]
        short_line=am.close_array[-1]<self.MA3[-1] and self.MAOBV[-1]<self.MAOBV[-2]

        #print(f'{self.MAOBV[-1]},MAOBV')
        #print(f'{Long_line},Long_line')
        #print(f'{short_line},short_line')
        if Long_line:
           if self.cond4[-1]>self.cond5[-2] and self.cond5[-1]>self.cond5[-2] and self.cond4[-1]>self.cond4[-2] and self.cond3[-1]>self.cond3[-2] and self.cond2[-1]>self.cond2[-2] and self.cond1[-1]>self.cond1[-2] :
                if self.pos!=1 and bar.high_price>=self.Highup :
                    if self.pos == 0:
                        self.buy(bar.close_price, self.lots)
                    elif self.pos < 0:
                        self.cover(bar.close_price, self.lots)
                        self.buy(bar.close_price, self.lots)
                    self.LowerAfterEntry = bar.close_price

        if short_line:
            if self.kong_cond4[-1]<self.kong_cond5[-1] and self.kong_cond5[-1]<self.kong_cond5[-2] and self.kong_cond4[-1]<self.kong_cond4[-2] and self.kong_cond3[-1]<self.kong_cond3[-2] and self.kong_cond2[-1]<self.kong_cond2[-2] and self.kong_cond1[-1]<self.kong_cond1[-2]:
                if self.pos!=-1 and bar.low_price<=self.Lowdown :
                    if self.pos == 0:
                        self.short(bar.close_price, self.lots)
                    elif self.pos > 0:
                        self.sell(bar.close_price, self.lots)
                        self.short(bar.close_price, self.lots)
                        
                    self.HigherAfterEntry = bar.close_price


        #记录入场后的最高价和最低价
        if self.pos>0:
            #self.HigherAfterEntry = HigherAfterEntry
            self.LowerAfterEntry = max(self.LowerAfterEntry, bar.low_price)
        elif self.pos<0:
            self.HigherAfterEntry = min(self.HigherAfterEntry, bar.high_price)	
            #self.LowerAfterEntry = LowerAfterEntry

        if self.pos == 0:
            self.liQKA = 1
        else:
            self.liQKA = self.liQKA - 0.1
            self.liQKA = max(self.liQKA, 0.5)
        
        if self.pos>0:
            DliqPoint = self.LowerAfterEntry - ((bar.open_price)*self.TS/1000)*self.liQKA
            #print(bar.datetime,"DliqPoint",DliqPoint,"low_price",bar.low_price,"self.LowerAfterEntry",self.LowerAfterEntry)
        if self.pos<0:
            KliqPoint = self.HigherAfterEntry + ((bar.open_price)*self.TS/1000)*self.liQKA
            #print(bar.datetime,"KliqPoint",KliqPoint,"high_price",bar.high_price,"self.HigherAfterEntry",self.HigherAfterEntry)
        
        if (self.pos>0 and bar.low_price <= DliqPoint and DliqPoint>0):
            #多头出场
            self.sell(bar.close_price, self.lots)
        elif (self.pos<0 and bar.high_price > KliqPoint and KliqPoint>0):
            #空头出场
            self.cover(bar.close_price, self.lots)

        self.current_bar += 1
        self.put_event()


    def on_trade(self, trade: TradeData):
        """
        Callback of new trade data update.
        """
        self.put_event()
        

    def on_stop_order(self, stop_order: StopOrder):
        """
        Callback of stop order update.
        """
        pass