"""
Object tracking script that reads predictions.json and outputs tracking.json with track IDs.

Usage:
    python track_objects.py --input runs/val_viz/exp/predictions.json --output runs/val_viz/exp/tracking.json
"""

import argparse
import csv
import json
import os
import re
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
import numpy as np
try:
    import yaml
except ImportError:
    yaml = None
try:
    from scipy.optimize import linear_sum_assignment
    SCIPY_AVAILABLE = True
except ImportError:
    SCIPY_AVAILABLE = False
    print("Warning: scipy not available, using greedy matching instead of Hungarian algorithm")

from merge_tracking_results import merge_case, normalize_image_name


DEFAULT_DATA_CONFIG_PATH = Path(__file__).resolve().parents[2] / 'ultralytics' / 'cfg' / 'datasets' / 'mono3d_ground.yaml'
DEFAULT_CLASS_MAP = {
    'car': 0,
    'suv': 1,
    'pickup': 2,
    'medium_car': 3,
    'van': 4,
    'bus': 5,
    'truck': 6,
    'tanker': 6,
    'large_truck': 6,
    'construction_vehicle': 6,
    'special_vehicle': 7,
    'unknown': 8,
    'pedestrian': 9,
    'bicyclist': 10,
    'motorcyclist': 10,
    'bicycle': 11,
    'motorcycle': 11,
    'tricycle': 12,
    'tricyclist': 12,
    'traffic_sign': 13,
    'wheel': 14,
    'plate': 15,
    'face': 16,
    'car_fake': 17,
    'bicyclist_fake': 18,
    'pedestrian_fake': 19,
}
DEFAULT_FACE_3D_CLASS_IDS = {0, 1, 2, 3, 4, 5, 6, 7, 8, 17}
DEFAULT_COMPLETE_3D_CLASS_IDS = {9, 10, 11, 12, 18, 19}
DEFAULT_TRACKED_CLASS_IDS = sorted(DEFAULT_FACE_3D_CLASS_IDS | DEFAULT_COMPLETE_3D_CLASS_IDS)
VEHICLE_CLASS_NAME_ALIASES = {
    'vehicle',
    'car',
    'car_fake',
    'suv',
    'pickup',
    'medium_car',
    'van',
    'bus',
    'truck',
    'tanker',
    'large_truck',
    'construction_vehicle',
    'special_vehicle',
    'unknown',
    'tricycle',
    'tricyclist',
}
PEDESTRIAN_CLASS_NAME_ALIASES = {
    'pedestrian',
    'pedestrian_fake',
    'person',
    'ped',
    'bicyclist',
    'bicyclist_fake',
    'motorcyclist',
    'cyclist',
    'rider',
}
DEFAULT_SUB_CLS = -1
DEFAULT_TIMESTAMP = 0
DEFAULT_MODEL_VERSION = '20260317'
TIMESTAMP_CSV_NAME = 'frame_timestamps_interp.csv'
TIMESTAMP_FRAME_ID_COLUMN = 'camera4_frame_id'
TIMESTAMP_VALUE_COLUMN = 'camera4'
DEFAULT_DUPLICATE_OVERLAP_THRESHOLD = 0.95
VRU_ASSOCIATION_CLASS_ID = -100
VRU_ASSOCIATION_TYPE_NAME = 'vru'


def parse_numeric_value(value):
    """Convert CSV / JSON scalar strings to int or float when possible."""
    if value is None:
        return None
    if isinstance(value, (int, float)):
        return value

    value_str = str(value).strip()
    if not value_str:
        return None

    # Preserve exact integer precision for large ids / timestamps such as
    # nanosecond log times written into exported filenames.
    if re.fullmatch(r'[+-]?\d+', value_str):
        try:
            return int(value_str)
        except ValueError:
            return value

    try:
        numeric = float(value_str)
    except ValueError:
        return value

    if numeric.is_integer():
        return int(numeric)
    return numeric


def normalize_class_name(name):
    """Normalize class-like strings to a lowercase underscore form."""
    normalized = re.sub(r'[^a-z0-9]+', '_', str(name or '').strip().lower())
    return normalized.strip('_')


def invert_class_map(class_map):
    """Convert class_map{name->id} into the first canonical name per class id."""
    class_names = {}
    for raw_name, class_id in class_map.items():
        class_names.setdefault(int(class_id), normalize_class_name(raw_name))
    return dict(sorted(class_names.items()))


def is_vehicle_name(normalized_name):
    """Check whether a normalized class name belongs to the vehicle attribute task."""
    return (
        normalized_name in VEHICLE_CLASS_NAME_ALIASES
        or normalized_name == 'unknown'
        or 'vehicle' in normalized_name
        or 'truck' in normalized_name
        or 'bus' in normalized_name
        or 'tanker' in normalized_name
        or 'pickup' in normalized_name
        or normalized_name in {'car', 'suv', 'van', 'tricycle'}
    )


def is_pedestrian_name(normalized_name):
    """Check whether a normalized class name belongs to the pedestrian attribute task."""
    return (
        normalized_name in PEDESTRIAN_CLASS_NAME_ALIASES
        or 'pedestrian' in normalized_name
        or 'bicyclist' in normalized_name
        or 'motorcyclist' in normalized_name
        or 'cyclist' in normalized_name
        or 'rider' in normalized_name
    )


def resolve_task_class_ids(class_names):
    """Resolve vehicle / pedestrian class-id groups from canonical class names."""
    vehicle_class_ids = set()
    pedestrian_class_ids = set()
    for cls_id, cls_name in class_names.items():
        normalized_name = normalize_class_name(cls_name)
        if is_pedestrian_name(normalized_name):
            pedestrian_class_ids.add(int(cls_id))
            continue
        if is_vehicle_name(normalized_name):
            vehicle_class_ids.add(int(cls_id))
    return vehicle_class_ids, pedestrian_class_ids


def load_tracking_metadata(data_config_path=DEFAULT_DATA_CONFIG_PATH, verbose=False):
    """Load Ground3D class metadata, falling back to synced in-script defaults."""
    class_map = DEFAULT_CLASS_MAP.copy()
    face_3d_classes = set(DEFAULT_FACE_3D_CLASS_IDS)
    complete_3d_classes = set(DEFAULT_COMPLETE_3D_CLASS_IDS)
    source = 'builtin defaults'

    config_path = Path(data_config_path) if data_config_path else None
    if config_path and config_path.is_file():
        if yaml is None:
            if verbose:
                print(f"Warning: PyYAML not available, falling back to built-in class map instead of {config_path}")
        else:
            try:
                with open(config_path, 'r', encoding='utf-8') as f:
                    config = yaml.safe_load(f) or {}
            except Exception as exc:
                if verbose:
                    print(f"Warning: Failed to parse data config {config_path}: {exc}. Using built-in class map.")
            else:
                raw_class_map = config.get('class_map')
                if isinstance(raw_class_map, dict) and raw_class_map:
                    class_map = {str(key): int(value) for key, value in raw_class_map.items()}
                raw_face_3d = config.get('face_3d_classes')
                if raw_face_3d:
                    face_3d_classes = {int(value) for value in raw_face_3d}
                raw_complete_3d = config.get('complete_3d_classes')
                if raw_complete_3d:
                    complete_3d_classes = {int(value) for value in raw_complete_3d}
                source = str(config_path)
    elif config_path and verbose:
        print(f"Warning: Data config not found: {config_path}. Using built-in class map.")

    class_names = invert_class_map(class_map)
    tracked_classes = sorted((face_3d_classes | complete_3d_classes) or set(class_names))
    vehicle_class_ids, pedestrian_class_ids = resolve_task_class_ids(class_names)
    return {
        'source': source,
        'class_names': class_names,
        'face_3d_classes': face_3d_classes,
        'complete_3d_classes': complete_3d_classes,
        'tracked_classes': tracked_classes,
        'vehicle_class_ids': vehicle_class_ids,
        'pedestrian_class_ids': pedestrian_class_ids,
    }


