import numpy as np
import matplotlib.collections
import matplotlib.pyplot as plt
import openseespy.opensees as ops


def ops_sec_vis(ele_tag, sec_num=1, rebar_d=None, color=None):
    """
    ele_tag: 要显示截面的所属单元号；
    sec_num: 显示单元上的哪个积分点截面，从i段到j段从1开始编号；
    rebar_d: 钢筋直径，若想突出显示钢筋，请传入一个标量或列表等，列表则表示有多种钢筋直径。
    color: 显示颜色，默认为渐变颜色，其值按与截面中心点的距离变化。
    """
    # 利用eleResponse()命令提取出纤维数据
    FiberData = ops.eleResponse(ele_tag, 'section', 'fiberData', sec_num)
    # FiberData从第1列至第5列分别为："yCoord", "zCoord", "area", "stress", "strain"
    FiberData = np.array(FiberData).reshape((-1, 5))   # 变形为5列数组
    ylocs, zlocs, areas = FiberData[:, 0], FiberData[:, 1], FiberData[:, 2]
    # 如下提取出边界与坐标中心
    ymin, ymax = np.min(ylocs), np.max(ylocs)
    zmin, zmax = np.min(zlocs), np.max(zlocs)
    ymean, zmean = np.mean(ylocs), np.mean(zlocs)

    # matplotlib 绘图
    # plt.style.use('fivethirtyeight')
    # plt.style.use('ggplot')
    fig, ax = plt.subplots(figsize=(8, 8))
    patches = [plt.Circle((yloc, zloc), np.sqrt(area / np.pi))
               for yloc, zloc, area in zip(ylocs, zlocs, areas)]
    coll = matplotlib.collections.PatchCollection(patches, alpha=0.75)
    if color is None:
        colors = (ylocs - ymean)**2 + (zlocs - zmean)**2
        coll.set_array(colors)
    else:
        coll.set_color(color)
    ax.add_collection(coll)

    # 如果包含钢筋
    if rebar_d is not None:
        rebar_d = np.atleast_1d(rebar_d)
        for d in rebar_d:
            rebar_area = np.pi * d ** 2 / 4
            idx = np.argwhere(np.abs(areas - rebar_area) < 1e-6)
            rebar_ys = ylocs[idx]
            rebar_zs = zlocs[idx]
            rebar_areas = areas[idx]
            patches_rebar = [plt.Circle((yloc, zloc), np.sqrt(area / np.pi))
                             for yloc, zloc, area in zip(rebar_ys, rebar_zs, rebar_areas)]
            coll_rebar = matplotlib.collections.PatchCollection(patches_rebar, color='black', alpha=1)
            ax.add_collection(coll_rebar)
    ax.set_aspect('equal')
    ax.set_xlim(ymin * 1.5, ymax * 1.5)
    ax.set_ylim(zmin * 1.5, zmax * 1.5)
    ax.set_xlabel('y', fontsize=20)
    ax.set_ylabel('z', fontsize=20)
    plt.xticks(fontsize=12)
    plt.yticks(fontsize=12)
    plt.show();