import pandas as pd
import numpy as np


# 创建BOLL线

def get_high_low(df: pd.DataFrame) -> list:
    # 用最高价和最低价构建新的数据，用于显示和判断是否突破BOLL线支撑位、压力位
    high_low = []
    for i in df.index:
        if df['High'][i] > df['阻力线'][i]:
            high_low.append(df['High'][i])
        elif df['Low'][i] < df['支撑线'][i]:
            high_low.append((df['Low'][i]))
        else:
            high_low.append(df['Close'][i])
    return high_low


def get_boll_lines(df: pd.DataFrame, tp: int, sf: int) -> pd.DataFrame:
    # SMA:简单移动平均(Simple Moving Average)
    time_period = 21  # SMA的计算周期，默认为20
    stdev_factor = 2  # 上下频带的标准偏差比例因子
    if tp is not None:
        time_period = tp
    if sf is not None:
        stdev_factor = sf

    history = []  # 每个计算周期所需的价格数据
    sma_values = []  # 初始化SMA值
    upper_band = []  # 初始化阻力线价格
    lower_band = []  # 初始化支撑线价格

    # 构造列表形式的绘图数据
    for close_price in df['Close']:
        #
        history.append(close_price)

        # 计算移动平均时先确保时间周期不大于20
        if len(history) > time_period:
            # 删除第0个元素
            del (history[0])

        # 将计算的SMA值存入列表
        # 求取history列表数据的平均值
        sma = np.mean(history)
        # 将平均值添加到平均值列表中
        sma_values.append(sma)
        # 计算标准差
        stdev = np.sqrt(np.sum((history - sma) ** 2) / len(history))
        upper_band.append(sma + stdev_factor * stdev)
        lower_band.append(sma - stdev_factor * stdev)

    # 将BOLL线加入到数据中
    df = df.assign(收盘价=pd.Series(df['Close'], index=df.index))
    df = df.assign(中界线=pd.Series(sma_values, index=df.index))
    df = df.assign(阻力线=pd.Series(upper_band, index=df.index))
    df = df.assign(支撑线=pd.Series(lower_band, index=df.index))
    # 将最高最低组合数据，加入到数据中
    # 用最高价和最低价构建新的数据，用于显示
    high_low = get_high_low(df)
    df = df.assign(hl=pd.Series(high_low, index=df.index))

    return df