#!/usr/bin/env python3
"""
Script to convert YOLO txt label format to LabelMe JSON format.

YOLO format: class_id x_center y_center width height (normalized 0.0-1.0)
LabelMe format: JSON with shapes containing rectangles with pixel coordinates
"""

import os
import sys
import argparse
import json
import shutil
from pathlib import Path

try:
    from PIL import Image
    HAS_PIL = True
except ImportError:
    HAS_PIL = False
    print("Warning: PIL/Pillow not installed. Image dimension detection required for conversion.")
    print("Install with: pip install pillow")


def get_image_dimensions(image_path):
    """Get image width and height."""
    if not HAS_PIL:
        return None, None
    try:
        with Image.open(image_path) as img:
            return img.size  # Returns (width, height)
    except Exception as e:
        print(f"Warning: Could not read image {image_path}: {e}")
        return None, None


def yolo_to_labelme_rectangle(x_center_norm, y_center_norm, width_norm, height_norm,
                              img_width, img_height):
    """
    Convert YOLO normalized bounding box to LabelMe rectangle coordinates.
    
    Args:
        x_center_norm, y_center_norm, width_norm, height_norm: Normalized coordinates (0.0-1.0)
        img_width, img_height: Image dimensions in pixels
    
    Returns:
        List of two points: [[x1, y1], [x2, y2]] for top-left and bottom-right corners
    """
    # Denormalize center coordinates and dimensions
    x_center = x_center_norm * img_width
    y_center = y_center_norm * img_height
    width = width_norm * img_width
    height = height_norm * img_height
    
    # Calculate top-left and bottom-right corners
    x1 = x_center - width / 2.0
    y1 = y_center - height / 2.0
    x2 = x_center + width / 2.0
    y2 = y_center + height / 2.0
    
    # Ensure coordinates are within image bounds
    x1 = max(0.0, min(img_width, x1))
    y1 = max(0.0, min(img_height, y1))
    x2 = max(0.0, min(img_width, x2))
    y2 = max(0.0, min(img_height, y2))
    
    return [[float(x1), float(y1)], [float(x2), float(y2)]]


def is_normalized(value):
    """Check if a coordinate value is normalized (0.0-1.0)."""
    return 0.0 <= float(value) <= 1.0


def find_image_file(txt_file, image_extensions=None):
    """
    Find corresponding image file for a txt annotation file.
    
    Args:
        txt_file: Path to txt annotation file
        image_extensions: List of image extensions to try (default: ['.jpg', '.jpeg', '.png', '.bmp'])
    
    Returns:
        Path to image file or None if not found
    """
    if image_extensions is None:
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff']
    
    txt_file = Path(txt_file)
    base_name = txt_file.stem
    txt_dir = txt_file.parent
    
    # First, check if txt_file is in a 'labels' directory
    # If so, look for corresponding 'images' directory
    if txt_dir.name.lower() == 'labels':
        # Try to find images directory at the same level
        images_dir = txt_dir.parent / 'images'
        if images_dir.exists():
            # Look for image in images directory
            for ext in image_extensions:
                potential_image = images_dir / f"{base_name}{ext}"
                if potential_image.exists():
                    return potential_image
            # Try case variations
            for ext in image_extensions:
                for case_ext in [ext, ext.upper(), ext.capitalize()]:
                    potential_image = images_dir / f"{base_name}{case_ext}"
                    if potential_image.exists():
                        return potential_image
    
    # Check in same directory as txt file
    for ext in image_extensions:
        potential_image = txt_dir / f"{base_name}{ext}"
        if potential_image.exists():
            return potential_image
    
    # Check with case variations in same directory
    for ext in image_extensions:
        for case_ext in [ext, ext.upper(), ext.capitalize()]:
            potential_image = txt_dir / f"{base_name}{case_ext}"
            if potential_image.exists():
                return potential_image
    
    return None


def find_images_directory_for_labels(labels_dir):
    """
    Find the corresponding images directory for a labels directory.
    
    Args:
        labels_dir: Path to labels directory
    
    Returns:
        Path to images directory or None if not found
    """
    labels_dir = Path(labels_dir)
    
    # If the directory name is 'labels', look for 'images' at the same level
    if labels_dir.name.lower() == 'labels':
        images_dir = labels_dir.parent / 'images'
        if images_dir.exists():
            return images_dir
    
    return None


