import matplotlib.pyplot as plt
import matplotlib.patches as patches
import json


def visualize_solution(row_boundaries, col_boundaries, car_paths_coords, W, H):
    plt.rcParams['font.family'] = ['sans-serif']
    plt.rcParams['font.sans-serif'] = ['SimHei']
    fig, ax = plt.subplots()
    ax.set_xlim(0, W)
    ax.set_ylim(H, 0)  # 调整y轴方向，原点在左上角
    ax.set_title("区域划分与车-机-巢系统覆盖")
    ax.set_xlabel("区域宽度")
    ax.set_ylabel("区域高度")

    # 定义若干颜色以区分不同系统（系统编号从0开始）
    colors = ['red', 'blue', 'green', 'orange', 'purple', 'cyan', 'magenta']

    # 绘制区域中心
    region_center = (H / 2.0, W / 2.0)
    ax.plot(region_center[1], region_center[0],
            'ko', markersize=8, label="区域中心")

    # 绘制行分割边界
    for row in row_boundaries:
        ax.axhline(y=row * H, color='black', linestyle='--')

    # 绘制列分割边界
    for col in col_boundaries:
        ax.axvline(x=col * W, color='black', linestyle='--')

    # 绘制每辆车的轨迹并标注区域序号
    for system_id, path in enumerate(car_paths_coords):
        path = [(region_center[0], region_center[1])] + \
            path + [(region_center[0], region_center[1])]
        y, x = zip(*path)
        ax.plot(x, y, marker='o', color=colors[int(
            system_id) % len(colors)], label=f"系统 {system_id}")

        # 标注每个区域的序号
        for idx, (px, py) in enumerate(zip(x[1:-1], y[1:-1])):  # 跳过起点和终点
            ax.text(px, py, str(idx), color='black', fontsize=8, ha='center', va='center')

    # 添加图例
    ax.legend()
    plt.show()


if __name__ == "__main__":
    import yaml

    # ---------------------------
    # 需要修改的超参数
    # ---------------------------
    params_file = 'params2'
    solution_file = r'solutions\mtkl_params2.json'

    with open(params_file + '.yml', 'r', encoding='utf-8') as file:
        params = yaml.safe_load(file)

    H = params['H']
    W = params['W']
    k = params['num_cars']

    # 读取最佳方案的JSON文件
    with open(solution_file, 'r', encoding='utf-8') as f:
        best_solution = json.load(f)

    row_boundaries = best_solution['row_boundaries']
    col_boundaries = best_solution['col_boundaries']
    car_paths = best_solution['car_paths']

    # 计算分块区域的中心点坐标
    rectangle_centers = []
    for i in range(len(row_boundaries) - 1):
        for j in range(len(col_boundaries) - 1):
            r1 = row_boundaries[i]
            r2 = row_boundaries[i + 1]
            c1 = col_boundaries[j]
            c2 = col_boundaries[j + 1]
            d = (r2 - r1) * H * (c2 - c1) * W  # 任务的照片数量（矩形面积）

            # 计算任务矩形中心，用于后续车辆移动时间计算
            center_r = (r1 + r2) / 2.0 * H
            center_c = (c1 + c2) / 2.0 * W

            rectangle_centers.append((center_r, center_c))
    
    # 将car_paths里的index换成坐标
    car_paths_coords = [[] for _ in range(k)]
    for car_idx in range(k):
        car_path = car_paths[car_idx]
        for point in car_path:
            car_paths_coords[car_idx].append(rectangle_centers[point])

    visualize_solution(row_boundaries, col_boundaries, car_paths_coords, W, H)