import pandas as pd
# 将项目根目录添加到 sys.path
#project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
#sys.path.append(project_root)

# 导入当站目录的config文件
try:
    # 尝试相对导入（作为包的一部分）
    from .config import TS_TOKEN, START_DATE, END_DATE, PRECISION_CONFIG
    from .stock_utils import *
except (ImportError, SystemError):  
    # 失败则使用绝对导入（直接运行脚本）
    from config import TS_TOKEN, START_DATE, END_DATE, PRECISION_CONFIG
    from stock_utils import *

from .data_source import get_tushare_pro
pro = get_tushare_pro()
def getStockEp(TS_CODE,start_date=START_DATE,end_date=END_DATE): 
    """
    从tushare接口获取单只股票的财务数据，计算并填充每日的每股收益(EP)指标。
    Parameters:
        TS_CODE (str): 股票代码,格式为 '股票代码.SZ' 或 '股票代码.SH',例如 '000001.SZ'
        START_DATE (str): 开始日期,格式为 'YYYYMMDD'
        END_DATE (str): 结束日期,格式为 'YYYYMMDD'

    Returns:
        pd.DataFrame: 包含以下字段的DataFrame:
            ts_code: 股票代码
            ann_date: 财报公告日期
            end_date: 财报结束日期/填充日期
            basic_eps: 基本每股收益
            diluted_eps: 稀释每股收益
            trade_date: 实际交易日(来自日线数据)
            close: 收盘价
            basic_ep: 基本EP值(basic_eps/close)
            diluted_ep: 稀释EP值(diluted_eps/close)

    Raises:
        Exception: 如果从Tushare接口获取数据时发生错误。
    """

    '''
    1. 调用get_trading_dates函数，填充完整df_income 变量内部end_date，规则为：以end_date为限，往前天从到上一个财报日期为止(不包含上一个财报日期)，期间填充basic_eps,diluted_eps值为报告日的相同值
    2. 新增一个basic_ep列，数据为:basic_eps/当天交易日期的股价，当天交易日期的股价来自df_daily的close列
    3. 新增一个diluted_ep列，数据为:diluted_eps/当天交易日期的股价，当天交易日期的股价来自df_daily的close列
    '''
    try:
        # 获取股票财务数据
        df_income = pro.income(ts_code=TS_CODE, start_date=start_date, end_date=end_date,
                                fields='ts_code,ann_date,end_date,basic_eps,diluted_eps')

        # 填充df_income的end_date区间数据
        if not df_income.empty:
            # 按end_date分组并排序财报数据
            df_income = df_income.sort_values('end_date')
            grouped = df_income.groupby('end_date')
            # 存储填充后的数据
            filled_data = []

            # 遍历每个财报期间
            for end_date, group in grouped:
                # 获取该日期到上一个财报日的所有交易日期
                trading_dates = get_trading_dates(end_date)
                # 填充数据
                for date in trading_dates:
                    # 将日期对象转换为字符串格式（YYYYMMDD）
                    # 如果date是日期对象（有strftime方法），则调用strftime格式化
                    # 否则直接使用原值（假设已经是字符串格式）
                    date_str = date.strftime('%Y%m%d') if hasattr(date, 'strftime') else date
                    filled_row = {
                        'ts_code': TS_CODE,
                        'ann_date': group['ann_date'].iloc[0],
                        'end_date': date_str,
                        'basic_eps': group['basic_eps'].iloc[0],
                        'diluted_eps': group['diluted_eps'].iloc[0]
                    }
                    filled_data.append(filled_row)

            # 合并填充后的数据
            df_income = pd.DataFrame(filled_data)
            # 获取股票日线数据以获取股价
            min_end_date = df_income['end_date'].min()
            df_daily = pro.daily(ts_code=TS_CODE, start_date=min_end_date, end_date=END_DATE)

        # 2. 计算basic_ep和diluted_ep
        if not df_income.empty and not df_daily.empty:
            # 修改原因：原代码使用left_on='end_date'和right_on='trade_date'导致匹配失败
            # 解决方案：将df_daily的trade_date转换为字符串格式再合并
            df_daily['trade_date_str'] = df_daily['trade_date'].astype(str)
            df_income = pd.merge(
                df_income,
                df_daily[['trade_date_str', 'close']],
                left_on='end_date',
                right_on='trade_date_str',
                how='left'
            )
            # 将合并后的trade_date_str重命名为trade_date
            df_income.rename(columns={'trade_date_str': 'trade_date'}, inplace=True)

            # 计算basic_ep和diluted_ep
            df_income['basic_ep'] = (df_income['basic_eps'] / df_income['close']).round(PRECISION_CONFIG)
            df_income['diluted_ep'] = (df_income['diluted_eps'] / df_income['close']).round(PRECISION_CONFIG)
            #print("df_income:",df_income)
            # 按日期截断
            #df_income = df_income[(df_income['trade_date'] >= START_DATE) & (df_income['trade_date'] <= END_DATE)]
            
            # 数据修正
            cols = ['basic_eps', 'diluted_eps', 'close', 'basic_ep', 'diluted_ep']
            df_income = dataCorrect(df_income, cols).sort_values('trade_date')

        return df_income

    except Exception as e:
        # 捕获并处理异常
        print(f"错误: 获取股票 {TS_CODE} 的EP数据时发生错误: {e}")
        return pd.DataFrame()

