import pandas as pd
import numpy as np
from datetime import datetime

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 *

'''
给定个股代码和起止日期：
1.  调用tushare接口查询个股分红数据，接口文档：https://tushare.pro/document/2?doc_id=103 查询日期范围内的所有分红记录
2.  返回调整后的分红记录，包含以下字段：
    - ts_code: 股票代码
    - end_date: 分红年度
    - ann_date: 公告日期
    - ex_date: 除权除息日
    - cash_div_tax: 每股现金分红（含税）
    - div_proc: 实施进度,仅筛选 div_proc='实施'的记录
3.  调用tushare 的trade_cal接口，查询起止日期内的所有交易日
4.  建一个新的df,在起止日期内填充，规则：
    - 下列运算过程中将起始时间往前推 gap_days,赋值360天。运算结束后，截取起止时间内的数据
    - trade_date: 以交易日历为准，填充所有交易日。
    - 填入当ex_date=trade_date的日期，填入：ex_date,cash_div_tax，其余日期ex_date留空，cash_div_tax 置 0
    - cash_div_year:逐行计算填入，以trade_date往前计算过去gap_days天内的cash_div_tax之和
5. 调用tushare的个股日线行情接口：https://tushare.pro/document/2?doc_id=27，查询起止日期内的个股日线行情数据，包含以下字段：
    - ts_code: 股票代码
    - trade_date: 交易日期
    - close: 收盘价
6. 将分红数据和日线行情数据按交易日期合并，得到最终结果，包含以下字段：
    - ts_code: 股票代码
    - trade_date: 交易日期
    - close: 收盘价
    - cach_div_tax: 每股现金分红（含税）
    - cach_div_year: 每股现金TTM年度分红（含税）
    - div_yield: 股息率，计算公式为 (cach_div_year / close) * 100，保留PRECISION_CONFIG位小数，如果cash_div_year为0则div_yield也为0
7. 返回最终结果的DataFrame
'''


# 初始化tushare
from .data_source import get_tushare_pro
pro = get_tushare_pro()

def get_dividend_yield(ts_code, start_date, end_date):
    # 检查日期范围是否超过当前日期，若超过则调整为当前日期
    today = datetime.now().strftime("%Y%m%d")
    if end_date > today:
        end_date = today
    if start_date > today:
        start_date = today
    
    gap_days = 360  # 定义TTM计算窗口期为360天

    # 1. 获取分红数据：查询指定股票的分红信息
    df_div = pro.dividend(ts_code=ts_code, fields='ts_code,end_date,ann_date,ex_date,cash_div_tax,div_proc')
    # 筛选实施状态的分红记录
    df_div = df_div[df_div['div_proc'] == '实施']
    print(f"分红原始数据：\n {df_div.to_string()} ")

    # 2. 获取交易日历：查询指定日期范围内的开盘日
    df_cal = pro.trade_cal(exchange='', start_date=start_date, end_date=end_date, is_open='1')
    trade_dates = df_cal['cal_date'].tolist()

    # 3. 创建基础DataFrame：将计算窗口期向前扩展gap_days天
    extended_start = pd.to_datetime(start_date) - pd.Timedelta(days=gap_days)
    extended_start_str = extended_start.strftime('%Y%m%d')
    # 获取扩展后的交易日历
    df_cal_ext = pro.trade_cal(exchange='', start_date=extended_start_str, end_date=end_date, is_open='1')
    df_base = pd.DataFrame({'trade_date': df_cal_ext['cal_date']})

    # 4. 合并分红数据：将分红信息按除权除息日合并到基础交易日历
    df_base = df_base.merge(df_div[['ex_date', 'cash_div_tax']],
                           left_on='trade_date', right_on='ex_date', how='left')
    # 填充空值：无分红日期现金分红设为0
    df_base['cash_div_tax'] = df_base['cash_div_tax'].fillna(0)
    
    # 5. 计算TTM年度分红：滚动计算过去gap_days天的现金分红总和
    # 修正滚动窗口计算：使用固定窗口大小，确保在窗口期内正确累加
    df_base['cash_div_year'] = df_base['cash_div_tax'].rolling(window=gap_days, min_periods=0).sum()


    print(f"分红填充数据S1：\n {df_base.to_string()} ")
    # 6. 截取指定日期范围：保留原始查询日期范围内的数据
    df_base = df_base[df_base['trade_date'] >= start_date]
    print(f"分红填充数据S2：\n {df_base.to_string()} ")
    # 7. 获取日线行情数据：查询指定股票的日线收盘价
    df_daily = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date,
                        fields='ts_code,trade_date,close')

    # 8. 合并数据：将行情数据与分红数据按交易日合并
    result = df_base.merge(df_daily, on='trade_date', how='left')
    # 填充股票代码：确保所有行都有股票代码
    result['ts_code'] = result['ts_code'].fillna(ts_code)

    # 9. 计算股息率：当TTM分红>0时计算（分红/收盘价）*100，否则为0
    result['div_yield'] = np.where(
        result['cash_div_year'] > 0,
        (result['cash_div_year'] / result['close']) * 100,
        0
    )
    # 精度处理：保留配置指定的小数位数
    result['div_yield'] = result['div_yield'].round(PRECISION_CONFIG)

    # 10. 整理列顺序：选择并排列最终输出的列
    result = result[['ts_code', 'trade_date', 'close', 'cash_div_tax', 'cash_div_year', 'div_yield']]

    return result

if __name__ == "__main__":
    # 测试代码
    test_code = "600900.SH"
    test_start = "20230101"
    test_end = "20251231"

    try:
        result = get_dividend_yield(test_code, test_start, test_end)
        print(f"股票 {test_code} 的分红股息率数据:")
        print(result.head(100))
        print(f"\n数据形状: {result.shape}")
        print(f"\n数据列名: {result.columns.tolist()}")

        # 检查是否有分红数据
        if not result.empty:
            print(f"\n股息率统计:")
            print(result['div_yield'].describe())
        else:
            print("未找到分红数据")

    except Exception as e:
        print(f"测试过程中出现错误: {e}")