"""
股票PE/PB估值分析模块

提供股票历史PE、PB分位数分析功能，包括：
1. 历史PE、PB数据获取
2. 分位数计算
3. 可视化展示
4. 行业和概念板块对比分析
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sqlalchemy import create_engine, text
import datetime
import logging
from typing import Tuple, Dict, List, Optional, Union
import os
import matplotlib.dates as mdates
from matplotlib.ticker import FuncFormatter
from pathlib import Path

from .config import DB_URL, OUTPUT_DIR, LOG_FILE

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(LOG_FILE),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("valuation_analysis")

# 设置matplotlib中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


class ValuationAnalyzer:
    """股票估值分析器类"""

    def __init__(self, db_url: str = DB_URL):
        """
        初始化估值分析器
        
        Args:
            db_url: 数据库连接URL
        """
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
        logger.info("估值分析器初始化完成")
        
    def get_stock_name(self, stock_code: str) -> str:
        """
        根据股票代码获取股票名称
        
        Args:
            stock_code: 股票代码
            
        Returns:
            股票名称
        """
        try:
            # 查询数据库获取股票名称
            query = text("""
                SELECT DISTINCT gp_name 
                FROM gp_gnbk 
                WHERE gp_code = :stock_code
                UNION
                SELECT DISTINCT gp_name 
                FROM gp_hybk 
                WHERE gp_code = :stock_code
                LIMIT 1
            """)
            
            with self.engine.connect() as conn:
                result = conn.execute(query, {"stock_code": stock_code}).fetchone()
                
            if result:
                return result[0]
            else:
                logger.warning(f"未找到股票 {stock_code} 的名称信息")
                return stock_code
        except Exception as e:
            logger.error(f"获取股票名称失败: {e}")
            return stock_code

    def get_industry_stocks(self, industry_name: str) -> List[str]:
        """
        获取指定行业的所有股票代码
        
        Args:
            industry_name: 行业名称
            
        Returns:
            股票代码列表
        """
        try:
            query = text("""
                SELECT DISTINCT gp_code 
                FROM gp_hybk 
                WHERE bk_name = :industry_name
            """)
            
            with self.engine.connect() as conn:
                result = conn.execute(query, {"industry_name": industry_name}).fetchall()
                
            if result:
                return [row[0] for row in result]
            else:
                logger.warning(f"未找到行业 {industry_name} 的股票")
                return []
        except Exception as e:
            logger.error(f"获取行业股票失败: {e}")
            return []

    def get_concept_stocks(self, concept_name: str) -> List[str]:
        """
        获取指定概念板块的所有股票代码
        
        Args:
            concept_name: 概念板块名称
            
        Returns:
            股票代码列表
        """
        try:
            query = text("""
                SELECT DISTINCT gp_code 
                FROM gp_gnbk 
                WHERE bk_name = :concept_name
            """)
            
            with self.engine.connect() as conn:
                result = conn.execute(query, {"concept_name": concept_name}).fetchall()
                
            if result:
                return [row[0] for row in result]
            else:
                logger.warning(f"未找到概念板块 {concept_name} 的股票")
                return []
        except Exception as e:
            logger.error(f"获取概念板块股票失败: {e}")
            return []

    def get_industry_avg_data(self, industry_name: str, start_date: str, metric: str) -> pd.DataFrame:
        """
        获取行业平均估值数据
        
        Args:
            industry_name: 行业名称
            start_date: 开始日期
            metric: 估值指标（pe或pb）
            
        Returns:
            包含行业平均估值数据的DataFrame
        """
        try:
            # 获取行业所有股票
            stock_codes = self.get_industry_stocks(industry_name)
            if not stock_codes:
                return pd.DataFrame()

            # 构建查询
            query = text(f"""
                WITH valid_data AS (
                    SELECT 
                        `timestamp`,
                        symbol,
                        {metric}
                    FROM 
                        gp_day_data 
                    WHERE 
                        symbol IN :stock_codes AND 
                        `timestamp` >= :start_date AND
                        {metric} > 0 AND
                        {metric} < 1000  -- 过滤掉极端异常值
                )
                SELECT 
                    `timestamp`,
                    AVG({metric}) as avg_{metric},
                    COUNT(*) as stock_count
                FROM 
                    valid_data
                GROUP BY 
                    `timestamp`
                ORDER BY 
                    `timestamp`
            """)
            
            with self.engine.connect() as conn:
                df = pd.read_sql(
                    query, 
                    conn, 
                    params={"stock_codes": tuple(stock_codes), "start_date": start_date}
                )
            
            if df.empty:
                logger.warning(f"未找到行业 {industry_name} 的估值数据")
                return pd.DataFrame()
                
            logger.info(f"成功获取行业 {industry_name} 的平均{metric.upper()}数据，共 {len(df)} 条记录")
            return df
            
        except Exception as e:
            logger.error(f"获取行业平均数据失败: {e}")
            return pd.DataFrame()

    def get_concept_avg_data(self, concept_name: str, start_date: str, metric: str) -> pd.DataFrame:
        """
        获取概念板块平均估值数据
        
        Args:
            concept_name: 概念板块名称
            start_date: 开始日期
            metric: 估值指标（pe或pb）
            
        Returns:
            包含概念板块平均估值数据的DataFrame
        """
        try:
            # 获取概念板块所有股票
            stock_codes = self.get_concept_stocks(concept_name)
            if not stock_codes:
                return pd.DataFrame()

            # 构建查询
            query = text(f"""
                WITH valid_data AS (
                    SELECT 
                        `timestamp`,
                        symbol,
                        {metric}
                    FROM 
                        gp_day_data 
                    WHERE 
                        symbol IN :stock_codes AND 
                        `timestamp` >= :start_date AND
                        {metric} > 0 AND
                        {metric} < 1000  -- 过滤掉极端异常值
                )
                SELECT 
                    `timestamp`,
                    AVG({metric}) as avg_{metric},
                    COUNT(*) as stock_count
                FROM 
                    valid_data
                GROUP BY 
                    `timestamp`
                ORDER BY 
                    `timestamp`
            """)
            
            with self.engine.connect() as conn:
                df = pd.read_sql(
                    query, 
                    conn, 
                    params={"stock_codes": tuple(stock_codes), "start_date": start_date}
                )
            
            if df.empty:
                logger.warning(f"未找到概念板块 {concept_name} 的估值数据")
                return pd.DataFrame()
                
            logger.info(f"成功获取概念板块 {concept_name} 的平均{metric.upper()}数据，共 {len(df)} 条记录")
            return df
            
        except Exception as e:
            logger.error(f"获取概念板块平均数据失败: {e}")
            return pd.DataFrame()
            
    def get_historical_data(self, stock_code: str, start_date: str = '2018-01-01') -> pd.DataFrame:
        """
        获取股票的历史PE、PB数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，默认为2018-01-01
            
        Returns:
            包含历史PE、PB和价格数据的DataFrame
        """
        try:
            query = text("""
                SELECT 
                    `timestamp`, `close`, `pe`, `pb` 
                FROM 
                    gp_day_data 
                WHERE 
                    symbol = :symbol AND 
                    `timestamp` >= :start_date AND
                    pe > 0 
                ORDER BY 
                    `timestamp`
            """)
            
            with self.engine.connect() as conn:
                df = pd.read_sql(
                    query, 
                    conn, 
                    params={"symbol": stock_code, "start_date": start_date}
                )
            
            if df.empty:
                logger.warning(f"未找到股票 {stock_code} 的历史数据")
                return pd.DataFrame()
                
            # 转换数据类型
            df['close'] = df['close'].astype(float)
            df['pe'] = df['pe'].astype(float)
            df['pb'] = df['pb'].astype(float)
            
            # 过滤异常值(可能有极端值导致图表不可读)
            pe_q1, pe_q3 = df['pe'].quantile(0.05), df['pe'].quantile(0.95)
            pe_iqr = pe_q3 - pe_q1
            pe_upper = pe_q3 + 1.5 * pe_iqr
            
            pb_q1, pb_q3 = df['pb'].quantile(0.05), df['pb'].quantile(0.95)
            pb_iqr = pb_q3 - pb_q1
            pb_upper = pb_q3 + 1.5 * pb_iqr
            
            # 过滤异常值但保留记录
            df['pe_filtered'] = np.where(df['pe'] > pe_upper, pe_upper, df['pe'])
            df['pb_filtered'] = np.where(df['pb'] > pb_upper, pb_upper, df['pb'])
            
            logger.info(f"成功获取到股票 {stock_code} 的历史数据，共 {len(df)} 条记录")
            return df
        except Exception as e:
            logger.error(f"获取历史数据失败: {e}")
            return pd.DataFrame()
    
    def calculate_percentiles(self, data: pd.DataFrame, metric: str = 'pe') -> Dict:
        """
        计算估值指标的分位数
        
        Args:
            data: 历史数据DataFrame
            metric: 估值指标，pe或pb
            
        Returns:
            包含分位数信息的字典
        """
        if data.empty:
            logger.warning(f"数据为空，无法计算{metric}分位数")
            return {}
            
        # 使用过滤后的数据计算分位数
        metric_filtered = f'{metric}_filtered'
        if metric_filtered not in data.columns:
            metric_filtered = metric
            
        # 计算各种分位数
        result = {
            'min': data[metric_filtered].min(),
            'max': data[metric_filtered].max(),
            'current': data[metric_filtered].iloc[-1],
            'mean': data[metric_filtered].mean(),
            'median': data[metric_filtered].median(),
            'q1': data[metric_filtered].quantile(0.25),  # 第一四分位数
            'q3': data[metric_filtered].quantile(0.75),  # 第三四分位数
        }
        
        # 计算当前值的百分位
        result['percentile'] = (data[metric_filtered] <= result['current']).mean() * 100
        
        logger.info(f"计算 {metric} 分位数完成: 当前{metric}={result['current']:.2f}, 百分位={result['percentile']:.2f}%")
        return result
    
    def plot_valuation_bands(
        self, 
        data: pd.DataFrame, 
        stock_code: str,
        metric: str = 'pe',
        industry_name: Optional[str] = None,
        concept_name: Optional[str] = None,
        save_path: Optional[str] = None
    ) -> str:
        """
        绘制估值分位数曲线图
        
        Args:
            data: 历史数据DataFrame
            stock_code: 股票代码
            metric: 估值指标，pe或pb
            industry_name: 行业名称（可选）
            concept_name: 概念板块名称（可选）
            save_path: 保存路径，如果为None则使用默认路径
            
        Returns:
            图表保存路径
        """
        if data.empty:
            logger.warning(f"数据为空，无法绘制{metric}分位数曲线图")
            return ""
            
        stock_name = self.get_stock_name(stock_code)
        
        # 使用过滤后的数据绘图
        metric_filtered = f'{metric}_filtered'
        if metric_filtered not in data.columns:
            metric_filtered = metric
            
        # 计算分位数
        percentiles = self.calculate_percentiles(data, metric)
        
        # 创建图形和子图
        fig, ax1 = plt.subplots(figsize=(12, 8))
        
        # 绘制估值曲线（左轴）
        ax1.plot(data['timestamp'], data[metric_filtered], color='black', linewidth=1.5, label=f'{stock_name} {metric.upper()}')
        
        # 获取并绘制行业平均数据
        industry_data = None
        if industry_name:
            industry_data = self.get_industry_avg_data(industry_name, data['timestamp'].min().strftime('%Y-%m-%d'), metric)
            if not industry_data.empty:
                ax1.plot(industry_data['timestamp'], industry_data[f'avg_{metric}'], 
                        color='blue', linewidth=1.5, label=f'{industry_name}行业平均{metric.upper()}')
                
                # 在图表中添加行业统计信息
                industry_min_count = industry_data['stock_count'].min()
                industry_max_count = industry_data['stock_count'].max()
                industry_avg_count = industry_data['stock_count'].mean()
                logger.info(f"行业 {industry_name} 股票数量统计: 最小={industry_min_count}, 最大={industry_max_count}, 平均={industry_avg_count:.1f}")
        
        # 获取并绘制概念板块平均数据
        concept_data = None
        if concept_name:
            concept_data = self.get_concept_avg_data(concept_name, data['timestamp'].min().strftime('%Y-%m-%d'), metric)
            if not concept_data.empty:
                ax1.plot(concept_data['timestamp'], concept_data[f'avg_{metric}'], 
                        color='red', linewidth=1.5, label=f'{concept_name}概念平均{metric.upper()}')
                
                # 在图表中添加概念板块统计信息
                concept_min_count = concept_data['stock_count'].min()
                concept_max_count = concept_data['stock_count'].max()
                concept_avg_count = concept_data['stock_count'].mean()
                logger.info(f"概念 {concept_name} 股票数量统计: 最小={concept_min_count}, 最大={concept_max_count}, 平均={concept_avg_count:.1f}")
        
        # 绘制分位数线
        ax1.axhline(y=percentiles['min'], color='green', linestyle='-', alpha=0.7, label=f'最小{metric.upper()}')
        ax1.axhline(y=percentiles['max'], color='red', linestyle='-', alpha=0.7, label=f'最大{metric.upper()}')
        ax1.axhline(y=percentiles['mean'], color='purple', linestyle='--', alpha=0.7, label=f'{metric.upper()}均值')
        ax1.axhline(y=percentiles['q1'], color='orange', linestyle='--', alpha=0.7, label=f'第一四分位数')
        ax1.axhline(y=percentiles['q3'], color='brown', linestyle='--', alpha=0.7, label=f'第三四分位数')
        
        # 设置图表标题和标签
        metric_name = "PE" if metric == "pe" else "PB"
        current_percentile = percentiles['percentile']
        title = f"{stock_code} {stock_name} 历史{metric_name}分位数分析"
        if industry_name:
            title += f" vs {industry_name}行业"
        if concept_name:
            title += f" vs {concept_name}概念"
        title += f" (当前{metric_name}百分位: {current_percentile:.2f}%)"
        
        plt.title(title, fontsize=14)
        ax1.set_xlabel('日期', fontsize=12)
        ax1.set_ylabel(f'{metric_name}值', fontsize=12)
        
        # 设置日期格式
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
        ax1.xaxis.set_major_locator(mdates.MonthLocator(interval=6))
        fig.autofmt_xdate()
        
        # 设置图例
        lines1, labels1 = ax1.get_legend_handles_labels()
        ax1.legend(lines1, labels1, loc='upper left')
        
        # 添加网格线
        ax1.grid(True, linestyle='--', alpha=0.6)
        
        # 添加当前估值的标注
        current_value = percentiles['current']
        info_text = f'当前{metric_name}: {current_value:.2f}\n' \
                   f'历史区间: [{percentiles["min"]:.2f}, {percentiles["max"]:.2f}]\n' \
                   f'均值: {percentiles["mean"]:.2f}\n' \
                   f'百分位: {current_percentile:.2f}%'
        
        # 添加行业/概念统计信息
        if industry_data is not None and not industry_data.empty:
            info_text += f'\n\n{industry_name}行业统计:\n' \
                        f'成分股数: {industry_min_count}-{industry_max_count}只\n' \
                        f'均值: {industry_data[f"avg_{metric}"].mean():.2f}'
                        
        if concept_data is not None and not concept_data.empty:
            info_text += f'\n\n{concept_name}概念统计:\n' \
                        f'成分股数: {concept_min_count}-{concept_max_count}只\n' \
                        f'均值: {concept_data[f"avg_{metric}"].mean():.2f}'
        
        ax1.text(
            0.02, 0.05, 
            info_text,
            transform=ax1.transAxes,
            fontsize=10,
            bbox=dict(facecolor='white', alpha=0.8, edgecolor='gray')
        )
        
        # 保存图表
        if save_path is None:
            filename = f"{stock_code}_{metric}_analysis_{datetime.datetime.now().strftime('%Y%m%d')}.png"
            save_path = os.path.join(OUTPUT_DIR, filename)
            
        plt.tight_layout()
        plt.savefig(save_path, dpi=300)
        plt.close()
        
        logger.info(f"估值分位数曲线图已保存至: {save_path}")
        return save_path
    
    def analyze_stock_valuation(
        self, 
        stock_code: str,
        start_date: str = '2018-01-01',
        metrics: List[str] = ['pe', 'pb'],
        industry_name: Optional[str] = None,
        concept_name: Optional[str] = None
    ) -> Dict:
        """
        分析股票估值情况
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期
            metrics: 估值指标列表
            industry_name: 行业名称（可选）
            concept_name: 概念板块名称（可选）
            
        Returns:
            分析结果字典
        """
        logger.info(f"开始分析股票 {stock_code} 的估值情况")
        
        # 获取历史数据
        data = self.get_historical_data(stock_code, start_date)
        if data.empty:
            logger.error(f"无法获取股票 {stock_code} 的历史数据，分析终止")
            return {
                'success': False,
                'message': f"无法获取股票 {stock_code} 的历史数据"
            }
            
        result = {
            'success': True,
            'stock_code': stock_code,
            'stock_name': self.get_stock_name(stock_code),
            'analysis_date': datetime.datetime.now().strftime('%Y-%m-%d'),
            'metrics': {}
        }
        
        # 分析各个估值指标
        for metric in metrics:
            if metric not in ['pe', 'pb']:
                logger.warning(f"不支持的估值指标: {metric}，跳过")
                continue
                
            # 计算分位数
            percentiles = self.calculate_percentiles(data, metric)
            if not percentiles:
                logger.warning(f"无法计算股票 {stock_code} 的 {metric} 分位数")
                continue
                
            # 绘制图表
            chart_path = self.plot_valuation_bands(
                data, 
                stock_code, 
                metric,
                industry_name,
                concept_name
            )
            
            # 添加到结果中
            result['metrics'][metric] = {
                'current': percentiles['current'],
                'percentile': percentiles['percentile'],
                'min': percentiles['min'],
                'max': percentiles['max'],
                'mean': percentiles['mean'],
                'median': percentiles['median'],
                'q1': percentiles['q1'],
                'q3': percentiles['q3'],
                'chart_path': chart_path
            }
            
        logger.info(f"股票 {stock_code} 的估值分析完成")
        return result


def analyze_stock(
    stock_code: str, 
    db_url: Optional[str] = None,
    industry_name: Optional[str] = None,
    concept_name: Optional[str] = None
) -> Dict:
    """
    分析单个股票的估值情况
    
    Args:
        stock_code: 股票代码
        db_url: 数据库连接URL
        industry_name: 行业名称（可选）
        concept_name: 概念板块名称（可选）
        
    Returns:
        分析结果字典
    """
    analyzer = ValuationAnalyzer(db_url) if db_url else ValuationAnalyzer()
    return analyzer.analyze_stock_valuation(stock_code, industry_name=industry_name, concept_name=concept_name) 