"""
Merge multiple per-source tracking JSON files into a single combined JSON.

Each source (roi0, roi1, merge) is the output of track_objects.py.  The script
aligns frames by image_name, adds a "source" field to every detection, and
writes a unified JSON where each frame contains detections from all sources.

Output format (list of frame dicts):
  [
    {
      "image_name": "1741824631_0",
      "detections": [
        { "lane_assignment": 0, "track_id": 1, "class_id": 0, "bbox": [...], ... },
        { "lane_assignment": 1, "track_id": 1, ... },
        { "lane_assignment": 2, "track_id": 2, ... },
        ...
      ],
      "tracking_stats": {
        "roi0":  { "total_tracks": 3, "active_tracks_by_class": {...} },
        "roi1":  { ... },
        "merge": { ... }
      }
    },
    ...
  ]

Notes:
  - track_id values are offset per source so they are globally unique:
    roi0 keeps original IDs (0–9999), roi1 adds 10000, merge adds 20000.
  - lane_assignment values: 0=roi0, 1=roi1, 2=merge (positional index in source_map).
  - Missing sources are silently skipped (partial merge is valid).
  - Frames present in some sources but not others are still included, with
    detections only from the sources that have that frame.

Usage:
    python merge_tracking_results.py \
        --roi0  {case_dir}/roi0.json \
        --roi1  {case_dir}/roi1.json \
        --merge {case_dir}/merge.json \
        --output {case_dir}/combined_tracking.json

    # Batch mode: scan all cases under results_root/model_name
    python merge_tracking_results.py \
        --results-root /data1/dongying/Mono3d/G1M3/cases_regular \
        --model-name model_20260228

    # Custom source names (optional):
    python merge_tracking_results.py \
        --roi0  roi0.json --roi1 roi1.json --merge merge.json \
        --source-names roi0_track roi1_track merge_track \
        --output combined_tracking.json
"""

import argparse
import json
import re
from pathlib import Path


def normalize_image_name(name):
    """Normalize image_name to a canonical form for cross-source frame alignment.

    merge_json filenames are saved as ``{saved_count:06d}_{img_name}_merged.json``
    while roi0/roi1 use plain ``{img_name}.json``.  This function strips both
    the leading 6-digit counter prefix and the trailing ``_merged`` suffix so
    all three sources resolve to the same key.

    Examples::

        '000001_1741824631_0_merged' -> '1741824631_0'
        '1741824631_0'               -> '1741824631_0'  (unchanged)
    """
    name = re.sub(r'^\d{6}_', '', name)  # strip leading counter, e.g. '000001_'
    name = re.sub(r'_merged$', '', name)  # strip trailing '_merged'
    return name


def load_tracking_json(path, verbose=True):
    """Load a tracking JSON file produced by track_objects.py.

    Args:
        path: Path to the JSON file (list of frame dicts).

    Returns:
        List of frame dicts, or empty list if file does not exist / is invalid.
    """
    p = Path(path)
    if not p.exists():
        if verbose:
            print(f"Warning: tracking file not found, skipping: {p}")
        return []
    with open(p, 'r', encoding='utf-8') as f:
        data = json.load(f)
    if not isinstance(data, list):
        if verbose:
            print(f"Warning: expected a list in {p}, got {type(data).__name__}. Skipping.")
        return []
    if verbose:
        print(f"Loaded {len(data)} frame(s) from {p}")
    return data


def build_source_map(roi0_path, roi1_path, merge_path, source_names, extra_sources=None, verbose=True):
    """Load all available tracking inputs for a single case."""
    roi0_name, roi1_name, merge_name = source_names

    source_map = []
    if roi0_path is not None:
        frames = load_tracking_json(roi0_path, verbose=verbose)
        if frames:
            source_map.append((roi0_name, frames))
    if roi1_path is not None:
        frames = load_tracking_json(roi1_path, verbose=verbose)
        if frames:
            source_map.append((roi1_name, frames))
    if merge_path is not None:
        frames = load_tracking_json(merge_path, verbose=verbose)
        if frames:
            source_map.append((merge_name, frames))
    for source_name, source_path in extra_sources or []:
        frames = load_tracking_json(source_path, verbose=verbose)
        if frames:
            source_map.append((source_name, frames))

    return source_map


