#!/usr/bin/env python3
"""
Export filtered badcase lists from analyze_2d_fp_fn.py results.

The script reads ``analysis_report.json`` and produces:
  - a filtered JSON file with matching examples
  - a plain-text case/frame list for downstream visualization
  - a compact text summary
"""

import argparse
import json
from collections import Counter, defaultdict
from pathlib import Path


def parse_args():
    parser = argparse.ArgumentParser(
        description="Filter and export 2D FP/FN badcases from analysis_report.json."
    )
    parser.add_argument(
        "--input",
        type=str,
        required=True,
        help="Path to analysis_report.json generated by analyze_2d_fp_fn.py",
    )
    parser.add_argument(
        "--mode",
        type=str,
        default="both",
        choices=["fp", "fn", "both"],
        help="Which badcase pool to export",
    )
    parser.add_argument(
        "--classes",
        nargs="+",
        default=None,
        help="Optional class-name filter, e.g. vehicle pedestrian rider",
    )
    parser.add_argument(
        "--error-types",
        nargs="+",
        default=None,
        help="Optional error-type filter, e.g. duplicate localization low_score",
    )
    parser.add_argument(
        "--min-confidence",
        type=float,
        default=None,
        help="Minimum detection confidence for FP examples or matched dets in FN examples",
    )
    parser.add_argument(
        "--max-confidence",
        type=float,
        default=None,
        help="Maximum detection confidence for FP examples or matched dets in FN examples",
    )
    parser.add_argument(
        "--min-distance",
        type=float,
        default=None,
        help="Minimum target distance in metres",
    )
    parser.add_argument(
        "--max-distance",
        type=float,
        default=None,
        help="Maximum target distance in metres",
    )
    parser.add_argument(
        "--min-best-iou",
        type=float,
        default=None,
        help="Minimum best IoU field to keep an example",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=None,
        help="Only keep the first K filtered examples after sorting",
    )
    parser.add_argument(
        "--dedup-frame",
        action="store_true",
        help="Keep at most one example per case/frame/error/class combination",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=None,
        help="Output directory. Defaults to sibling folder next to the input report.",
    )
    return parser.parse_args()


def normalize_tokens(values):
    if not values:
        return None
    return {str(value).strip().lower() for value in values if str(value).strip()}


def get_confidence(item):
    if "confidence" in item and item["confidence"] is not None:
        return float(item["confidence"])
    if "best_det_confidence" in item and item["best_det_confidence"] is not None:
        return float(item["best_det_confidence"])
    return None


def get_best_iou(item):
    if "best_det_iou" in item and item["best_det_iou"] is not None:
        return float(item["best_det_iou"])
    return max(
        float(item.get("best_same_class_iou", 0.0)),
        float(item.get("best_other_class_iou", 0.0)),
    )


def passes_filters(item, class_filter, error_filter, args):
    class_name = str(item.get("class_name", "")).lower()
    error_type = str(item.get("error_type", "")).lower()

    if class_filter and class_name not in class_filter:
        return False
    if error_filter and error_type not in error_filter:
        return False

    confidence = get_confidence(item)
    if args.min_confidence is not None and (confidence is None or confidence < args.min_confidence):
        return False
    if args.max_confidence is not None and (confidence is None or confidence > args.max_confidence):
        return False

    distance = item.get("distance_m")
    if args.min_distance is not None and (distance is None or float(distance) < args.min_distance):
        return False
    if args.max_distance is not None and (distance is None or float(distance) > args.max_distance):
        return False

    best_iou = get_best_iou(item)
    if args.min_best_iou is not None and best_iou < args.min_best_iou:
        return False

    return True


def rank_key(item):
    confidence = get_confidence(item) or 0.0
    best_iou = get_best_iou(item)
    distance = item.get("distance_m")
    distance = float(distance) if distance is not None else -1.0
    return (confidence, best_iou, distance)


def summarize(items):
    by_error = Counter()
    by_class = Counter()
    by_class_error = defaultdict(Counter)

    for item in items:
        class_name = item.get("class_name", "unknown")
        error_type = item.get("error_type", "unknown")
        by_error[error_type] += 1
        by_class[class_name] += 1
        by_class_error[class_name][error_type] += 1

    return {
        "total_examples": len(items),
        "by_error_type": dict(sorted(by_error.items())),
        "by_class": dict(sorted(by_class.items())),
        "by_class_and_error": {
            class_name: dict(sorted(counter.items()))
            for class_name, counter in sorted(by_class_error.items())
        },
    }


def ensure_output_dir(input_path, output_dir):
    if output_dir:
        path = Path(output_dir)
    else:
        path = Path(input_path).resolve().parent / "exported_badcases"
    path.mkdir(parents=True, exist_ok=True)
    return path


