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

json_path = Path("/home/jovyan/AURAD_dataset/train_prompt_layout2image_multi_total_device_fracture_w_demo.json")

# ---------- 1. 读入 + 去重 ----------
seen = set()
items = []
total_raw = 0
dup_count = 0

with json_path.open("r", encoding="utf-8") as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        total_raw += 1
        item = json.loads(line)
        # 规范化序列化作为去重键：内容完全相同(忽略 key 顺序/空白)即视为重复
        canonical = json.dumps(item, sort_keys=True, ensure_ascii=False)
        if canonical in seen:
            dup_count += 1
            continue
        seen.add(canonical)
        items.append(item)

total_lines = len(items)
print(f"原始记录数 : {total_raw}")
print(f"重复删除   : {dup_count}")
print(f"去重后记录 : {total_lines}\n")


# ---------- 字段探测（性别/年龄）----------
SEX_KEYS = ["sex"]
AGE_KEYS = ["age"]
TEXT_KEYS = ["prompt"]

def find_key(item, candidates):
    for k in candidates:
        if k in item and item[k] not in (None, "", []):
            return k
    return None

sample = items[0] if items else {}
sex_key = find_key(sample, SEX_KEYS)
age_key = find_key(sample, AGE_KEYS)
text_key = find_key(sample, TEXT_KEYS)
print("字段探测:")
print(f"  性别字段 = {sex_key}")
print(f"  年龄字段 = {age_key}")
print(f"  文本字段 = {text_key}  (探测不到性别/年龄时从这里正则提取)")
print(f"  所有可用 key 示例 = {list(sample.keys())}\n")


# ---------- 从文本兜底提取 ----------
def norm_sex(v):
    if v is None:
        return "Unknown"
    s = str(v).strip().lower()
    if s in ("m", "male", "man", "男", "1"):
        return "Male"
    if s in ("f", "female", "woman", "女", "0", "2"):
        return "Female"
    return "Unknown" if s == "" else s

def extract_sex_from_text(t):
    if not t:
        return "Unknown"
    t = t.lower()
    # 常见模式: "55-year-old male", "F,", "sex: male" 等
    if re.search(r"\b(female|woman|girl)\b", t) or re.search(r"\bsex[:\s]+f\b", t):
        return "Female"
    if re.search(r"\b(male|man|boy)\b", t) or re.search(r"\bsex[:\s]+m\b", t):
        return "Male"
    return "Unknown"

def extract_age_from_text(t):
    if not t:
        return None
    m = re.search(r"(\d{1,3})\s*[- ]?\s*year[- ]?old", t.lower())
    if not m:
        m = re.search(r"\bage[:\s]+(\d{1,3})\b", t.lower())
    if m:
        a = int(m.group(1))
        if 0 <= a <= 120:
            return a
    return None

def get_sex(item):
    if sex_key:
        return norm_sex(item.get(sex_key))
    return extract_sex_from_text(item.get(text_key, "")) if text_key else "Unknown"

def get_age(item):
    if age_key:
        v = item.get(age_key)
        try:
            return int(float(v)) if v not in (None, "") else None
        except (ValueError, TypeError):
            return extract_age_from_text(str(v))
    return extract_age_from_text(item.get(text_key, "")) if text_key else None


# ---------- 2. 统计 ----------
disease_mask_count = defaultdict(int)
disease_line_count = defaultdict(int)
dataset_count = defaultdict(int)
sex_count = Counter()
age_values = []
age_bins = Counter()

def age_bin(a):
    if a is None:
        return "Unknown"
    lo = (a // 10) * 10
    return f"{lo}-{lo+9}"

for item in items:
    dataset = item["file_name"].split("/")[0]
    dataset_count[dataset] += 1

    diseases_in_line = set()
    for disease, mask_path in item.get("attn_list", []):
        disease_mask_count[disease] += 1
        diseases_in_line.add(disease)
    for disease in diseases_in_line:
        disease_line_count[disease] += 1

    sex_count[get_sex(item)] += 1
    a = get_age(item)
    age_bins[age_bin(a)] += 1
    if a is not None:
        age_values.append(a)


# ---------- 3. 输出 ----------
print(f"{'Disease':35s} {'Mask_Count':>10s} {'Patient_Count':>14s}")
print("-" * 63)
for disease in sorted(disease_mask_count):
    print(f"{disease:35s} {disease_mask_count[disease]:10d} {disease_line_count[disease]:14d}")
print("-" * 63)
print(f"{'TOTAL':35s} {sum(disease_mask_count.values()):10d} {total_lines:14d}  (人数=去重后总行数)")

print(f"\n{'Dataset':35s} {'Record_Count':>14s}")
print("-" * 51)
for ds in sorted(dataset_count, key=lambda k: -dataset_count[k]):
    print(f"{ds:35s} {dataset_count[ds]:14d}")
print("-" * 51)
print(f"{'TOTAL':35s} {total_lines:14d}")

print(f"\n{'Sex':12s} {'Count':>8s} {'Pct':>8s}")
print("-" * 30)
for s, c in sex_count.most_common():
    print(f"{s:12s} {c:8d} {c/total_lines*100:7.1f}%")

print(f"\n{'Age_Group':12s} {'Count':>8s} {'Pct':>8s}")
print("-" * 30)
def bin_sort(k):
    return (1, 0) if k == "Unknown" else (0, int(k.split("-")[0]))
for grp in sorted(age_bins, key=bin_sort):
    c = age_bins[grp]
    print(f"{grp:12s} {c:8d} {c/total_lines*100:7.1f}%")

if age_values:
    import statistics
    print(f"\n年龄 (有效 {len(age_values)} 条): "
          f"min={min(age_values)}, max={max(age_values)}, "
          f"mean={statistics.mean(age_values):.1f}, median={statistics.median(age_values)}")

from collections import defaultdict

# 年龄分箱：按十岁一段，与之前保持一致
def age_decade(a):
    if a is None:
        return None
    lo = (a // 10) * 10
    return f"{lo}-{lo+9}"

# (age_group, sex) -> count
agesex = defaultdict(int)
age_order = []                       # 保持出现顺序里有效的箱
for item in items:                   # items = 去重后的记录
    a = get_age(item)
    s = get_sex(item)                # "Male" / "Female" / "Unknown"
    grp = age_decade(a)
    if grp is None or s not in ("Male", "Female"):
        continue                     # 年龄或性别缺失的跳过(也可单列统计)
    agesex[(grp, s)] += 1

# 整理成有序表
all_groups = sorted({g for (g, _) in agesex}, key=lambda x: int(x.split("-")[0]))

print(f"{'Age_Group':10s} {'Male':>8s} {'Female':>8s} {'Total':>8s}")
print("-" * 38)
tot_m = tot_f = 0
for g in all_groups:
    m = agesex[(g, "Male")]
    f = agesex[(g, "Female")]
    tot_m += m
    tot_f += f
    print(f"{g:10s} {m:8d} {f:8d} {m+f:8d}")
print("-" * 38)
print(f"{'TOTAL':10s} {tot_m:8d} {tot_f:8d} {tot_m+tot_f:8d}")

# ---- 直接可粘贴回来的字面量(把这段输出发我，我把精确数字填进 PPTX) ----
print("\n# copy below")
print("age_male   =", [agesex[(g, 'Male')]   for g in all_groups])
print("age_female =", [agesex[(g, 'Female')] for g in all_groups])
print("age_labels =", all_groups)