import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

# 获取股票名称（通过tushare查询或使用预设映射）
# 尝试获取股票名称
try:
    stock_name_map = {}
    for code in returns.columns:
        stock_info = pro.stock_basic(ts_code=code, fields='ts_code,name')
        if len(stock_info) > 0:
            stock_name_map[code] = stock_info.iloc[0]['name']
        else:
            stock_name_map[code] = code
except:
    # 如果API调用失败，使用代码作为名称
    stock_name_map = {code: code for code in returns.columns}

# 从精度矩阵计算偏相关系数矩阵
# 正确公式: ρ_ij|rest = -Θ_ij / √(Θ_ii × Θ_jj)
d = np.sqrt(np.diag(precision))
partial_corr = -precision / np.outer(d, d)  # 注意负号！
np.fill_diagonal(partial_corr, 1)  # 对角线为1（自相关）

# 创建图结构
G = nx.Graph()

# 使用股票名称作为节点标签
stock_codes = returns.columns.tolist()
stock_names_display = [stock_name_map.get(code, code) for code in stock_codes]
G.add_nodes_from(stock_names_display)

# 添加边（设定阈值，只显示较强的关联）
threshold = 0.15 
edges = []
weights = []
edge_colors_list = []

n_features = len(stock_codes)
for i in range(n_features):
    for j in range(i + 1, n_features):
        coef = partial_corr[i, j]
        if abs(coef) > threshold:
            name_i = stock_names_display[i]
            name_j = stock_names_display[j]
            G.add_edge(name_i, name_j, weight=coef)
            edges.append((name_i, name_j))
            weights.append(coef)
            # 根据偏相关系数正负确定颜色
            if coef > 0:
                edge_colors_list.append('#e74c3c')  # 红色：正相关
            else:
                edge_colors_list.append('#27ae60')  # 绿色：负相关

# 绘图 - 使用更紧凑的尺寸
fig, ax = plt.subplots(figsize=(10, 7))

# 布局算法 - 减小k值使节点更紧凑
pos = nx.spring_layout(G, k=0.15, seed=42, iterations=100)

# 绘制节点
nx.draw_networkx_nodes(G, pos, node_size=400, node_color='#3498db', alpha=0.85, ax=ax)

# 绘制边（使用正确的颜色列表）
if edges:  # 只有当有边时才绘制
    edge_widths = [abs(w) * 5 for w in weights]
    nx.draw_networkx_edges(G, pos, edgelist=edges, width=edge_widths, 
                           edge_color=edge_colors_list, alpha=0.7, ax=ax)

# 绘制标签（使用较小字体以适应中文名称）
nx.draw_networkx_labels(G, pos, font_size=5, font_family='sans-serif', ax=ax)

# 添加图例
pos_patch = mpatches.Patch(color='#e74c3c', label='正偏相关（协同波动）')
neg_patch = mpatches.Patch(color='#27ae60', label='负偏相关（反向波动）')
ax.legend(handles=[pos_patch, neg_patch], loc='upper left', fontsize=10, frameon=True)

# 统计边的数量
n_pos_edges = sum(1 for w in weights if w > 0)
n_neg_edges = sum(1 for w in weights if w <= 0)

ax.set_title(f'上证50成分股条件依赖网络\n(阈值={threshold}, 正边数={n_pos_edges}, 负边数={n_neg_edges})', 
             fontsize=14, fontweight='bold')
ax.axis('off')
plt.tight_layout()
plt.show()

# 打印边统计信息
print(f"\n=== 网络统计 ===")
print(f"节点数: {G.number_of_nodes()}")
print(f"边数: {G.number_of_edges()}")
print(f"  - 正偏相关边: {n_pos_edges}")
print(f"  - 负偏相关边: {n_neg_edges}")

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import networkx as nx
from sklearn.covariance import GraphicalLassoCV
from sklearn.preprocessing import StandardScaler


