#!/usr/bin/env python3
"""Analyze mono_target_label vs mono_audio_labels in ov1_foa.jsonl"""
import json
from collections import Counter

JSONL = "/apdcephfs_cq10/share_1603164/user/schmittzhu/data/metadata/ov1_foa.jsonl"

train_samples = []
with open(JSONL) as f:
    for line in f:
        d = json.loads(line)
        if d["split"] == "train":
            train_samples.append(d)

print(f"Total train samples: {len(train_samples)}\n")

# Class distribution
label_counts = Counter(s["mono_target_label"] for s in train_samples)
print("=== mono_target_label distribution (train) ===")
for label, cnt in label_counts.most_common():
    print(f"  {label}: {cnt}")
print(f"Total unique labels: {len(label_counts)}\n")

# Specific labels
targets = ["guitar", "string_instrument", "musical_instrument", "singing", "male_singing", "female_singing"]
for target in targets:
    matches = [s for s in train_samples if s["mono_target_label"] == target]
    print(f'=== mono_target_label = "{target}" ({len(matches)} samples) ===')
    if not matches:
        print("  (no samples found)")
    else:
        combos = Counter(tuple(s["mono_audio_labels"]) for s in matches)
        for combo, cnt in combos.most_common(20):
            print(f"  [{cnt}x] {list(combo)}")
    print()

# primary labels
print("=== mono_primary_label for targets of interest ===")
for target in targets:
    matches = [s for s in train_samples if s["mono_target_label"] == target]
    if matches:
        primaries = Counter(s["mono_primary_label"] for s in matches)
        print(f"  {target}: {dict(primaries.most_common(20))}")
print()

# Reverse: what target labels contain Musical_instrument
print('=== Samples with "Musical_instrument" in mono_audio_labels ===')
mi_labels = Counter()
for s in train_samples:
    if "Musical_instrument" in s["mono_audio_labels"]:
        mi_labels[s["mono_target_label"]] += 1
for label, cnt in mi_labels.most_common():
    print(f"  {label}: {cnt}")