import bisect
from functools import reduce

import pandas as pd
import numpy as np

from ldnode.calculator.schemas import GridBacktestParameters, GridType


def cb_grid_backtest(
    df: pd.DataFrame,
    strategy_config: GridBacktestParameters,
) -> pd.DataFrame:
    """
    网格回测
    """
    df.set_index('datetime', inplace=True)
    df.dropna(inplace=True)
    # 初始化网格
    if strategy_config['grid_type'] == GridType.LINE:
        grid = list(np.linspace(strategy_config['grid_range'][0], strategy_config['grid_range'][1], strategy_config['grid_count'] + 1).round(3))
    elif strategy_config['grid_type'] == GridType.GEOM:
        grid = list(np.geomspace(strategy_config['grid_range'][0], strategy_config['grid_range'][1], strategy_config['grid_count'] + 1).round(3))
    else:
        raise ValueError(f"网格类型错误: {strategy_config['grid_type']}")
    # 生成所需
    df['preclose'] = df.close.shift(1)
    df['inc'] = df.close - df.preclose
    df['grid_index'] = None
    df['pos_chg'] = None
    df['pos'] = None
    df['pos_val'] = None
    df['cash'] = None
    df['grid_return'] = None
    df['grid_return_acc'] = None
    df['asset'] = None
    # 定位网格初始位置 投入底仓
    grid_index = df.columns.get_loc('grid_index')
    inc = df.columns.get_loc('inc')
    open_price = df.columns.get_loc('open')
    close = df.columns.get_loc('close')
    pos = df.columns.get_loc('pos')
    cash = df.columns.get_loc('cash')
    grid_return = df.columns.get_loc('grid_return')
    df.iat[0, grid_index] = grid.index(min(grid, key=lambda x: abs(x - df.iat[0, 0]))) # 当前位置
    df.iat[0, pos] = ((len(grid) - 1) - df.iat[0, grid_index]) * 10 # 投入底仓
    df.iat[0, cash] = -(df.iat[0, pos] * df.iat[0, open_price])# 借入现金

    # 计算网格位置
    for i in range(1, df.shape[0]):
        top_flag = False # 触碰网格顶部
        bottom_flag = False # 触碰网格底部
        pre_index = df.iat[i - 1, grid_index] # 上周期网格所处位置
        if pre_index == (len(grid) - 1):
            # print('向上破网\n', df.iloc[[i]])
            if strategy_config['stop_on_top']:
                break
            top_flag = True
        elif pre_index == 0:
            # print('向下破网，不再买入\n', df.iloc[[i]].to_string())
            bottom_flag = True
        pre_cash = df.iat[i - 1, cash] # 上周期现金
        upper_con = (df.iat[i, inc] > 0) & (df.iat[i, close] >= grid[pre_index + 1]) if not top_flag else False
        lower_con = (df.iat[i, inc] < 0) & (df.iat[i, close] <= grid[pre_index - 1]) if not bottom_flag else False
        # 网格上移
        if upper_con:
            index = (bisect.bisect_right(grid, df.iat[i, close]) - 1) # 上移后index位置
            df.iat[i, grid_index] = index # 写入当前index位置
            l = grid[pre_index + 1 : index + 1] # 上移经过位置
            cost = reduce(lambda x, y: x + y, l)
            df.iat[i, cash] = pre_cash + 10 * cost * (1 - strategy_config['grid_c_rate']) # 卖出变现
            # 纯网格收益
            df.iat[i, grid_return] = sum(map(lambda x : (grid[grid.index(x)] - grid[grid.index(x) - 1]) , l)) * 10 * (1 - strategy_config['grid_c_rate']) ** 2
        # 网格下移
        elif lower_con:
            index = bisect.bisect_left(grid ,df.iat[i, close]) # 下移后index位置
            df.iat[i, grid_index] = index # 写入当前index位置
            l = grid[index:pre_index] # 下移经过位置
            cost = reduce(lambda x, y: (x + y), l)
            df.iat[i, cash] = pre_cash - 10 * cost * (1 + strategy_config['grid_c_rate'])# 借入现金
        # 网格不变
        else:
            df.iat[i, grid_index] = pre_index
            df.iat[i, cash] = pre_cash # 现金保持不变

    # 生成资产数据
    df.dropna(subset=['cash'], inplace=True) # 删除停止之后的数据
    df['pos_chg'] = (df.grid_index.shift(1) - df.grid_index) * 10 # 持仓变化
    df.loc[:,'pos'] = df.iat[0, pos] + df.pos_chg.expanding().sum() # 累计持仓
    df['pos_val'] = df.close * df.pos
    df.grid_return = df.grid_return.fillna(0)
    df['grid_return_acc'] = df.grid_return.expanding().sum()
    df['asset'] = df.pos_val + df.cash
    return df
