import os
import json
from itertools import combinations
import cv2
import numpy as np
import skimage
from scipy.ndimage import zoom
from tqdm import tqdm  # 添加 tqdm 库

def to_int(x):
    return tuple(map(int, x))

def convert_labelme_to_custom_json(labelme_json_path):
    """
    Convert LabelMe JSON to custom JSON format.
    """
    try:
        # Read LabelMe JSON file
        with open(labelme_json_path, 'r', encoding='utf-8') as f:
            labelme_data = json.load(f)
        # 检查必要的键是否存在
        if 'imageHeight' not in labelme_data:
            raise KeyError("Missing 'imageHeight' in JSON")
        if 'imageWidth' not in labelme_data:
            raise KeyError("Missing 'imageWidth' in JSON")
        # Extract image information
        image_height = labelme_data['imageHeight']
        image_width = labelme_data['imageWidth']
        # Extract line segment information
        lines = []
        for shape in labelme_data.get('shapes', []):
            if shape.get('label') == 'dseam1':
                if len(shape.get('points', [])) != 2:
                    print(f"Warning: Skipping invalid line in {labelme_json_path}")
                    continue
                start_point = shape['points'][0]
                end_point = shape['points'][1]
                lines.append([start_point, end_point])
        # 检查是否有线段
        if not lines:
            raise ValueError("No line segments found in the JSON")
        # Create custom JSON data
        custom_json_data = {
            "image_id": os.path.splitext(os.path.basename(labelme_json_path))[0],  # 使用 JSON 文件名作为 image_id
            # "height": image_height,
            # "width": image_width,
            "lines": lines
        }
        return custom_json_data
    except json.JSONDecodeError:
        print(f"JSON Decode Error: Invalid JSON file - {labelme_json_path}")
    except KeyError as e:
        print(f"Key Error in {labelme_json_path}: {e}")
    except ValueError as e:
        print(f"Value Error in {labelme_json_path}: {e}")
    except Exception as e:
        print(f"Unexpected error processing {labelme_json_path}: {e}")
    return None

def find_image_file(json_path, image_extensions=('.jpg', '.jpeg', '.png', '.bmp')):
    """
    Find the corresponding image file for a given JSON file.
    """
    base_name = os.path.splitext(json_path)[0]
    for ext in image_extensions:
        image_path = f"{base_name}{ext}"
        if os.path.exists(image_path):
            return image_path
    return None

def generate_heatmap_data(image, lines):
    """
    Generate heatmap data from wireframe lines (in memory).
    """
    try:
        im_rescale = (512, 512)
        heatmap_scale = (128, 128)
        fy, fx = heatmap_scale[1] / image.shape[0], heatmap_scale[0] / image.shape[1]
        jmap = np.zeros((1,) + heatmap_scale, dtype=np.float32)
        joff = np.zeros((1, 2) + heatmap_scale, dtype=np.float32)
        lmap = np.zeros(heatmap_scale, dtype=np.float32)
        lines[:, :, 0] = np.clip(lines[:, :, 0] * fx, 0, heatmap_scale[0] - 1e-4)
        lines[:, :, 1] = np.clip(lines[:, :, 1] * fy, 0, heatmap_scale[1] - 1e-4)
        lines = lines[:, :, ::-1]
        junc = []
        jids = {}
        def jid(jun):
            jun = tuple(jun[:2])
            if jun in jids:
                return jids[jun]
            jids[jun] = len(junc)
            junc.append(np.array(jun + (0,)))
            return len(junc) - 1
        lnid = []
        lpos, lneg = [], []
        for v0, v1 in lines:
            lnid.append((jid(v0), jid(v1)))
            lpos.append([junc[jid(v0)], junc[jid(v1)]])
            vint0, vint1 = to_int(v0), to_int(v1)
            jmap[0][vint0] = 1
            jmap[0][vint1] = 1
            rr, cc, value = skimage.draw.line_aa(*to_int(v0), *to_int(v1))
            lmap[rr, cc] = np.maximum(lmap[rr, cc], value)
        for v in junc:
            vint = to_int(v[:2])
            joff[0, :, vint[0], vint[1]] = v[:2] - vint - 0.5
        llmap = zoom(lmap, [0.5, 0.5])
        lineset = set([frozenset(l) for l in lnid])
        # Generate negative samples
        for i0, i1 in combinations(range(len(junc)), 2):
            if frozenset([i0, i1]) not in lineset:
                v0, v1 = junc[i0], junc[i1]
                vint0, vint1 = to_int(v0[:2] / 2), to_int(v1[:2] / 2)
                rr, cc, value = skimage.draw.line_aa(*vint0, *vint1)
                lneg.append([v0, v1, i0, i1, np.average(np.minimum(value, llmap[rr, cc]))])
        # 如果没有负样本，模拟生成一些负样本
        if len(lneg) == 0:
            print("Warning: No negative samples found. Simulating artificial negative samples.")
            for i0, i1 in combinations(range(len(junc)), 2):
                v0, v1 = junc[i0], junc[i1]
                vint0, vint1 = to_int(v0[:2] / 2), to_int(v1[:2] / 2)
                rr, cc, value = skimage.draw.line_aa(*vint0, *vint1)
                simulated_value = np.random.uniform(0.01, 0.1)  # 模拟低置信度值
                lneg.append([v0, v1, i0, i1, simulated_value])
        assert len(lneg) != 0, "Failed to generate negative samples even after simulation."
        lneg.sort(key=lambda l: -l[-1])
        junc = np.array(junc, dtype=np.float32)
        Lpos = np.array(lnid, dtype=np.int32)
        Lneg = np.array([l[2:4] for l in lneg][:4000], dtype=np.int32)
        lpos = np.array(lpos, dtype=np.float32)
        lneg = np.array([l[:2] for l in lneg[:2000]], dtype=np.float32)
        image = cv2.resize(image, im_rescale)
        # Organize data into the target format
        wires_data = {
            "junc_map": {
                "name": "junc_map",
                "shape": list(jmap.shape),
                "content": jmap.tolist()
            },
            "junc_offset": {
                "name": "junc_offset",
                "shape": list(joff.shape),
                "content": joff.tolist()
            },
            "line_map": {
                "name": "line_map",
                "shape": list(lmap.shape),
                "content": lmap.tolist()
            },
            "junc_coords": {
                "name": "junc_coords",
                "shape": list(junc.shape),
                "content": junc.tolist()
            },
            "line_pos_idx": {
                "name": "line_pos_idx",
                "shape": list(Lpos.shape),
                "content": Lpos.tolist()
            },
            "line_neg_idx": {
                "name": "line_neg_idx",
                "shape": list(Lneg.shape),
                "content": Lneg.tolist()
            },
            "line_pos_coords": {
                "name": "line_pos_coords",
                "shape": list(lpos.shape),
                "content": lpos.tolist()
            },
            "line_neg_coords": {
                "name": "line_neg_coords",
                "shape": list(lneg.shape),
                "content": lneg.tolist()
            }
        }
        return wires_data
    except Exception as e:
        print(f"Error generating heatmap data: {e}")
        raise