def convert_yolo_to_labelme(txt_file, image_file=None, class_names=None, 
                            image_extensions=None, include_image_data=False):
    """
    Convert a single YOLO txt annotation file to LabelMe JSON format.
    
    Args:
        txt_file: Path to YOLO txt annotation file
        image_file: Path to corresponding image file (optional, will be searched if not provided)
        class_names: Dictionary mapping class_id to class name (optional)
        image_extensions: List of image extensions to search (default: ['.jpg', '.jpeg', '.png', '.bmp'])
        include_image_data: Whether to include base64-encoded image data in JSON
    
    Returns:
        Dictionary with LabelMe JSON structure
    """
    txt_file = Path(txt_file)
    
    if not txt_file.exists():
        raise FileNotFoundError(f"Annotation file not found: {txt_file}")
    
    # Find image file if not provided
    if image_file is None:
        image_file = find_image_file(txt_file, image_extensions)
    
    if image_file is None:
        raise FileNotFoundError(
            f"Image file not found for {txt_file}. "
            f"Please provide image_file or ensure image exists in same directory."
        )
    
    image_file = Path(image_file)
    if not image_file.exists():
        raise FileNotFoundError(f"Image file not found: {image_file}")
    
    # Get image dimensions
    img_width, img_height = get_image_dimensions(image_file)
    if img_width is None or img_height is None:
        raise ValueError(
            f"Could not determine image dimensions for {image_file}. "
            f"PIL/Pillow is required for this operation."
        )
    
    # Read YOLO annotations
    shapes = []
    with open(txt_file, 'r') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:  # Skip empty lines
                continue
            
            parts = line.split()
            if len(parts) < 5:
                print(f"Warning: Invalid YOLO format in {txt_file} line {line_num}: {line}")
                continue
            
            try:
                class_id = int(parts[0])
                x_center = float(parts[1])
                y_center = float(parts[2])
                width = float(parts[3])
                height = float(parts[4])
                
                # Check if coordinates are normalized
                if not (is_normalized(x_center) and is_normalized(y_center) and
                       is_normalized(width) and is_normalized(height)):
                    print(f"Warning: Coordinates in {txt_file} line {line_num} may not be normalized. "
                          f"Assuming normalized format.")
                
                # Convert to LabelMe rectangle format
                points = yolo_to_labelme_rectangle(
                    x_center, y_center, width, height, img_width, img_height
                )
                
                # Get class name
                if class_names and class_id in class_names:
                    label = class_names[class_id]
                else:
                    label = str(class_id)  # Use class_id as label if no mapping provided
                
                # Create shape annotation
                shape = {
                    "label": label,
                    "points": points,
                    "group_id": None,
                    "shape_type": "rectangle",
                    "flags": {}
                }
                shapes.append(shape)
                
            except (ValueError, IndexError) as e:
                print(f"Warning: Could not parse line {line_num} in {txt_file}: {line} - {e}")
                continue
    
    # Get image data if requested
    image_data = None
    if include_image_data:
        try:
            with open(image_file, 'rb') as f:
                import base64
                image_data = base64.b64encode(f.read()).decode('utf-8')
        except Exception as e:
            print(f"Warning: Could not encode image data: {e}")
    
    # Create LabelMe JSON structure
    labelme_json = {
        "version": "5.0.1",
        "flags": {},
        "shapes": shapes,
        "imagePath": image_file.name,
        "imageData": image_data,
        "imageHeight": img_height,
        "imageWidth": img_width
    }
    
    return labelme_json


