#!/usr/bin/env python3
"""
STRING PPI网络分析 + 可视化
1. 查询STRING API获取PPI网络
2. 绘制富集分析柱状图
3. 绘制Hub基因图
4. 绘制成分-靶点网络图
"""

import os
import json
import requests
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib import font_manager
import warnings
warnings.filterwarnings('ignore')

# 尝试设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['figure.dpi'] = 150

OUTPUT_DIR = r"E:\workbuddy\Claw\发物网络药理学"
TABLES_DIR = os.path.join(OUTPUT_DIR, "results", "tables")
FIGURES_DIR = os.path.join(OUTPUT_DIR, "results", "figures")
os.makedirs(FIGURES_DIR, exist_ok=True)

def query_string_ppi(gene_list, species=9606, score_threshold=400):
    """查询STRING API获取PPI网络"""
    print("查询STRING PPI网络...")
    
    string_api_url = "https://version-11-5.string-db.org/api"
    output_format = "json"
    method = "network"
    
    params = {
        "identifiers": "%0d".join(gene_list),
        "species": species,
        "required_score": score_threshold,
        "caller_identity": "fawu_research"
    }
    
    url = f"{string_api_url}/{output_format}/{method}"
    
    try:
        resp = requests.post(url, data=params, timeout=60)
        if resp.status_code == 200:
            data = resp.json()
            print(f"  STRING返回 {len(data)} 条互作")
            return data
        else:
            print(f"  STRING API返回 {resp.status_code}")
            # 尝试v12
            url2 = f"https://string-db.org/api/{output_format}/{method}"
            resp2 = requests.post(url2, data=params, timeout=60)
            if resp2.status_code == 200:
                data = resp2.json()
                print(f"  STRING v12返回 {len(data)} 条互作")
                return data
    except Exception as e:
        print(f"  STRING查询失败: {e}")
    
    return []

def analyze_ppi(ppi_data, gene_list):
    """分析PPI网络，计算节点度"""
    print("\nPPI网络分析...")
    
    if not ppi_data:
        print("  无PPI数据")
        return pd.DataFrame(), pd.DataFrame()
    
    # 构建边表
    edges = []
    for item in ppi_data:
        edges.append({
            'source': item.get('preferredName_A', item.get('stringId_A', '')),
            'target': item.get('preferredName_B', item.get('stringId_B', '')),
            'score': item.get('score', 0),
            'interaction': 'PPI'
        })
    
    edges_df = pd.DataFrame(edges)
    edges_df.to_csv(os.path.join(TABLES_DIR, "PPI_edges.csv"), index=False)
    
    # 计算节点度
    degree = {}
    for _, row in edges_df.iterrows():
        degree[row['source']] = degree.get(row['source'], 0) + 1
        degree[row['target']] = degree.get(row['target'], 0) + 1
    
    # 确保所有基因都在（即使度为0）
    for gene in gene_list:
        if gene not in degree:
            degree[gene] = 0
    
    degree_df = pd.DataFrame([
        {'gene': k, 'degree': v} for k, v in degree.items()
    ]).sort_values('degree', ascending=False)
    
    degree_df.to_csv(os.path.join(TABLES_DIR, "PPI_node_degree.csv"), index=False)
    
    print(f"  PPI节点: {len(degree_df)}, 边: {len(edges_df)}")
    print(f"  Top 15 Hub基因 (按度排序):")
    print(degree_df.head(15).to_string(index=False))
    
    return edges_df, degree_df

def plot_enrichment_bar():
    """绘制KEGG富集柱状图"""
    print("\n绘制KEGG富集柱状图...")
    
    kegg_path = os.path.join(TABLES_DIR, "intersection_KEGG.csv")
    if not os.path.exists(kegg_path):
        print("  KEGG结果文件不存在")
        return
    
    df = pd.read_csv(kegg_path)
    df = df.sort_values('Adjusted P-value')
    
    # 取Top 20
    top_n = min(20, len(df))
    top = df.head(top_n).copy()
    
    # 处理0值
    top['neg_log_p'] = -np.log10(top['Adjusted P-value'].clip(lower=1e-50))
    
    # 截断过长的通路名
    top['Term_short'] = top['Term'].apply(lambda x: x[:50] + '...' if len(x) > 50 else x)
    
    fig, ax = plt.subplots(figsize=(10, 8))
    
    colors = plt.cm.RdYlBu_r(np.linspace(0.2, 0.9, top_n))
    y_pos = np.arange(top_n)
    
    bars = ax.barh(y_pos, top['neg_log_p'], color=colors, edgecolor='grey', linewidth=0.5)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(top['Term_short'], fontsize=8)
    ax.invert_yaxis()
    ax.set_xlabel('-log10(Adjusted P-value)', fontsize=11)
    ax.set_title('KEGG Pathway Enrichment Analysis', fontsize=13, fontweight='bold')
    
    # 添加p值标注
    for i, (bar, pval) in enumerate(zip(bars, top['Adjusted P-value'])):
        ax.text(bar.get_width() + 0.5, bar.get_y() + bar.get_height()/2,
                f'{pval:.2e}', va='center', fontsize=7, color='grey')
    
    ax.axvline(x=-np.log10(0.05), color='red', linestyle='--', linewidth=0.8, alpha=0.7)
    ax.text(-np.log10(0.05)+0.5, top_n-1, 'P=0.05', fontsize=8, color='red', va='top')
    
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, "KEGG_enrichment_barplot.png"), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  已保存: KEGG_enrichment_barplot.png")