TRACKING_METADATA = load_tracking_metadata(verbose=False)
CLASS_NAME_LOOKUP = TRACKING_METADATA['class_names']
VEHICLE_CLASS_IDS = TRACKING_METADATA['vehicle_class_ids']
PEDESTRIAN_CLASS_IDS = TRACKING_METADATA['pedestrian_class_ids']
TRACKED_CLASS_IDS = TRACKING_METADATA['tracked_classes']


def safe_int(value):
    """Best-effort integer conversion."""
    numeric = parse_numeric_value(value)
    if numeric is None:
        return None
    try:
        return int(numeric)
    except (TypeError, ValueError):
        return None


def normalize_heading_debug(heading_debug):
    """Normalize per-detection heading debug data to numeric Python types."""
    if not isinstance(heading_debug, dict):
        return None

    normalized = {}

    yaw_bin = safe_int(heading_debug.get('yaw_bin'))
    if yaw_bin is not None:
        normalized['yaw_bin'] = yaw_bin

    yaw_delta = parse_numeric_value(heading_debug.get('yaw_delta'))
    if isinstance(yaw_delta, (int, float)):
        normalized['yaw_delta'] = float(yaw_delta)

    rot_y_decoded = parse_numeric_value(heading_debug.get('rot_y_decoded'))
    if isinstance(rot_y_decoded, (int, float)):
        normalized['rot_y_decoded'] = float(rot_y_decoded)

    yaw_probs_raw = heading_debug.get('yaw_probs')
    if isinstance(yaw_probs_raw, (list, tuple)):
        yaw_probs = []
        valid = True
        for value in yaw_probs_raw:
            numeric = parse_numeric_value(value)
            if not isinstance(numeric, (int, float)):
                valid = False
                break
            yaw_probs.append(float(numeric))
        if valid:
            normalized['yaw_probs'] = yaw_probs

    return normalized or None


def resolve_output_version(model_version=None, existing_version=None):
    """Resolve the version string written into tracking detections."""
    if model_version is not None:
        return str(model_version)
    if existing_version is not None:
        return str(existing_version)
    return DEFAULT_MODEL_VERSION


def resolve_detection_type_name(det):
    """Resolve the best available normalized type name for one detection."""
    for key in ('type_name', 'cls_name', 'class_name'):
        value = det.get(key)
        normalized = normalize_class_name(value)
        if normalized:
            return normalized

    class_id = safe_int(det.get('class_id'))
    if class_id is None:
        return ''
    return normalize_class_name(CLASS_NAME_LOOKUP.get(class_id, ''))


def extract_frame_metadata_from_image_name(image_name):
    """Extract frame_id and optional timestamp from an image / JSON stem.

    Preferred filename format:
        camera4_<frame_id>_<timestamp>.json

    Exported case / clip filename format also supported:
        ..._<frame_id>_<timestamp>.json
        e.g. G1M3_xxx_<saved_idx>_<frame_id>_<timestamp>.json
             019dxxxx-clip_uuid_<frame_id>_<timestamp>.json

    Legacy fallback:
        use the last numeric token as frame_id and leave timestamp unset.
    """
    if not image_name:
        return None, None

    normalized = normalize_image_name(Path(image_name).stem)

    # Prefer explicit frame_id + timestamp carried in the filename, e.g.
    # camera4_36607_1760795616410702.json
    match = re.search(r'(?:^|_)camera\d+_(\d+)_(\d+)(?:$|_)', normalized)
    if match is not None:
        return int(match.group(1)), parse_numeric_value(match.group(2))

    # Exported case / clip files often end with `_<frame_id>_<timestamp>`, while
    # earlier tokens may include scene ids, datetimes, UUIDs, or saved indices.
    # Split on `_` rather than matching arbitrary digits so UUID-internal digits
    # do not interfere with the parsed tail fields.
    parts = normalized.split('_')
    if len(parts) >= 3 and parts[-1].isdigit() and parts[-2].isdigit():
        return int(parts[-2]), parse_numeric_value(parts[-1])

    match = re.search(r'(\d+)(?!.*\d)', normalized)
    if match is None:
        return None, None
    return int(match.group(1)), None


def extract_frame_id_from_image_name(image_name):
    """Extract frame_id from an image / JSON stem."""
    frame_id, _ = extract_frame_metadata_from_image_name(image_name)
    return frame_id


def resolve_frame_metadata(image_name=None, frame_info=None, timestamp_lookup=None):
    """Resolve frame_id / timestamp with frame_info overrides when available.

    Priority:
      1) frame_info.original_frame_id / frame_info.frame_id / frame_info.frameId
      2) frame_id parsed from image_name

    Timestamp priority:
      1) frame_info.timestamp
      2) timestamp parsed from image_name
      3) timestamp_lookup[frame_id]
    """
    parsed_frame_id, parsed_timestamp = extract_frame_metadata_from_image_name(image_name)

    frame_info_frame_id = None
    frame_info_timestamp = None
    if isinstance(frame_info, dict):
        for key in ('original_frame_id', 'frame_id', 'frameId'):
            frame_info_frame_id = safe_int(frame_info.get(key))
            if frame_info_frame_id is not None:
                break
        frame_info_timestamp = parse_numeric_value(frame_info.get('timestamp'))

    resolved_frame_id = frame_info_frame_id if frame_info_frame_id is not None else parsed_frame_id

    if frame_info_timestamp is not None:
        resolved_timestamp = frame_info_timestamp
    elif parsed_timestamp is not None:
        resolved_timestamp = parsed_timestamp
    elif timestamp_lookup and resolved_frame_id in timestamp_lookup:
        resolved_timestamp = timestamp_lookup[resolved_frame_id]
    else:
        resolved_timestamp = None

    return resolved_frame_id, resolved_timestamp


def find_timestamp_csv(search_path):
    """Find frame_timestamps_interp.csv near an input file or directory."""
    current = Path(search_path)
    if current.is_file():
        current = current.parent

    for candidate_dir in [current, *current.parents]:
        candidate = candidate_dir / TIMESTAMP_CSV_NAME
        if candidate.is_file():
            return candidate
    return None


def load_frame_timestamp_lookup(search_path, verbose=True):
    """Load a frame_id -> timestamp lookup from frame_timestamps_interp.csv."""
    csv_path = find_timestamp_csv(search_path)
    if csv_path is None:
        if verbose:
            print(f"Warning: {TIMESTAMP_CSV_NAME} not found near {search_path}")
        return {}

    with open(csv_path, 'r', encoding='utf-8-sig', newline='') as f:
        reader = csv.DictReader(f)
        if reader.fieldnames is None:
            if verbose:
                print(f"Warning: Empty timestamp CSV: {csv_path}")
            return {}

        field_map = {name.strip(): name for name in reader.fieldnames if name is not None}
        frame_id_key = field_map.get(TIMESTAMP_FRAME_ID_COLUMN)
        timestamp_key = field_map.get(TIMESTAMP_VALUE_COLUMN)
        if frame_id_key is None or timestamp_key is None:
            if verbose:
                print(
                    f"Warning: Missing required columns in {csv_path}. "
                    f"Expected '{TIMESTAMP_FRAME_ID_COLUMN}' and '{TIMESTAMP_VALUE_COLUMN}'."
                )
            return {}

        timestamp_lookup = {}
        for row in reader:
            frame_id = safe_int(row.get(frame_id_key))
            if frame_id is None:
                continue
            timestamp_value = parse_numeric_value(row.get(timestamp_key))
            if timestamp_value is None:
                continue
            timestamp_lookup[frame_id] = timestamp_value

    if verbose:
        print(f"Loaded {len(timestamp_lookup)} frame timestamp(s) from {csv_path}")
    return timestamp_lookup


