"""
美股PE估值分析模块

提供美股历史PE分位数分析功能，包括：
1. 美股历史PE数据获取
2. 分位数计算
3. 可视化展示
"""

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

# 导入配置
try:
    from src.valuation_analysis.config import DB_URL, OUTPUT_DIR, LOG_FILE
except ImportError:
    # 如果导入失败，使用默认配置
    DB_URL = 'mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'
    OUTPUT_DIR = 'results/valuation_analysis'
    LOG_FILE = 'logs/us_valuation_analysis.log'

# 确保输出目录存在
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(os.path.dirname(LOG_FILE), exist_ok=True)

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

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


class USValuationAnalyzer:
    """美股估值分析器类"""

    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_us_stock_name(self, stock_code: str) -> str:
        """
        根据美股代码获取股票名称
        
        Args:
            stock_code: 美股代码 (如 AAPL, GOOGL)
            
        Returns:
            股票名称
        """
        try:
            # 查询数据库获取美股名称
            query = text("""
                SELECT DISTINCT stock_name 
                FROM us_stock_price_data 
                WHERE stock_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:
                # 如果实时数据表中没有，尝试从日线数据表获取
                query2 = text("""
                    SELECT DISTINCT symbol 
                    FROM us_day_data 
                    WHERE symbol = :stock_code
                    LIMIT 1
                """)
                
                with self.engine.connect() as conn:
                    result2 = conn.execute(query2, {"stock_code": stock_code}).fetchone()
                
                if result2:
                    return stock_code  # 返回股票代码作为名称
                else:
                    logger.warning(f"未找到美股 {stock_code} 的名称信息")
                    return stock_code
        except Exception as e:
            logger.error(f"获取美股名称失败: {e}")
            return stock_code

    def get_us_historical_data(self, stock_code: str, start_date: str = '2018-01-01') -> pd.DataFrame:
        """
        获取美股的历史PE数据
        
        Args:
            stock_code: 美股代码 (如 AAPL, GOOGL)
            start_date: 开始日期，默认为2018-01-01
            
        Returns:
            包含历史PE和价格数据的DataFrame
        """
        try:
            # 美股数据查询，只包含PE字段，不包含PB
            query = text("""
                SELECT 
                    `timestamp`, `close`, `pe` 
                FROM 
                    us_day_data 
                WHERE 
                    symbol = :symbol AND 
                    `timestamp` >= :start_date AND
                    pe > 0 AND
                    pe IS NOT NULL
                ORDER BY 
                    `timestamp` ASC
            """)
            
            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 df
            
            # 转换时间戳列
            df['timestamp'] = pd.to_datetime(df['timestamp'])
            
            # 确保数值列是正确类型
            df['close'] = pd.to_numeric(df['close'], errors='coerce')
            df['pe'] = pd.to_numeric(df['pe'], errors='coerce')
            
            # 过滤异常值
            df = self._filter_extreme_values(df, 'pe')
            
            logger.info(f"获取美股 {stock_code} 历史数据成功，共 {len(df)} 条记录")
            return df
            
        except Exception as e:
            logger.error(f"获取美股历史数据失败: {e}")
            return pd.DataFrame()

    def _filter_extreme_values(self, data: pd.DataFrame, metric: str) -> pd.DataFrame:
        """
        过滤极端值
        
        Args:
            data: 数据DataFrame
            metric: 指标名称
            
        Returns:
            过滤后的DataFrame
        """
        if data.empty or metric not in data.columns:
            return data
            
        # 过滤负值和异常大的值
        if metric == 'pe':
            # PE通常在0-100之间，过滤掉大于1000的异常值
            data = data[(data[metric] > 0) & (data[metric] <= 1000)]
        elif metric == 'pb':
            # PB通常在0-50之间，过滤掉大于100的异常值
            data = data[(data[metric] > 0) & (data[metric] <= 100)]
        
        # 创建过滤后的列
        data[f'{metric}_filtered'] = data[metric]
        
        return data

    def calculate_us_percentiles(self, data: pd.DataFrame, metric: str = 'pe') -> Dict:
        """
        计算美股估值指标的分位数
        
        Args:
            data: 历史数据DataFrame
            metric: 估值指标，只支持pe
            
        Returns:
            包含分位数信息的字典
        """
        if data.empty:
            logger.warning(f"数据为空，无法计算{metric}分位数")
            return {}
            
        # 美股只支持PE分析
        if metric != 'pe':
            logger.warning(f"美股估值分析只支持PE指标，不支持{metric}")
            return {}
            
        # 使用过滤后的数据计算分位数
        metric_filtered = f'{metric}_filtered'
        if metric_filtered not in data.columns:
            metric_filtered = metric
            
        # 计算各种分位数
        percentiles = {
            'min': data[metric_filtered].min(),
            'max': data[metric_filtered].max(),
            'mean': data[metric_filtered].mean(),
            'median': data[metric_filtered].median(),
            'q1': data[metric_filtered].quantile(0.25),
            'q3': data[metric_filtered].quantile(0.75),
            'p10': data[metric_filtered].quantile(0.1),
            'p90': data[metric_filtered].quantile(0.9),
            'p5': data[metric_filtered].quantile(0.05),
            'p95': data[metric_filtered].quantile(0.95)
        }
        
        # 获取当前值（最新值）
        if not data.empty:
            percentiles['current'] = data[metric_filtered].iloc[-1]
        else:
            percentiles['current'] = None
            
        # 计算当前值在历史分位数中的位置
        if percentiles['current'] is not None:
            current_value = percentiles['current']
            sorted_values = sorted(data[metric_filtered].dropna())
            if sorted_values:
                percentile_rank = (np.searchsorted(sorted_values, current_value) / len(sorted_values)) * 100
                percentiles['current_percentile'] = percentile_rank
            else:
                percentiles['current_percentile'] = None
        else:
            percentiles['current_percentile'] = None
            
        logger.info(f"美股{metric}分位数计算完成")
        return percentiles

    def get_us_industry_avg_data(self, industry_name: str, start_date: str, metric: str = 'pe') -> Optional[pd.DataFrame]:
        """
        获取美股行业平均数据（暂不支持，返回None）
        
        Args:
            industry_name: 行业名称
            start_date: 开始日期
            metric: 估值指标
            
        Returns:
            None (美股暂不支持行业分析)
        """
        logger.info("美股暂不支持行业平均数据分析")
        return None

    def get_us_concept_avg_data(self, concept_name: str, start_date: str, metric: str = 'pe') -> Optional[pd.DataFrame]:
        """
        获取美股概念板块平均数据（暂不支持，返回None）
        
        Args:
            concept_name: 概念板块名称
            start_date: 开始日期
            metric: 估值指标
            
        Returns:
            None (美股暂不支持概念板块分析)
        """
        logger.info("美股暂不支持概念板块平均数据分析")
        return None

    def create_us_valuation_chart(self, stock_code: str, start_date: str = '2018-01-01', 
                                 metric: str = 'pe', save_path: Optional[str] = None) -> str:
        """
        创建美股估值分析图表
        
        Args:
            stock_code: 美股代码
            start_date: 开始日期
            metric: 估值指标（只支持pe）
            save_path: 保存路径
            
        Returns:
            图表文件路径
        """
        try:
            # 获取历史数据
            data = self.get_us_historical_data(stock_code, start_date)
            if data.empty:
                raise ValueError(f"未找到美股 {stock_code} 的历史数据")
            
            # 计算分位数
            percentiles = self.calculate_us_percentiles(data, metric)
            if not percentiles:
                raise ValueError(f"无法计算美股 {stock_code} 的{metric}分位数")
            
            # 获取股票名称
            stock_name = self.get_us_stock_name(stock_code)
            
            # 创建图表
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(15, 12))
            
            # 上图：股价走势
            ax1.plot(data['timestamp'], data['close'], linewidth=2, color='blue', label='股价')
            ax1.set_title(f'{stock_name} ({stock_code}) 股价走势', fontsize=16, fontweight='bold')
            ax1.set_ylabel('股价 (USD)', fontsize=12)
            ax1.legend()
            ax1.grid(True, alpha=0.3)
            
            # 下图：PE估值分析
            metric_filtered = f'{metric}_filtered' if f'{metric}_filtered' in data.columns else metric
            
            # 绘制PE曲线
            ax2.plot(data['timestamp'], data[metric_filtered], linewidth=2, color='red', label=f'{metric.upper()}估值')
            
            # 绘制分位数线
            ax2.axhline(y=percentiles['p95'], color='purple', linestyle='--', alpha=0.7, label='95%分位线')
            ax2.axhline(y=percentiles['q3'], color='orange', linestyle='--', alpha=0.7, label='75%分位线')
            ax2.axhline(y=percentiles['median'], color='green', linestyle='-', alpha=0.7, label='中位数')
            ax2.axhline(y=percentiles['q1'], color='orange', linestyle='--', alpha=0.7, label='25%分位线')
            ax2.axhline(y=percentiles['p5'], color='purple', linestyle='--', alpha=0.7, label='5%分位线')
            
            # 标记当前值
            if percentiles['current'] is not None:
                current_date = data['timestamp'].iloc[-1]
                ax2.scatter([current_date], [percentiles['current']], color='red', s=100, zorder=5)
                ax2.annotate(f'当前{metric.upper()}: {percentiles["current"]:.2f}', 
                           xy=(current_date, percentiles['current']),
                           xytext=(10, 10), textcoords='offset points',
                           bbox=dict(boxstyle='round,pad=0.3', facecolor='yellow', alpha=0.7),
                           arrowprops=dict(arrowstyle='->', connectionstyle='arc3,rad=0'))
            
            ax2.set_title(f'{stock_name} ({stock_code}) {metric.upper()}估值分析', fontsize=16, fontweight='bold')
            ax2.set_xlabel('日期', fontsize=12)
            ax2.set_ylabel(f'{metric.upper()}', fontsize=12)
            ax2.legend()
            ax2.grid(True, alpha=0.3)
            
            # 格式化x轴日期
            for ax in [ax1, ax2]:
                ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m'))
                ax.xaxis.set_major_locator(mdates.MonthLocator(interval=6))
                plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
            
            plt.tight_layout()
            
            # 保存图表
            if save_path is None:
                save_path = os.path.join(OUTPUT_DIR, f'{stock_code}_{metric}_analysis_{datetime.datetime.now().strftime("%Y%m%d")}.png')
            
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            plt.close()
            
            logger.info(f"美股估值分析图表保存成功: {save_path}")
            return save_path
            
        except Exception as e:
            logger.error(f"创建美股估值分析图表失败: {e}")
            raise


# 创建全局实例
us_valuation_analyzer = USValuationAnalyzer()
