# vizpaint/sankey_diagram.py
"""
桑基图（Sankey Diagram）模块
用于展示流量、能量或资源流动的图表
"""

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.sankey import Sankey


def sankey_diagram(flows, labels=None, orientations=0,
                   pathlengths=0.25, trunklength=1.0,
                   head_angle=100, margin=0.4,
                   title="Sankey Diagram", color_palette='Set3'):
    """
    绘制桑基图（Sankey Diagram）

    参数:
        flows: 流量列表，正值为向右，负值为向左
              例如: [0.5, 0.3, -0.2, 0.1, -0.7]
        labels: 标签列表，长度应为 flows.length + 1
        orientations: 流向方向列表（0: 右，1: 上，-1: 下）
        pathlengths: 路径长度
        trunklength: 主干长度
        head_angle: 箭头角度
        margin: 边距
        title: 图表标题
        color_palette: 颜色调色板名称

    返回:
        fig, ax, sankey: 图形、坐标轴和桑基图对象
    """
    # 创建图形
    fig, ax = plt.subplots(figsize=(14, 10))

    # 设置标签（如果未提供）
    if labels is None:
        labels = [f'Node {i}' for i in range(len(flows) + 1)]

    # 设置方向（如果未提供）
    if isinstance(orientations, int):
        orientations = [orientations] * (len(flows) - 1)

    # 创建桑基图对象
    sankey = Sankey(ax=ax, scale=1.0, offset=0.1,
                    head_angle=head_angle,
                    trunklength=trunklength,
                    margin=margin)

    # 添加主要流动
    sankey.add(flows=flows,
               labels=labels,
               orientations=orientations,
               pathlengths=pathlengths,
               facecolor=plt.cm.get_cmap(color_palette)(0.3))

    # 完成图表
    diagrams = sankey.finish()

    # 自定义每个箭头的颜色
    cmap = plt.cm.get_cmap(color_palette)
    for i, diagram in enumerate(diagrams):
        for j, patch in enumerate(diagram.patches):
            # 为每个路径分配不同的颜色
            color = cmap(j / max(len(diagram.patches) - 1, 1))
            patch.set_facecolor(color)
            patch.set_alpha(0.8)
            patch.set_edgecolor('white')
            patch.set_linewidth(1)

    # 添加标签到每个节点
    for diagram in diagrams:
        for text in diagram.texts:
            text.set_fontsize(10)
            text.set_fontweight('bold')

    # 设置标题
    ax.set_title(title, fontsize=18, fontweight='bold', pad=20)

    # 移除坐标轴
    ax.axis('off')

    # 添加图例说明
    total_flow = sum(abs(f) for f in flows if f > 0)
    info_text = f'总流量: {total_flow:.2f} 单位\n箭头宽度表示流量大小'

    props = dict(boxstyle='round', facecolor='lightgray', alpha=0.7)
    ax.text(0.02, 0.02, info_text, transform=ax.transAxes,
            fontsize=10, verticalalignment='bottom',
            bbox=props)

    plt.tight_layout()
    return fig, ax, sankey


def create_simple_sankey():
    """创建示例桑基图数据"""
    # 示例：能源流动
    flows = [0.5, 0.3, -0.2, 0.1, -0.7, 0.2, -0.2]
    labels = ['能源生产', '发电', '输电损耗',
              '工业用电', '居民用电', '商业用电', '损耗']

    fig, ax, sankey = sankey_diagram(
        flows=flows,
        labels=labels,
        orientations=[0, 1, 0, -1, 1, 0],  # 方向控制
        title="能源分配桑基图",
        color_palette='Set2'
    )

    return fig, ax, sankey


def create_budget_sankey():
    """创建预算分配桑基图"""
    flows = [1.0, -0.4, -0.3, -0.2, -0.1,  # 第一层分配
             0.4, -0.2, -0.1, -0.1,  # 运营费用细分
             0.3, -0.15, -0.10, -0.05]  # 研发费用细分

    labels = ['总预算', '运营', '研发', '市场', '其他',
              '运营预算', '人力', '设备', '办公',
              '研发预算', '新产品', '技术升级', '测试']

    fig, ax, sankey = sankey_diagram(
        flows=flows,
        labels=labels,
        title="年度预算分配桑基图",
        color_palette='tab20c',
        margin=0.3
    )

    return fig, ax, sankey