def categorize_detection(det):
    """Map a detection to a coarse attribute category."""
    attribute = det.get('attribute')
    if isinstance(attribute, dict):
        task = normalize_class_name(attribute.get('task', ''))
        if task == 'vehicle':
            return 'vehicle'
        if task == 'pedestrian':
            return 'pedestrian'

    type_name = resolve_detection_type_name(det)
    if type_name and is_vehicle_name(type_name):
        return 'vehicle'
    if type_name and is_pedestrian_name(type_name):
        return 'pedestrian'

    class_id = safe_int(det.get('class_id'))
    if class_id in VEHICLE_CLASS_IDS:
        return 'vehicle'
    if class_id in PEDESTRIAN_CLASS_IDS:
        return 'pedestrian'
    return None


def class_id_is_vru(class_id):
    """Return whether a class id belongs to the pedestrian/rider association group."""
    if class_id in PEDESTRIAN_CLASS_IDS:
        return True
    cls_name = normalize_class_name(CLASS_NAME_LOOKUP.get(class_id, ''))
    return bool(cls_name and is_pedestrian_name(cls_name))


def resolve_detection_association_class_id(det, association_mode='class'):
    """Resolve the class key used only for temporal association."""
    class_id = det.get('class_id')
    if association_mode != 'vru':
        return class_id

    association_group = normalize_class_name(det.get('association_group') or det.get('association_type_name'))
    if association_group == VRU_ASSOCIATION_TYPE_NAME:
        return VRU_ASSOCIATION_CLASS_ID

    type_name = resolve_detection_type_name(det)
    if type_name and is_pedestrian_name(type_name):
        return VRU_ASSOCIATION_CLASS_ID
    if class_id_is_vru(class_id):
        return VRU_ASSOCIATION_CLASS_ID
    return class_id


def association_type_name_for_class(association_class_id):
    if association_class_id == VRU_ASSOCIATION_CLASS_ID:
        return VRU_ASSOCIATION_TYPE_NAME
    return CLASS_NAME_LOOKUP.get(association_class_id, str(association_class_id))


def resolve_tracker_class_ids(target_classes, association_mode='class'):
    """Map target classes to tracker buckets without changing output class ids."""
    tracker_class_ids = []
    seen = set()
    for cls_id in target_classes:
        assoc_id = VRU_ASSOCIATION_CLASS_ID if association_mode == 'vru' and class_id_is_vru(cls_id) else cls_id
        if assoc_id in seen:
            continue
        tracker_class_ids.append(assoc_id)
        seen.add(assoc_id)
    return tracker_class_ids


def compute_sub_cls(det):
    """Derive sub_cls from attribute outputs for vehicles and pedestrians."""
    category = categorize_detection(det)
    attribute = det.get('attribute')
    if category is None or not isinstance(attribute, dict):
        return DEFAULT_SUB_CLS

    attr_cls = safe_int(attribute.get('attr_cls'))
    if attr_cls is None:
        return DEFAULT_SUB_CLS

    is_fake = safe_int(attribute.get('is_fake')) or 0
    if category == 'vehicle':
        if is_fake == 1:
            return 26
        if attr_cls <= 11:
            return attr_cls
        if attr_cls == 23:
            return 12
        return attr_cls + 3

    if category == 'pedestrian':
        return attr_cls

    return DEFAULT_SUB_CLS


def enrich_detection_metadata(detection, image_name=None, timestamp_lookup=None, model_version=None, frame_info=None):
    """Inject tracking metadata and normalize optional heading debug info."""
    type_name = resolve_detection_type_name(detection)
    if type_name:
        detection['type_name'] = type_name

    detection['version'] = resolve_output_version(
        model_version=model_version,
        existing_version=detection.get('version'),
    )
    detection['sub_cls'] = compute_sub_cls(detection)

    frame_id, resolved_timestamp = resolve_frame_metadata(
        image_name=image_name,
        frame_info=frame_info,
        timestamp_lookup=timestamp_lookup,
    )
    if frame_id is not None:
        detection['frameId'] = str(frame_id)
    else:
        detection.setdefault('frameId', None)

    if resolved_timestamp is not None:
        detection['timestamp'] = resolved_timestamp
    elif detection.get('timestamp') is None:
        detection['timestamp'] = DEFAULT_TIMESTAMP
    else:
        detection.setdefault('timestamp', DEFAULT_TIMESTAMP)

    heading_debug = normalize_heading_debug(detection.get('heading_debug'))
    if heading_debug is not None:
        detection['heading_debug'] = heading_debug
    elif 'heading_debug' in detection:
        detection.pop('heading_debug', None)

    return detection


def enrich_predictions_data(predictions_data, timestamp_lookup=None, model_version=None):
    """Inject sub_cls / timestamp fields into existing frame predictions."""
    enriched = []
    for frame_data in predictions_data:
        image_name = frame_data.get('image_name')
        frame_info = frame_data.get('frame_info')
        detections = frame_data.get('detections', [])
        for det in detections:
            enrich_detection_metadata(
                det,
                image_name=image_name,
                timestamp_lookup=timestamp_lookup,
                model_version=model_version,
                frame_info=frame_info,
            )
        enriched.append(frame_data)
    return enriched


def limit_predictions_data(predictions_data, max_frames=None, verbose=True, source_name='input'):
    """Optionally keep only the first N frames in temporal order."""
    if max_frames is None:
        return predictions_data

    if max_frames <= 0:
        return []

    if len(predictions_data) <= max_frames:
        return predictions_data

    if verbose:
        print(
            f"Limiting {source_name} from {len(predictions_data)} frame(s) "
            f"to first {max_frames} frame(s)"
        )
    return predictions_data[:max_frames]


def compute_iou(box1, box2):
    """Compute IoU between two bounding boxes.
    
    Args:
        box1: [x1, y1, x2, y2]
        box2: [x1, y1, x2, y2]
        
    Returns:
        IoU value (0-1)
    """
    x1_1, y1_1, x2_1, y2_1 = box1
    x1_2, y1_2, x2_2, y2_2 = box2
    
    # Compute intersection area
    x1_inter = max(x1_1, x1_2)
    y1_inter = max(y1_1, y1_2)
    x2_inter = min(x2_1, x2_2)
    y2_inter = min(y2_1, y2_2)
    
    if x2_inter < x1_inter or y2_inter < y1_inter:
        return 0.0
    
    inter_area = (x2_inter - x1_inter) * (y2_inter - y1_inter)
    
    # Compute union area
    box1_area = (x2_1 - x1_1) * (y2_1 - y1_1)
    box2_area = (x2_2 - x1_2) * (y2_2 - y1_2)
    union_area = box1_area + box2_area - inter_area
    
    if union_area == 0:
        return 0.0
    
    return inter_area / union_area


def compute_intersection_over_min_area(box1, box2):
    """Measure how fully the smaller box is covered by the larger one."""
    x1_1, y1_1, x2_1, y2_1 = box1
    x1_2, y1_2, x2_2, y2_2 = box2

    x1_inter = max(x1_1, x1_2)
    y1_inter = max(y1_1, y1_2)
    x2_inter = min(x2_1, x2_2)
    y2_inter = min(y2_1, y2_2)

    if x2_inter < x1_inter or y2_inter < y1_inter:
        return 0.0

    inter_area = (x2_inter - x1_inter) * (y2_inter - y1_inter)
    box1_area = max(0.0, (x2_1 - x1_1) * (y2_1 - y1_1))
    box2_area = max(0.0, (x2_2 - x1_2) * (y2_2 - y1_2))
    min_area = min(box1_area, box2_area)

    if min_area <= 0.0:
        return 0.0

    return inter_area / min_area


