import requests
import pandas as pd

CN_CODE_TYPE = {
    "600":'sh',
    "601":'sh',
    "603":'sh',
    "605":'sh',
    "000":'sz',
    "300":'sz',
    "688":'sh',
    "002":'sz',
    "50":"sh",
    "51":"sh",
    "52":"sh",
    "16":"sz",
    "15":"sz",
    "18":"sz"
    }


def calculate_brar(df: pd.DataFrame, N: int = 26) -> pd.DataFrame:
    """
    计算AR和BR情绪指标

    参数:
        df (pd.DataFrame): 包含 'high', 'low', 'open', 'close' 列的DataFrame
        N (int): 计算周期，默认通常为26
    返回:
        pd.DataFrame: 包含原始数据及计算出的 'ar' 和 'br' 列的DataFrame
    """
    # 创建副本，避免修改原始数据
    data = df.copy()

    # --- 计算 AR (人气指标) ---
    # 公式: AR = (N日内 (最高价 - 开盘价) 之和) / (N日内 (开盘价 - 最低价) 之和) * 100
    high_open = data['high'] - data['open']
    open_low = data['open'] - data['low']

    data['ar'] = (high_open.rolling(window=N, min_periods=1).sum() /
                  open_low.rolling(window=N, min_periods=1).sum()) * 100

    # --- 计算 BR (意愿指标) ---
    # 公式: BR = (N日内 MAX(0, 最高价 - 前一日收盘价) 之和) /
    #           (N日内 MAX(0, 前一日收盘价 - 最低价) 之和) * 100
    prev_close = data['close'].shift(1)

    # 计算多头意愿：当日最高价 - 前一日收盘价，若小于0则记为0
    high_prev_close = data['high'] - prev_close
    high_prev_close = high_prev_close.where(high_prev_close > 0, 0)

    # 计算空头意愿：前一日收盘价 - 当日最低价，若小于0则记为0
    prev_close_low = prev_close - data['low']
    prev_close_low = prev_close_low.where(prev_close_low > 0, 0)

    data['br'] = (high_prev_close.rolling(window=N, min_periods=1).sum() /
                  prev_close_low.rolling(window=N, min_periods=1).sum()) * 100

    return data


def get_history_k(code, start_date='19700101', end_date=None):

    if '.' in code:
        code = code.split('.')[0]

    stock_code = ''

    if code[0:2] in CN_CODE_TYPE:
        stock_code = code + '.' + CN_CODE_TYPE[code[0:2]]
    elif code[0:3] in CN_CODE_TYPE:
        stock_code = code + '.' + CN_CODE_TYPE[code[0:3]]
    if stock_code != '':

        url = "https://api.zhituapi.com/hs/history/%s/d/f?token=6E0E86BC-15AD-4275-8A95-B02D168D63C1&st=%s"%(stock_code, start_date)

        if end_date is not None:
            url += "&et=%s"%(end_date)
        response = requests.get(url)

        data = response.json()
        df = pd.DataFrame(data)

        df['t'] = df['t'].apply(pd.to_datetime)
        df.rename({'t': 'date', 'o': 'open', 'h': 'high', 'l': 'low', 'c': 'close', 'v': 'volume', 'a': 'amount'},
                  axis='columns', inplace=True)
        df = df[['date', 'open', 'high', 'low', 'close', 'volume', 'amount']].set_index('date')

    else:
        df = pd.DataFrame()


    return df


def get_5min_k(code, start_date='19700101', end_date=None):
    stock_code = ''

    if code[0:2] in CN_CODE_TYPE:
        stock_code = code + '.' + CN_CODE_TYPE[code[0:2]]
    elif code[0:3] in CN_CODE_TYPE:
        stock_code = code + '.' + CN_CODE_TYPE[code[0:3]]

    if stock_code != '':

        url = "https://api.zhituapi.com/hs/history/%s/5/n?token=6E0E86BC-15AD-4275-8A95-B02D168D63C1&st=%s"%(stock_code, start_date)

        if end_date is not None:
            url += "&end_date=%s"%(end_date)

        response = requests.get(url)

        data = response.json()

        df = pd.DataFrame(data)
        df['t'] = df['t'].apply(pd.to_datetime)
        df.rename({'t': 'date', 'o': 'open', 'h': 'high', 'l': 'low', 'c': 'close', 'v': 'volume', 'a': 'amount'},
                  axis='columns', inplace=True)
        df = df[['date','close', 'volume']].set_index('date')

    else:
        df = pd.DataFrame()


    return df