from __future__ import annotations

import talib.abstract as ta
from . import state

def add_indicators(df):
	macd = ta.MACD(df, fastperiod=state.macd_fast_period, slowperiod=state.macd_slow_period, signalperiod=state.macd_signal_period)

	df['macd'] = macd['macd']
	df['macdsignal'] = macd['macdsignal']
	df['macdhist'] = macd['macdhist']
	df['ma5'] = (ta.MA(df, timeperiod=5)).fillna(0)
	df['ma10'] = (ta.MA(df, timeperiod=10)).fillna(0)
	df['ma30'] = (ta.EMA(df, timeperiod=30)).fillna(0)
	df['ma250'] = (ta.MA(df, timeperiod=250)).fillna(0)
	# 新增 EMA 指标
	df['ema5'] = (ta.EMA(df, timeperiod=5)).fillna(0)
	df['ema10'] = (ta.EMA(df, timeperiod=10)).fillna(0)
	df['ema24'] = (ta.EMA(df, timeperiod=24)).fillna(0)
	df['ema52'] = (ta.EMA(df, timeperiod=52)).fillna(0)
	df['ema26'] = (ta.EMA(df, timeperiod=26)).fillna(0)
	df['ema13'] = (ta.EMA(df, timeperiod=13)).fillna(0)
	df['ema7'] = (ta.EMA(df, timeperiod=7)).fillna(0)
	df['ema104'] = (ta.EMA(df, timeperiod=104)).fillna(0)
	df['ema156'] = (ta.EMA(df, timeperiod=156)).fillna(0)
	df['ema208'] = (ta.EMA(df, timeperiod=208)).fillna(0)
	# 常用SMA 24/52
	try:
		df['sma24'] = (ta.SMA(df, timeperiod=24)).fillna(0)
		df['sma52'] = (ta.SMA(df, timeperiod=52)).fillna(0)
	except Exception:
		df['sma24'] = 0
		df['sma52'] = 0
	df['rsi'] = ta.RSI(df, timeperiod=14)
	
	# 计算布林带 (当前周期 - 20周期，2标准差)
	bb = ta.BBANDS(df, timeperiod=365, nbdevup=3.0, nbdevdn=3.0, matype=0)
	df['bb_upper'] = bb['upperband'].fillna(0)
	df['bb_middle'] = bb['middleband'].fillna(0)
	df['bb_lower'] = bb['lowerband'].fillna(0)
	bb30 = ta.BBANDS(df, timeperiod=41, nbdevup=2.3, nbdevdn=2.3, matype=0)
	#bb30 = ta.BBANDS(df, timeperiod=20, nbdevup=2.0, nbdevdn=2.0, matype=0)
	df['bbup30'] = bb30['upperband'].fillna(0)
	df['bblow30'] = bb30['lowerband'].fillna(0)
	bb302 = ta.BBANDS(df, timeperiod=41, nbdevup=2.0, nbdevdn=2.0, matype=0)
	#bb302 = ta.BBANDS(df, timeperiod=20, nbdevup=2.0, nbdevdn=2.0, matype=0)
	df['bbup302'] = bb302['upperband'].fillna(0)
	df['bblow302'] = bb302['lowerband'].fillna(0)
	# 计算次周期布林带 (14周期，2标准差)
	bb_element = ta.BBANDS(df, timeperiod=14, nbdevup=2.0, nbdevdn=2.0, matype=0)
	df['element_bb_upper'] = bb_element['upperband'].fillna(0)
	df['element_bb_middle'] = bb_element['middleband'].fillna(0)
	df['element_bb_lower'] = bb_element['lowerband'].fillna(0)
	
	df['macd'] = df['macd'].fillna(0)
	df['macdsignal'] = df['macdsignal'].fillna(0)
	df['macdhist'] = df['macdhist'].fillna(0)
	df['ma5'] = df['ma5'].fillna(0)
	df['ma10'] = df['ma10'].fillna(0)
	df['ma30'] = df['ma30'].fillna(0)
	df['ma250'] = df['ma250'].fillna(0)
	df['ema5'] = df['ema5'].fillna(0)
	df['ema10'] = df['ema10'].fillna(0)
	df['ema24'] = df['ema24'].fillna(0)
	df['ema52'] = df['ema52'].fillna(0)
	df['sma24'] = df['sma24'].fillna(0)
	df['sma52'] = df['sma52'].fillna(0)
	df['rsi'] = df['rsi'].fillna(0)
	df['avg_volume'] = df['volume'].rolling(10).mean()
	# 计算量比，避免产生Infinity值
	df['volume_ratio'] = df['volume'] / df['avg_volume']
	# 填充缺失值（前N根K线）
	df['volume_ratio'] = df['volume_ratio'].fillna(1.0)
	df['avg_volume'] = df['avg_volume'].fillna(0)
	
	# 处理Infinity和-Infinity值
	df['volume_ratio'] = df['volume_ratio'].replace([float('inf'), float('-inf')], 1.0)
	
	# 计算ATR (Average True Range) - 14周期
	df['atr'] = ta.ATR(df, timeperiod=14)
	df['atr'] = df['atr'].fillna(0)
	bb2633 = ta.BBANDS(df, timeperiod=26, nbdevup=3.0, nbdevdn=3.0, matype=0)
	bbp2633 = (df['close'] - bb2633['lowerband']) / (bb2633['upperband'] - bb2633['lowerband'])
	df['bb2633upper'] = bb2633['upperband'].fillna(0)
	df['bb2633lower'] = bb2633['lowerband'].fillna(0)
	df['bbp2633'] = bbp2633.fillna(0)
	df['bb2633middle'] = bb2633['middleband'].fillna(0)
	return df

def calculate_macd(df):
	"""计算MACD指标"""
	exp1 = df['close'].ewm(span=state.macd_fast_period, adjust=False).mean()
	exp2 = df['close'].ewm(span=state.macd_slow_period, adjust=False).mean()
	macd = exp1 - exp2
	signal = macd.ewm(span=state.macd_signal_period, adjust=False).mean()
	histogram = macd - signal
	
	return {
		'macd': macd.tolist(),
		'signal': signal.tolist(),
		'histogram': histogram.tolist()
	}