def compute_distance(box1, box2):
    """Compute center distance between two bounding boxes.
    
    Args:
        box1: [x1, y1, x2, y2]
        box2: [x1, y1, x2, y2]
        
    Returns:
        Euclidean distance between centers
    """
    cx1 = (box1[0] + box1[2]) / 2
    cy1 = (box1[1] + box1[3]) / 2
    cx2 = (box2[0] + box2[2]) / 2
    cy2 = (box2[1] + box2[3]) / 2
    
    return np.sqrt((cx1 - cx2)**2 + (cy1 - cy2)**2)


def get_detection_confidence(det):
    """Return a sortable confidence score for one detection."""
    confidence = parse_numeric_value(det.get('confidence', det.get('score')))
    if isinstance(confidence, (int, float)):
        return float(confidence)
    return 0.0


def suppress_near_duplicate_detections(
    detections,
    overlap_threshold=DEFAULT_DUPLICATE_OVERLAP_THRESHOLD,
    class_key='class_id',
):
    """Remove same-association detections that almost completely overlap a higher-score box."""
    if overlap_threshold <= 0 or len(detections) < 2:
        return list(detections), 0

    ranked = sorted(
        enumerate(detections),
        key=lambda item: (-get_detection_confidence(item[1]), item[0]),
    )

    kept = []
    suppressed_count = 0
    for det_idx, det in ranked:
        bbox = det.get('bbox')
        class_id = det.get(class_key, det.get('class_id'))
        if not isinstance(bbox, (list, tuple)) or len(bbox) < 4:
            kept.append((det_idx, det))
            continue

        is_duplicate = False
        for _, kept_det in kept:
            if kept_det.get(class_key, kept_det.get('class_id')) != class_id:
                continue

            kept_bbox = kept_det.get('bbox')
            if not isinstance(kept_bbox, (list, tuple)) or len(kept_bbox) < 4:
                continue

            overlap = compute_intersection_over_min_area(kept_bbox, bbox)
            if overlap >= overlap_threshold:
                is_duplicate = True
                break

        if is_duplicate:
            suppressed_count += 1
            continue

        kept.append((det_idx, det))

    kept_indices = sorted(det_idx for det_idx, _ in kept)
    return [detections[det_idx] for det_idx in kept_indices], suppressed_count


def compute_3d_distance(det1, det2):
    """Compute 3D spatial distance between two detections.
    
    Args:
        det1, det2: Detection dicts with 'object_3d' containing 'location' [x, y, z]
        
    Returns:
        3D Euclidean distance in meters (or None if 3D info not available)
    """
    if 'object_3d' not in det1 or 'object_3d' not in det2:
        return None
    
    obj1 = det1['object_3d']
    obj2 = det2['object_3d']
    
    # Support both dict and list formats
    if isinstance(obj1, dict) and isinstance(obj2, dict):
        if 'location' not in obj1 or 'location' not in obj2:
            return None
        loc1 = np.array(obj1['location'][:3])  # [x, y, z]
        loc2 = np.array(obj2['location'][:3])
    elif isinstance(obj1, (list, tuple)) and isinstance(obj2, (list, tuple)):
        if len(obj1) < 3 or len(obj2) < 3:
            return None
        loc1 = np.array(obj1[:3])  # First 3 elements are [x, y, z]
        loc2 = np.array(obj2[:3])
    else:
        return None
    
    return np.linalg.norm(loc1 - loc2)


def compute_size_similarity(det1, det2):
    """Compute size similarity between two detections based on 3D dimensions.
    
    Args:
        det1, det2: Detection dicts with 'object_3d' containing 'dimensions' [l, h, w]
        
    Returns:
        Similarity score (0-1), higher is more similar (or None if 3D info not available)
    """
    if 'object_3d' not in det1 or 'object_3d' not in det2:
        return None
    
    obj1 = det1['object_3d']
    obj2 = det2['object_3d']
    
    # Support both dict and list formats
    if isinstance(obj1, dict) and isinstance(obj2, dict):
        if 'dimensions' not in obj1 or 'dimensions' not in obj2:
            return None
        dim1 = np.array(obj1['dimensions'][:3])  # [l, h, w]
        dim2 = np.array(obj2['dimensions'][:3])
    elif isinstance(obj1, (list, tuple)) and isinstance(obj2, (list, tuple)):
        if len(obj1) < 6 or len(obj2) < 6:
            return None
        dim1 = np.array(obj1[3:6])  # Elements [3,4,5] are [l, h, w]
        dim2 = np.array(obj2[3:6])
    else:
        return None
    
    # Compute relative difference for each dimension
    diff = np.abs(dim1 - dim2) / (np.maximum(dim1, dim2) + 1e-6)
    similarity = 1.0 - np.mean(diff)
    
    return max(0.0, similarity)


