import math
from util.tushare_dec import *

#计算等比例线
def calGridlines(hightN,lowN,lines=7):
    #7条水平线位置
    y = math.pow((lowN/hightN), 1/(lines - 1))
    r = []
    for i in range(lines):
        r.append(round(hightN*y**i, 3))
    return r

#计算价格区间有哪些等比例线
def calPriceReachedLines(gridlines=None,high=None,low=None):
    if gridlines is None or high is None or low is None:
        return None
    highline, lowline = gridlines[0], gridlines[-1:][0]
    y = math.pow((lowline / highline), 1 / (len(gridlines) - 1))
    x1 = math.log(high/highline)/math.log(y)
    x2 = math.log(low/highline)/math.log(y)
    # print(x1, x2)
    x1 = math.ceil(x1)
    x2 = math.floor(x2)
    # return (x1, x2)
    r = []
    x = x1
    while x <= x2:
        r.append(x)
        x += 1
    return r
#区分哪些等比例线对应的买卖操作
def calsplitBSLines(gridlines,position,positiondict,open,high, low):
    r, Blist, Slist = {}, [], []
    info = ""
    profit = 0
    #持仓对应的股价
    posprice = gridlines[position+1]
    #计算open,high, low 相对lines的位置
    highline, lowline = gridlines[0], gridlines[-1:][0]
    print("H L O:", high, low, open," UPer Lower:", highline, lowline ,"pos:", position, "posprice:", posprice)
    y = math.pow((lowline / highline), 1 / (len(gridlines) - 1))
    relativehigh= math.log(high / highline) / math.log(y)
    relativelow = math.log(low / highline) / math.log(y)
    # print(relativehigh,relativelow)
    relativehigh = math.ceil(relativehigh)
    relativelow = math.floor(relativelow)
    # print("relative high low:",relativehigh,relativelow)
    buyFirst = True if open < posprice else False
    for loop in range(2):
        if buyFirst:
            buyFirst = False
            # buy
            for i in range(position + 1, relativelow):
                # print("buy:", i, gridlines[i+1])
                if position > 4 :
                    break
                position += 1
                Blist.append(gridlines[i+1])
                positiondict[i]=gridlines[i+1]
            r["B"] = Blist
        else:
            buyFirst = True
            # sell
            for i in range(relativehigh, position + 1):
                if i < 1:
                    continue
                # print("sell:", i, gridlines[i+1], "->", gridlines[i])
                # info += "sell:{} {}->{}  ".format( i, gridlines[i+1],  gridlines[i])
                info += "sell:{} {}->{}  ".format(i, positiondict[i],  gridlines[i])
                profit += gridlines[i] - positiondict[i]
                position -= 1
                Slist.append(gridlines[i+1])
                positiondict[i] = 0
            r["S"] = Slist
    r["I"] = info
    print("r:", r , "profit:", profit)
    return r,position,positiondict,profit

#账户
class account():
    positiondict = {}
    def __init__(self, position=0, profit=0, split=4):
        self.position = position  # 分四份
        self.split = split
        self.profit = profit  # 利润
        self.amplitude = 1  #追求一格利润
        # if "1" in self.positiondict.keys():
        #     x = self.positiondict["1"]
        #     print("has ",x)
        #     self.positiondict["1"] = x+1
        # else:
        #     self.positiondict["1"] = 1
        #     print("init d")

    def cal(self, high, low, open,
            hightN, lowN, lines=7):
        # print(high, low, open)
        gridlines = calGridlines(hightN, lowN, lines)
        print(gridlines)

        # 计算在[h l]区间的Gridlines(等比例线)
        p = calPriceReachedLines(gridlines, high, low)
        # print("lines in range :", p, "pos:",self.position)

        #需要加\减仓的等比例线
        v, self.position,self.positiondict,profit = calsplitBSLines(
            gridlines, self.position,self.positiondict,
            open, high, low)

        print(self.positiondict,"profit:",profit)
        return self.positiondict,profit
        pass

# ac = account()
if __name__ == '__main__':
    hightN, lowN = 0, 0
    ac = account(position=0)
    # high, low, open, hightN, lowN, = 6.14, 5.46, 5.94, 6.83, 5.20
    # high, low, open, hightN, lowN, = 3.032, 0.76, 3.0, 6, 0.1
    # high, low, open, hightN, lowN, = 0.78, 0.76, 0.77, 6, 0.1
    # ac.cal(high, low, open, hightN, lowN)
    # print("="*30)

    # high, low, open = 5.97,5.23,5.47
    # high, low, open = 0.39, 0.197,0.39
    # ac.cal(high, low, open, hightN, lowN)
    # print("=" * 30)
    #
    # high, low, open = 0.78, 0.76, 0.77
    # ac.cal(high, low, open, hightN, lowN)
    # print("=" * 30)

    # high, low, open = 3.04, 0.197, 3.04
    # ac.cal(high, low, open, hightN, lowN)
    # print("=" * 30)

    # df = get_symbol_close_volum(symbol, begin_date, end_date, clear=clear)
    symbol =  ["603605","000156","688036"]
    df = get_symbol_close_volum(symbol[0], "2020-01-18", "2021-02-20", clear=False)
    cycle = 5
    df['hightN'] = df['high'].rolling(window=cycle, center=False).max().shift(1)
    df['lowN'] = df['low'].rolling(window=cycle, center=False).min().shift(1)
    df = df.dropna()
    profitsum = 0
    tradecnt = 0
    close = 0
    for i, v in df.iterrows():
        hightN, lowN = v['hightN'], v['lowN'],
        high, low, open = v['high'], v['low'], v['open'],
        print("H L O HN LN:",high, low, open, hightN, lowN ,v['date'])
        _,profit = ac.cal(high, low, open, hightN, lowN)
        if profit!=0:
            tradecnt += 1
        profitsum += profit
        print("=" * 30)
        close = open
    print("profit:",profitsum,tradecnt,len(df),tradecnt/len(df),profitsum/4/close)