# 选取跨行业代表性股票（去除金融和地产，加入农业和消费）
# 样本区间扩展为3年，增加观察样本量
cross_industry_stocks = {
    # 科技成长
    '海康威视': '002415.SZ',
    '立讯精密': '002475.SZ',
    '中兴通讯': '000063.SZ',
    # 医药（防御性）
    '恒瑞医药': '600276.SH',
    '片仔癀': '600436.SH',
    '云南白药': '000538.SZ',
    # 白酒消费（防御+成长）
    '贵州茅台': '600519.SH',
    '五粮液': '000858.SZ',
    '泸州老窖': '000568.SZ',
    # 日常消费
    '伊利股份': '600887.SH',
    '海天味业': '603288.SH',
    '双汇发展': '000895.SZ',
    # 农业
    '牧原股份': '002714.SZ',
    '温氏股份': '300498.SZ',
    '新希望': '000876.SZ',
    # 新能源（高成长）
    '宁德时代': '300750.SZ',
    '隆基绿能': '601012.SH',
    # 传统制造
    '三一重工': '600031.SH',
    '美的集团': '000333.SZ',
}

# 获取数据（3年期间）
data_cross = {}
for name, code in cross_industry_stocks.items():
    try:
        df = pro.daily(ts_code=code, start_date='20210101', end_date='20231231')  # 3年数据
        df = df.set_index('trade_date').sort_index()
        data_cross[name] = df['close']
    except Exception as e:
        print(f"获取 {name} 失败: {e}")

# 合并数据
prices_cross = pd.DataFrame(data_cross)
prices_cross.index = pd.to_datetime(prices_cross.index)

# 处理缺失值
prices_cross = prices_cross.dropna(axis=1, how='any')  # 删除有缺失的股票
prices_cross = prices_cross.dropna()  # 删除有缺失的日期

print(f"成功获取 {len(prices_cross.columns)} 只股票，{len(prices_cross)} 个交易日的数据")

# 计算对数收益率
returns_cross = np.log(prices_cross / prices_cross.shift(1)).dropna()

# 标准化
scaler_cross = StandardScaler()
X_cross = scaler_cross.fit_transform(returns_cross)

# 训练 Graphical Lasso
model_cross = GraphicalLassoCV(cv=5)
model_cross.fit(X_cross)
precision_cross = model_cross.precision_

print(f"Graphical Lasso 最优 alpha: {model_cross.alpha_:.4f}")

# 计算偏相关系数（注意负号！）
d_cross = np.sqrt(np.diag(precision_cross))
partial_corr_cross = -precision_cross / np.outer(d_cross, d_cross)
np.fill_diagonal(partial_corr_cross, 1)

# 创建网络
G_cross = nx.Graph()
stock_names_cross = returns_cross.columns.tolist()
G_cross.add_nodes_from(stock_names_cross)

# 添加边（使用较低阈值以显示更多边）
threshold_cross = 0.15  # 与上证50样本使用相同阈值，便于对比
edges_cross = []
weights_cross = []
colors_cross = []

n_cross = len(stock_names_cross)
for i in range(n_cross):
    for j in range(i + 1, n_cross):
        coef = partial_corr_cross[i, j]
        if abs(coef) > threshold_cross:
            G_cross.add_edge(stock_names_cross[i], stock_names_cross[j], weight=coef)
            edges_cross.append((stock_names_cross[i], stock_names_cross[j]))
            weights_cross.append(coef)
            if coef > 0:
                colors_cross.append('#e74c3c')  # 红色：正偏相关
            else:
                colors_cross.append('#27ae60')  # 绿色：负偏相关

# 统计
n_pos_cross = sum(1 for w in weights_cross if w > 0)
n_neg_cross = sum(1 for w in weights_cross if w <= 0)

# 按行业分配颜色（与stock list对应）
industry_colors = {
    # 科技：青色
    '海康威视': '#1abc9c', '立讯精密': '#1abc9c', '中兴通讯': '#1abc9c',
    # 医药：红色
    '恒瑞医药': '#e74c3c', '片仔癀': '#e74c3c', '云南白药': '#e74c3c',
    # 白酒：金色
    '贵州茅台': '#f1c40f', '五粮液': '#f1c40f', '泸州老窖': '#f1c40f',
    # 日常消费：橙色
    '伊利股份': '#e67e22', '海天味业': '#e67e22', '双汇发展': '#e67e22',
    # 农业：棕色
    '牧原股份': '#8B4513', '温氏股份': '#8B4513', '新希望': '#8B4513',
    # 新能源：绿色
    '宁德时代': '#2ecc71', '隆基绿能': '#2ecc71',
    # 传统制造：蓝色
    '三一重工': '#3498db', '美的集团': '#3498db',
}
node_colors = [industry_colors.get(n, '#3498db') for n in G_cross.nodes()]