class SimpleTracker:
    """IoU-based object tracker with optional 3D distance matching."""
    
    def __init__(self, iou_threshold=0.3, max_age=5, min_hits=1, distance_threshold=100,
                 use_3d=False, max_3d_distance=10.0, w_iou=1.0, w_2d_dist=0.5, w_3d_dist=1.0, w_size=0.3,
                 shared_id_counter=None):
        """Initialize tracker.
        
        Args:
            iou_threshold: Minimum IoU to match detections with tracks
            max_age: Maximum frames to keep track without detection
            min_hits: Minimum detections before track is confirmed
            distance_threshold: Maximum center distance for matching (pixels)
            use_3d: Whether to use 3D distance for matching
            max_3d_distance: Maximum 3D distance for matching (meters)
            w_iou: Weight for IoU in cost function
            w_2d_dist: Weight for 2D distance in cost function
            w_3d_dist: Weight for 3D distance in cost function (if use_3d=True)
            w_size: Weight for size similarity in cost function (if use_3d=True)
            shared_id_counter: Optional mutable list [next_id] shared across trackers
                to ensure globally unique track IDs across classes.
        """
        self.iou_threshold = iou_threshold
        self.max_age = max_age
        self.min_hits = min_hits
        self.distance_threshold = distance_threshold
        self.use_3d = use_3d
        self.max_3d_distance = max_3d_distance
        self.w_iou = w_iou
        self.w_2d_dist = w_2d_dist
        self.w_3d_dist = w_3d_dist
        self.w_size = w_size
        
        # Use shared counter if provided, otherwise use a private one
        self._shared_id_counter = shared_id_counter if shared_id_counter is not None else [1]
        self.tracks = {}  # track_id -> track_info
        
    def update(self, detections):
        """Update tracks with new detections.
        
        Args:
            detections: List of detection dicts with 'bbox', 'confidence', 'class_id', etc.
            
        Returns:
            List of detections with added 'track_id' field
        """
        if not detections:
            # Age all tracks
            for track_id in list(self.tracks.keys()):
                self.tracks[track_id]['age'] += 1
                if self.tracks[track_id]['age'] > self.max_age:
                    del self.tracks[track_id]
            return []
        
        # Match detections to existing tracks
        matched_detections = []
        unmatched_detections = list(range(len(detections)))
        unmatched_tracks = list(self.tracks.keys())
        
        # Build cost matrix (IoU + distance)
        if unmatched_tracks and unmatched_detections:
            cost_matrix = np.zeros((len(unmatched_tracks), len(unmatched_detections)))
            
            for i, track_id in enumerate(unmatched_tracks):
                track_box = self.tracks[track_id]['bbox']
                track_class = self.tracks[track_id].get('association_class_id', self.tracks[track_id]['class_id'])
                track_obj3d = self.tracks[track_id].get('object_3d', None)
                
                for j, det_idx in enumerate(unmatched_detections):
                    det_box = detections[det_idx]['bbox']
                    det_class = detections[det_idx].get('association_class_id', detections[det_idx]['class_id'])
                    
                    # Only match same class
                    if track_class != det_class:
                        cost_matrix[i, j] = 0.0
                        continue
                    
                    # Compute 2D IoU
                    iou = compute_iou(track_box, det_box)
                    
                    # Compute 2D center distance
                    distance_2d = compute_distance(track_box, det_box)
                    
                    # Initialize cost components
                    cost = 0.0
                    
                    # IoU component (primary)
                    if iou > self.iou_threshold:
                        cost += self.w_iou * iou
                    else:
                        cost_matrix[i, j] = 0.0
                        continue
                    
                    # 2D distance component
                    if distance_2d < self.distance_threshold:
                        cost += self.w_2d_dist * (1.0 - distance_2d / self.distance_threshold)
                    else:
                        cost_matrix[i, j] = 0.0
                        continue
                    
                    # 3D components (if enabled and available)
                    if self.use_3d and track_obj3d is not None:
                        # Create temporary track detection dict for 3D functions
                        track_det = {'object_3d': track_obj3d}
                        det_dict = detections[det_idx]
                        
                        # 3D distance component
                        dist_3d = compute_3d_distance(track_det, det_dict)
                        if dist_3d is not None and dist_3d < self.max_3d_distance:
                            cost += self.w_3d_dist * (1.0 - dist_3d / self.max_3d_distance)
                        elif dist_3d is not None and dist_3d >= self.max_3d_distance:
                            # 3D distance too large, reject match
                            cost_matrix[i, j] = 0.0
                            continue
                        
                        # Size similarity component
                        size_sim = compute_size_similarity(track_det, det_dict)
                        if size_sim is not None:
                            cost += self.w_size * size_sim
                    
                    cost_matrix[i, j] = cost
            
            # Hungarian algorithm matching (optimal assignment)
            if SCIPY_AVAILABLE and cost_matrix.max() > 0:
                # Use Hungarian algorithm for optimal matching
                row_ind, col_ind = linear_sum_assignment(-cost_matrix)  # Maximize (negative for minimize)
                matches = [(unmatched_tracks[i], unmatched_detections[j]) 
                          for i, j in zip(row_ind, col_ind) if cost_matrix[i, j] > 0]
            else:
                # Fallback to greedy matching if scipy not available
                matches = []
                cost_copy = cost_matrix.copy()
                while True:
                    max_val = cost_copy.max()
                    if max_val == 0:
                        break
                    
                    i, j = np.unravel_index(cost_copy.argmax(), cost_copy.shape)
                    track_id = unmatched_tracks[i]
                    det_idx = unmatched_detections[j]
                    
                    matches.append((track_id, det_idx))
                    
                    # Mark as matched
                    cost_copy[i, :] = 0
                    cost_copy[:, j] = 0
            
            # Update matched tracks
            for track_id, det_idx in matches:
                self.tracks[track_id]['bbox'] = detections[det_idx]['bbox']
                self.tracks[track_id]['age'] = 0
                self.tracks[track_id]['hits'] += 1
                
                # Update 3D info if available
                if 'object_3d' in detections[det_idx]:
                    self.tracks[track_id]['object_3d'] = detections[det_idx]['object_3d']
                
                # Add track_id to detection
                detections[det_idx]['track_id'] = track_id
                matched_detections.append(det_idx)
                
                unmatched_tracks.remove(track_id)
                unmatched_detections.remove(det_idx)
        
        # Age unmatched tracks
        for track_id in unmatched_tracks:
            self.tracks[track_id]['age'] += 1
            if self.tracks[track_id]['age'] > self.max_age:
                del self.tracks[track_id]
        
        # Create new tracks for unmatched detections
        for det_idx in unmatched_detections:
            track_id = self._shared_id_counter[0]
            self._shared_id_counter[0] += 1
            
            track_info = {
                'bbox': detections[det_idx]['bbox'],
                'class_id': detections[det_idx]['class_id'],
                'association_class_id': detections[det_idx].get('association_class_id', detections[det_idx]['class_id']),
                'age': 0,
                'hits': 1
            }
            
            # Store 3D info if available
            if 'object_3d' in detections[det_idx]:
                track_info['object_3d'] = detections[det_idx]['object_3d']
            
            self.tracks[track_id] = track_info
            
            # Add track_id to detection
            detections[det_idx]['track_id'] = track_id
            matched_detections.append(det_idx)
        
        # Return detections with track IDs (only confirmed tracks)
        result = []
        for det_idx in matched_detections:
            det = detections[det_idx]
            track_id = det['track_id']
            
            # Only include if track has enough hits
            if self.tracks[track_id]['hits'] >= self.min_hits:
                result.append(det)
        
        return result


