import requests
import pandas as pd
import numpy as np

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"
    }

## arbr指标
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

## kdj计算
def calc_kdj_tdx(df, n=9, m1=3, m2=3):
    # 1. 计算 RSV
    low_min = df['low'].rolling(n).min()
    high_max = df['high'].rolling(n).max()
    rsv = (df['close'] - low_min) / (high_max - low_min) * 100
    rsv = rsv.fillna(0)  # 或者处理 NaN

    # 2. 自定义 SMA 函数（通达信风格）
    def sma_tdx(series, period, weight):
        # period=N, weight=M。公式：Y = (X*M + PREV_Y*(N-M))/N
        result = np.zeros(len(series))
        # 初始化第一个有效值。通常，如果数据足够，初始值为 series[0]
        result[0] = series[0] if not np.isnan(series[0]) else 0
        for i in range(1, len(series)):
            # 如果 series[i] 是 NaN（例如，前 n 天），则保持 0 或向前填充，但通常 RSV 已填充。
            val = series[i] if not np.isnan(series[i]) else 0
            result[i] = (val * weight + result[i - 1] * (period - weight)) / period
        return result

    # 3. 计算 K 和 D
    k_series = sma_tdx(rsv.values, m1, 1)  # K = SMA(RSV, 3, 1)
    d_series = sma_tdx(k_series, m2, 1)  # D = SMA(K, 3, 1)

    df['K'] = k_series
    df['D'] = d_series
    df['J'] = 3 * df['K'] - 2 * df['D']

    return df

## 获取历史k线
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


## 获取当日5分钟的k线
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