def getStockEp_ttm(TS_CODE,start_date=START_DATE,end_date=END_DATE): 
    """
    获取股票的TTM(最近12个月)每股收益与股价比率(EP)数据

    Parameters:
        TS_CODE (str): 股票代码,格式为 '股票代码.SZ' 或 '股票代码.SH',例如 '000001.SZ'
        START_DATE (str): 开始日期,格式为 'YYYYMMDD'
        END_DATE (str): 结束日期,格式为 'YYYYMMDD'

    Returns:
        pd.DataFrame: 包含以下字段的DataFrame:
            ts_code: 股票代码
            trade_date: 交易日
            close: 收盘价
            basic_ep_ttm: 基本EP_TTM值(basic_eps_q_ttm/close)
            diluted_ep_ttm: 稀释EP_TTM值(diluted_eps_q_ttm/close)
            basic_eps_q_ttm: 基本每股收益TTM值
            diluted_eps_q_ttm: 稀释每股收益TTM值

    Raises:
        Exception: 如果从Tushare接口获取数据时发生错误
    """
    try:
        start_date=date_format_correction(start_date) #confirm date as yyyymmdd
        end_date=date_format_correction(end_date) #confirm date as yyyymmdd
        df_eps = get_quarterly_eps(tscodeCheck(TS_CODE), start_date, end_date)
        df_eps_ttm = calculate_ttm_eps(df_eps)
        df_eps_ttm=fill_trading_dates_with_eps(df_eps_ttm,start_date=start_date,end_date=end_date)
        df_daily = pro.daily(ts_code=TS_CODE, start_date=start_date, end_date=end_date)
        df_eps_ttm = pd.merge(
            df_eps_ttm,
            df_daily[['trade_date', 'close']],
            on='trade_date',
            how='left'
        )
        df_eps_ttm['basic_ep_ttm'] = (100*df_eps_ttm['basic_eps_q_ttm'] / df_eps_ttm['close']).round(PRECISION_CONFIG)
        df_eps_ttm['diluted_ep_ttm'] = (100*df_eps_ttm['diluted_eps_q_ttm'] / df_eps_ttm['close']).round(PRECISION_CONFIG)
        # 比较最大日期并补充数据, 解决当ep不存在，close值也不会显示的问题
        if not df_eps_ttm.empty and not df_daily.empty:
            max_eps_date = df_eps_ttm['trade_date'].max()
            max_daily_date = df_daily['trade_date'].max()
            print(f"max_eps_date:{max_eps_date}")
            print(f"max_daily_date:{max_daily_date}")
            if max_eps_date != max_daily_date:
                # 获取需要补充的日期范围
                mask = (df_daily['trade_date'] > max_eps_date) & (df_daily['trade_date'] <= max_daily_date)
                additional_data = df_daily.loc[mask, ['ts_code', 'trade_date', 'close']].copy()

                # 补充空列
                for col in df_eps_ttm.columns:
                    if col not in ['ts_code', 'trade_date', 'close']:
                        additional_data[col] = ''
                print(additional_data)
                # 合并数据
                df_eps_ttm = pd.concat([df_eps_ttm, additional_data], ignore_index=True)
                
        cols=["close","basic_ep_ttm","diluted_ep_ttm","basic_eps_q_ttm","diluted_eps_q_ttm"]
        df_eps_ttm = dataCorrect(df_eps_ttm,cols)
        return df_eps_ttm[['ts_code', 'trade_date', 'close', 'basic_ep_ttm', 'diluted_ep_ttm', 'basic_eps_q_ttm', 'diluted_eps_q_ttm']]
    except Exception as e:
        # 捕获并处理异常
        print(f"错误: 获取股票 {TS_CODE} 的EP_TTM数据时发生错误: {e}")
        return pd.DataFrame()