def parse_det_format(det_dict, image_name=None, timestamp_lookup=None, model_version=None, frame_info=None):
    """Parse a single-frame detection result in det_format into internal frame data format.

    Args:
        det_dict: Either a flat dict of detections keyed by string index (as in
                  det_format.json reference), OR a dict with a top-level
                  "detections" key wrapping that flat dict (actual merge_json format).
                  Each detection entry has: type, score, roi_id, box2d, xyzlhwyaw,
                  face_cls, cut_cls.  type_name and cut_cls_name are optional.
        image_name: Optional frame name / image stem.
        frame_info: Optional frame metadata dict. When it contains
                    original_frame_id/frame_id/frameId, that value is preferred
                    over any frame id parsed from image_name.

    Returns:
        Frame data dict with keys 'image_name' and 'detections', compatible with
        the list expected by track_objects().
    """
    # Support both flat format and wrapped {"detections": {...}} format
    if 'detections' in det_dict and isinstance(det_dict['detections'], dict):
        raw_detections = det_dict['detections']
    else:
        raw_detections = det_dict

    face_map = {
        'front': 'kMonocular3DFront',
        'tail': 'kMonocular3DRear',
        'back': 'kMonocular3DRear',
        'left': 'kMonocular3DLeft',
        'right': 'kMonocular3DRight',
        'center': 'kMonocular3DCenter',
        'none': 'kMonocular3DCenter'
    }

    resolved_frame_id, resolved_timestamp = resolve_frame_metadata(
        image_name=image_name,
        frame_info=frame_info,
        timestamp_lookup=timestamp_lookup,
    )

    detections = []
    for det in raw_detections.values():
        class_id = int(det['type'])
        bbox = [float(v) for v in det['box2d']]
        score = float(det['score'])

        # Parse 3D info from xyzlhwyaw; sentinel value -1 means no 3D
        xyzlhwyaw_raw = det.get('xyzlhwyaw', [])
        object_3d = None
        if xyzlhwyaw_raw and float(xyzlhwyaw_raw[0]) != -1:
            object_3d = [float(v) for v in xyzlhwyaw_raw]  # [x, y, z, l, h, w, yaw]

        # Parse 3D info from xyzlhwyaw_ego; sentinel value -1 means no 3D
        xyzlhwyaw_ego_raw = det.get('xyzlhwyaw_ego', [])
        object_3d_ego = None
        if xyzlhwyaw_ego_raw and float(xyzlhwyaw_ego_raw[0]) != -1:
            object_3d_ego = [float(v) for v in xyzlhwyaw_ego_raw]  # [x, y, z, l, h, w, yaw]

        # Parse 3D box center (separate from face-center xyzlhwyaw)
        box_center_xyz_raw = det.get('box_center_xyz', [])
        box_center_3d = None
        if box_center_xyz_raw and len(box_center_xyz_raw) >= 3:
            box_center_3d = [float(v) for v in box_center_xyz_raw[:3]]  # [x, y, z]

        box_center_xyz_ego_raw = det.get('box_center_xyz_ego', [])
        box_center_3d_ego = None
        if box_center_xyz_ego_raw and len(box_center_xyz_ego_raw) >= 3:
            box_center_3d_ego = [float(v) for v in box_center_xyz_ego_raw[:3]]  # [x, y, z]

        detection = {
            'bbox': bbox,
            'confidence': score,
            'class_id': class_id,
            'type_name': det.get('type_name', ''),
            'face_cls': det.get('face_cls', 'none'),
            'cut_cls': int(det.get('cut_cls', -1)),
            'cut_cls_name': det.get('cut_cls_name', 'none'),
            'frameId': str(resolved_frame_id) if resolved_frame_id is not None else None,
            'version': resolve_output_version(model_version=model_version, existing_version=det.get('version')),
            'timestamp': resolved_timestamp if resolved_timestamp is not None else DEFAULT_TIMESTAMP,
            'roi_id': int(det.get('roi_id', -1)),
        }

        detection['anchor'] = face_map.get(detection['face_cls'], 'kMonocular3DCenter')
        if 'attribute' in det:
            detection['attribute'] = det.get('attribute')
        if 'heading_debug' in det:
            detection['heading_debug'] = det.get('heading_debug')
        for field_name in (
            'difficulty_logit',
            'difficulty_prob',
            'difficulty_label',
            'difficulty_name',
        ):
            if field_name in det:
                detection[field_name] = det.get(field_name)
        for field_name in (
            'association_group',
            'association_class_id',
            'association_type_name',
            'raw_cls_id',
            'raw_cls_name',
        ):
            if field_name in det:
                detection[field_name] = det.get(field_name)

        if object_3d is not None:
            detection['object_3d'] = object_3d
        if box_center_3d is not None:
            detection['box_center_3d'] = box_center_3d
        if object_3d_ego is not None:
            detection['object_3d_ego'] = object_3d_ego
        if box_center_3d_ego is not None:
            detection['box_center_3d_ego'] = box_center_3d_ego
        enrich_detection_metadata(
            detection,
            image_name=image_name,
            timestamp_lookup=timestamp_lookup,
            model_version=model_version,
            frame_info=frame_info,
        )
        detections.append(detection)

    return {
        'image_name': image_name,
        'detections': detections,
    }


def load_predictions_from_dir(input_dir, pattern='*.json', verbose=True,
                              timestamp_lookup=None, max_frames=None, model_version=None):
    """Load per-frame detection results from a directory of det_format JSON files.

    Files are sorted lexicographically (i.e. by filename) to preserve temporal order.

    Args:
        input_dir: Directory containing per-frame JSON files in det_format.
        pattern: Glob pattern used to match JSON files (default: '*.json').

    Returns:
        List of frame data dicts compatible with track_objects().
    """
    input_dir = Path(input_dir)
    json_files = sorted(input_dir.glob(pattern))

    if not json_files:
        if verbose:
            print(f"Warning: No files matching '{pattern}' found in {input_dir}")
        return []

    if verbose:
        print(f"Found {len(json_files)} frame file(s) in {input_dir}")

    if max_frames is not None and max_frames > 0 and len(json_files) > max_frames:
        if verbose:
            print(f"Limiting {input_dir} to first {max_frames} frame file(s)")
        json_files = json_files[:max_frames]

    if timestamp_lookup is None:
        timestamp_lookup = load_frame_timestamp_lookup(input_dir, verbose=verbose)

    predictions_data = []
    for json_file in json_files:
        with open(json_file, 'r', encoding='utf-8') as f:
            det_dict = json.load(f)
        frame_data = parse_det_format(
            det_dict,
            image_name=json_file.stem,
            timestamp_lookup=timestamp_lookup,
            model_version=model_version,
        )
        predictions_data.append(frame_data)

    return predictions_data


def track_objects(predictions_data, target_classes=None,
                  iou_threshold=0.3, max_age=5, min_hits=1, distance_threshold=100,
                  use_3d=False, max_3d_distance=10.0, verbose=True,
                  duplicate_overlap_threshold=DEFAULT_DUPLICATE_OVERLAP_THRESHOLD,
                  association_mode='class'):
    """Track objects across frames.
    
    Args:
        predictions_data: List of frame predictions from predictions.json
        target_classes: List of class IDs to track. Defaults to Ground3D
            face_3d_classes ∪ complete_3d_classes.
        iou_threshold: Minimum IoU for matching
        max_age: Maximum frames without detection
        min_hits: Minimum detections before confirmed
        distance_threshold: Maximum center distance for matching (pixels)
        use_3d: Whether to use 3D distance for matching
        max_3d_distance: Maximum 3D distance for matching (meters)
        duplicate_overlap_threshold: Suppress same-class detections when the
            smaller box is almost fully covered by a higher-score box.
        
    Returns:
        List of frame predictions with track IDs added
    """
    if target_classes is None:
        target_classes = list(TRACKED_CLASS_IDS)
    target_classes = [int(cls_id) for cls_id in target_classes]
    target_class_set = set(target_classes)
    tracker_class_ids = resolve_tracker_class_ids(target_classes, association_mode=association_mode)

    # Shared counter ensures track IDs are unique across all classes
    shared_id_counter = [1]
    trackers = {cls_id: SimpleTracker(iou_threshold, max_age, min_hits, distance_threshold,
                                     use_3d=use_3d, max_3d_distance=max_3d_distance,
                                     shared_id_counter=shared_id_counter) 
                for cls_id in tracker_class_ids}
    
    tracking_results = []
    
    for frame_idx, frame_data in enumerate(predictions_data):
        if verbose:
            print(f"Processing frame {frame_idx + 1}/{len(predictions_data)}: {frame_data.get('image_name', 'unknown')}")
        
        # Group detections by class
        detections_by_class = {cls_id: [] for cls_id in tracker_class_ids}
        non_tracked_detections = []
        for det in frame_data.get('detections', []):
            cls_id = safe_int(det.get('class_id'))
            if cls_id in target_class_set:
                det['class_id'] = cls_id
                association_class_id = resolve_detection_association_class_id(det, association_mode=association_mode)
                det['association_class_id'] = association_class_id
                det['association_type_name'] = association_type_name_for_class(association_class_id)
                if association_class_id == VRU_ASSOCIATION_CLASS_ID:
                    det['association_group'] = VRU_ASSOCIATION_TYPE_NAME
                detections_by_class.setdefault(association_class_id, []).append(det)
            else:
                non_tracked_detections.append(det)
        
        # Track each class separately
        tracked_detections = []
        frame_duplicate_count = 0
        for cls_id in tracker_class_ids:
            if cls_id in detections_by_class:
                deduped_detections, suppressed_count = suppress_near_duplicate_detections(
                    detections_by_class[cls_id],
                    overlap_threshold=duplicate_overlap_threshold,
                    class_key='association_class_id' if association_mode == 'vru' else 'class_id',
                )
                frame_duplicate_count += suppressed_count
                tracked = trackers[cls_id].update(deduped_detections)
                tracked_detections.extend(tracked)

        if verbose and frame_duplicate_count > 0:
            print(
                f"  Suppressed {frame_duplicate_count} near-duplicate detection(s) "
                f"before association"
            )
        
        # Append non-tracked classes directly (no tracking needed)
        tracked_detections.extend(non_tracked_detections)
        
        # Create output frame data
        result_frame = {
            'image_name': frame_data.get('image_name'),
            'detections': tracked_detections
        }
        
        # Copy frame_info if present
        if 'frame_info' in frame_data:
            result_frame['frame_info'] = frame_data['frame_info']
        
        # Add tracking statistics
        result_frame['tracking_stats'] = {
            'total_tracks': sum(len(tracker.tracks) for tracker in trackers.values()),
            'active_tracks_by_class': {cls_id: len([t for t in tracker.tracks.values() if t['age'] == 0])
                                      for cls_id, tracker in trackers.items()}
        }
        
        tracking_results.append(result_frame)
    
    return tracking_results