def convert_dataset(input_dir, output_dir=None, class_names_file=None, 
                   image_extensions=None, include_image_data=False, 
                   copy_images=False, recursive=False):
    """
    Convert a directory of YOLO txt annotations to LabelMe JSON format.
    
    Args:
        input_dir: Input directory containing txt files and images
        output_dir: Output directory for LabelMe JSON files (optional, if None, JSON files are placed next to images)
        class_names_file: Path to file with class names (one per line, optional)
        image_extensions: List of image extensions to search
        include_image_data: Whether to include base64-encoded image data
        copy_images: Whether to copy images to output directory (only used if output_dir is specified)
        recursive: Whether to process subdirectories recursively
    
    Returns:
        Dictionary with conversion statistics
    """
    input_dir = Path(input_dir)
    
    if not input_dir.exists():
        raise FileNotFoundError(f"Input directory not found: {input_dir}")
    
    # Load class names if provided
    class_names = None
    if class_names_file:
        class_names_file = Path(class_names_file)
        if class_names_file.exists():
            class_names = {}
            with open(class_names_file, 'r') as f:
                for idx, line in enumerate(f):
                    class_name = line.strip()
                    if class_name:
                        class_names[idx] = class_name
            print(f"Loaded {len(class_names)} class names from {class_names_file}")
        else:
            print(f"Warning: Class names file not found: {class_names_file}")
    
    # Find all txt files (recursive or not)
    if recursive:
        txt_files = list(input_dir.rglob('*.txt'))
    else:
        txt_files = list(input_dir.glob('*.txt'))
    
    if not txt_files:
        search_type = "recursively" if recursive else "in"
        raise ValueError(f"No .txt files found {search_type} {input_dir}")
    
    stats = {
        'files_processed': 0,
        'total_annotations': 0,
        'errors': []
    }
    
    # Process each txt file
    for txt_file in txt_files:
        try:
            # Find corresponding image
            image_file = find_image_file(txt_file, image_extensions)
            
            if not image_file:
                error_msg = f"Image file not found for {txt_file}"
                stats['errors'].append(error_msg)
                print(f"ERROR: {error_msg}")
                continue
            
            # Convert to LabelMe format
            labelme_json = convert_yolo_to_labelme(
                txt_file, image_file, class_names, image_extensions, include_image_data
            )
            
            # Determine where to place the JSON file
            if output_dir:
                # If output_dir is specified, preserve relative path structure when recursive
                output_dir = Path(output_dir)
                if recursive:
                    # Preserve relative path from input_dir
                    relative_path = txt_file.relative_to(input_dir)
                    output_subdir = output_dir / relative_path.parent
                    output_subdir.mkdir(parents=True, exist_ok=True)
                    json_file = output_subdir / f"{image_file.stem}.json"
                    
                    # Copy image if requested, preserving directory structure
                    if copy_images:
                        output_image = output_subdir / image_file.name
                        if not output_image.exists():
                            shutil.copy2(image_file, output_image)
                else:
                    # Non-recursive: just use output_dir
                    output_dir.mkdir(parents=True, exist_ok=True)
                    json_file = output_dir / f"{image_file.stem}.json"
                    
                    # Copy image if requested
                    if copy_images:
                        output_image = output_dir / image_file.name
                        if not output_image.exists():
                            shutil.copy2(image_file, output_image)
            else:
                # Check if txt_file is in a 'labels' directory
                # If so, place JSON in corresponding 'images' directory
                txt_dir = txt_file.parent
                if txt_dir.name.lower() == 'labels':
                    images_dir = find_images_directory_for_labels(txt_dir)
                    if images_dir:
                        # Place JSON in images directory
                        json_file = images_dir / f"{image_file.stem}.json"
                    else:
                        # Fallback: place next to image file
                        json_file = image_file.parent / f"{image_file.stem}.json"
                else:
                    # Otherwise, place JSON file next to the image file
                    json_file = image_file.parent / f"{image_file.stem}.json"
            
            json_file.parent.mkdir(parents=True, exist_ok=True)
            with open(json_file, 'w') as f:
                json.dump(labelme_json, f, indent=2)
            
            stats['files_processed'] += 1
            stats['total_annotations'] += len(labelme_json['shapes'])
            
            print(f"Processed: {txt_file} -> {json_file} ({len(labelme_json['shapes'])} annotations)")
            
        except Exception as e:
            error_msg = f"Error processing {txt_file}: {str(e)}"
            stats['errors'].append(error_msg)
            print(f"ERROR: {error_msg}")
    
    return stats


