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

class vip13(CtaTemplate):
    author = "松鼠Quant"

    boll_n = 20
    offset = 1.5
    trailing_stop_rate =60
    x = 30
    rsi_n = 14
    over_sold = 30
    over_bought = 70
    
    parameters = ["boll_n", "offset", "trailing_stop_rate", "x"]
    variables = ["lots", "up_line", "down_line", "mid_line", "band", "rsi_value", "kg", "highest_low_after_entry", "lowest_high_after_entry", "rsi_n", "over_sold", "over_bought"]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        """"""
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        
        self.lots = 0
        self.up_line = 0
        self.down_line = 0
        self.mid_line = 0
        self.band = 0
        self.rsi_value = 0
        self.kaiduo = False
        self.kaikong = False
        self.kg = 0
        self.highest_low_after_entry = 0
        self.lowest_high_after_entry = 0
        self.bg = BarGenerator(self.on_bar)
        self.am = ArrayManager(300)
        self.pos=0
        
    def on_init(self):
        self.write_log("策略初始化")
        self.load_bar(10)
        
    def on_start(self):
        self.write_log("策略启动")
        
    def on_stop(self):
        self.write_log("策略停止")
        
    def on_tick(self, tick: TickData):
        pass
    
    def on_bar(self, bar: BarData):
        self.am.update_bar(bar)
        if not self.am.inited:
            return
        
        self.lots = 1
        
        self.calculate_boll()
        self.calculate_rsi()
        
        self.kaikong = self.rsi_value > self.over_bought and self.am.close[-1] < self.up_line
        self.kaiduo = self.rsi_value < self.over_sold and self.am.close[-1] > self.down_line
        
        if self.kaiduo:
            self.kg = 1
            self.hd = self.am.high[-self.x:-1].max()
        
        if self.kaikong:
            self.kg = -1
            self.ld = self.am.low[-self.x:-1].min()
        

        if  self.pos != 1 and self.kg == 1 and bar.close_price >= self.hd :
            self.buy(bar.close_price, self.lots)
            self.highest_low_after_entry = bar.close_price
            self.pos=1
        if  self.pos != -1 and self.kg == -1 and bar.close_price <= self.ld:
            self.short(bar.close_price, self.lots)
            self.lowest_high_after_entry = bar.close_price
            self.pos=-1
        
        self.trailing_stop(bar)
        
        self.put_event()
    
    def calculate_boll(self):
        self.mid_line = self.am.sma(self.boll_n)
        self.band = self.am.std(self.boll_n)
        self.up_line = self.mid_line + self.offset * self.band
        self.down_line = self.mid_line - self.offset * self.band
    
    def calculate_rsi(self):
        self.rsi_value = self.am.rsi(self.rsi_n)
    
    def trailing_stop(self, bar: BarData):
        if self.pos > 0:
            self.highest_low_after_entry = max(self.highest_low_after_entry, bar.low_price)
            dliq_point = self.highest_low_after_entry - bar.open_price * (self.trailing_stop_rate/1000)
            if bar.close_price < dliq_point:
                self.sell(bar.close_price, self.lots)
                self.pos=0
        if self.pos < 0:
            self.lowest_high_after_entry = min(self.lowest_high_after_entry, bar.high_price)
            kliq_point = self.lowest_high_after_entry + bar.open_price * (self.trailing_stop_rate/1000)
            if bar.close_price > kliq_point:
                self.cover(bar.close_price, self.lots)
                self.pos=0