def save_tracking_results(tracking_results, output_path):
    """Save tracking results to a JSON file."""
    output_path = Path(output_path)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(tracking_results, f, indent=2, ensure_ascii=False)


def count_unique_tracks(tracking_results):
    """Count unique track ids in tracking output."""
    total_tracks = set()
    for frame in tracking_results:
        for det in frame.get('detections', []):
            if 'track_id' in det:
                total_tracks.add(det['track_id'])
    return len(total_tracks)


def run_tracking_job(input_dir, output_path, file_pattern, classes, iou_threshold,
                     max_age, min_hits, distance_threshold, use_3d,
                     max_3d_distance, verbose=True, timestamp_lookup=None,
                     max_frames=None, model_version=None, association_mode='class'):
    """Run one tracking job from a directory of per-frame JSON files."""
    predictions_data = load_predictions_from_dir(
        input_dir,
        file_pattern,
        verbose=verbose,
        timestamp_lookup=timestamp_lookup,
        max_frames=max_frames,
        model_version=model_version,
    )
    if not predictions_data:
        return {
            'ok': False,
            'input_dir': str(input_dir),
            'output_path': str(output_path),
            'frames': 0,
            'unique_tracks': 0,
            'reason': 'no_frames_loaded',
        }

    tracking_results = track_objects(
        predictions_data,
        target_classes=classes,
        iou_threshold=iou_threshold,
        max_age=max_age,
        min_hits=min_hits,
        distance_threshold=distance_threshold,
        use_3d=use_3d,
        max_3d_distance=max_3d_distance,
        verbose=verbose,
        association_mode=association_mode,
    )
    save_tracking_results(tracking_results, output_path)
    return {
        'ok': True,
        'input_dir': str(input_dir),
        'output_path': str(output_path),
        'frames': len(predictions_data),
        'unique_tracks': count_unique_tracks(tracking_results),
    }


def find_case_dirs(results_root, model_name):
    """Scan the model directory once and return case dirs containing tracking inputs."""
    search_root = Path(results_root) / model_name
    if not search_root.exists():
        print(f"Error: model directory does not exist: {search_root}")
        return []

    case_dirs = set()
    target_dirs = {'roi0', 'roi1', 'merge_json'}
    for current_dir, dirnames, _ in os.walk(search_root):
        if target_dirs.intersection(dirnames):
            case_dirs.add(Path(current_dir))

    return sorted(case_dirs)


def process_case(case_dir, file_pattern, classes, iou_threshold, max_age, min_hits,
                 distance_threshold, use_3d, max_3d_distance, merge_output_name,
                 max_frames=None, model_version=None, association_mode='class'):
    """Track all sources for one case and write the merged output."""
    case_dir = Path(case_dir)
    timestamp_lookup = load_frame_timestamp_lookup(case_dir, verbose=False)
    source_specs = [
        ('roi0', 'roi0.json'),
        ('roi1', 'roi1.json'),
        ('merge_json', 'merge.json'),
    ]

    tracking_outputs = {}
    source_summaries = []
    missing_sources = []

    for input_name, output_name in source_specs:
        input_dir = case_dir / input_name
        output_path = case_dir / output_name
        if not input_dir.is_dir():
            missing_sources.append(input_name)
            continue

        result = run_tracking_job(
            input_dir=input_dir,
            output_path=output_path,
            file_pattern=file_pattern,
            classes=classes,
            iou_threshold=iou_threshold,
            max_age=max_age,
            min_hits=min_hits,
            distance_threshold=distance_threshold,
            use_3d=use_3d,
            max_3d_distance=max_3d_distance,
            verbose=False,
            timestamp_lookup=timestamp_lookup,
            max_frames=max_frames,
            model_version=model_version,
            association_mode=association_mode,
        )
        result['source'] = input_name
        source_summaries.append(result)
        if result['ok']:
            tracking_outputs[input_name] = output_path

    merge_ok = merge_case(
        roi0_path=tracking_outputs.get('roi0'),
        roi1_path=tracking_outputs.get('roi1'),
        merge_path=tracking_outputs.get('merge_json'),
        output_path=case_dir / merge_output_name,
        source_names=('roi0', 'roi1', 'merge'),
        verbose=False,
    )

    ok_sources = [item['source'] for item in source_summaries if item['ok']]
    failed_sources = [item['source'] for item in source_summaries if not item['ok']]
    return {
        'case_dir': str(case_dir),
        'ok': bool(ok_sources) and merge_ok and not failed_sources,
        'tracked_sources': ok_sources,
        'failed_sources': failed_sources,
        'missing_sources': missing_sources,
        'merge_ok': merge_ok,
        'source_summaries': source_summaries,
        'merge_output': str(case_dir / merge_output_name),
    }


def print_case_summary(result):
    """Print a clean per-case summary after batch processing."""
    print("")
    print("==========================================")
    print(f"Case   : {result['case_dir']}")
    print(f"Merge  : {'ok' if result['merge_ok'] else 'failed'}")
    print(f"Output : {result['merge_output']}")
    if result['tracked_sources']:
        print(f"Tracked: {', '.join(result['tracked_sources'])}")
    if result['failed_sources']:
        print(f"Failed : {', '.join(result['failed_sources'])}")
    if result['missing_sources']:
        print(f"Missing: {', '.join(result['missing_sources'])}")
    for item in result['source_summaries']:
        status = 'ok' if item['ok'] else item.get('reason', 'failed')
        print(
            f"  - {item['source']}: {status}, frames={item['frames']}, "
            f"tracks={item['unique_tracks']}"
        )
    print("==========================================")


def track_cases_in_batch(results_root, model_name, file_pattern, classes,
                         iou_threshold, max_age, min_hits, distance_threshold,
                         use_3d, max_3d_distance, merge_output_name,
                         num_workers, max_frames=None, model_version=None,
                         association_mode='class'):
    """Run tracking and merge end-to-end for every case under one model directory."""
    case_dirs = find_case_dirs(results_root, model_name)
    if not case_dirs:
        print("Error: No case directories containing roi0/roi1/merge_json were found.")
        return False

    print(f"Found {len(case_dirs)} case(s) under {Path(results_root) / model_name}")
    print(f"Using {num_workers} worker(s)")

    results = []
    if num_workers <= 1:
        for case_dir in case_dirs:
            result = process_case(
                case_dir=case_dir,
                file_pattern=file_pattern,
                classes=classes,
                iou_threshold=iou_threshold,
                max_age=max_age,
                min_hits=min_hits,
                distance_threshold=distance_threshold,
                use_3d=use_3d,
                max_3d_distance=max_3d_distance,
                merge_output_name=merge_output_name,
                max_frames=max_frames,
                model_version=model_version,
                association_mode=association_mode,
            )
            results.append(result)
            print_case_summary(result)
    else:
        with ProcessPoolExecutor(max_workers=num_workers) as executor:
            future_to_case = {
                executor.submit(
                    process_case,
                    case_dir,
                    file_pattern,
                    classes,
                    iou_threshold,
                    max_age,
                    min_hits,
                    distance_threshold,
                    use_3d,
                    max_3d_distance,
                    merge_output_name,
                    max_frames,
                    model_version,
                    association_mode,
                ): case_dir
                for case_dir in case_dirs
            }
            for future in as_completed(future_to_case):
                result = future.result()
                results.append(result)
                print_case_summary(result)

    succeeded = sum(1 for item in results if item['ok'])
    failed = len(results) - succeeded
    print("")
    print("Batch tracking complete")
    print(f"  succeeded_cases: {succeeded}")
    print(f"  failed_cases: {failed}")
    return failed == 0