def write_summary(path, mode, filters, summary, items):
    with open(path, "w") as file:
        file.write("=" * 90 + "\n")
        file.write("2D FP/FN BADCASE EXPORT SUMMARY\n")
        file.write("=" * 90 + "\n\n")
        file.write(f"Mode: {mode}\n")
        file.write(f"Classes: {filters['classes']}\n")
        file.write(f"Error types: {filters['error_types']}\n")
        file.write(f"Min confidence: {filters['min_confidence']}\n")
        file.write(f"Max confidence: {filters['max_confidence']}\n")
        file.write(f"Min distance: {filters['min_distance']}\n")
        file.write(f"Max distance: {filters['max_distance']}\n")
        file.write(f"Min best IoU: {filters['min_best_iou']}\n")
        file.write(f"Top K: {filters['top_k']}\n")
        file.write(f"Dedup frame: {filters['dedup_frame']}\n\n")

        file.write("SUMMARY\n")
        file.write("-" * 90 + "\n")
        file.write(f"Total examples: {summary['total_examples']}\n\n")

        file.write("BY ERROR TYPE\n")
        file.write("-" * 90 + "\n")
        for key, value in summary["by_error_type"].items():
            file.write(f"{key:<24} {value}\n")
        file.write("\n")

        file.write("BY CLASS\n")
        file.write("-" * 90 + "\n")
        for key, value in summary["by_class"].items():
            file.write(f"{key:<24} {value}\n")
        file.write("\n")

        file.write("TOP EXAMPLES\n")
        file.write("-" * 90 + "\n")
        for item in items[:50]:
            confidence = get_confidence(item)
            best_iou = get_best_iou(item)
            file.write(
                f"{item.get('case_name')}/{item.get('frame_name')}  "
                f"{item.get('class_name')}  {item.get('error_type')}  "
                f"conf={confidence if confidence is not None else '-'}  "
                f"best_iou={best_iou:.4f}\n"
            )


def write_case_frame_list(path, items):
    seen = set()
    with open(path, "w") as file:
        for item in items:
            key = (item.get("case_name"), item.get("frame_name"))
            if key in seen:
                continue
            seen.add(key)
            file.write(f"{key[0]}\t{key[1]}\n")


def main():
    args = parse_args()

    with open(args.input, "r") as file:
        data = json.load(file)

    pools = []
    if args.mode in ("fp", "both"):
        pools.extend(data.get("false_positive_examples", []))
    if args.mode in ("fn", "both"):
        pools.extend(data.get("false_negative_examples", []))

    class_filter = normalize_tokens(args.classes)
    error_filter = normalize_tokens(args.error_types)

    filtered = [
        item for item in pools if passes_filters(item, class_filter, error_filter, args)
    ]

    filtered.sort(key=rank_key, reverse=True)

    if args.dedup_frame:
        deduped = []
        seen = set()
        for item in filtered:
            key = (
                item.get("case_name"),
                item.get("frame_name"),
                item.get("class_name"),
                item.get("error_type"),
            )
            if key in seen:
                continue
            seen.add(key)
            deduped.append(item)
        filtered = deduped

    if args.top_k is not None:
        filtered = filtered[: args.top_k]

    summary = summarize(filtered)
    output_dir = ensure_output_dir(args.input, args.output_dir)

    base_name = f"{args.mode}_badcases"
    json_path = output_dir / f"{base_name}.json"
    txt_path = output_dir / f"{base_name}.txt"
    list_path = output_dir / f"{base_name}_case_frame_list.txt"

    with open(json_path, "w") as file:
        json.dump(
            {
                "source": str(Path(args.input).resolve()),
                "mode": args.mode,
                "filters": {
                    "classes": sorted(class_filter) if class_filter else None,
                    "error_types": sorted(error_filter) if error_filter else None,
                    "min_confidence": args.min_confidence,
                    "max_confidence": args.max_confidence,
                    "min_distance": args.min_distance,
                    "max_distance": args.max_distance,
                    "min_best_iou": args.min_best_iou,
                    "top_k": args.top_k,
                    "dedup_frame": args.dedup_frame,
                },
                "summary": summary,
                "examples": filtered,
            },
            file,
            indent=2,
        )

    write_summary(
        txt_path,
        args.mode,
        {
            "classes": sorted(class_filter) if class_filter else None,
            "error_types": sorted(error_filter) if error_filter else None,
            "min_confidence": args.min_confidence,
            "max_confidence": args.max_confidence,
            "min_distance": args.min_distance,
            "max_distance": args.max_distance,
            "min_best_iou": args.min_best_iou,
            "top_k": args.top_k,
            "dedup_frame": args.dedup_frame,
        },
        summary,
        filtered,
    )
    write_case_frame_list(list_path, filtered)

    print(f"Filtered JSON saved to: {json_path}")
    print(f"Summary text saved to:  {txt_path}")
    print(f"Case/frame list saved to: {list_path}")


if __name__ == "__main__":
    main()
