# coding:utf-8
#计算股价平均距离因子-行业个股列表来计算

import pandas as pd
import numpy as np
from sqlalchemy import create_engine, text
from datetime import datetime, timedelta
from sklearn.preprocessing import StandardScaler
from sklearn.metrics.pairwise import euclidean_distances
import warnings
warnings.filterwarnings('ignore')

class AverageDistanceFactor:
    """平均距离因子计算器"""
    
    def __init__(self, db_url):
        self.engine = create_engine(
            db_url,
            pool_size=5,
            max_overflow=10,
            pool_recycle=3600
        )
    
    def get_industry_stocks(self, industry_name=None, concept_name=None):
        """获取指定行业或概念的股票列表"""
        if industry_name:
            query = text("""
                SELECT DISTINCT gp_code as symbol
                FROM gp_hybk 
                WHERE bk_name = :name
            """)
            params = {"name": industry_name}
        elif concept_name:
            query = text("""
                SELECT DISTINCT gp_code as symbol
                FROM gp_gnbk 
                WHERE bk_name = :name
            """)
            params = {"name": concept_name}
        else:
            raise ValueError("必须提供 industry_name 或 concept_name 之一")
        
        try:
            with self.engine.connect() as conn:
                result = conn.execute(query, params).fetchall()
            
            if result:
                symbols = [row[0] for row in result]
                print(f"获取到 {len(symbols)} 只股票")
                return symbols
            else:
                print(f"未找到{'行业' if industry_name else '概念'} {industry_name or concept_name} 的股票")
                return []
        except Exception as e:
            print(f"获取股票列表失败: {e}")
            return []
    
    def get_stock_data(self, symbols, days=20, end_date=None):
        """
        获取股票的历史数据
        
        Args:
            symbols: 股票代码列表
            days: 需要获取的天数
            end_date: 结束日期（datetime对象或字符串），如果为None则使用今天
        """
        if not symbols:
            return pd.DataFrame()
        
        # 计算结束日期（及时数据使用目标日期当天）
        if end_date is None:
            end_date = datetime.now()
        elif isinstance(end_date, str):
            end_date = datetime.strptime(end_date, '%Y-%m-%d')
        
        # 及时数据（日线数据）应该包含目标日期当天的数据
        # 例如：如果目标日期是2025-11-01，则使用2025-11-01及之前的数据
        query_end_date = end_date
        
        # 计算开始日期（多取一些数据以防节假日）
        start_date = query_end_date - timedelta(days=days * 2)
        
        # 构建SQL查询（使用 <= 包含目标日期当天）
        symbols_str = "', '".join(symbols)
        query = f"""
        SELECT symbol, timestamp, volume, open, high, low, close, 
               chg, percent, turnoverrate, amount, pb, pe, ps
        FROM gp_day_data 
        WHERE symbol IN ('{symbols_str}')
        AND timestamp >= '{start_date.strftime('%Y-%m-%d')}'
        AND timestamp <= '{end_date.strftime('%Y-%m-%d')}'
        ORDER BY symbol, timestamp DESC
        """
        
        try:
            df = pd.read_sql(query, self.engine)
            return df
        except Exception as e:
            return pd.DataFrame()
    
    def calculate_technical_indicators(self, df, days=20):
        """
        计算技术指标
        
        Args:
            df: 股票历史数据DataFrame（应已按目标日期筛选）
            days: 计算指标使用的天数（取目标日期前N天）
        """
        result_data = []
        
        for symbol in df['symbol'].unique():
            stock_data = df[df['symbol'] == symbol].copy()
            # 按时间升序排序，确保时间顺序正确
            stock_data = stock_data.sort_values('timestamp')
            
            # 只取目标日期前N个交易日的数据（包括目标日期当天）
            # tail取最后N行，即最接近目标日期的N个交易日（包括目标日期当天）
            stock_data = stock_data.tail(days)
            
            # 检查数据是否足够（需要至少有days个交易日的数据）
            if len(stock_data) < days:
                continue  # 数据不足，跳过
            
            # 转换数据类型
            for col in ['open', 'high', 'low', 'close', 'chg']:
                stock_data[col] = pd.to_numeric(stock_data[col], errors='coerce')
            
            # 计算各种技术指标
            indicators = self._compute_indicators(stock_data)
            indicators['symbol'] = symbol
            result_data.append(indicators)
        
        return pd.DataFrame(result_data)
    
    def _compute_indicators(self, data):
        """计算具体的技术指标"""
        indicators = {}
        
        # 1. 收益率指标
        data['returns'] = data['close'].pct_change()
        indicators['return_5d'] = data['returns'].tail(5).sum()  # 5日累计收益率
        indicators['return_10d'] = data['returns'].tail(10).sum()  # 10日累计收益率
        indicators['return_20d'] = data['returns'].tail(20).sum()  # 20日累计收益率
        
        # 2. 波动率指标
        indicators['volatility_5d'] = data['returns'].tail(5).std()  # 5日波动率
        indicators['volatility_10d'] = data['returns'].tail(10).std()  # 10日波动率
        indicators['volatility_20d'] = data['returns'].tail(20).std()  # 20日波动率
        
        # 3. 价格相对位置
        indicators['price_position_5d'] = (data['close'].iloc[-1] - data['low'].tail(5).min()) / (data['high'].tail(5).max() - data['low'].tail(5).min())
        indicators['price_position_10d'] = (data['close'].iloc[-1] - data['low'].tail(10).min()) / (data['high'].tail(10).max() - data['low'].tail(10).min())
        indicators['price_position_20d'] = (data['close'].iloc[-1] - data['low'].tail(20).min()) / (data['high'].tail(20).max() - data['low'].tail(20).min())
        
        # 4. 移动平均偏离度
        ma_5 = data['close'].tail(5).mean()
        ma_10 = data['close'].tail(10).mean()
        ma_20 = data['close'].tail(20).mean()
        current_price = data['close'].iloc[-1]
        
        indicators['ma_deviation_5d'] = (current_price - ma_5) / ma_5
        indicators['ma_deviation_10d'] = (current_price - ma_10) / ma_10
        indicators['ma_deviation_20d'] = (current_price - ma_20) / ma_20
        
        # 5. 成交量相关指标
        indicators['volume_ratio_5d'] = data['volume'].tail(5).mean() / data['volume'].mean()
        indicators['volume_ratio_10d'] = data['volume'].tail(10).mean() / data['volume'].mean()
        indicators['turnover_avg_5d'] = data['turnoverrate'].tail(5).mean()
        indicators['turnover_avg_10d'] = data['turnoverrate'].tail(10).mean()
        
        # 6. 价格振幅指标
        data['amplitude'] = (data['high'] - data['low']) / data['close']
        indicators['amplitude_avg_5d'] = data['amplitude'].tail(5).mean()
        indicators['amplitude_avg_10d'] = data['amplitude'].tail(10).mean()
        
        # 7. 趋势强度（连续涨跌）
        indicators['consecutive_up'] = self._count_consecutive(data['percent'] > 0)
        indicators['consecutive_down'] = self._count_consecutive(data['percent'] < 0)
        
        # 8. 估值动量（如果有PE、PB、PS数据）
        if 'pe' in data.columns and not data['pe'].isna().all():
            pe_change = data['pe'].pct_change().tail(5).mean()
            indicators['pe_momentum'] = pe_change if not np.isnan(pe_change) else 0
        else:
            indicators['pe_momentum'] = 0
            
        if 'pb' in data.columns and not data['pb'].isna().all():
            pb_change = data['pb'].pct_change().tail(5).mean()
            indicators['pb_momentum'] = pb_change if not np.isnan(pb_change) else 0
        else:
            indicators['pb_momentum'] = 0
        
        # 处理NaN值
        for key, value in indicators.items():
            if np.isnan(value) or np.isinf(value):
                indicators[key] = 0
        
        return indicators
    
    def _count_consecutive(self, condition_series):
        """计算连续满足条件的天数"""
        if len(condition_series) == 0:
            return 0
        
        count = 0
        for value in reversed(condition_series.tolist()):
            if value:
                count += 1
            else:
                break
        return count
    
    def calculate_distance_factor(self, indicators_df):
        """计算平均距离因子"""
        if len(indicators_df) < 2:
            print("股票数量不足，无法计算距离因子")
            return pd.DataFrame()
        
        # 准备特征矩阵
        feature_columns = [col for col in indicators_df.columns if col != 'symbol']
        X = indicators_df[feature_columns].values
        
        # 标准化
        scaler = StandardScaler()
        X_scaled = scaler.fit_transform(X)
        
        # 计算距离矩阵
        distances = euclidean_distances(X_scaled)
        
        # 计算每只股票的平均距离
        n_stocks = len(indicators_df)
        avg_distances = []
        
        for i in range(n_stocks):
            # 排除自己与自己的距离（对角线元素）
            other_distances = np.concatenate([distances[i, :i], distances[i, i+1:]])
            avg_distance = np.mean(other_distances)
            avg_distances.append(avg_distance)
        
        # 创建结果DataFrame
        result_df = pd.DataFrame({
            'symbol': indicators_df['symbol'],
            'avg_distance_factor': avg_distances
        })
        
        # 按距离因子降序排列
        result_df = result_df.sort_values('avg_distance_factor', ascending=False)
        result_df['rank'] = range(1, len(result_df) + 1)
        
        return result_df
    
    def analyze_industry(self, industry_name=None, concept_name=None, days=20):
        """分析指定行业或概念的平均距离因子"""
        print(f"开始分析{'行业' if industry_name else '概念'}: {industry_name or concept_name}")
        
        # 1. 获取股票列表
        symbols = self.get_industry_stocks(industry_name, concept_name)
        if not symbols:
            return pd.DataFrame()
        
        # 2. 获取历史数据
        stock_data = self.get_stock_data(symbols, days)
        if stock_data.empty:
            return pd.DataFrame()
        
        # 3. 计算技术指标
        print("计算技术指标...")
        indicators_df = self.calculate_technical_indicators(stock_data, days)
        if indicators_df.empty:
            return pd.DataFrame()
        
        print(f"成功计算了 {len(indicators_df)} 只股票的技术指标")
        
        # 4. 计算平均距离因子
        print("计算平均距离因子...")
        distance_df = self.calculate_distance_factor(indicators_df)
        
        if not distance_df.empty:
            print(f"计算完成，共 {len(distance_df)} 只股票")
            print(distance_df.to_string(index=False))
        
        return distance_df
    
    def get_available_industries(self):
        """获取可用的行业列表"""
        query = "SELECT DISTINCT bk_name FROM gp_hybk ORDER BY bk_name"
        try:
            df = pd.read_sql(query, self.engine)
            return df['bk_name'].tolist()
        except Exception as e:
            print(f"获取行业列表失败: {e}")
            return []
    
    def get_available_concepts(self):
        """获取可用的概念列表"""
        query = "SELECT DISTINCT bk_name FROM gp_gnbk ORDER BY bk_name"
        try:
            df = pd.read_sql(query, self.engine)
            return df['bk_name'].tolist()
        except Exception as e:
            print(f"获取概念列表失败: {e}")
            return []
    
    def __del__(self):
        if hasattr(self, 'engine'):
            self.engine.dispose()

def main():
    """主函数示例"""
    db_url = 'mysql+pymysql://root:Chlry#$.8@192.168.18.199:3306/db_gp_cj'
    
    # 创建分析器
    analyzer = AverageDistanceFactor(db_url)
    
    # 示例1: 分析特定行业
    result = analyzer.analyze_industry(industry_name="军工电子")
    
    # 示例2: 分析特定概念
    # result = analyzer.analyze_industry(concept_name="固态电池")
    
    # 示例3: 查看可用的行业列表
    # industries = analyzer.get_available_industries()
    # print("可用行业列表（前10个）:")
    # for industry in industries[:10]:
    #     print(f"  - {industry}")
    
    # 示例4: 查看可用的概念列表
    # concepts = analyzer.get_available_concepts()
    # print("\n可用概念列表（前10个）:")
    # for concept in concepts[:10]:
    #     print(f"  - {concept}")

if __name__ == "__main__":
    main()