def main():
    parser = argparse.ArgumentParser(description='Track objects from predictions.json')
    parser.add_argument('--input', type=str, default=None,
                        help='Input predictions.json file (multi-frame list format) '
                             'or a single-frame det_format JSON file')
    parser.add_argument('--input-dir', type=str, default=None,
                        help='Directory containing per-frame JSON files in det_format '
                             '(used when inference results are saved as individual files)')
    parser.add_argument('--results-root', type=str, default=None,
                        help='Batch mode root directory that contains per-model case results')
    parser.add_argument('--model-name', type=str, default=None,
                        help='Batch mode model directory name under results-root')
    parser.add_argument('--file-pattern', type=str, default='*.json',
                        help='Glob pattern for JSON files inside --input-dir (default: %(default)s)')
    parser.add_argument('--output', type=str, default=None, help='Output tracking.json file path (default: same dir as input)')
    parser.add_argument('--merge-output-name', type=str, default='combined_tracking.json',
                        help='Batch mode output filename written inside each case directory')
    parser.add_argument('--num-workers', type=int, default=max(1, min(8, os.cpu_count() or 1)),
                        help='Batch mode worker count for case-level parallelism (default: %(default)s)')
    parser.add_argument(
        '--classes',
        type=int,
        nargs='+',
        default=None,
        help='Class IDs to track (default: Ground3D face_3d_classes + complete_3d_classes)',
    )
    parser.add_argument('--iou-threshold', type=float, default=0.3, help='IoU threshold for matching (default: 0.3)')
    parser.add_argument('--max-age', type=int, default=5, help='Maximum frames without detection (default: 5)')
    parser.add_argument('--min-hits', type=int, default=1, help='Minimum detections before confirmed (default: 1)')
    parser.add_argument('--distance-threshold', type=float, default=100, help='Maximum center distance for matching in pixels (default: 100)')
    parser.add_argument('--max-frames', type=int, default=None,
                        help='Only process the first N frames in temporal order (default: all frames)')
    parser.add_argument('--model-version', type=str, default=None,
                        help=f"Version string written into detections (default: keep existing or '{DEFAULT_MODEL_VERSION}' when missing)")
    parser.add_argument('--use-3d', action='store_true', help='Use 3D distance for matching (requires object_3d in detections)')
    parser.add_argument('--max-3d-distance', type=float, default=10.0, help='Maximum 3D distance for matching in meters (default: 10.0)')
    parser.add_argument(
        '--association-mode',
        type=str,
        choices=('class', 'vru'),
        default='class',
        help='Temporal association class mapping. "vru" tracks pedestrian/rider classes as one VRU group while preserving raw classes.',
    )
    
    args = parser.parse_args()
    if args.classes is None:
        args.classes = list(TRACKED_CLASS_IDS)
    if args.max_frames is not None and args.max_frames <= 0:
        parser.error('--max-frames must be a positive integer')

    if args.results_root or args.model_name:
        if not args.results_root or not args.model_name:
            parser.error('--results-root and --model-name must be provided together in batch mode')
        track_cases_in_batch(
            results_root=args.results_root,
            model_name=args.model_name,
            file_pattern=args.file_pattern,
            classes=args.classes,
            iou_threshold=args.iou_threshold,
            max_age=args.max_age,
            min_hits=args.min_hits,
            distance_threshold=args.distance_threshold,
            use_3d=args.use_3d,
            max_3d_distance=args.max_3d_distance,
            merge_output_name=args.merge_output_name,
            num_workers=args.num_workers,
            max_frames=args.max_frames,
            model_version=args.model_version,
            association_mode=args.association_mode,
        )
        return
    
    # Read input predictions
    if args.input_dir is not None:
        # New format: directory of per-frame det_format JSON files
        predictions_data = load_predictions_from_dir(
            args.input_dir,
            args.file_pattern,
            max_frames=args.max_frames,
            model_version=args.model_version,
        )
        if not predictions_data:
            print("Error: No frame files loaded. Check --input-dir and --file-pattern.")
            return
        input_path = Path(args.input_dir)  # used for default output path
    elif args.input is not None:
        input_path = Path(args.input)
        if not input_path.exists():
            print(f"Error: Input file not found: {input_path}")
            return

        print(f"Loading predictions from {input_path}")
        timestamp_lookup = load_frame_timestamp_lookup(input_path, verbose=False)
        with open(input_path, 'r', encoding='utf-8') as f:
            raw_data = json.load(f)

        # Auto-detect format: dict → single-frame det_format; list → multi-frame predictions.json
        if isinstance(raw_data, dict):
            print("Detected single-frame det_format JSON. Wrapping into one-frame sequence.")
            predictions_data = [
                parse_det_format(
                    raw_data,
                    image_name=input_path.stem,
                    timestamp_lookup=timestamp_lookup,
                    model_version=args.model_version,
                )
            ]
        else:
            predictions_data = raw_data
            predictions_data = enrich_predictions_data(
                predictions_data,
                timestamp_lookup=timestamp_lookup,
                model_version=args.model_version,
            )
            predictions_data = limit_predictions_data(
                predictions_data,
                max_frames=args.max_frames,
                source_name=str(input_path),
            )
    else:
        print("Error: Either --input or --input-dir must be specified.")
        return

    print(f"Loaded {len(predictions_data)} frame(s)")
    
    # Perform tracking
    print(f"Tracking objects of classes: {args.classes}")
    resolved_class_names = [CLASS_NAME_LOOKUP.get(int(cls_id), str(cls_id)) for cls_id in args.classes]
    print(f"Tracking class names: {', '.join(resolved_class_names)}")
    print(f"Tracking class config source: {TRACKING_METADATA['source']}")
    if args.max_frames is not None:
        print(f"Tracking only first {args.max_frames} frame(s)")
    if args.model_version is not None:
        print(f"Tracking model version: {args.model_version}")
    if args.use_3d:
        print(f"Using 3D distance matching (max 3D distance: {args.max_3d_distance}m)")
    print(f"Association mode: {args.association_mode}")
    tracking_results = track_objects(
        predictions_data,
        target_classes=args.classes,
        iou_threshold=args.iou_threshold,
        max_age=args.max_age,
        min_hits=args.min_hits,
        distance_threshold=args.distance_threshold,
        use_3d=args.use_3d,
        max_3d_distance=args.max_3d_distance,
        association_mode=args.association_mode,
    )
    
    # Determine output path
    if args.output is None:
        out_base = input_path if input_path.is_dir() else input_path.parent
        output_path = out_base / 'tracking.json'
    else:
        output_path = Path(args.output)
    
    # Save results
    print(f"Saving tracking results to {output_path}")
    save_tracking_results(tracking_results, output_path)
    
    # Print statistics
    total_track_count = count_unique_tracks(tracking_results)

    print(f"\nTracking completed!")
    print(f"Total unique tracks: {total_track_count}")
    print(f"Output saved to: {output_path}")


if __name__ == '__main__':
    main()