# Offset added to track_id per source so IDs are globally unique:
#   source_idx 0 (roi0)       -> no offset
#   source_idx 1 (roi1)       -> +10000
#   source_idx 2 (merge)      -> +20000
#   source_idx 3+ (extra src) -> keep incrementing by +10000
TRACK_ID_OFFSET_PER_SOURCE = 10000


def merge_tracking_results(source_map):
    """Merge tracking results from multiple sources into a single per-frame list.

    Args:
        source_map: Ordered dict / list of (source_name, frames_list) pairs.
                    Each frames_list is the output of track_objects.py.

    Returns:
        List of merged frame dicts, one per unique image_name found across all
        sources.  Frames are ordered by the first source that contains them;
        remaining sources are appended in order.
    """
    # Build per-source index: image_name -> frame_dict
    source_indices = {}
    all_image_names_ordered = []
    seen_names = set()

    for source_name, frames in source_map:
        idx = {}
        for frame in frames:
            name = normalize_image_name(frame.get('image_name') or '')
            idx[name] = frame
            if name not in seen_names:
                all_image_names_ordered.append(name)
                seen_names.add(name)
        source_indices[source_name] = idx

    merged_frames = []
    for image_name in all_image_names_ordered:
        merged_detections = []
        merged_stats = {}

        for source_idx, (source_name, _) in enumerate(source_map):
            frame = source_indices[source_name].get(image_name)
            if frame is None:
                continue

            # Tag every detection with its lane_assignment (0=roi0, 1=roi1, 2=merge)
            # and offset track_id so it is globally unique across sources.
            id_offset = source_idx * TRACK_ID_OFFSET_PER_SOURCE
            for det in frame.get('detections', []):
                tagged = dict(det)
                tagged['lane_assignment'] = source_idx
                if 'track_id' in tagged and tagged['track_id'] is not None:
                    tagged['track_id'] = tagged['track_id'] + id_offset
                merged_detections.append(tagged)

            # Collect per-source tracking_stats
            if 'tracking_stats' in frame:
                merged_stats[source_name] = frame['tracking_stats']

        merged_frame = {
            'image_name': image_name,
            'detections': merged_detections,
        }
        if merged_stats:
            merged_frame['tracking_stats'] = merged_stats

        merged_frames.append(merged_frame)

    return merged_frames


def write_merged_output(merged, output_path):
    """Persist merged frames to disk."""
    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(merged, f, indent=2, ensure_ascii=False)


def print_merge_summary(merged, source_map, output_path):
    """Print aggregate detection counts per source for one merged output."""
    idx_to_name = {i: s for i, (s, _) in enumerate(source_map)}
    total_by_source = {s: 0 for s, _ in source_map}
    for frame in merged:
        for det in frame.get('detections', []):
            lane = det.get('lane_assignment')
            src = idx_to_name.get(lane)
            if src is not None:
                total_by_source[src] += 1

    print(f"\nMerge complete: {len(merged)} frame(s) written to {output_path}")
    for i, (src, _) in enumerate(source_map):
        print(f"  lane_assignment={i} ({src}): {total_by_source[src]} detection(s) total")


def merge_case(roi0_path, roi1_path, merge_path, output_path, source_names, extra_sources=None, verbose=True):
    """Merge one case from explicit input paths."""
    source_map = build_source_map(
        roi0_path,
        roi1_path,
        merge_path,
        source_names,
        extra_sources=extra_sources,
        verbose=verbose,
    )
    if not source_map:
        if verbose:
            print("Error: No valid tracking files were loaded. Nothing to merge.")
        return False

    if verbose:
        print(f"Merging {len(source_map)} source(s): {[s for s, _ in source_map]}")
    merged = merge_tracking_results(source_map)
    write_merged_output(merged, output_path)
    if verbose:
        print_merge_summary(merged, source_map, output_path)
    return True