# 绘图 - 使用更紧凑的尺寸
fig, ax = plt.subplots(figsize=(10, 7))

# 布局 - 减小k值使节点更紧凑
pos_cross = nx.spring_layout(G_cross, k=0.25, seed=42, iterations=100)

# 绘制节点
nx.draw_networkx_nodes(G_cross, pos_cross, node_size=500, node_color=node_colors, alpha=0.9, ax=ax)

# 绘制边
if edges_cross:
    edge_widths = [abs(w) * 6 for w in weights_cross]
    nx.draw_networkx_edges(G_cross, pos_cross, edgelist=edges_cross, width=edge_widths, 
                           edge_color=colors_cross, alpha=0.7, ax=ax)

# 绘制标签
nx.draw_networkx_labels(G_cross, pos_cross, font_size=7, font_family='sans-serif', ax=ax)

# 图例
pos_patch = mpatches.Patch(color='#e74c3c', label=f'正偏相关（协同波动）: {n_pos_cross} 条')
neg_patch = mpatches.Patch(color='#27ae60', label=f'负偏相关（反向波动）: {n_neg_cross} 条')
ax.legend(handles=[pos_patch, neg_patch], loc='upper left', fontsize=11, frameon=True)

ax.set_title(f'跨行业股票条件依赖网络\n(阈值={threshold_cross}, 正边={n_pos_cross}, 负边={n_neg_cross})', 
             fontsize=14, fontweight='bold')
ax.axis('off')
plt.tight_layout()
plt.show()

# 打印统计
print(f"\n=== 跨行业网络统计 ===")
print(f"节点数: {G_cross.number_of_nodes()}")
print(f"边数: {G_cross.number_of_edges()}")
print(f"  - 正偏相关边（红色）: {n_pos_cross}")
print(f"  - 负偏相关边（绿色）: {n_neg_cross}")

# 展示部分偏相关系数
partial_corr_df_cross = pd.DataFrame(partial_corr_cross, 
                                      index=stock_names_cross, 
                                      columns=stock_names_cross)
print("\n=== 偏相关系数矩阵（部分展示）===")
print(partial_corr_df_cross.round(2))

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.covariance import GraphicalLassoCV
from sklearn.preprocessing import StandardScaler



# 获取三只股票的数据
# 招商银行: 600036.SH, 平安银行: 000001.SZ, 中国平安: 601318.SH
three_stocks = {
    '招商银行': '600036.SH',
    '平安银行': '000001.SZ',
    '中国平安': '601318.SH'
}

# 获取数据
data_3stocks = {}
for name, code in three_stocks.items():
    try:
        df = pro.daily(ts_code=code, start_date='20210101', end_date='20231231')  # 3年数据
        df = df.set_index('trade_date').sort_index()
        data_3stocks[name] = df['close']
    except Exception as e:
        print(f"Error fetching {name}: {e}")

# 合并数据
prices_3 = pd.DataFrame(data_3stocks)
prices_3.index = pd.to_datetime(prices_3.index)

# 计算对数收益率
returns_3 = np.log(prices_3 / prices_3.shift(1)).dropna()

# 标准化
scaler_3 = StandardScaler()
X_3 = scaler_3.fit_transform(returns_3)

# ====== 1. 计算边际相关系数（皮尔逊相关） ======
corr_matrix = returns_3.corr()

print("=" * 60)
print("【边际相关系数矩阵（皮尔逊相关）】")
print("=" * 60)
print(corr_matrix.round(3))

# ====== 2. 计算条件相关（偏相关系数） ======
# 使用精度矩阵计算
cov_matrix = returns_3.cov().values
precision_3 = np.linalg.inv(cov_matrix)

# 计算偏相关系数
d_3 = np.sqrt(np.diag(precision_3))
partial_corr_3 = -precision_3 / np.outer(d_3, d_3)
np.fill_diagonal(partial_corr_3, 1)

