import os
import json
import numpy as np
import matplotlib.pyplot as plt
import cv2
from tifffile import tifffile

def draw_lines_on_resized_tiff(json_path, tiff_dir, output_dir=None):
    with open(json_path, 'r', encoding='utf-8') as f:
        data = json.load(f)

    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    target_width, target_height = 826, 818  # 目标尺寸

    for item in data:
        original_img_name = item["image"]
        base_name = os.path.splitext(original_img_name)[0].replace("_color", "_depth")
        tiff_name = "filled_" + base_name + ".tiff"
        tiff_path = os.path.join(tiff_dir, tiff_name)

        if not os.path.exists(tiff_path):
            print(f"[!] TIFF file not found: {tiff_path}")
            continue

        tiff_np = tifffile.imread(tiff_path)
        orig_height, orig_width = tiff_np.shape[:2]

        resized_img = cv2.resize(tiff_np, (target_width, target_height), interpolation=cv2.INTER_LINEAR)

        fig, ax = plt.subplots(figsize=(target_width / 100, target_height / 100), dpi=100)

        if resized_img.ndim == 2:
            ax.imshow(resized_img, cmap='gray')
        else:
            ax.imshow(resized_img)

        scale_x = target_width / orig_width
        scale_y = target_height / orig_height

        for line in item["lines"]:
            start_2000 = np.array(line["start"])
            end_2000 = np.array(line["end"])

            scale_line_x = orig_width / 2000
            scale_line_y = orig_height / 2000

            start_orig = start_2000 * np.array([scale_line_x, scale_line_y])
            end_orig = end_2000 * np.array([scale_line_x, scale_line_y])

            start = start_orig * np.array([scale_x, scale_y])
            end = end_orig * np.array([scale_x, scale_y])

            x0, y0 = int(round(start[0])), int(round(start[1]))
            x1, y1 = int(round(end[0])), int(round(end[1]))

            x0 = np.clip(x0, 0, target_width - 1)
            y0 = np.clip(y0, 0, target_height - 1)
            x1 = np.clip(x1, 0, target_width - 1)
            y1 = np.clip(y1, 0, target_height - 1)

            val_start = resized_img[y0, x0]
            val_end = resized_img[y1, x1]

            ax.plot([x0, x1], [y0, y1], 'r-', linewidth=1)
            ax.scatter([x0, x1], [y0, y1], c='blue', s=10)
            ax.text(x0, y0, f"{val_start}", color='white', fontsize=8,
                    bbox=dict(facecolor='black', alpha=0.5))
            ax.text(x1, y1, f"{val_end}", color='white', fontsize=8,
                    bbox=dict(facecolor='black', alpha=0.5))

        ax.set_title(original_img_name)
        ax.axis('off')

        if output_dir:
            save_name = base_name + "_vis_resized.jpg"
            save_path = os.path.join(output_dir, save_name)
            plt.savefig(save_path, bbox_inches='tight', dpi=100)
            plt.close(fig)
        else:
            plt.show()

    print("✅ 所有图像处理完成。")

if __name__ == '__main__':
    json_path = r"G:\python_ws_g\data\pcd2color_result\a_predict_restnet50\predictions.json"
    tiff_dir = r"G:\python_ws_g\data\pcd2color_result\dilated_tiff_filled"
    output_dir = r"G:\python_ws_g\data\pcd2color_result\line_vis_resized"

    draw_lines_on_resized_tiff(json_path, tiff_dir, output_dir)