def find_case_dirs(results_root, model_name):
    """Return all case directories that contain roi1.json under results_root/model_name."""
    search_root = Path(results_root) / model_name
    if not search_root.exists():
        print(f"Error: model directory does not exist: {search_root}")
        return []

    return sorted(path.parent for path in search_root.rglob('roi1.json'))


def merge_cases_in_batch(results_root, model_name, source_names, output_name):
    """Scan case directories in Python and merge each one sequentially."""
    case_dirs = find_case_dirs(results_root, model_name)
    if not case_dirs:
        print("Error: No roi1.json files were found. Nothing to merge.")
        return False

    print(f"Found {len(case_dirs)} case(s) under {Path(results_root) / model_name}")
    merged_cases = 0
    failed_cases = 0

    for case_dir in case_dirs:
        print("")
        print("==========================================")
        print(f"Case   : {case_dir}")
        print(f"Output : {case_dir / output_name}")
        print("==========================================")

        ok = merge_case(
            roi0_path=case_dir / 'roi0.json',
            roi1_path=case_dir / 'roi1.json',
            merge_path=case_dir / 'merge.json',
            output_path=case_dir / output_name,
            source_names=source_names,
        )
        if ok:
            merged_cases += 1
        else:
            failed_cases += 1

    print("")
    print("Batch merge complete")
    print(f"  merged_cases: {merged_cases}")
    print(f"  failed_cases: {failed_cases}")
    return failed_cases == 0


def parse_extra_source_specs(specs):
    """Parse CLI specs of the form NAME=PATH for additional tracking sources."""
    parsed = []
    for spec in specs or []:
        if '=' not in spec:
            raise ValueError(f"Invalid --extra-source {spec!r}; expected NAME=PATH")
        source_name, source_path = spec.split('=', 1)
        source_name = source_name.strip()
        source_path = source_path.strip()
        if not source_name or not source_path:
            raise ValueError(f"Invalid --extra-source {spec!r}; expected NAME=PATH")
        parsed.append((source_name, source_path))
    return parsed


def main():
    parser = argparse.ArgumentParser(
        description='Merge per-source tracking JSONs into a single combined JSON.'
    )
    parser.add_argument('--roi0', type=str, default=None,
                        help='Path to ROI0 tracking JSON (output of track_objects.py)')
    parser.add_argument('--roi1', type=str, default=None,
                        help='Path to ROI1 tracking JSON (output of track_objects.py)')
    parser.add_argument('--merge', type=str, default=None,
                        help='Path to merged-ROI tracking JSON (output of track_objects.py)')
    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('--output-name', type=str, default='combined_tracking.json',
                        help='Batch mode output filename written inside each case directory')
    parser.add_argument('--source-names', type=str, nargs=3,
                        default=['roi0', 'roi1', 'merge'],
                        metavar=('ROI0_NAME', 'ROI1_NAME', 'MERGE_NAME'),
                        help='Custom labels for the three sources (default: roi0 roi1 merge)')
    parser.add_argument('--extra-source', type=str, action='append', default=[],
                        metavar='NAME=PATH',
                        help='Additional tracking source to append in single-case mode. Can be passed multiple times.')
    parser.add_argument('--output', type=str, default=None,
                        help='Output path for the combined tracking JSON in single-case mode')
    args = parser.parse_args()

    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')
        merge_cases_in_batch(
            results_root=args.results_root,
            model_name=args.model_name,
            source_names=args.source_names,
            output_name=args.output_name,
        )
        return

    if args.output is None:
        parser.error('--output is required in single-case mode')

    try:
        extra_sources = parse_extra_source_specs(args.extra_source)
    except ValueError as exc:
        parser.error(str(exc))

    merge_case(
        roi0_path=args.roi0,
        roi1_path=args.roi1,
        merge_path=args.merge,
        output_path=args.output,
        source_names=args.source_names,
        extra_sources=extra_sources,
    )


if __name__ == '__main__':
    main()