def get_quarterly_eps(TS_CODE, start_date=START_DATE, end_date=END_DATE):
    """
    获取单季度每股收益数据
    Parameters:
        TS_CODE (str): 股票代码
        START_DATE (str): 开始日期(YYYYMMDD)
        END_DATE (str): 结束日期(YYYYMMDD)

    Returns:
        pd.DataFrame: 包含以下字段的DataFrame:
            ts_code: 股票代码
            ann_date: 财报公告日期
            end_date: 财报结束日期(YYYYMMDD格式)
            basic_eps: 累计基本每股收益
            diluted_eps: 累计稀释每股收益
            basic_eps_q: 单季度基本每股收益
            diluted_eps_q: 单季度稀释每股收益
    """
    
    try:
        # 扩展日期范围往前三个季度
        extended_start = (pd.to_datetime(start_date) - pd.DateOffset(months=12)).strftime('%Y%m%d')

        # 获取原始财务数据
        df = pro.income(ts_code=TS_CODE, start_date=extended_start, end_date=end_date,
                       fields='ts_code,ann_date,end_date,basic_eps,diluted_eps')

        if df.empty:
            return pd.DataFrame()

        # 按财报日期排序并去重（解决重复数据问题）
        df = df.sort_values('end_date').drop_duplicates(subset=['end_date'], keep='last')  # 修改原因：确保每个end_date只保留最新数据
        # 计算单季度数据
        df['basic_eps_q'] = df['basic_eps']
        df['diluted_eps_q'] = df['diluted_eps']

        # 非第一季度数据需要减去上季度数据
        mask = ~df['end_date'].str.endswith('0331')
        df.loc[mask, 'basic_eps_q'] = df['basic_eps'].diff()
        df.loc[mask, 'diluted_eps_q'] = df['diluted_eps'].diff()

        # 过滤掉非季报数据（保留3/6/9/12月数据）
        df = df[df['end_date'].str.endswith(('0331', '0630', '0930', '1231'))]

        # 删除最小日期的数据行
        if not df.empty:
            df = df[df['end_date'] != df['end_date'].min()]
        
        df['basic_eps_q'] = df['basic_eps_q'].round(PRECISION_CONFIG)
        df['diluted_eps_q'] = df['diluted_eps_q'].round(PRECISION_CONFIG)
        # 重置行号并保留原始EPS值
        #return df[['ts_code', 'ann_date', 'end_date', 'basic_eps', 'diluted_eps', 'basic_eps_q', 'diluted_eps_q']].reset_index(drop=True)
        #重置行号，去掉原始EPS值
        return df[['ts_code', 'ann_date', 'end_date', 'basic_eps_q', 'diluted_eps_q']].reset_index(drop=True)

    except Exception as e:
        print(f"获取季度EPS数据出错: {e}")
        return pd.DataFrame()

def calculate_ttm_eps(df):
    """
    计算EPS指标的TTM(最近12个月)值

    Parameters:
        df (pd.DataFrame): get_quarterly_eps函数返回的DataFrame，包含以下列:
            - ts_code: 股票代码
            - ann_date: 公告日期
            - end_date: 财报结束日期
            - basic_eps: 基本每股收益(累计)
            - diluted_eps: 稀释每股收益(累计)
            - basic_eps_q: 单季度基本每股收益
            - diluted_eps_q: 单季度稀释每股收益

    Returns:
        pd.DataFrame: 包含原始数据和TTM计算结果的DataFrame，新增以下列:
            - basic_eps_ttm: 基本每股收益TTM值
            - diluted_eps_ttm: 稀释每股收益TTM值
            - basic_eps_q_ttm: 单季度基本每股收益TTM值
            - diluted_eps_q_ttm: 单季度稀释每股收益TTM值
    """
    if df.empty:
        return df

    try:
        # 确保数据按end_date降序排列
        df = df.sort_values('end_date', ascending=False).reset_index(drop=True)

        # 初始化TTM结果列
        #df['basic_eps_ttm'] = None
        #df['diluted_eps_ttm'] = None
        df['basic_eps_q_ttm'] = None
        df['diluted_eps_q_ttm'] = None

        # 遍历每一行数据计算TTM
        for i in range(len(df)):
            # 检查是否有足够的后续数据(至少3个季度)
            if i + 3 >= len(df):
                continue  # 数据不足，跳过计算

            # 计算TTM值(当前季度+后续3个季度)
            #df.at[i, 'basic_eps_ttm'] = df.loc[i:i+3, 'basic_eps'].sum()
            # df.at[i, 'diluted_eps_ttm'] = df.loc[i:i+3, 'diluted_eps'].sum()
            df.at[i, 'basic_eps_q_ttm'] = df.loc[i:i+3, 'basic_eps_q'].sum()
            df.at[i, 'diluted_eps_q_ttm'] = df.loc[i:i+3, 'diluted_eps_q'].sum()

        # 四舍五入保留指定小数位数
        ttm_cols = [ 'basic_eps_q_ttm', 'diluted_eps_q_ttm']
        df[ttm_cols] = df[ttm_cols].round(PRECISION_CONFIG)
        # 删除最后三行数据，因为最后三行ttm数据为None
        df = df.iloc[:-3]

        return df

    except Exception as e:
        print(f"计算TTM值时出错: {e}")
        return pd.DataFrame()