def main():
    parser = argparse.ArgumentParser(
        description='Convert YOLO txt label format to LabelMe JSON format',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Convert single file (JSON placed in images folder if txt is in labels folder)
  python convert_yolo_to_labelme.py train/labels/x.txt --image train/images/x.jpg
  # Output: train/images/x.json
  
  # Convert directory (JSON files placed in images folders when txt files are in labels folders)
  python convert_yolo_to_labelme.py --input-dir ./train/labels
  # Converts train/labels/x.txt -> train/images/x.json
  
  # Convert directory recursively (processes all subdirectories)
  python convert_yolo_to_labelme.py --input-dir ./dataset --recursive
  # Converts train/labels/x.txt -> train/images/x.json
  # Converts val/labels/y.txt -> val/images/y.json
  
  # Convert directory with custom output directory
  python convert_yolo_to_labelme.py --input-dir ./labels --output-dir ./labelme_annotations
  
  # Convert recursively with custom output directory (preserves directory structure)
  python convert_yolo_to_labelme.py --input-dir ./labels --output-dir ./labelme_annotations --recursive
  
  # Convert with class names file
  python convert_yolo_to_labelme.py --input-dir ./labels --class-names classes.txt
  
  # Convert and include image data in JSON
  python convert_yolo_to_labelme.py --input-dir ./labels --include-image-data
        """
    )
    
    parser.add_argument(
        'input',
        nargs='?',
        help='Input YOLO txt file (if converting single file)'
    )
    
    parser.add_argument(
        '--image',
        type=str,
        help='Image file path (required for single file conversion)'
    )
    
    parser.add_argument(
        '--output', '-o',
        type=str,
        help='Output JSON file path (for single file conversion). If not specified, JSON is placed in images folder when txt is in labels folder, otherwise next to image file.'
    )
    
    parser.add_argument(
        '--input-dir',
        type=str,
        help='Input directory containing txt files and images (for batch conversion)'
    )
    
    parser.add_argument(
        '--output-dir',
        type=str,
        default=None,
        help='Output directory for LabelMe JSON files (for batch conversion). If not specified, JSON files are placed in the images folder when txt files are in a labels folder (e.g., train/labels/x.txt -> train/images/x.json), otherwise next to image files.'
    )
    
    parser.add_argument(
        '--class-names',
        type=str,
        dest='class_names_file',
        help='File with class names (one per line, line number = class_id)'
    )
    
    parser.add_argument(
        '--image-extensions',
        nargs='+',
        default=['.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff'],
        help='Image file extensions to search for (default: .jpg .jpeg .png .bmp .tif .tiff)'
    )
    
    parser.add_argument(
        '--include-image-data',
        action='store_true',
        help='Include base64-encoded image data in JSON (increases file size)'
    )
    
    parser.add_argument(
        '--copy-images',
        action='store_true',
        help='Copy images to output directory (for batch conversion)'
    )
    
    parser.add_argument(
        '--recursive', '-r',
        action='store_true',
        help='Process subdirectories recursively'
    )
    
    args = parser.parse_args()
    
    # Determine mode: single file or batch
    if args.input:
        # Single file mode
        if not args.image:
            parser.error("--image is required for single file conversion")
        
        # Load class names if provided
        class_names = None
        if args.class_names_file:
            class_names_file = Path(args.class_names_file)
            if class_names_file.exists():
                class_names = {}
                with open(class_names_file, 'r') as f:
                    for idx, line in enumerate(f):
                        class_name = line.strip()
                        if class_name:
                            class_names[idx] = class_name
            else:
                print(f"Warning: Class names file not found: {class_names_file}")
        
        try:
            labelme_json = convert_yolo_to_labelme(
                args.input,
                args.image,
                class_names,
                args.image_extensions,
                args.include_image_data
            )
            
            # Determine output file path
            if args.output:
                output_file = Path(args.output)
            else:
                # Check if txt file is in a 'labels' directory
                # If so, place JSON in corresponding 'images' directory
                txt_file = Path(args.input)
                txt_dir = txt_file.parent
                image_file = Path(args.image)
                
                if txt_dir.name.lower() == 'labels':
                    images_dir = find_images_directory_for_labels(txt_dir)
                    if images_dir:
                        # Place JSON in images directory
                        output_file = images_dir / f"{image_file.stem}.json"
                    else:
                        # Fallback: place next to image file
                        output_file = image_file.parent / f"{image_file.stem}.json"
                else:
                    # Place JSON file next to the image file
                    output_file = image_file.parent / f"{image_file.stem}.json"
            
            output_file.parent.mkdir(parents=True, exist_ok=True)
            
            with open(output_file, 'w') as f:
                json.dump(labelme_json, f, indent=2)
            
            print(f"Successfully converted {args.input} to {output_file}")
            print(f"  Annotations: {len(labelme_json['shapes'])}")
            print(f"  Image: {labelme_json['imagePath']} ({labelme_json['imageWidth']}x{labelme_json['imageHeight']})")
            
        except Exception as e:
            print(f"ERROR: {e}", file=sys.stderr)
            sys.exit(1)
    
    elif args.input_dir:
        # Batch mode - output_dir is optional
        
        try:
            stats = convert_dataset(
                args.input_dir,
                args.output_dir,
                args.class_names_file,
                args.image_extensions,
                args.include_image_data,
                args.copy_images,
                args.recursive
            )
            
            print("\n" + "="*50)
            print("Conversion Summary:")
            print(f"  Files processed: {stats['files_processed']}")
            print(f"  Total annotations: {stats['total_annotations']}")
            if stats['errors']:
                print(f"  Errors: {len(stats['errors'])}")
                for error in stats['errors']:
                    print(f"    - {error}")
            print("="*50)
            
        except Exception as e:
            print(f"ERROR: {e}", file=sys.stderr)
            sys.exit(1)
    
    else:
        parser.error("Either provide input file or --input-dir for batch conversion")


if __name__ == '__main__':
    main()