def plot_go_bp():
    """绘制GO BP富集柱状图"""
    print("绘制GO BP富集柱状图...")
    
    path = os.path.join(TABLES_DIR, "intersection_GO_BP.csv")
    if not os.path.exists(path):
        print("  GO BP结果文件不存在")
        return
    
    df = pd.read_csv(path)
    df = df.sort_values('Adjusted P-value')
    
    top_n = min(15, len(df))
    top = df.head(top_n).copy()
    top['neg_log_p'] = -np.log10(top['Adjusted P-value'].clip(lower=1e-50))
    top['Term_short'] = top['Term'].apply(lambda x: x[:45] + '...' if len(x) > 45 else x)
    
    fig, ax = plt.subplots(figsize=(10, 7))
    colors = plt.cm.YlGnBu_r(np.linspace(0.3, 0.9, top_n))
    y_pos = np.arange(top_n)
    
    ax.barh(y_pos, top['neg_log_p'], color=colors, edgecolor='grey', linewidth=0.5)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(top['Term_short'], fontsize=8)
    ax.invert_yaxis()
    ax.set_xlabel('-log10(Adjusted P-value)', fontsize=11)
    ax.set_title('GO Biological Process Enrichment', fontsize=13, fontweight='bold')
    
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, "GO_BP_enrichment_barplot.png"), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  已保存: GO_BP_enrichment_barplot.png")

def plot_hub_genes(degree_df):
    """绘制Hub基因度柱状图"""
    print("绘制Hub基因图...")
    
    if degree_df is None or len(degree_df) == 0:
        print("  无PPI度数据")
        return
    
    top_n = min(20, len(degree_df))
    top = degree_df.head(top_n).copy()
    
    fig, ax = plt.subplots(figsize=(10, 7))
    
    colors = plt.cm.viridis(np.linspace(0.2, 0.9, top_n))
    y_pos = np.arange(top_n)
    
    bars = ax.barh(y_pos, top['degree'], color=colors, edgecolor='grey', linewidth=0.5)
    ax.set_yticks(y_pos)
    ax.set_yticklabels(top['gene'], fontsize=9, fontweight='bold')
    ax.invert_yaxis()
    ax.set_xlabel('Degree (Number of PPI Interactions)', fontsize=11)
    ax.set_title('Hub Genes in PPI Network', fontsize=13, fontweight='bold')
    
    # 标注度数
    for i, (bar, deg) in enumerate(zip(bars, top['degree'])):
        ax.text(bar.get_width() + 0.3, bar.get_y() + bar.get_height()/2,
                str(int(deg)), va='center', fontsize=9, fontweight='bold')
    
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, "hub_genes_PPI_degree.png"), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  已保存: hub_genes_PPI_degree.png")

def plot_compound_target_network():
    """绘制成分-靶点网络"""
    print("绘制成分-靶点网络图...")
    
    edges_path = os.path.join(TABLES_DIR, "network_compound_target_edges.csv")
    if not os.path.exists(edges_path):
        print("  网络边文件不存在")
        return
    
    edges_df = pd.read_csv(edges_path)
    
    # 获取节点
    compounds = list(edges_df['source'].unique())
    targets = list(edges_df['target'].unique())
    
    # 颜色映射
    cat_colors = {
        '辛辣类': '#FF6347',   # tomato
        '海鲜类': '#1E90FF',   # dodger blue
        '肉类': '#DAA520',     # goldenrod
    }
    
    fig, ax = plt.subplots(figsize=(14, 10))
    ax.set_xlim(-1.5, 1.5)
    ax.set_ylim(-1.5, 1.5)
    ax.set_aspect('equal')
    ax.axis('off')
    
    # 成分放在左侧，靶点放在右侧
    n_compounds = len(compounds)
    n_targets = len(targets)
    
    compound_pos = {}
    for i, comp in enumerate(compounds):
        y = 1.4 - (i / max(n_compounds - 1, 1)) * 2.8
        compound_pos[comp] = (-0.8, y)
    
    target_pos = {}
    for i, gene in enumerate(targets):
        y = 1.4 - (i / max(n_targets - 1, 1)) * 2.8
        target_pos[gene] = (0.8, y)
    
    # 绘制边
    for _, row in edges_df.iterrows():
        src = row['source']
        tgt = row['target']
        cat = row['category']
        
        x1, y1 = compound_pos.get(src, (0, 0))
        x2, y2 = target_pos.get(tgt, (0, 0))
        
        color = cat_colors.get(cat, 'grey')
        ax.plot([x1, x2], [y1, y2], color=color, alpha=0.3, linewidth=0.8, zorder=1)
    
    # 绘制成分节点
    for comp, (x, y) in compound_pos.items():
        cat = edges_df[edges_df['source'] == comp]['category'].iloc[0]
        color = cat_colors.get(cat, 'grey')
        circle = plt.Circle((x, y), 0.08, color=color, ec='black', linewidth=1.2, zorder=3)
        ax.add_patch(circle)
        ax.text(x, y, comp[:8], ha='center', va='center', fontsize=6, fontweight='bold', zorder=4, color='white')
        ax.text(x - 0.15, y, comp, ha='right', va='center', fontsize=7, fontweight='bold')
    
    # 绘制靶点节点
    for gene, (x, y) in target_pos.items():
        # 根据连接的成分数调整大小
        n_compounds_connected = len(edges_df[edges_df['target'] == gene]['source'].unique())
        size = 0.04 + (n_compounds_connected / 8) * 0.06
        
        circle = plt.Circle((x, y), size, color='lightblue', ec='navy', linewidth=1, zorder=3)
        ax.add_patch(circle)
        ax.text(x, y, gene, ha='center', va='center', fontsize=5.5, fontweight='bold', zorder=4)
    
    # 图例
    handles = [mpatches.Patch(color=color, label=cat) for cat, color in cat_colors.items()]
    handles.append(mpatches.Patch(color='lightblue', ec='navy', label='Target Gene'))
    ax.legend(handles=handles, loc='lower left', fontsize=9, framealpha=0.9)
    
    ax.set_title('Compound-Target Network\n"发物" Active Compounds vs. Intersection Targets',
                 fontsize=13, fontweight='bold', pad=20)
    
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, "compound_target_network.png"), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  已保存: compound_target_network.png")