def fill_trading_dates_with_eps(df_ttm,start_date=START_DATE,end_date=END_DATE):
    """
    填充交易日期并保留EPS值

    参数:
        df_ttm (pd.DataFrame): calculate_ttm_eps函数返回的DataFrame，包含以下列:
            - end_date: 财报结束日期(YYYYMMDD格式)
            - basic_eps_ttm: 基本每股收益TTM值
            - diluted_eps_ttm: 稀释每股收益TTM值
            - basic_eps_q_ttm: 单季度基本每股收益TTM值
            - diluted_eps_q_ttm: 单季度稀释每股收益TTM值

    返回:
        pd.DataFrame: 包含填充后的交易日期和对应EPS值的DataFrame

    异常处理:
        - 输入为空DataFrame时直接返回
        - Tushare接口调用失败时返回原始数据
    """
    if df_ttm.empty:
        return df_ttm

    try:
        # 1. 准备数据: 按end_date排序并转换为datetime格式
        df_ttm = df_ttm.sort_values('end_date')
        df_ttm['end_date_dt'] = pd.to_datetime(df_ttm['end_date'])
        # 2. 获取所有需要填充的日期区间
        date_ranges = []
        for i in range(len(df_ttm)-1):
            s_date = df_ttm['end_date_dt'].iloc[i]
            e_date = df_ttm['end_date_dt'].iloc[i+1]
            date_ranges.append((s_date, e_date))
        # 检查是否需要添加最后一个区间
        if df_ttm['end_date_dt'].max() < pd.to_datetime(end_date):
            next_report_date = pd.to_datetime(get_next_report_date(df_ttm['end_date_dt'].max()))
            next_report_date = min(next_report_date, pd.to_datetime(end_date))
            
            date_ranges.append((df_ttm['end_date_dt'].max(), next_report_date))
        # 3. 获取交易所交易日历
        exchange = 'SZSE' if df_ttm['ts_code'].iloc[0].endswith('SZ') else 'SSE' # 可以不用考虑SH,SZ,BJ, 理论上日历应该是一样的
        trade_cal = pro.trade_cal(exchange=exchange,
                                start_date=start_date,
                                end_date=end_date)

        # 过滤出交易日
        trade_cal = trade_cal[trade_cal['is_open'] == 1]
        trade_cal['cal_date_dt'] = pd.to_datetime(trade_cal['cal_date'])
        
        # 4. 填充每个区间内的交易日
        filled_data = []
        #eps_cols = ['basic_eps_ttm', 'diluted_eps_ttm', 'basic_eps_q_ttm', 'diluted_eps_q_ttm']

        # 首先添加原始数据
        for _, row in df_ttm.iterrows():
            filled_data.append(row.to_dict())

        # 然后填充区间数据
        for start_date, end_date in date_ranges:
            # 获取该区间内的所有交易日
            mask = (trade_cal['cal_date_dt'] > start_date) & (trade_cal['cal_date_dt'] < end_date)
            dates_in_range = trade_cal[mask]['cal_date_dt']

            # 使用较小的日期(即start_date)的EPS值填充
            ref_row = df_ttm[df_ttm['end_date_dt'] == start_date].iloc[0]

            for date in dates_in_range:
                new_row = ref_row.copy()
                new_row['end_date'] = date.strftime('%Y%m%d')
                new_row['end_date_dt'] = date
                filled_data.append(new_row)

        # 5. 转换为DataFrame并整理
        result = pd.DataFrame(filled_data)
        result = result.sort_values('end_date_dt')

        # 删除临时列并重置索引
        result = result.drop(columns=['end_date_dt']).reset_index(drop=True)
        result = result.drop(columns=['ann_date']) #去掉ann_date列
        result = result.rename(columns={'end_date': 'trade_date'}) #重命名 end_date 为trade_date
        # 过滤掉非交易日
        result = result[result['trade_date'].isin(trade_cal['cal_date'])]
        return result

    except Exception as e:
        print(f"填充交易日期时出错: {e}, 返回原始数据")
        return df_ttm
if __name__ == '__main__':
    #df = getStockEp(tscodeCheck('688469'))
    '''df = get_quarterly_eps(tscodeCheck('688556'), '20230101', END_DATE)
    print(df)
    df2 = calculate_ttm_eps(df)
    print(df2)
    df3=fill_trading_dates_with_eps(df2)
    print(df3)'''
    df4= getStockEp_ttm(tscodeCheck('300316'), '20230101', END_DATE)
    print(df4)
