import numpy as np
import math
import akshare as ak
import os
from datetime import datetime, timedelta, date
import pandas as pd
import mplfinance as mpf

def calc_sma_atr_pd(kdf,period):
    """计算TR与ATR

    Args:
        kdf (_type_): 历史数据
        period (_type_): ATR周期

    Returns:
        _type_: 返回kdf，增加TR与ATR列
    """
    kdf['HL'] = kdf['最高'] - kdf['最低']
    kdf['HC'] = np.abs(kdf['最高'] - kdf['收盘'].shift(1))
    kdf['LC'] = np.abs(kdf['最低'] - kdf['收盘'].shift(1))
    kdf['TR'] = np.round(kdf[['HL','HC','LC']].max(axis=1), 3)
    # ranges = pd.concat([high_low, high_close, low_close], axis=1)
    # true_range = np.max(ranges, axis=1)
    kdf['ATR'] = np.round(kdf['TR'].rolling(period).mean(), 3)
    
    return kdf.drop(['HL','HC','LC'], axis = 1)

class TurtleTrading(object):
    """对象范围较小，对某一个标的创建一个海龟，如513300，
    计算ATR、
    Position Size，
    买入、卖出、加仓等行为
    Args:
        object (_type_): _description_
    """
    def __init__(self, TradeCode) -> None:
        self.TradeCode = TradeCode
    
    def GetRecentData(self):
        """获取某个标的的最近数据，从两年前到今天, 计算后的数据保存在self.CurrentData

        Returns:
            _type_: _description_
        """
        Today = datetime.today()
        # print(Today)
        formatted_date = Today.strftime("%Y%m%d")
        two_years_ago = date.today() - timedelta(days=365*2).strftime("%Y%m%d")
        # print(formatted_date)
        Code = f"{self.TradeCode}"
        self.CurrentData = ak.fund_etf_hist_em(symbol=Code, period="daily", start_date=two_years_ago, end_date=formatted_date, adjust="")
        # return CurrentData
        
    def CalATR(self, data, ATRday, SaveOrNot):
        """计算某个标的的ATR，从上市日到今天, 计算后的数据保存在self.CurrentData

        Args:
            ATRday: 多少日ATR
            SaveOrNot (_type_): 是否保存.csv数据
        """

        self.CurrentData = calc_sma_atr_pd(data, ATRday)
        self.N = self.CurrentData['ATR']
        if SaveOrNot:
            self.CurrentData.to_csv('513300data-N.csv', index=False)
            print("csv保存成功")
        
        return self.N
    
    
    def CalPositionSize(self, RiskCoef, Capital):
        """计算PosizionSize 持有的单位，该单位某标的，1N波动对应RiskCoef * Capital资金

        Args:
            RiskCoef (_type_): 风险系数
            Capital (_type_): 资金
        """
        
        N = self.CurrentData.iloc[-1]['ATR']
        # N = 0.473
        Price = self.CurrentData.iloc[-1]['收盘']
        # Price = 5.60
        self.PositionSize = RiskCoef * Capital /( N*100*Price) # 默认用股票形式了 100
        
        
        return self.PositionSize
        
       
    def ReadExistData(self, data):
        """除了通过发请求获取数据，也可以读本地的数据库，赋值给self.CurrentData

        Args:
            data (_type_): 本地csv名称
        """
        self.CurrentData = pd.read_csv(data)

    def DrawKLine(self, days):
        """画出k线图看看,画出最近days天的K线图
        """
        
        # 日期部分

        dates = pd.to_datetime(self.CurrentData['日期'][-days:])
        # Klinedf['Data'] = pd.to_datetime(self.CurrentData['日期'])
        Klinedf = pd.DataFrame()
        # Klinedf.set_index = Klinedf['Data']
        
        # 其他数据
        Klinedf['Open'] = self.CurrentData['开盘'][-days:]
        Klinedf['High'] = self.CurrentData['最高'][-days:]
        Klinedf['Low'] = self.CurrentData['最低'][-days:]
        Klinedf['Close'] = self.CurrentData['收盘'][-days:]
        Klinedf['Volume'] = self.CurrentData['成交量'][-days:]
        
        Klinedf.set_index(dates, inplace=True)
        # 画图
        mpf.plot(Klinedf, type='candle', style='yahoo', volume=False, mav=(5,), addplot=[mpf.make_addplot(self.Donchian[['Upper', 'Lower']])])
    
    def calculate_donchian_channel(self, days, n):
        """
        计算唐奇安通道days一共多少日， n多少日唐奇安

        参数:
        self.CurrentData (DataFrame): 包含价格数据的Pandas DataFrame，至少包含"High"和"Low"列
        n (int): 时间周期

        返回:self.Donchian
        DataFrame: 唐奇安通道的DataFrame，包含"Upper", "Lower", 和 "Middle"列
        """
        Donchian = pd.DataFrame()
        # 计算最高价和最低价的N日移动平均线
        Donchian['Upper'] = self.CurrentData['最高'][-days:].rolling(n).max()
        Donchian['Lower'] = self.CurrentData['最低'][-days:].rolling(n).min()

        # # 计算中间线
        # Donchian['Middle'] = (self.Donchian['Upper'] + self.Donchian['Lower']) / 2

        return Donchian