"""Explore the current dataset to understand lesion_id, target3, and NROI mixing."""
import json
from datasets import load_from_disk
from collections import defaultdict, Counter

dd = load_from_disk('sandbox/datasets/F_hf_dataset_lesion')
ALL = []
for split in ['train', 'validation', 'test']:
    ds = dd[split]
    for idx, r in enumerate(ds):
        ALL.append({
            'split': split, 'idx': idx,
            'lesion_id_auto': r['lesion_id'],
            'patient_id': r['patient_id'],
            'source_dataset': r['source_dataset'],
            'track_id': r['track_id'],
            'target3': r['target3'],
            'filename': r['original_filename'],
        })

with open('sandbox/tools/lesion_editor/edits.json') as f:
    EDITS = json.load(f)

def edit_key(split, idx):
    return f'{split}:{idx}'

def eff_lesion(rec):
    k = edit_key(rec['split'], rec['idx'])
    if k in EDITS:
        return EDITS[k]['lesion_id']
    return rec['lesion_id_auto']

# target3 mapping: 0=MT, 1=NML, 2=NROI
TARGET3_NAMES = {0: 'MT', 1: 'NML', 2: 'NROI'}

# 1. What does lesion_id look like for NROI images?
nroi_records = [r for r in ALL if r['target3'] == 2]
nroi_lesion_ids = Counter()
for r in nroi_records:
    lid = eff_lesion(r)
    nroi_lesion_ids[lid] += 1

print(f'=== NROI (target3=2) images: {len(nroi_records)} ===')
print(f'How many have non-null lesion_id: {sum(1 for r in nroi_records if eff_lesion(r) is not None and eff_lesion(r) != "")}')
print(f'How many have null/empty lesion_id: {sum(1 for r in nroi_records if eff_lesion(r) is None or eff_lesion(r) == "")}')
print(f'Unique lesion_ids among NROI: {len(nroi_lesion_ids)}')
print(f'Sample NROI lesion_ids: {list(nroi_lesion_ids.most_common(10))}')
print()

# 2. Check the raw lesion_id for NROI (before edits)
nroi_auto_lids = Counter(r['lesion_id_auto'] for r in nroi_records)
print(f'NROI auto lesion_ids (top 10): {list(nroi_auto_lids.most_common(10))}')
print()

# 3. Find groups that mix NROI with lesions (MT/NML)
# Build group -> target3 counts
group_targets = defaultdict(Counter)
group_size = defaultdict(int)
for r in ALL:
    lid = eff_lesion(r)
    group_targets[lid][r['target3']] += 1
    group_size[lid] += 1

# Groups with both NROI and non-NROI
mixed_groups = {}
for lid, counts in group_targets.items():
    if counts[2] > 0 and (counts[0] > 0 or counts[1] > 0):
        mixed_groups[lid] = dict(counts)

print(f'=== Groups mixing NROI with lesions (MT/NML): {len(mixed_groups)} ===')
for lid, counts in sorted(mixed_groups.items()):
    total = sum(counts.values())
    names = {TARGET3_NAMES[k]: v for k, v in counts.items()}
    print(f'  {lid}: {total} imgs — {names}')
print()

# 4. Groups that are ALL NROI
all_nroi_groups = {}
for lid, counts in group_targets.items():
    if counts[2] > 0 and counts[0] == 0 and counts[1] == 0:
        all_nroi_groups[lid] = dict(counts)

print(f'=== Groups that are ALL NROI: {len(all_nroi_groups)} ===')
for lid, counts in sorted(all_nroi_groups.items()):
    total = sum(counts.values())
    print(f'  {lid}: {total} imgs — all NROI')
print()

# 5. Distribution of target3 overall
print(f'=== Overall target3 distribution ===')
overall = Counter(r['target3'] for r in ALL)
for k in sorted(overall):
    print(f'  {TARGET3_NAMES[k]}: {overall[k]}')
print()

# 6. What are the lesion_id values like? Show some samples
print(f'=== Sample lesion_id values (first 20 unique) ===')
unique_lids = sorted(set(eff_lesion(r) for r in ALL))
print(f'Total unique lesion_ids: {len(unique_lids)}')
for lid in unique_lids[:20]:
    counts = group_targets[lid]
    names = {TARGET3_NAMES[k]: v for k, v in counts.items()}
    print(f'  {lid}: {group_size[lid]} imgs — {names}')