#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
行业重叠度和概念重叠度分析工具
用于分析某股票与相似股票的重叠度，识别滞涨投资机会
"""

import sys
import os
import logging
from typing import Dict, List, Optional, Tuple
from pathlib import Path
from sqlalchemy import create_engine, text
import pandas as pd
from datetime import datetime, timedelta

# 添加项目根路径到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.append(str(project_root))

# 导入配置
try:
    from src.valuation_analysis.config import DB_URL
except ImportError:
    # 如果上面的导入失败，尝试直接导入
    import importlib.util
    config_path = os.path.join(project_root, 'src', 'valuation_analysis', 'config.py')
    spec = importlib.util.spec_from_file_location("config", config_path)
    config_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(config_module)
    DB_URL = config_module.DB_URL

# 导入股票代码格式转换工具
try:
    from tools.stock_code_formatter import StockCodeFormatter
except ImportError:
    # 如果上面的导入失败，尝试直接导入
    import importlib.util
    formatter_path = os.path.join(project_root, 'tools', 'stock_code_formatter.py')
    spec = importlib.util.spec_from_file_location("stock_code_formatter", formatter_path)
    formatter_module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(formatter_module)
    StockCodeFormatter = formatter_module.StockCodeFormatter

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class OverlapAnalyzer:
    """行业重叠度和概念重叠度分析器"""
    
    def __init__(self):
        """初始化"""
        # MySQL连接
        self.mysql_engine = None
        
        # 股票代码格式转换器
        self.code_formatter = StockCodeFormatter()
        
        # 重叠度评分权重
        self.INDUSTRY_WEIGHT = 3  # 行业重叠权重
        self.CONCEPT_WEIGHT = 1   # 概念重叠权重
        self.MIN_OVERLAP_SCORE = 5  # 最小重叠分数阈值
        
        self.connect_mysql()
    
    def connect_mysql(self):
        """连接MySQL数据库"""
        try:
            self.mysql_engine = create_engine(
                DB_URL,
                pool_size=5,
                max_overflow=10,
                pool_recycle=3600
            )
            
            # 测试连接
            with self.mysql_engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            
            logger.info("MySQL数据库连接成功")
            
        except Exception as e:
            logger.error(f"MySQL数据库连接失败: {str(e)}")
            raise
    
    def normalize_stock_code(self, stock_code: str) -> str:
        """
        标准化股票代码格式，转换为数据库中使用的格式
        
        Args:
            stock_code: 输入的股票代码，支持多种格式
            
        Returns:
            str: 标准化后的股票代码
        """
        return self.code_formatter.to_prefix_format(stock_code)
    
    def get_stock_industries(self, stock_code: str) -> List[Tuple[str, str]]:
        """
        获取股票所属的行业板块
        
        Args:
            stock_code: 股票代码
            
        Returns:
            List[Tuple[str, str]]: [(板块代码, 板块名称)] 列表
        """
        try:
            normalized_code = self.normalize_stock_code(stock_code)
            
            query = text("""
                SELECT DISTINCT bk_code, bk_name
                FROM gp_hybk 
                WHERE gp_code = :stock_code
                AND bk_code IS NOT NULL
                AND bk_name IS NOT NULL
            """)
            
            with self.mysql_engine.connect() as conn:
                result = conn.execute(query, {"stock_code": normalized_code}).fetchall()
                
            industries = [(str(row[0]), str(row[1])) for row in result]
            logger.info(f"股票 {stock_code} 找到 {len(industries)} 个行业板块")
            return industries
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 行业板块失败: {str(e)}")
            return []
    
    def get_stock_concepts(self, stock_code: str) -> List[Tuple[str, str]]:
        """
        获取股票所属的概念板块
        
        Args:
            stock_code: 股票代码
            
        Returns:
            List[Tuple[str, str]]: [(板块代码, 板块名称)] 列表
        """
        try:
            normalized_code = self.normalize_stock_code(stock_code)
            
            query = text("""
                SELECT DISTINCT bk_code, bk_name
                FROM gp_gnbk 
                WHERE gp_code = :stock_code
                AND bk_code IS NOT NULL
                AND bk_name IS NOT NULL
            """)
            
            with self.mysql_engine.connect() as conn:
                result = conn.execute(query, {"stock_code": normalized_code}).fetchall()
                
            concepts = [(str(row[0]), str(row[1])) for row in result]
            logger.info(f"股票 {stock_code} 找到 {len(concepts)} 个概念板块")
            return concepts
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 概念板块失败: {str(e)}")
            return []
    
    def get_similar_stocks_by_industry(self, industries: List[Tuple[str, str]]) -> Dict[str, List[str]]:
        """
        根据行业板块获取相似股票
        
        Args:
            industries: 行业板块列表 [(板块代码, 板块名称)]
            
        Returns:
            Dict[str, List[str]]: {板块代码: [股票代码列表]}
        """
        try:
            similar_stocks = {}
            
            for bk_code, bk_name in industries:
                query = text("""
                    SELECT DISTINCT gp_code
                    FROM gp_hybk 
                    WHERE bk_code = :bk_code
                    AND gp_code IS NOT NULL
                """)
                
                with self.mysql_engine.connect() as conn:
                    result = conn.execute(query, {"bk_code": bk_code}).fetchall()
                    
                stock_codes = [row[0] for row in result if row[0]]
                similar_stocks[bk_code] = stock_codes
                logger.info(f"行业板块 {bk_name}({bk_code}) 包含 {len(stock_codes)} 只股票")
            
            return similar_stocks
            
        except Exception as e:
            logger.error(f"获取行业相似股票失败: {str(e)}")
            return {}
    
    def get_similar_stocks_by_concept(self, concepts: List[Tuple[str, str]]) -> Dict[str, List[str]]:
        """
        根据概念板块获取相似股票
        
        Args:
            concepts: 概念板块列表 [(板块代码, 板块名称)]
            
        Returns:
            Dict[str, List[str]]: {板块代码: [股票代码列表]}
        """
        try:
            similar_stocks = {}
            
            for bk_code, bk_name in concepts:
                query = text("""
                    SELECT DISTINCT gp_code
                    FROM gp_gnbk 
                    WHERE bk_code = :bk_code
                    AND gp_code IS NOT NULL
                """)
                
                with self.mysql_engine.connect() as conn:
                    result = conn.execute(query, {"bk_code": bk_code}).fetchall()
                    
                stock_codes = [row[0] for row in result if row[0]]
                similar_stocks[bk_code] = stock_codes
                logger.info(f"概念板块 {bk_name}({bk_code}) 包含 {len(stock_codes)} 只股票")
            
            return similar_stocks
            
        except Exception as e:
            logger.error(f"获取概念相似股票失败: {str(e)}")
            return {}
    
    def calculate_overlap_scores(self, target_stock: str, industry_stocks: Dict[str, List[str]], concept_stocks: Dict[str, List[str]]) -> Dict[str, float]:
        """
        计算重叠度分数
        
        Args:
            target_stock: 目标股票代码
            industry_stocks: 行业相似股票字典 {板块代码: [股票代码列表]}
            concept_stocks: 概念相似股票字典 {板块代码: [股票代码列表]}
            
        Returns:
            Dict[str, float]: {股票代码: 重叠分数}
        """
        try:
            overlap_scores = {}
            target_stock_normalized = self.normalize_stock_code(target_stock)
            
            # 计算行业重叠分数（每个行业+3分）
            for bk_code, stock_list in industry_stocks.items():
                for stock_code in stock_list:
                    if stock_code == target_stock_normalized:
                        continue  # 跳过目标股票本身
                    
                    if stock_code not in overlap_scores:
                        overlap_scores[stock_code] = 0
                    
                    overlap_scores[stock_code] += self.INDUSTRY_WEIGHT
            
            # 计算概念重叠分数（每个概念+1分）
            for bk_code, stock_list in concept_stocks.items():
                for stock_code in stock_list:
                    if stock_code == target_stock_normalized:
                        continue  # 跳过目标股票本身
                    
                    if stock_code not in overlap_scores:
                        overlap_scores[stock_code] = 0
                    
                    overlap_scores[stock_code] += self.CONCEPT_WEIGHT
            
            logger.info(f"计算出 {len(overlap_scores)} 只股票的重叠分数")
            return overlap_scores
            
        except Exception as e:
            logger.error(f"计算重叠分数失败: {str(e)}")
            return {}
    
    def filter_high_overlap_stocks(self, overlap_scores: Dict[str, float]) -> List[Tuple[str, float]]:
        """
        筛选高重叠度股票
        
        Args:
            overlap_scores: 重叠分数字典
            
        Returns:
            List[Tuple[str, float]]: [(股票代码, 分数)] 按分数降序排列
        """
        try:
            # 筛选分数大于阈值的股票
            high_overlap_stocks = [
                (stock_code, score) for stock_code, score in overlap_scores.items() 
                if score > self.MIN_OVERLAP_SCORE
            ]
            
            # 按分数降序排列
            high_overlap_stocks.sort(key=lambda x: x[1], reverse=True)
            
            logger.info(f"筛选出 {len(high_overlap_stocks)} 只高重叠度股票（分数>{self.MIN_OVERLAP_SCORE}）")
            return high_overlap_stocks
            
        except Exception as e:
            logger.error(f"筛选高重叠度股票失败: {str(e)}")
            return []
    
    def get_recent_price_changes(self, stock_codes: List[str], days: int = 3) -> Dict[str, Dict]:
        """
        获取股票的近期涨跌幅（默认近3个交易日）
        
        Args:
            stock_codes: 股票代码列表
            days: 统计天数，默认3个交易日
            
        Returns:
            Dict[str, Dict]: {股票代码: {price_change, change_pct, latest_price}}
        """
        try:
            price_changes = {}
            
            for stock_code in stock_codes:
                # 获取最近N个交易日的价格数据
                query = text("""
                    SELECT close, percent, timestamp
                    FROM gp_day_data 
                    WHERE symbol = :stock_code 
                    ORDER BY timestamp DESC 
                    LIMIT :days
                """)
                
                with self.mysql_engine.connect() as conn:
                    result = conn.execute(query, {"stock_code": stock_code, "days": days}).fetchall()
                
                if len(result) >= 2:  # 至少需要2个交易日的数据
                    # 最新价格（第一个）
                    latest_price = float(result[0][0]) if result[0][0] else None
                    latest_change_pct = float(result[0][1]) if result[0][1] else None
                    
                    # N天前价格（最后一个）
                    historical_price = float(result[-1][0]) if result[-1][0] else None
                    
                    if latest_price and historical_price:
                        price_change = latest_price - historical_price
                        change_pct = (price_change / historical_price) * 100
                        
                        price_changes[stock_code] = {
                            'latest_price': latest_price,
                            'price_change': price_change,
                            'change_pct': change_pct,
                            'latest_change_pct': latest_change_pct,
                            'trading_days': len(result)
                        }
            
            logger.info(f"获取到 {len(price_changes)} 只股票的近{days}个交易日涨跌幅数据")
            return price_changes
            
        except Exception as e:
            logger.error(f"获取近期涨跌幅失败: {str(e)}")
            return {}
    
    def analyze_lag_performance(self, target_change: Dict, top5_stocks: List[Tuple[str, float]], similar_changes: Dict[str, Dict]) -> Dict:
        """
        分析目标股票的滞涨情况（基于近3个交易日）
        
        Args:
            target_change: 目标股票涨跌幅数据
            top5_stocks: 重叠度最高的5只股票 [(股票代码, 重叠分数)]
            similar_changes: 相似股票涨跌幅数据
            
        Returns:
            Dict: 滞涨分析结果
        """
        try:
            if not target_change or 'change_pct' not in target_change:
                logger.warning("目标股票涨跌幅数据不完整，无法进行滞涨分析")
                return {
                    'target_lag_vs_top5': None,
                    'target_lag_vs_top1': None,
                    'top5_avg_change': None,
                    'top1_change': None
                }
            
            target_change_pct = target_change['change_pct']
            
            # 获取前5只股票的有效涨跌幅数据
            top5_changes = []
            for stock_code, _ in top5_stocks:
                if stock_code in similar_changes and 'change_pct' in similar_changes[stock_code]:
                    change_pct = similar_changes[stock_code]['change_pct']
                    if change_pct is not None:
                        top5_changes.append(change_pct)
            
            if len(top5_changes) < 1:
                logger.warning("没有有效的前5只股票涨跌幅数据")
                return {
                    'target_lag_vs_top5': None,
                    'target_lag_vs_top1': None,
                    'top5_avg_change': None,
                    'top1_change': None
                }
            
            # 计算前5只股票的平均涨幅
            top5_avg_change = sum(top5_changes) / len(top5_changes)
            
            # 获取重叠度最高的1只股票涨幅
            top1_change = top5_changes[0] if top5_changes else None
            
            # 计算滞涨分数
            lag_vs_top5 = target_change_pct - top5_avg_change if top5_avg_change is not None else None
            lag_vs_top1 = target_change_pct - top1_change if top1_change is not None else None
            
            # 滞涨等级判断
            def get_lag_level(lag_score):
                if lag_score is None:
                    return "无数据"
                elif lag_score < -8:
                    return "严重滞涨"
                elif lag_score < -4:
                    return "明显滞涨"
                elif lag_score < -1.5:
                    return "轻微滞涨"
                elif lag_score > 4:
                    return "跑赢同行"
                elif lag_score > 1.5:
                    return "略胜同行"
                else:
                    return "正常水平"
            result = {
                'target_lag_vs_top5': {
                    'lag_score': lag_vs_top5,
                    'lag_level': get_lag_level(lag_vs_top5),
                    'target_change_pct': target_change_pct,
                    'top5_avg_change': top5_avg_change
                },
                'target_lag_vs_top1': {
                    'lag_score': lag_vs_top1,
                    'lag_level': get_lag_level(lag_vs_top1),
                    'target_change_pct': target_change_pct,
                    'top1_change': top1_change
                },
                'top5_avg_change': top5_avg_change,
                'top1_change': top1_change
            }
            
            logger.info(f"滞涨分析完成 - 目标股票: {target_change_pct:.2f}%, 前5平均: {top5_avg_change:.2f}%, 最高1只: {top1_change:.2f}%")
            return result
            
        except Exception as e:
            logger.error(f"滞涨分析失败: {str(e)}")
            return {
                'target_lag_vs_top5': None,
                'target_lag_vs_top1': None,
                'top5_avg_change': None,
                'top1_change': None
            }
    
    def get_stock_name(self, stock_code: str) -> Optional[str]:
        """
        获取股票名称
        
        Args:
            stock_code: 股票代码
            
        Returns:
            str: 股票名称
        """
        try:
            normalized_code = self.normalize_stock_code(stock_code)
            
            query = text("""
                SELECT gp_name
                FROM gp_code_all 
                WHERE gp_code = :stock_code
                LIMIT 1
            """)
            
            with self.mysql_engine.connect() as conn:
                result = conn.execute(query, {"stock_code": normalized_code}).fetchone()
                
            if result and result[0]:
                return result[0]
            else:
                return None
                
        except Exception as e:
            logger.error(f"获取股票名称失败 {stock_code}: {str(e)}")
            return None
    
    def analyze_stock_overlap(self, stock_code: str, days: int = 20) -> Dict:
        """
        分析股票重叠度的主函数
        
        Args:
            stock_code: 目标股票代码
            days: 统计涨跌幅的天数
            
        Returns:
            Dict: 分析结果
        """
        try:
            logger.info(f"开始分析股票 {stock_code} 的重叠度")
            
            # 1. 获取目标股票的行业和概念板块
            industries = self.get_stock_industries(stock_code)
            concepts = self.get_stock_concepts(stock_code)
            
            if not industries and not concepts:
                logger.warning(f"股票 {stock_code} 没有找到行业或概念板块数据")
                return {
                    'target_stock': stock_code,
                    'error': '未找到行业或概念板块数据'
                }
            
            # 2. 获取相似股票
            industry_stocks = self.get_similar_stocks_by_industry(industries)
            concept_stocks = self.get_similar_stocks_by_concept(concepts)
            
            # 3. 计算重叠分数
            overlap_scores = self.calculate_overlap_scores(stock_code, industry_stocks, concept_stocks)
            
            # 4. 筛选高重叠度股票
            high_overlap_stocks = self.filter_high_overlap_stocks(overlap_scores)
            
            if not high_overlap_stocks:
                logger.info(f"股票 {stock_code} 没有找到高重叠度股票（分数>{self.MIN_OVERLAP_SCORE}）")
                return {
                    'target_stock': stock_code,
                    'similar_stocks_count': 0,
                    'message': f'没有找到重叠度超过{self.MIN_OVERLAP_SCORE}分的股票'
                }
            
            # 5. 获取前5只重叠度最高的股票
            top5_stocks = high_overlap_stocks[:5]  # 取前5只
            top5_stock_codes = [stock[0] for stock in top5_stocks]
            all_stock_codes = [self.normalize_stock_code(stock_code)] + top5_stock_codes
            
            price_changes = self.get_recent_price_changes(all_stock_codes, days)
            
            # 6. 滞涨分析
            target_change = price_changes.get(self.normalize_stock_code(stock_code))
            similar_changes = {code: data for code, data in price_changes.items() 
                             if code != self.normalize_stock_code(stock_code)}
            
            lag_analysis = self.analyze_lag_performance(target_change, top5_stocks, similar_changes)
            
            # 7. 组装前5只股票结果
            top5_results = []
            for similar_stock, overlap_score in top5_stocks:
                stock_name = self.get_stock_name(similar_stock)
                price_data = price_changes.get(similar_stock, {})
                
                top5_results.append({
                    'stock_code': similar_stock,
                    'stock_name': stock_name,
                    'overlap_score': overlap_score,
                    'latest_price': price_data.get('latest_price'),
                    'price_change': price_data.get('price_change'),
                    'change_pct': price_data.get('change_pct')
                })
            
            return {
                'target_stock': stock_code,
                'target_stock_name': self.get_stock_name(stock_code),
                'target_price_data': target_change,
                'analysis_date': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'analysis_period_days': days,
                'industries': [{'code': code, 'name': name} for code, name in industries],
                'concepts': [{'code': code, 'name': name} for code, name in concepts],
                'overlap_threshold': self.MIN_OVERLAP_SCORE,
                'top5_stocks': top5_results,
                'lag_analysis': lag_analysis
            }
            
        except Exception as e:
            logger.error(f"分析股票重叠度失败 {stock_code}: {str(e)}")
            return {
                'target_stock': stock_code,
                'error': str(e)
            }
    
    def close_connection(self):
        """关闭数据库连接"""
        try:
            if self.mysql_engine:
                self.mysql_engine.dispose()
                logger.info("MySQL连接已关闭")
        except Exception as e:
            logger.error(f"关闭MySQL连接失败: {str(e)}")


def main():
    """主函数 - 示例用法"""
    analyzer = None
    try:
        # 创建分析器实例
        analyzer = OverlapAnalyzer()
        
        # 示例：分析某只股票的重叠度
        stock_code = "301308.SZ"  # 圣邦股份
        
        print(f"=== 股票重叠度分析：{stock_code} ===")
        
        # 执行分析
        result = analyzer.analyze_stock_overlap(stock_code, days=3)
        
        if 'error' in result:
            print(f"分析失败: {result['error']}")
            return
        
        # 输出结果
        print(f"\n目标股票: {result['target_stock_name']} ({result['target_stock']})")
        print(f"分析日期: {result['analysis_date']}")
        print(f"统计周期: {result['analysis_period_days']}天")
        
        if 'target_price_data' in result and result['target_price_data']:
            target_data = result['target_price_data']
            print(f"目标股票涨跌幅: {target_data.get('change_pct', 0):.2f}%")
        
        print(f"\n行业板块: {len(result['industries'])}个")
        for industry in result['industries'][:3]:  # 显示前3个
            print(f"  - {industry['name']} ({industry['code']})")
        
        print(f"\n概念板块: {len(result['concepts'])}个")
        for concept in result['concepts'][:3]:  # 显示前3个
            print(f"  - {concept['name']} ({concept['code']})")
        
        print(f"\n重叠度最高的5只股票:")
        print("-" * 80)
        print(f"{'股票代码':<12} {'股票名称':<15} {'重叠分数':<8} {'最新价格':<10} {'涨跌幅':<10}")
        print("-" * 80)
        
        for stock in result['top5_stocks']:
            print(f"{stock['stock_code']:<12} "
                  f"{stock['stock_name'] or '未知':<15} "
                  f"{stock['overlap_score']:<8.1f} "
                  f"{stock['latest_price'] or 0:<10.2f} "
                  f"{stock['change_pct'] or 0:<10.2f}%")
        
        # 显示滞涨分析结果
        lag_analysis = result.get('lag_analysis', {})
        if lag_analysis:
            print(f"\n滞涨分析结果:")
            print("-" * 50)
            
            # vs 前5平均
            vs_top5 = lag_analysis.get('target_lag_vs_top5', {})
            if vs_top5:
                print(f"vs 前5平均: {vs_top5.get('lag_score', 0):.2f}% ({vs_top5.get('lag_level', '未知')})")
                print(f"  目标股票: {vs_top5.get('target_change_pct', 0):.2f}%")
                print(f"  前5平均: {vs_top5.get('top5_avg_change', 0):.2f}%")
            
            # vs 最高1只
            vs_top1 = lag_analysis.get('target_lag_vs_top1', {})
            if vs_top1:
                print(f"vs 最高1只: {vs_top1.get('lag_score', 0):.2f}% ({vs_top1.get('lag_level', '未知')})")
                print(f"  目标股票: {vs_top1.get('target_change_pct', 0):.2f}%")
                print(f"  最高1只: {vs_top1.get('top1_change', 0):.2f}%")
        
    except Exception as e:
        logger.error(f"程序执行失败: {str(e)}")
    finally:
        if analyzer:
            analyzer.close_connection()


if __name__ == "__main__":
    main()
