# backend/app/data_collection/market_classifier_etl.py
"""
Created on 2025/11/12
---------
@summary: 市场阶段分片 ,多指数综合市场分类器
---------
@author: NBR
"""
import os
from datetime import datetime
from typing import List, Dict

import pandas as pd
import tushare as ts
from sqlalchemy import create_engine, text
from sqlalchemy.exc import SQLAlchemyError


class _MultiIndexMarketClassifier:
    """
    多指数综合市场分类器 - 纯ETL版本
    职责：从Tushare获取数据，计算市场阶段，并将结果写入数据库。
    """

    # --- 配置 ---
    PROXY_DB_URL = os.environ.get('PROXY_DB_URL')
    if not PROXY_DB_URL:
        raise ValueError("PROXY_DB_URL environment variable must be set.")

    # 采用更健壮的连接池配置
    engine = create_engine(PROXY_DB_URL, pool_pre_ping=True, pool_recycle=3600)

    MARKET_INDICES = {
        '000001.SH': {'name': '上证指数', 'weight': 0.3},
        '000300.SH': {'name': '沪深300', 'weight': 0.3},
        '399006.SZ': {'name': '创业板指', 'weight': 0.2},
        '000905.SH': {'name': '中证500', 'weight': 0.2}
    }

    def __init__(self, token: str, db_engine):
        ts.set_token(token)
        self.pro = ts.pro_api()
        self.engine = db_engine
        self._init_database()

    def _init_database(self):
        """通过Proxy确保目标表存在。"""
        create_table_sql = """
        CREATE TABLE IF NOT EXISTS gp_market_category (
            id INT AUTO_INCREMENT PRIMARY KEY,
            market_type VARCHAR(20) NOT NULL,
            start_date DATE NOT NULL,
            end_date DATE NOT NULL,
            duration_days INT NOT NULL,
            analysis_date DATE,
            created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            UNIQUE KEY uk_market_period (market_type, start_date, end_date)
        ) ENGINE=InnoDB COMMENT='市场阶段分类表'
        """
        try:
            with self.engine.connect() as connection:
                with connection.begin() as trans:
                    connection.execute(text(create_table_sql))
                    trans.commit()
            print("✅ 数据库表'gp_market_category'检查/创建完成。")
        except SQLAlchemyError as e:
            print(f"❌ 数据库表初始化失败: {e}")
            raise

    def get_all_index_data(self, start_date: str, end_date: str) -> Dict[str, pd.DataFrame]:
        """获取所有指数的日线数据"""
        print("📈 获取多指数数据...")

        all_data = {}
        for ts_code, info in self.MARKET_INDICES.items():
            try:
                df = self.pro.index_daily(ts_code=ts_code,
                                          start_date=start_date,
                                          end_date=end_date)
                if not df.empty:
                    df['trade_date'] = pd.to_datetime(df['trade_date'])
                    df = df.sort_values('trade_date').reset_index(drop=True)

                    # 计算技术指标
                    df['ma20'] = df['close'].rolling(window=20).mean()
                    df['ma60'] = df['close'].rolling(window=60).mean()
                    df['ma120'] = df['close'].rolling(window=120).mean()
                    df['high_60'] = df['close'].rolling(window=60).max()
                    df['low_60'] = df['close'].rolling(window=60).min()
                    df['volume_ma20'] = df['vol'].rolling(window=20).mean()

                    all_data[ts_code] = df
                    print(f"   ✅ 获取 {info['name']} 数据: {len(df)} 个交易日")
                else:
                    print(f"   ❌ 获取 {info['name']} 数据为空")

            except Exception as e:
                print(f"   ❌ 获取 {info['name']} 数据失败: {e}")

        return all_data

    def calculate_composite_indicators(self, all_data: Dict[str, pd.DataFrame]) -> pd.DataFrame:
        """计算综合市场指标"""
        print("📊 计算综合市场指标...")

        # 获取所有指数的共同交易日
        common_dates = None
        for ts_code, df in all_data.items():
            if common_dates is None:
                common_dates = set(df['trade_date'])
            else:
                common_dates = common_dates.intersection(set(df['trade_date']))

        if not common_dates:
            raise Exception("没有共同的交易日数据")

        common_dates = sorted(common_dates)
        composite_data = []

        for date in common_dates:
            date_metrics = {'trade_date': date}

            # 计算每个指数在该日期的状态
            bull_signals = 0
            bear_signals = 0
            total_weight = 0

            for ts_code, df in all_data.items():
                weight = self.MARKET_INDICES[ts_code]['weight']
                row = df[df['trade_date'] == date]

                if not row.empty:
                    close = row['close'].iloc[0]
                    ma20 = row['ma20'].iloc[0]
                    ma60 = row['ma60'].iloc[0]
                    ma120 = row['ma120'].iloc[0]
                    high_60 = row['high_60'].iloc[0]
                    low_60 = row['low_60'].iloc[0]
                    volume_ratio = row['vol'].iloc[0] / row['volume_ma20'].iloc[0] if row['volume_ma20'].iloc[
                                                                                          0] > 0 else 1

                    # 判断单个指数的状态
                    index_bull_signals = 0
                    index_bear_signals = 0

                    # 牛市信号
                    if close > ma60 > ma120:  # 多头排列
                        index_bull_signals += 1
                    if close > high_60 * 0.95:  # 接近近期高点
                        index_bull_signals += 1
                    if volume_ratio > 1.2:  # 放量
                        index_bull_signals += 1
                    if (close - low_60) / low_60 > 0.2:  # 从低点上涨超过20%
                        index_bull_signals += 1

                    # 熊市信号
                    if close < ma60 < ma120:  # 空头排列
                        index_bear_signals += 1
                    if close < low_60 * 1.05:  # 接近近期低点
                        index_bear_signals += 1
                    if volume_ratio > 1.1 and close < ma20:  # 放量下跌
                        index_bear_signals += 1
                    if (high_60 - close) / high_60 > 0.15:  # 从高点下跌超过15%
                        index_bear_signals += 1

                    # 加权累计信号
                    if index_bull_signals >= 2:
                        bull_signals += weight
                    if index_bear_signals >= 2:
                        bear_signals += weight

                    total_weight += weight

            # 计算综合市场状态
            if total_weight > 0:
                bull_ratio = bull_signals / total_weight
                bear_ratio = bear_signals / total_weight

                if bull_ratio > 0.6:  # 60%以上的指数显示牛市信号
                    market_type = 'bull'
                elif bear_ratio > 0.6:  # 60%以上的指数显示熊市信号
                    market_type = 'bear'
                else:
                    market_type = 'consolidation'
            else:
                market_type = 'consolidation'

            date_metrics.update({
                'market_type': market_type,
                'bull_ratio': bull_ratio,
                'bear_ratio': bear_ratio
            })

            composite_data.append(date_metrics)

        return pd.DataFrame(composite_data)

    def detect_market_phases(self, composite_df: pd.DataFrame) -> List[Dict]:
        """检测市场阶段"""
        print("🔍 检测市场阶段...")

        phases = []
        current_phase = None

        for i, row in composite_df.iterrows():
            date = row['trade_date'].date()
            market_type = row['market_type']

            if current_phase is None:
                current_phase = {
                    'market_type': market_type,
                    'start_date': date,
                    'end_date': date,
                    'dates': [date]
                }
            elif current_phase['market_type'] == market_type:
                # 相同阶段，更新结束日期
                current_phase['end_date'] = date
                current_phase['dates'].append(date)
            else:
                # 阶段转换，保存当前阶段
                phase_duration = (current_phase['end_date'] - current_phase['start_date']).days
                if phase_duration >= 30:  # 至少30天的阶段才保留
                    phases.append(current_phase)

                # 开始新阶段
                current_phase = {
                    'market_type': market_type,
                    'start_date': date,
                    'end_date': date,
                    'dates': [date]
                }

        # 处理最后一个阶段
        if current_phase:
            phase_duration = (current_phase['end_date'] - current_phase['start_date']).days
            if phase_duration >= 30:
                phases.append(current_phase)

        return phases

    def merge_same_type_phases(self, phases: List[Dict]) -> List[Dict]:
        """合并相同类型的相邻阶段"""
        if len(phases) <= 1:
            return phases

        merged_phases = []
        current_phase = phases[0]

        for i in range(1, len(phases)):
            next_phase = phases[i]

            # 检查是否相邻且类型相同
            gap_days = (next_phase['start_date'] - current_phase['end_date']).days
            if (current_phase['market_type'] == next_phase['market_type'] and
                    gap_days <= 60):  # 间隔小于60天可以合并

                # 合并阶段
                current_phase['end_date'] = next_phase['end_date']
                current_phase['dates'].extend(next_phase['dates'])
            else:
                merged_phases.append(current_phase)
                current_phase = next_phase

        merged_phases.append(current_phase)

        print(f"🔄 阶段合并: {len(phases)} → {len(merged_phases)}")
        return merged_phases

    def format_phases_for_storage(self, phases: List[Dict]) -> List[Dict]:
        """格式化阶段数据用于存储"""
        formatted_phases = []

        for phase in phases:
            duration_days = (phase['end_date'] - phase['start_date']).days + 1

            formatted_phase = {
                'market_type': phase['market_type'],
                'start_date': phase['start_date'],
                'end_date': phase['end_date'],
                'duration_days': duration_days,
                # 如果需要存储日期集合，可以在这里添加
                # 'date_set': json.dumps([d.strftime('%Y-%m-%d') for d in sorted(set(phase['dates']))], ensure_ascii=False)
            }

            formatted_phases.append(formatted_phase)

        return formatted_phases

    def save_phases_to_mysql(self, phases: List[Dict]):
        """将市场阶段通过Proxy以事务方式写入MySQL。"""
        if not phases:
            print("ℹ️ 没有市场阶段数据需要保存。")
            return

        print("💾 保存市场阶段数据到MySQL...")
        phases_df = pd.DataFrame(phases)
        phases_df['analysis_date'] = datetime.now().date()
        table_name = 'gp_market_category'

        with self.engine.connect() as connection:
            with connection.begin() as transaction:
                try:
                    connection.execute(text(f"TRUNCATE TABLE {table_name}"))
                    phases_df.to_sql(
                        name=table_name,
                        con=connection,
                        if_exists='append',
                        index=False,
                        chunksize=1000
                    )
                    transaction.commit()
                    print(f"✅ 保存完成: 成功写入 {len(phases_df)} 个市场阶段。")
                except SQLAlchemyError as e:
                    print(f"❌ 保存数据到 {table_name} 失败: {e}")
                    transaction.rollback()
                    raise

    def run_comprehensive_analysis(self, start_date: str, end_date: str):
        """执行完整的ETL流程。"""
        print("🎯 开始多指数市场分析 (ETL)...")
        print("=" * 60)
        try:
            all_data = self.get_all_index_data(start_date, end_date)
            if not all_data:
                print("❌ 无法获取指数数据，ETL终止。")
                return

            composite_df = self.calculate_composite_indicators(all_data)
            if composite_df.empty:
                print("❌ 无法计算综合指标，ETL终止。")
                return

            raw_phases = self.detect_market_phases(composite_df)
            merged_phases = self.merge_same_type_phases(raw_phases)
            final_phases = self.format_phases_for_storage(merged_phases)

            self.save_phases_to_mysql(final_phases)
            self._print_analysis_results(final_phases)

            print(f"\n🎉 ETL分析完成: 共处理了 {len(final_phases)} 个市场阶段。")
        except Exception as e:
            print(f"❌ ETL分析过程中出错: {e}")
            import traceback
            traceback.print_exc()
            raise

    def _print_analysis_results(self, phases: List[Dict]):
        """打印分析结果"""
        print("\n📋 市场阶段分析结果:")
        print("=" * 70)
        print(f"{'序号':<3} {'市场类型':<10} {'开始日期':<12} {'结束日期':<12} {'持续时间':<8}")
        print("-" * 70)

        bull_count = bear_count = consolidation_count = 0

        for i, phase in enumerate(phases, 1):
            market_type = phase['market_type']
            if market_type == 'bull':
                bull_count += 1
                type_display = '🐂 牛市'
            elif market_type == 'bear':
                bear_count += 1
                type_display = '🐻 熊市'
            else:
                consolidation_count += 1
                type_display = '➰ 震荡市'

            print(f"{i:<3} {type_display:<10} {phase['start_date']} {phase['end_date']} "
                  f"{phase['duration_days']:>6} 天")

        print("-" * 70)
        print(f"📈 统计: 牛市{bull_count}个, 熊市{bear_count}个, 震荡市{consolidation_count}个")


