"""Verify all lesion groups are within a single patient."""
import json
from datasets import load_from_disk
from collections import defaultdict

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'],
            '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']

lesion_patients = defaultdict(set)
lesion_sources = defaultdict(set)
lesion_count = defaultdict(int)
for rec in ALL:
    lid = eff_lesion(rec)
    lesion_patients[lid].add(rec['patient_id'])
    lesion_sources[lid].add(rec['source_dataset'])
    lesion_count[lid] += 1

multi_patient = {lid: pts for lid, pts in lesion_patients.items() if len(pts) > 1}
multi_source = {lid: srcs for lid, srcs in lesion_sources.items() if len(srcs) > 1}

print(f'Total lesion groups: {len(lesion_patients)}')
print(f'Total images: {sum(lesion_count.values())}')
print(f'Edits applied: {len(EDITS)}')
print()
print(f'=== Lesion groups spanning MULTIPLE patients: {len(multi_patient)} ===')
if multi_patient:
    for lid, pts in sorted(multi_patient.items()):
        n = lesion_count[lid]
        srcs = lesion_sources[lid]
        print(f'  {lid}: {n} imgs, patients={sorted(pts)}, sources={sorted(srcs)}')
else:
    print('  NONE - all lesion groups are within a single patient.')
print()
print(f'=== Lesion groups spanning MULTIPLE sources: {len(multi_source)} ===')
if multi_source:
    for lid, srcs in sorted(multi_source.items()):
        pts = lesion_patients[lid]
        n = lesion_count[lid]
        print(f'  {lid}: {n} imgs, sources={sorted(srcs)}, patients={sorted(pts)}')
else:
    print('  NONE - all lesion groups are within a single source.')