partial_corr_df = pd.DataFrame(partial_corr_3, 
                                index=returns_3.columns, 
                                columns=returns_3.columns)

print("=" * 60)
print("【偏相关系数矩阵（条件相关）】")
print("=" * 60)
print(partial_corr_df.round(3))

# ====== 3. 使用 Graphical Lasso 验证 ======
model_3 = GraphicalLassoCV(cv=5)
model_3.fit(X_3)
precision_glasso = model_3.precision_

d_glasso = np.sqrt(np.diag(precision_glasso))
partial_corr_glasso = -precision_glasso / np.outer(d_glasso, d_glasso)
np.fill_diagonal(partial_corr_glasso, 1)

partial_corr_glasso_df = pd.DataFrame(partial_corr_glasso, 
                                       index=returns_3.columns, 
                                       columns=returns_3.columns)

print("=" * 60)
print("【Graphical Lasso 偏相关系数矩阵】")
print(f"（最优正则化参数 α = {model_3.alpha_:.4f}）")
print("=" * 60)
print(partial_corr_glasso_df.round(3))

# ====== 4. 可视化对比 ======
# 使用 constrained_layout 自动调整布局，防止重叠
fig, axes = plt.subplots(1, 3, figsize=(15, 5), constrained_layout=True)

# 热力图函数
def plot_heatmap(ax, matrix, title, cmap='RdBu_r'):
    im = ax.imshow(matrix, cmap=cmap, vmin=-1, vmax=1)
    ax.set_xticks(range(len(matrix.columns)))
    ax.set_yticks(range(len(matrix.index)))
    ax.set_xticklabels(matrix.columns, rotation=45, ha='right', fontsize=10)
    ax.set_yticklabels(matrix.index, fontsize=10)
    ax.set_title(title, fontsize=12, fontweight='bold')
    
    # 添加数值标签
    for i in range(len(matrix.index)):
        for j in range(len(matrix.columns)):
            val = matrix.iloc[i, j]
            color = 'white' if abs(val) > 0.5 else 'black'
            ax.text(j, i, f'{val:.2f}', ha='center', va='center', 
                   fontsize=11, color=color, fontweight='bold')
    return im

plot_heatmap(axes[0], corr_matrix, '边际相关系数\n（皮尔逊相关）')
plot_heatmap(axes[1], partial_corr_df, '偏相关系数\n（精度矩阵）')
im = plot_heatmap(axes[2], partial_corr_glasso_df, '偏相关系数\n（Graphical Lasso）')

# 添加 colorbar
fig.colorbar(im, ax=axes, shrink=0.8, label='相关系数', location='right')

# plt.tight_layout() # constrained_layout 已启用，不需要 tight_layout
plt.show()

# ====== 5. 结论分析 ======
print("=" * 60)
print("【结论分析】")
print("=" * 60)

# 提取关键数值
marginal_zs_zp = corr_matrix.loc['招商银行', '中国平安']
partial_zs_zp = partial_corr_df.loc['招商银行', '中国平安']
glasso_zs_zp = partial_corr_glasso_df.loc['招商银行', '中国平安']

print(f"招商银行 与 中国平安：")
print(f"  - 边际相关系数:   {marginal_zs_zp:.3f}")
print(f"  - 偏相关系数:     {partial_zs_zp:.3f}")
print(f"  - Glasso偏相关:   {glasso_zs_zp:.3f}")
print()

if abs(partial_zs_zp) < 0.3 and abs(marginal_zs_zp) > 0.5:
    print("✓ 验证成功！招商银行与中国平安的高边际相关主要是间接关联，")
    print("  控制平安银行后，两者的条件相关性大幅下降。")
    verification_result = "consistent"
elif abs(partial_zs_zp) >= 0.3:
    print("⚠ 结果与预期不完全一致：招商银行与中国平安之间存在一定的直接联系，")
    print("  但偏相关系数仍低于边际相关系数，说明部分相关来自间接联系。")
    verification_result = "partial"
else:
    print("ℹ 边际相关本身不高，条件独立性分析意义有限。")
    verification_result = "low_marginal"