import pyodbc
import numpy
import pandas as pd
import matplotlib.pyplot as plt
import xlwings as xw

def brinson(portfolioID, beginDate, endDate, analyst=None, industry=None):
    try:        
    # 建立与Excel的互动
        wb = xw.Book.caller()
        sht = wb.sheets['sheet1']
        sht2 = wb.sheets['sheet2']
    
        # 设置连接字符串和数据库连接
        conn_str = r"DRIVER={Microsoft Access Driver (*.mdb, *.accdb)};DBQ=C:\Users\LIQI\Documents\portfolio.accdb"
        conn = pyodbc.connect(conn_str)
    
        portfolio_benchmark_mapping = {
        '300Active': 'CSI300',
        '300Standard': 'CSI300',
        '500Active': 'CSI500',
        '500Standard': 'CSI500'}

        # 获取给定portfolioID的对应benchID
        if portfolioID not in portfolio_benchmark_mapping:
            print(f"未找到投资组合ID的基准: {portfolioID}")
        else:
            benchID = portfolio_benchmark_mapping[portfolioID]
        
        # 查询数据库
        portfolio_query = "SELECT * FROM Mockportfolio WHERE ID = ? AND TRADEDATE >= ? AND tradedate <= ?"
        portfolio_data = pd.read_sql(portfolio_query, conn, params=[portfolioID, beginDate, endDate])
        # 修改 benchmark_query，使用参数化查询
        benchmark_query = "SELECT * FROM Benchmark WHERE BenchID = ? AND TRADEDATE >= ? AND tradedate <= ?"
        benchmark_data = pd.read_sql(benchmark_query, conn, params=[benchID, beginDate, endDate])
    
        portfolio_data['RETURN'] = portfolio_data['RETURN'].astype(float)
        benchmark_data['Weight'] = benchmark_data['Weight'].astype(float)
        benchmark_data['Return'] = benchmark_data['Return'].astype(float)
    
        portfolio_data['CONTRIBUTION'] = portfolio_data['WEIGHT'] * portfolio_data['RETURN']
        benchmark_data['CONTRIBUTION'] = benchmark_data['Weight'] * benchmark_data['Return']
    
        if analyst is not None:
            # 为指定的分析师过滤投资组合数据
            portfolio_data = portfolio_data[portfolio_data['Analyst'] == analyst]
            benchmark_data = benchmark_data[benchmark_data['Analyst'] == analyst]
        
        benchmark_data_grouped = benchmark_data.groupby(['TradeDate', 'Industry I'])[['Weight', 'Return', 'CONTRIBUTION']].sum().reset_index()
        portfolio_data_grouped = portfolio_data.groupby(['TRADEDATE', 'Industry I'])[['WEIGHT', 'RETURN', 'CONTRIBUTION']].sum().reset_index()
        benchmark_data_grouped['Return'] = benchmark_data_grouped['CONTRIBUTION'] / benchmark_data_grouped['Weight']   
        portfolio_data_grouped['RETURN'] = portfolio_data_grouped['CONTRIBUTION'] / portfolio_data_grouped['WEIGHT']
        #portfolio_data_grouped['RETURN'] = portfolio_data_grouped['RETURN'].fillna(0)
        
        benchmark_data_grouped.rename(columns={'TradeDate': 'TRADEDATE','Weight': 'WEIGHT','Return': 'RETURN'}, inplace=True)
        portfolio_df = pd.merge(portfolio_data_grouped, benchmark_data_grouped, on=['TRADEDATE', 'Industry I'], suffixes=('_P', '_B'))
        
        portfolio_df['Q1_daily'] = portfolio_df['WEIGHT_B'] * portfolio_df['RETURN_B']
        portfolio_df['Q2_daily'] = portfolio_df['WEIGHT_P'] * portfolio_df['RETURN_B']
        portfolio_df['Q3_daily'] = portfolio_df['WEIGHT_B'] * portfolio_df['RETURN_P']
        portfolio_df['Q4_daily'] = portfolio_df['WEIGHT_P'] * portfolio_df['RETURN_P']
    
        portfolio_df['excess_return_daily'] = portfolio_df['Q4_daily'] - portfolio_df['Q1_daily']
        portfolio_df['allocation_effect_daily'] = portfolio_df['Q2_daily'] - portfolio_df['Q1_daily']
        portfolio_df['selection_effect_daily'] = portfolio_df['Q3_daily'] - portfolio_df['Q1_daily']
        portfolio_df['interaction_effect_daily'] = portfolio_df['Q4_daily'] - portfolio_df['Q3_daily'] - portfolio_df['Q2_daily'] + portfolio_df['Q1_daily']
    
        # 计算各行业的累积超额收益,每日收益率累乘-1
        portfolio_df['RETURN_P_compound']=(1 + portfolio_df['RETURN_P']).groupby([portfolio_df.iloc[:, 1]]).cumprod()-1
        portfolio_df['RETURN_B_compound']=(1 + portfolio_df['RETURN_B']).groupby([portfolio_df.iloc[:, 1]]).cumprod()-1
        portfolio_df['Q1_compound']=(1 + portfolio_df['Q1_daily']).groupby([portfolio_df.iloc[:, 1]]).cumprod()-1
        portfolio_df['Q2_compound']=(1 + portfolio_df['Q2_daily']).groupby([portfolio_df.iloc[:, 1]]).cumprod()-1
        portfolio_df['Q3_compound']=(1 + portfolio_df['Q3_daily']).groupby([portfolio_df.iloc[:, 1]]).cumprod()-1
        portfolio_df['Q4_compound']=(1 + portfolio_df['Q4_daily']).groupby([portfolio_df.iloc[:, 1]]).cumprod()-1
    
        #四种效应的累计收益
        portfolio_df['excess_return_compound'] = portfolio_df['Q4_compound'] - portfolio_df['Q1_compound']
        portfolio_df['allocation_effect_compound'] = portfolio_df['Q2_compound'] - portfolio_df['Q1_compound']
        portfolio_df['selection_effect_compound'] = portfolio_df['Q3_compound'] - portfolio_df['Q1_compound']
        portfolio_df['interaction_effect_compound'] = portfolio_df['Q4_compound'] - portfolio_df['Q3_compound'] - portfolio_df['Q2_compound'] + portfolio_df['Q1_compound']
    
        #对 portfolio_df 数据框进行了索引重置，以确保索引重新从零开始，以便于后续处理和分析。
        portfolio_df=portfolio_df.reset_index()
    
        # 计算整个投资组合的累积超额收益
        #四种效应下的累计收益按交易日期分组求和然后索引重置
        #excess_return_compound = portfolio_df.groupby('TRADEDATE')['excess_return_compound'].sum().reset_index()
        #allocation_effect_compound = portfolio_df.groupby('TRADEDATE')['allocation_effect_compound'].sum().reset_index()
        #selection_effect_compound = portfolio_df.groupby('TRADEDATE')['selection_effect_compound'].sum().reset_index()
        #interaction_effect_compound = portfolio_df.groupby('TRADEDATE')['interaction_effect_compound'].sum().reset_index()
    
        # 为特定行业绘制图表
        if industry is not None:
            industry_data = portfolio_df[portfolio_df['Industry I'] == industry]
            x_data = industry_data['TRADEDATE'].to_numpy() 
            y1_data = industry_data['excess_return_compound'].to_numpy() 
            y2_data = industry_data['allocation_effect_compound'].to_numpy() 
            y3_data = industry_data['selection_effect_compound'].to_numpy() 
            y4_data = industry_data['interaction_effect_compound'].to_numpy()
            
            # 创建一个图形和坐标轴
            plt.figure(figsize=(12, 9))
            fig, ax = plt.subplots()
            
            # 绘制折线图
            ax.plot(x_data, y1_data, color='blue', marker='o', label='excess_return_compound')
            ax.plot(x_data, y2_data, color='red', marker='o', label='allocation_effect_compound')
            ax.plot(x_data, y3_data, color='black', marker='o', label='selection_effect_compoundt')
            ax.plot(x_data, y4_data, color='purple', marker='o', label='interaction_effect_compound')
            
            plt.xticks(rotation=45)
            plt.title(f"Excess Return in {industry}")
            plt.xlabel('TRADEDATE')
            plt.ylabel('excess_return_compound')
            plt.legend()
            # Show the plot
            left = sht2.range('G3').left
            top = sht2.range('G3').top
            sht2.pictures.add(fig, name='Industry', update=True, left=left, top=top)
            
            #生成单行业时间序列
            industry_data = industry_data[['TRADEDATE','excess_return_compound','allocation_effect_compound','selection_effect_compound','interaction_effect_compound']]
            industry_data.rename(columns={'TRADEDATE':'交易日'}, inplace=True)
            industry_data.rename(columns={'excess_return_compound': '超额收益','allocation_effect_compound': '配置效应'}, inplace=True)
            industry_data.rename(columns={'selection_effect_compound': '选股效应','interaction_effect_compound': '交叉效应'}, inplace=True)
            # 将后四列的数据格式化为 #.###### 格式
            decimal_format = lambda x: f'{x:.6f}'
            columns_to_format = ['超额收益', '配置效应', '选股效应', '交叉效应']
            #写入excel
            sht2.range('A2').options(index=False, header=True, numbers=str).value = industry_data
            
            
        
    
        # 为特定组合绘制图表
        excess_return_compound = portfolio_df.groupby('TRADEDATE')['excess_return_compound'].sum().reset_index()
        allocation_effect_compound = portfolio_df.groupby('TRADEDATE')['allocation_effect_compound'].sum().reset_index()
        selection_effect_compound = portfolio_df.groupby('TRADEDATE')['selection_effect_compound'].sum().reset_index()
        interaction_effect_compound = portfolio_df.groupby('TRADEDATE')['interaction_effect_compound'].sum().reset_index()
    
        x21_data = excess_return_compound['TRADEDATE'].to_numpy() 
        y21_data = excess_return_compound['excess_return_compound'].to_numpy() 
        y22_data = allocation_effect_compound['allocation_effect_compound'].to_numpy()
        y23_data = selection_effect_compound['selection_effect_compound'].to_numpy()
        y24_data = interaction_effect_compound['interaction_effect_compound'].to_numpy() 
    
        # 创建一个图形和坐标轴
        plt.figure(figsize=(12, 9))
        fig, ax = plt.subplots()

        ax.plot(x21_data, y21_data, color='blue', marker='o', label='excess_return_compound')
        ax.plot(x21_data, y22_data, color='red', marker='o', label='allocation_effect_compound')
        ax.plot(x21_data, y23_data, color='black', marker='o', label='selection_effect_compoundt')
        ax.plot(x21_data, y24_data, color='purple', marker='o', label='interaction_effect_compound')

        plt.title('Entire Portfolio Compound Excess Return')
        plt.xlabel('Trade Date')
        plt.ylabel('Excess Return')
    
        # 添加图例
        ax.legend()
        plt.xticks(rotation=45)
    
           
        # 插入第二个图表到 sheet1
        left = sht.range('J3').left
        top = sht.range('J3').top
        sht.pictures.add(fig, name='PortfolioChart2', update=True, left=left, top=top)
        
        
        #生成投资组合快照
        portfolio_snap = portfolio_df[portfolio_df['TRADEDATE'] == endDate]
        portfolio_snap = portfolio_snap[['Industry I','WEIGHT_P','WEIGHT_B','RETURN_P_compound','RETURN_B_compound','excess_return_compound','allocation_effect_compound','selection_effect_compound','interaction_effect_compound']]
        portfolio_snap = portfolio_snap.sort_values('selection_effect_compound', ascending=False)
        portfolio_snap.rename(columns={'Industry I':'行业','WEIGHT_P':'组合权重','WEIGHT_B':'指数权重','RETURN_P_compound': '组合回报','RETURN_B_compound': '指数回报'}, inplace=True)
        portfolio_snap.rename(columns={'excess_return_compound': '超额收益','allocation_effect_compound': '配置效应'}, inplace=True)
        portfolio_snap.rename(columns={'selection_effect_compound': '选股效应','interaction_effect_compound': '交叉效应'}, inplace=True)
        # 将后四列的数据格式化为 #.###### 格式
        decimal_format = lambda x: f'{x:.6f}'
        columns_to_format = ['组合权重','指数权重','组合回报','超额收益', '配置效应', '选股效应', '交叉效应']
        #写入excel
        sht.range('A2').options(index=False, header=True, numbers=str).value = portfolio_snap
        
    except Exception as e:
        print(f"Error: {e}")

    finally:
        conn.close()

def generate():              
    # 建立与Excel的互动
    wb = xw.Book.caller()
    sht = wb.sheets['sheet1']
    portfolioID = sht.range('b1').value#300Active，300Standard……
    beginDate = sht.range('j1').value
    endDate = sht.range('l1').value
    analyst = sht.range('h1').value
    industry = sht.range('e1').value#一级行业
    
    brinson(portfolioID, beginDate, endDate, analyst, industry)  
                
if __name__ == "__main__":
    generate()     我写了一个brinson归因分析的脚本，现在遇到一个问题。如果从参数beginDate到endDate，某个analyst在某个industry的weight_p上一直>0，那么 在#生成投资组合快照这一段会出现有关该industry的数据，但一旦endDate那天，该industry的weight_p=0，数据就不显示了。问题出在哪里呢？因为即便在endDate，不再持有该industry的股票，在期间还是会有累积超额收益数据啊