断面图绘制

# -*- coding: utf-8 -*-
#
# @Author: CPS
# @email: 373704015@qq.com
# @Date:
# @Last Modified by: CPS
# @Last Modified time: 2022-12-07 00:42:18.589644
# @file_path "D:\CPS\IDE\JS_SublmieText\Data\Packages\cps-fileheader"
# @Filename "main.py"
# @Description: 处理指定目录下面所有.xlsx文件为折线断面图，xlsx文件需要固定格式

import os, sys
import glob

sys.path.append("..")

import matplotlib
import pandas as pd
import numpy as np
from pandas import DataFrame

import matplotlib.pyplot as plt
import matplotlib.patheffects as path_effects
from matplotlib.patches import Polygon


# class CoordsData(BaseModel):
#     x_coords: list[float]
#     y_coords: list[float]


# def test(line: list[tuple[float, float]], polygon: list[tuple[float, float]]):
#     # 检查直线与多边形相交的点
#     intersect_points = []
#     for i in range(len(polygon)):
#         p1 = polygon[i]
#         p2 = polygon[(i + 1) % len(polygon)]
#         x1, y1 = line[0]
#         x2, y2 = line[1]
#         x3, y3 = p1
#         x4, y4 = p2
#         den = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
#         if den != 0:
#             t = ((x1 - x3) * (y3 - y4) - (y1 - y3) * (x3 - x4)) / den
#             u = -((x1 - x2) * (y1 - y3) - (y1 - y2) * (x1 - x3)) / den
#             if 0 <= t <= 1 and 0 <= u <= 1:
#                 intersect_x = x1 + t * (x2 - x1)
#                 intersect_y = y1 + t * (y2 - y1)
#                 intersect_points.append((intersect_x, intersect_y))

#     intersect_x = [point[0] for point in intersect_points]
#     intersect_y = [point[1] for point in intersect_points]

#     return intersect_x, intersect_y


def extract_coordinates_from_excel(xls_dataframe: DataFrame):
    # 读取Excel文件
    df = xls_dataframe

    # 获取前三列有数据的内容
    data = df.iloc[:, :3].dropna()  # 通过dropna 过滤掉missing的空数据

    # 通过dropna 过滤掉missing的空数据
    # 使用how=all 处理了如果第一个数据是nan还是会往下取数据
    data2 = df.iloc[:, 4:6].dropna(how="all").replace(np.nan, 0)

    tags = data2.iloc[:, 0]
    water = data2.iloc[:, 1]

    # 提取点名、x坐标和y坐标
    names = data.iloc[:, 0]
    x_coords = data.iloc[:, 1]
    y_coords = data.iloc[:, 2]

    # 返回字典
    coordinates_dict = {
        "names": names,
        "x_coords": x_coords,
        "y_coords": y_coords,
        "tags": tags,
        "water_flow": water,
    }

    return coordinates_dict


def plot_data_with_triangles(
    plt, data, triangle_side_length=0.1, rotation_angle=0, vertical_offset=None
):
    offset = 30
    rotation_angle += offset

    # 计算等边三角形的外接圆半径
    radius = triangle_side_length / (2 * np.sin(np.pi / 3))

    # 计算垂直偏移量
    if vertical_offset is None:
        vertical_offset = radius

    # 绘制数据点和等边三角形
    for point in data:
        x, y = point

        # 计算等边三角形的顶点
        angle = np.radians(rotation_angle)
        triangle_x = [
            x + radius * np.cos(angle),
            x + radius * np.cos(angle + 2 * np.pi / 3),
            x + radius * np.cos(angle + 4 * np.pi / 3),
        ]
        triangle_y = [
            y + vertical_offset + radius * np.sin(angle),
            y + vertical_offset + radius * np.sin(angle + 2 * np.pi / 3),
            y + vertical_offset + radius * np.sin(angle + 4 * np.pi / 3),
        ]

        # 绘制等边三角形
        plt.fill(triangle_x, triangle_y, color="blue")

    # plt.axis("equal")


