import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from yfinance import download

# 获取基金数据示例
tickers = ['VFIAX', 'VTSAX', 'VBTLX', 'VGSLX', 'VINEX']  # 示例基金代码
start_date = '2020-01-01'
end_date = '2023-01-01'

data = download(tickers, start=start_date, end=end_date)['Adj Close']

# 计算每日收益率
returns = data.pct_change().dropna()

# 计算相关系数矩阵
correlation_matrix = returns.corr()

# 可视化相关性热力图
plt.figure(figsize=(10, 8))
sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0)
plt.title('基金收益率相关性矩阵')
plt.show()

# 基于相关性的简单推荐算法
def recommend_funds(corr_matrix, threshold=0.7):
    """
    推荐相关性低于阈值的基金组合
    """
    funds = corr_matrix.columns
    recommendations = []
    
    for i in range(len(funds)):
        for j in range(i+1, len(funds)):
            if abs(corr_matrix.iloc[i, j]) < threshold:
                recommendations.append((funds[i], funds[j], corr_matrix.iloc[i, j]))
    
    return sorted(recommendations, key=lambda x: abs(x[2]))

# 获取推荐组合
recommendations = recommend_funds(correlation_matrix)
print("推荐的低相关性基金组合:")
for fund1, fund2, corr in recommendations:
    print(f"{fund1} 和 {fund2}: 相关性 {corr:.2f}")