# --- 唯一的公共入口函数 ---
def run_market_classifier_etl(start_date: str, end_date: str):
    """
    ETL入口函数，供Celery任务调用。
    """
    print(f"🚀 [ETL START] Market Classifier for period {start_date} to {end_date}")

    # 1. 从环境变量获取配置
    tushare_token = os.environ.get('TUSHARE_TOKEN')
    proxy_db_url = os.environ.get('PROXY_DB_URL')

    if not tushare_token or not proxy_db_url:
        raise ValueError("TUSHARE_TOKEN and PROXY_DB_URL environment variables must be set.")

    # 2. 创建数据库引擎
    engine = create_engine(proxy_db_url, pool_pre_ping=True)

    # 3. 实例化并运行
    try:
        classifier = _MultiIndexMarketClassifier(tushare_token, engine)
        classifier.run_comprehensive_analysis(start_date, end_date)
        print(f"✅ [ETL SUCCESS] Market Classifier finished.")
        return {"message": f"Market classifier ETL completed for period {start_date} to {end_date}."}
    except Exception as e:
        print(f"❌ [ETL FAILED] Market Classifier failed: {e}")
        # 在Celery任务中，重新抛出异常很重要，这样任务状态才会变为FAILURE
        raise


# --- 用于独立测试的入口 ---
if __name__ == "__main__":

    # 为了方便测试，从.env文件加载环境变量
    try:
        from dotenv import load_dotenv

        # 假设.env文件在项目根目录，即上一级的上一级
        dotenv_path = os.path.join(os.path.dirname(__file__), '..', '..', '.env')
        load_dotenv(dotenv_path=dotenv_path)
        print(f"Loaded .env file from: {dotenv_path}")
    except ImportError:
        print("Warning: dotenv not installed. Please set environment variables manually for standalone testing.")

    # 定义测试运行的时间范围
    test_start_date = '20140101'
    test_end_date = '20251130'

    run_market_classifier_etl(test_start_date, test_end_date)