def batch_process(labelme_json_dir, output_dir):
    """
    Batch process LabelMe JSON files to generate custom JSONs and copy images.
    """
    json_files = [f for f in os.listdir(labelme_json_dir) if f.endswith('.json')]
    total_files = len(json_files)
    processed_files = 0
    error_files = 0
    os.makedirs(output_dir, exist_ok=True)
    for filename in tqdm(json_files, desc="Processing JSON files", unit="file"):
        labelme_json_path = os.path.join(labelme_json_dir, filename)
        try:
            # Step 1: Convert LabelMe JSON to custom JSON
            custom_json_data = convert_labelme_to_custom_json(labelme_json_path)
            if custom_json_data is None:
                error_files += 1
                continue
            # Step 2: Find the corresponding image file
            image_path = find_image_file(labelme_json_path)
            if not image_path or not os.path.exists(image_path):
                print(f"Error: Image file not found for {labelme_json_path}")
                error_files += 1
                continue
            # Step 2.5: Copy image to output directory
            image_filename = os.path.basename(image_path)
            output_image_path = os.path.join(output_dir, image_filename)
            if not os.path.exists(output_image_path):
                try:
                    import shutil
                    shutil.copy(image_path, output_image_path)
                except Exception as e:
                    print(f"Error copying image file {image_path} to {output_image_path}: {e}")
                    error_files += 1
                    continue
            # Step 3: Read and process image
            image = cv2.imread(image_path)
            if image is None:
                print(f"Error: Unable to read image file - {image_path}")
                error_files += 1
                continue
            # Step 4: Generate heatmap data (in memory)
            lines = np.array(custom_json_data['lines']).reshape(-1, 2, 2)
            wires_data = generate_heatmap_data(image, lines)
            custom_json_data.pop("lines", None)  # Remove 'lines' field
            # Step 5: Combine custom JSON and heatmap data
            custom_json_data.update({
                "image_id": custom_json_data["image_id"],
                "boxes": [],
                "segmentations": [],
                "wires": [wires_data]
            })
            # Step 6: Save the final custom JSON (compact format)
            output_json_path = os.path.join(output_dir, os.path.splitext(filename)[0] + '.json')
            with open(output_json_path, 'w', encoding='utf-8') as f:
                json.dump(custom_json_data, f, separators=(',', ':'))
            print(f"Saved compact JSON file: {output_json_path}")
            processed_files += 1
        except Exception as e:
            print(f"Unexpected error processing {labelme_json_path}: {e}")
            error_files += 1
    print(f"\nProcessing Complete:")
    print(f"Total Files: {total_files}")
    print(f"Processed Files: {processed_files}")
    print(f"Error Files: {error_files}")





if __name__ == "__main__":
    labelme_json_dir = r"G:\python_ws_g\data\a_strong_output"
    parent_dir = os.path.dirname(labelme_json_dir)
    custom_json_output_dir = os.path.join(parent_dir, "a_strong_json_output")

    batch_process(labelme_json_dir, custom_json_output_dir)