def plot_coordinates(
    coordinates_dict, output_name: str = "", title: str = "", debug: bool = False
):

    # 提取点名、x坐标和y坐标
    names = coordinates_dict["names"]
    x_coords = coordinates_dict["x_coords"]
    y_coords = coordinates_dict["y_coords"]

    # 如果是右岸开始的序列，进行镜像
    if names[0] == "右岸":
        x_coords = x_coords.max() - x_coords

    # 绘制折线图
    plt.plot(x_coords, y_coords, linestyle="-", color="black", zorder=2)
    plt.fill_between(x_coords, y_coords, color="white", zorder=1)

    # plt.xlabel(x_coords.name)
    # plt.ylabel(y_coords.name)
    plt.xlabel("起点距(m)")
    plt.ylabel("高程(m)")

    xlim_min = 0
    xlim_max = x_coords.max()
    ylim_min = y_coords.min() - 1
    ylim_max = max(coordinates_dict["water_flow"].max(), y_coords.max()) + 1
    plt.xlim(xlim_min, xlim_max)
    plt.ylim(ylim_min, ylim_max)  # 根据最大水位或者最大的高程来定y轴

    # 绘制历史水位线
    x = []
    y = []
    white_border = path_effects.Stroke(linewidth=3, foreground="white")  # 添加白边效果
    annotate_font_size = 12

    full_polygon_x = [*x_coords] + [x_coords.max(), 0]
    full_polygon_y = [*y_coords] + [y_coords.max(), y_coords.max()]

    minIndex = np.argmin(y_coords)  # 获取最低的点坐标索引

    # 排序根据规律可以一个坐标出现标签，右边出现标签
    # 最终格式: [['标签1','水深1'], ['标签2','水深2'], ....]
    water_data = sorted(
        [
            [coordinates_dict["tags"][i], e]
            for i, e in enumerate(coordinates_dict["water_flow"])
        ],
        key=(lambda x: x[1]),
    )

    for index, each_water in enumerate(water_data):
        # 不同高度绘制文字注释
        if each_water[1] == coordinates_dict["water_flow"].max():
            xytext_y = each_water[1] + annotate_font_size / 100
        elif each_water[1] == coordinates_dict["water_flow"].min():
            xytext_y = each_water[1] - annotate_font_size / 50
        else:
            xytext_y = each_water[1] + (annotate_font_size / 150)

        each_label = each_water[0]
        y.append(each_water[1] + 0.05)
        x.append(x_coords[minIndex])

        if index % 2 == 1:
            label = f"  {round(each_water[1],2)}m {each_label}"
            ha = "left"
        else:
            label = f"{round(each_water[1],2)}m {each_label}  "
            ha = "right"

        # 文字标注
        plt.annotate(
            label,
            (x_coords[minIndex], each_water[1]),
            xytext=(x_coords[minIndex], xytext_y),
            ha=ha,
            fontsize=annotate_font_size,
        ).set_path_effects(
            [white_border, path_effects.Normal()]
        )  # 添加白边

        # 水位线
        plt.axhline(y=each_water[1], color="r", linestyle="-", zorder=0)

        # 提取所有交点的线段
        # line = [(0, each_y), (x_coords.max(), each_y)]
        # draw_x, draw_y = test(line, polygon)
        # plt.scatter(draw_x, draw_y, color="b")

    # 【绘制】 倒三角
    plt.scatter(x, y, marker="v", color="b", s=50, zorder=5)

    # 【绘制】 左岸标签
    text_xy_offset = 15
    left_text_x = xlim_min + (xlim_max - xlim_min) / text_xy_offset
    left_text_y = ylim_max - (ylim_max - ylim_min) / text_xy_offset
    left_text_xy = [left_text_x, left_text_y]

    plt.annotate(
        "左岸",
        left_text_xy,
        xytext=left_text_xy,
        ha="center",
        fontsize=annotate_font_size,
        zorder=5,
    )

    # 【绘制】 右岸标签
    right_text_x = xlim_max - (xlim_max - xlim_min) / text_xy_offset
    right_text_y = ylim_max - (ylim_max - ylim_min) / text_xy_offset
    right_text_xy = [right_text_x, right_text_y]
    plt.annotate(
        "右岸",
        right_text_xy,
        xytext=right_text_xy,
        ha="center",
        fontsize=annotate_font_size,
        zorder=5,
    )

    # plot_data_with_triangles(plt, xy, left_text_y / 50)

    plt.xticks()
    plt.yticks()
    plt.grid(True, zorder=2)

    if title:
        plt.title(title)
    else:
        plt.title(os.path.basename(output_name))

    if debug:
        plt.show()
    else:
        plt.savefig(output_name)
        # print("output_name: ", output_name)


def main(target_dir: str, width: int = 1200, height: int = 768, output_dir: str = None):
    """
    @Description 处理目录下面的所有.xlsx文件

    - param target_dir :{str} 目标目录，绝对路径

    """
    xlsx_list = glob.glob(f"{target_dir}/*.xlsx")
    xls_list = glob.glob(f"{target_dir}/*.xls")
    target_list = xlsx_list + xls_list

    for each_xlsx in target_list:
        xls = pd.ExcelFile(each_xlsx)
        sheet_names = xls.sheet_names
        basename = os.path.splitext(os.path.basename(each_xlsx))[0]

        for each_sheet_name in sheet_names:
            try:
                plt.figure(figsize=(width / 100, height / 100))  # 设置画布
                df = xls.parse(each_sheet_name)  # 读取指定sheet
                data = extract_coordinates_from_excel(df)  # 提取数据

                output_name = f"{basename}_{each_sheet_name}"
                output_path = os.path.join(
                    output_dir or target_dir, f"{output_name}.png"
                )
                plot_coordinates(data, title=output_name, output_name=output_path)
            except Exception as e:
                print("文件处理错误，请检查: ")
                print(each_xlsx)
                print("sheet: ", each_sheet_name)
                print("【错误信息】: ", e)
                continue


if __name__ == "__main__":
    width = 1200
    height = 768
    matplotlib.rcParams["font.family"] = "SimHei"

    input_path = os.path.abspath("./data")
    output_path = os.path.abspath("./jpg")
    file_path = os.path.join(input_path, "抚河.xlsx")  # 替换成你的Excel文件路径

    # 单文件示例用法
    # plt.figure(figsize=(width / 100, height / 100))
    # xls = pd.ExcelFile(file_path)
    # sheet_names = xls.sheet_names
    # sheet_index = 1

    # output_dir = os.path.dirname(file_path)
    # basename = os.path.splitext(os.path.basename(file_path))[0]
    # output_path = os.path.join(output_dir, f"{basename}_{sheet_names[sheet_index]}")

    # df = xls.parse(sheet_names[sheet_index])
    # data = extract_coordinates_from_excel(df)

    # plot_coordinates(
    #     data, output_name=output_path, title=sheet_names[sheet_index], debug=True
    # )

    # 目录批处理
    main(input_path, width, height, output_dir=output_path)