import json
import argparse
import os
from collections import defaultdict
from sklearn.metrics import cohen_kappa_score
from statsmodels.stats.inter_rater import fleiss_kappa

# Compute IoU
def compute_iou(boxA, boxB):
    """
    Compute Intersection over Union between two bboxes in COCO [x, y, w, h] format.
    """
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[0] + boxA[2], boxB[0] + boxB[2])
    yB = min(boxA[1] + boxA[3], boxB[1] + boxB[3])

    interW = max(0, xB - xA)
    interH = max(0, yB - yA)
    interArea = interW * interH

    boxAArea = boxA[2] * boxA[3]
    boxBArea = boxB[2] * boxB[3]
    unionArea = boxAArea + boxBArea - interArea

    return interArea / unionArea if unionArea > 0 else 0.0

# Load annotations
def load_annotations(source):
    """
    Load annotations from a COCO-style JSON file or directory of such files.
    """
    anns = []
    if os.path.isdir(source):
        for fname in os.listdir(source):
            if fname.lower().endswith('.json'):
                path = os.path.join(source, fname)
                with open(path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                anns.extend(data.get('annotations', []))
    else:
        with open(source, 'r', encoding='utf-8') as f:
            data = json.load(f)
        anns = data.get('annotations', [])
    return anns

# Main entrypoint
def main(gt_path, user_paths, iou_threshold=0.5):
    # Load GT and annotator annotations
    gt_annos = load_annotations(gt_path)
    users_annos = [load_annotations(p) for p in user_paths]
    num_annotators = len(users_annos)

    # Align pose and visibility labels per GT
    categories_pose = set()
    categories_vis = set()
    aligned_pose = []  # each row: labels from all annotators for one GT box (pose)
    aligned_vis = []   # same for visibility

    for gt in gt_annos:
        row_pose = []
        row_vis = []
        for annos in users_annos:
            best_iou, best_usr = 0.0, None
            for usr in annos:
                if usr['image_id'] != gt['image_id']:
                    continue
                iou = compute_iou(gt['bbox'], usr['bbox'])
                if iou > best_iou:
                    best_iou, best_usr = iou, usr
            if best_iou >= iou_threshold and best_usr:
                pose_label = str(best_usr.get('attributes', {}).get('pose', '')).strip().lower()
                vis_label = str(best_usr.get('attributes', {}).get('visibility_level', '')).strip()
            else:
                pose_label = 'none'
                vis_label = 'none'
            row_pose.append(pose_label)
            row_vis.append(vis_label)
            categories_pose.add(pose_label)
            categories_vis.add(vis_label)
        aligned_pose.append(row_pose)
        aligned_vis.append(row_vis)

    # Build Fleiss matrices
    cats_pose = sorted(categories_pose)
    matrix_pose = [[row.count(c) for c in cats_pose] for row in aligned_pose]

    cats_vis = sorted(categories_vis)
    matrix_vis = [[row.count(c) for c in cats_vis] for row in aligned_vis]

    # Compute Fleiss' kappa
    fleiss_pose = fleiss_kappa(matrix_pose)
    fleiss_vis = fleiss_kappa(matrix_vis)

    # Print results
    print(f"Fleiss' κ (Pose):       {fleiss_pose:.4f}")
    print(f"Fleiss' κ (Visibility): {fleiss_vis:.4f}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Compute Fleiss kappa for pose and visibility across multiple annotators')
    parser.add_argument('--gt', help='Path to GT JSON or folder', required=True)
    parser.add_argument('--users', nargs='+', help='Paths to annotator JSONs/folders', required=True)
    parser.add_argument('--iou-th', type=float, default=0.5, help='IoU threshold for matching')
    args = parser.parse_args()
    main(args.gt, args.users, args.iou_th)

