"""
COCO Category Combiner

This script combines multiple COCO categories into a single category.
It updates all annotation references and remaps category IDs to be contiguous.

Key Features:
- Combines any number of categories into one new category
- Updates all annotation category_id references
- Remaps category IDs to contiguous values after combining
- Preserves all non-standard fields on categories, annotations, images, and top-level keys
- Configurable starting ID (0 or 1) for different framework requirements
- Optional supercategory override

Common Use Cases:
- Consolidating fine-grained labels (e.g., "sedan" + "SUV" + "truck" -> "vehicle")
- Merging overlapping categories after dataset combination
- Simplifying category hierarchies for training

The tool will:
1. Load the input COCO JSON file
2. Validate that all specified categories exist
3. Replace the specified categories with a single combined category
4. Remap all category IDs to be contiguous
5. Update all annotation category_id references
6. Save the updated COCO JSON

Usage:
    poetry run coco-combine-categories \\
        --input-coco-json ./annotations/coco.json \\
        --output-coco-json ./annotations/coco_combined.json \\
        --category-names "sedan" "SUV" "truck" \\
        --new-name "vehicle"

Example combining into an existing category name:
    poetry run coco-combine-categories \\
        -i ./annotations/coco.json \\
        -o ./annotations/coco_combined.json \\
        -c "cat" "dog" "bird" \\
        -n "animal" \\
        -s "living_thing"
"""

import argparse
import json
import sys
from pathlib import Path


def combine_categories(
    input_json_path, output_json_path, category_names, new_name, supercategory=None, start_id=1
):
    """Combine multiple COCO categories into one and remap IDs to be contiguous."""

    with open(input_json_path, "r") as f:
        coco = json.load(f)

    categories = coco.get("categories", [])
    annotations = coco.get("annotations", [])

    # Build name -> category lookup
    name_to_cat = {}
    for cat in categories:
        name_to_cat[cat["name"]] = cat

    # Validate all requested category names exist
    missing = [name for name in category_names if name not in name_to_cat]
    if missing:
        print(f"Error: The following categories were not found: {missing}", file=sys.stderr)
        print(f"Available categories: {list(name_to_cat.keys())}", file=sys.stderr)
        sys.exit(1)

    # Validate new_name doesn't collide with a non-combined category
    names_to_combine = set(category_names)
    if new_name in name_to_cat and new_name not in names_to_combine:
        print(
            f"Error: Category '{new_name}' already exists and is not in the combine list. "
            f"Use a different name or include it in --category-names.",
            file=sys.stderr,
        )
        sys.exit(1)

    if len(category_names) == 1:
        print(f"Warning: Only one category provided; this is equivalent to a rename.")

    # Collect old IDs that will be merged
    old_ids_to_merge = {name_to_cat[name]["id"] for name in category_names}

    # Build the combined category from a shallow copy of the first source category
    first_source = name_to_cat[category_names[0]]
    combined_cat = first_source.copy()
    combined_cat["name"] = new_name
    combined_cat["supercategory"] = supercategory or new_name

    # Build new category list, inserting combined category at position of first source
    new_categories = []
    combined_inserted = False
    for cat in categories:
        if cat["id"] in old_ids_to_merge:
            if not combined_inserted:
                new_categories.append(combined_cat)
                combined_inserted = True
            # Skip other source categories
        else:
            new_categories.append(cat)

    # Assign contiguous IDs and build old->new mapping
    old_id_to_new_id = {}
    for i, cat in enumerate(new_categories):
        new_id = i + start_id
        if cat is combined_cat:
            # Map all merged old IDs to this new ID
            for old_id in old_ids_to_merge:
                old_id_to_new_id[old_id] = new_id
        else:
            old_id_to_new_id[cat["id"]] = new_id
        cat["id"] = new_id

    # Update annotations - only mutate category_id
    for ann in annotations:
        old_id = ann["category_id"]
        if old_id in old_id_to_new_id:
            ann["category_id"] = old_id_to_new_id[old_id]
        else:
            print(
                f"Warning: Annotation {ann.get('id', '?')} references unknown category ID {old_id}",
                file=sys.stderr,
            )

    # Update the coco dict in place (preserves all top-level keys)
    coco["categories"] = new_categories

    # Save
    Path(output_json_path).parent.mkdir(parents=True, exist_ok=True)
    with open(output_json_path, "w") as f:
        json.dump(coco, f, indent=2)

    # Print summary
    print(f"Combined {list(category_names)} -> '{new_name}'")
    print(f"Categories: {len(categories)} -> {len(new_categories)}")

    cat_counts = {}
    for ann in annotations:
        cat_counts[ann["category_id"]] = cat_counts.get(ann["category_id"], 0) + 1
    print(f"\nCategory distribution:")
    for cat in new_categories:
        count = cat_counts.get(cat["id"], 0)
        print(f"  {cat['name']} (ID {cat['id']}): {count} annotations")

    print(f"\n✅ Saved to: {output_json_path}")


def main():
    parser = argparse.ArgumentParser(
        description="Combine multiple COCO categories into one."
    )
    parser.add_argument(
        "--input-coco-json",
        "-i",
        default="coco.json",
        help="Path to input COCO JSON file",
    )
    parser.add_argument(
        "--output-coco-json",
        "-o",
        default="coco_combined.json",
        help="Path to output COCO JSON file",
    )
    parser.add_argument(
        "--category-names",
        "-c",
        nargs="+",
        required=True,
        help="Category names to combine",
    )
    parser.add_argument(
        "--new-name",
        "-n",
        required=True,
        help="Name for the combined category",
    )
    parser.add_argument(
        "--supercategory",
        "-s",
        default=None,
        help="Supercategory for the combined category (defaults to new-name)",
    )
    parser.add_argument(
        "--start-id",
        type=int,
        choices=[0, 1],
        default=1,
        help="Starting ID for category mapping (0 or 1)",
    )

    args = parser.parse_args()

    combine_categories(
        input_json_path=args.input_coco_json,
        output_json_path=args.output_coco_json,
        category_names=args.category_names,
        new_name=args.new_name,
        supercategory=args.supercategory,
        start_id=args.start_id,
    )


if __name__ == "__main__":
    main()
