"""
Script Name: 
Description: 根据yahoo提供的不复权数据，和分红及拆股数据，来计算前复权和后复权数据。
    注意：
    结果对不上！
    按照yahoo规则，不复权数据已经处理了拆股，所以只要把分红加上去就行，但处理结果与它返回的前复权数据，仍然对不对上！
    有些比如AAPL，差不多可以对得上；但对于KDP等，差异甚大，找不到原因。。
    所以，这个程序暂时无法使用。。。

Author: [Your Name]
Created Date: YYYY-MM-DD
Last Modified: YYYY-MM-DD
Version: 1.0

Modification History:
    - YYYY-MM-DD [Your Name]: 
    - YYYY-MM-DD [Your Name]: 
    - YYYY-MM-DD [Your Name]: 
"""


import pymysql
import pandas as pd
import logging
import os
import time
import config

# 设置日志
filename = os.path.splitext(os.path.basename(__file__))[0]
config.setup_logging(f"./log/{filename}.log")
logger = logging.getLogger()

# 数据库连接函数
def connect_to_db():
    return pymysql.connect(**config.db_config)

# 读取 sp500 表中的所有行，获取 code 和 name
def fetch_sp500_codes(connection):
    query = "SELECT code, code_name as name FROM sp500 "
    return pd.read_sql(query, connection)

# 读取 sp500_his_kline_none 表中的数据并按 time_key 降序排列
def fetch_sp500_his_kline_none(connection, code):
    query = f"SELECT * FROM sp500_his_kline_none WHERE code = '{code}' ORDER BY time_key DESC"
    return pd.read_sql(query, connection)

# 将计算结果插入到 sp500_ajust_kline_202410 表中
def insert_adjusted_kline_data(connection, data):
    try:
        with connection.cursor() as cursor:
            insert_query = """
            INSERT INTO sp500_ajust_kline_202410 (code, name, time_key, hfq_open, hfq_close, qfq_open, qfq_close, none_open, none_close)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
            ON DUPLICATE KEY UPDATE
                hfq_open = VALUES(hfq_open),
                hfq_close = VALUES(hfq_close),
                qfq_open = VALUES(qfq_open),
                qfq_close = VALUES(qfq_close),
                none_open = VALUES(none_open),
                none_close = VALUES(none_close)
            """
            cursor.executemany(insert_query, data)
        connection.commit()
    except Exception as e:
        logger.error(f"Error inserting data: {e}")

# 计算前复权和后复权的价格，并插入到 sp500_ajust_kline_202410
def process_and_insert_adjusted_kline(connection, code, name, result_none):
    dividends_qfq = 0
    dividends_hfq = 0
    dividends_total = result_none['dividends'].sum()

    adjusted_data = []

    for index, row in result_none.iterrows():
        # 计算前复权和后复权的开盘价和收盘价
        qfq_close = row['close'] - dividends_qfq
        qfq_open = row['open'] - dividends_qfq
        hfq_close = row['close'] + dividends_hfq
        hfq_open = row['open'] + dividends_hfq

        adjusted_data.append((
            row['code'], row['name'], row['time_key'], 
            hfq_open, hfq_close, qfq_open, qfq_close, 
            row['open'], row['close']
        ))

        dividends_qfq += row['dividends']
        dividends_hfq = dividends_total - dividends_qfq

    # 插入到 sp500_ajust_kline_202410 表中
    insert_adjusted_kline_data(connection, adjusted_data)
    logger.info(f"Successfully processed and inserted data for code {code}")

# 主函数
def main():
    try:
        connection = connect_to_db()

        # 读取 sp500 表中的所有行，得到 code 和 name 字段
        sp500_codes = fetch_sp500_codes(connection)

        for index, row in sp500_codes.iterrows():
            code = row['code']
            name = row['name']
            logger.info(f"Processing data for code: {code}, name: {name}")

            # 读取 sp500_his_kline_none 表中的数据并按 time_key 降序排列
            result_none = fetch_sp500_his_kline_none(connection, code)

            if result_none.empty:
                logger.warning(f"No data found for code: {code}")
                continue

            # 处理并插入调整后的 K 线数据
            process_and_insert_adjusted_kline(connection, code, name, result_none)

    except Exception as e:
        logger.error(f"Error occurred: {e}")
    finally:
        if connection:
            connection.close()

if __name__ == "__main__":
    main()