def plot_target_compound_heatmap():
    """绘制靶点-成分热图"""
    print("绘制靶点-成分热图...")
    
    edges_path = os.path.join(TABLES_DIR, "network_compound_target_edges.csv")
    if not os.path.exists(edges_path):
        return
    
    edges_df = pd.read_csv(edges_path)
    
    # 构建矩阵
    matrix = pd.crosstab(edges_df['target'], edges_df['source'])
    
    # 按连接数排序
    matrix['total'] = matrix.sum(axis=1)
    matrix = matrix.sort_values('total', ascending=False)
    matrix = matrix.drop('total', axis=1)
    
    fig, ax = plt.subplots(figsize=(10, 12))
    
    im = ax.imshow(matrix.values, cmap='YlOrRd', aspect='auto', vmin=0, vmax=1)
    
    ax.set_xticks(np.arange(len(matrix.columns)))
    ax.set_yticks(np.arange(len(matrix.index)))
    ax.set_xticklabels(matrix.columns, rotation=45, ha='right', fontsize=8)
    ax.set_yticklabels(matrix.index, fontsize=8)
    
    # 标注
    for i in range(len(matrix.index)):
        for j in range(len(matrix.columns)):
            if matrix.values[i, j] > 0:
                ax.text(j, i, '●', ha='center', va='center', fontsize=10, color='darkred')
    
    ax.set_xlabel('Compounds', fontsize=11)
    ax.set_ylabel('Target Genes', fontsize=11)
    ax.set_title('Compound-Target Interaction Matrix', fontsize=13, fontweight='bold')
    
    plt.colorbar(im, ax=ax, label='Interaction', shrink=0.5)
    plt.tight_layout()
    plt.savefig(os.path.join(FIGURES_DIR, "compound_target_heatmap.png"), dpi=200, bbox_inches='tight')
    plt.close()
    print(f"  已保存: compound_target_heatmap.png")

def main():
    print("=" * 60)
    print("STRING PPI分析 + 可视化")
    print("=" * 60)
    
    # 1. 加载交集基因
    inter_path = os.path.join(TABLES_DIR, "intersection_targets.csv")
    if not os.path.exists(inter_path):
        print("交集文件不存在，请先运行04_intersection_enrichment.py")
        return
    
    inter_df = pd.read_csv(inter_path)
    gene_list = inter_df['gene'].tolist()
    print(f"交集基因: {len(gene_list)} 个")
    
    # 2. STRING PPI查询
    ppi_data = query_string_ppi(gene_list)
    edges_df, degree_df = analyze_ppi(ppi_data, gene_list)
    
    # 3. 可视化
    print(f"\n{'='*60}")
    print("生成图表")
    print(f"{'='*60}")
    
    plot_enrichment_bar()
    plot_go_bp()
    plot_hub_genes(degree_df)
    plot_compound_target_network()
    plot_target_compound_heatmap()
    
    print(f"\n{'='*60}")
    print("完成! 图表保存在:", FIGURES_DIR)
    print(f"{'='*60}")
    
    # 列出所有生成的文件
    print("\n生成的文件:")
    for f in sorted(os.listdir(FIGURES_DIR)):
        fpath = os.path.join(FIGURES_DIR, f)
        size = os.path.getsize(fpath) / 1024
        print(f"  {f} ({size:.1f} KB)")

if __name__ == "__main__":
    main()