# paint/treemap.py
"""
树状图可视化
"""
import matplotlib as plt
import squarify


def treemap(sizes=None, labels=None, colors=None, title="Treemap"):
    """绘制树状图"""
    fig, ax = plt.subplots(figsize=(12, 8))

    if sizes is None:
        sizes = [40, 30, 20, 10]
        labels = ['Group A', 'Group B', 'Group C', 'Group D']

    if colors is None:
        colors = plt.cm.Set3(np.linspace(0, 1, len(sizes)))

    squarify.plot(sizes=sizes, label=labels, color=colors, alpha=0.8, ax=ax)
    ax.set_title(title, fontsize=14, pad=20)
    ax.axis('off')

    return fig, ax