#!/usr/bin/env python3
"""Deep audit of PINO v6 dataset label quality and coverage gaps."""
import json
import numpy as np
from collections import Counter, defaultdict

DATA_PATH = "data/empirical_dataset_v6.jsonl"
VOCAB_PATH = "data/pyrfume_vocabulary.json"

with open(VOCAB_PATH) as f:
    vocab = json.load(f)
DESCRIPTOR_NAMES = vocab["vocabulary"]
N_DESCRIPTORS = vocab["dimension"]

records = []
with open(DATA_PATH) as f:
    for line in f:
        line = line.strip()
        if line:
            records.append(json.loads(line))

N = len(records)
print(f"=" * 70)
print(f"PINO v6 DATASET AUDIT — {N} total records")
print(f"=" * 70)

controls = [r for r in records if r.get("is_control")]
non_controls = [r for r in records if not r.get("is_control")]
print(f"\nControls: {len(controls)}, Non-controls: {len(non_controls)}")

genre_counts = Counter(r.get("genre") for r in records)
print(f"Genres: {dict(genre_counts)}")

# =====================================================================
# (1) PER-FORMULA INGREDIENT COUNT DISTRIBUTION
# =====================================================================
print(f"\n{'=' * 70}")
print("(1) PER-FORMULA INGREDIENT COUNT DISTRIBUTION")
print(f"{'=' * 70}")

counts_all = [len(r["formula"]) for r in records]
counts_nc = [len(r["formula"]) for r in non_controls]
counts_c = [len(r["formula"]) for r in controls]

for label, counts in [("ALL", counts_all), ("NON-CONTROL", counts_nc), ("CONTROL", counts_c)]:
    arr = np.array(counts)
    dist = Counter(counts)
    print(f"\n  [{label}] n={len(arr)}")
    print(f"    min={arr.min()}, max={arr.max()}, mean={arr.mean():.2f}, median={np.median(arr):.1f}, std={arr.std():.2f}")
    print(f"    Distribution:")
    for k in sorted(dist.keys()):
        pct = dist[k] / len(arr) * 100
        print(f"      {k} ingredients: {dist[k]} ({pct:.1f}%)")

# =====================================================================
# (2) WEIGHT FRACTION DISTRIBUTION — ETHANOL DOMINANCE
# =====================================================================
print(f"\n{'=' * 70}")
print("(2) WEIGHT FRACTION DISTRIBUTION — ETHANOL DOMINANCE")
print(f"{'=' * 70}")

ETHANOL_CAS = {"64-17-5", "25265-71-8"}

ethanol_fracs = []
non_ethanol_fracs = []
ethanol_85plus = 0
ethanol_90plus = 0
formulas_with_ethanol = 0

for r in records:
    formula = r["formula"]
    eth_frac = 0.0
    non_eth_total = 0.0
    has_eth = False
    for ing in formula:
        cas = ing["cas"]
        wf = ing["weight_fraction"]
        if cas in ETHANOL_CAS:
            eth_frac += wf
            has_eth = True
        else:
            non_eth_total += wf
    if has_eth:
        formulas_with_ethanol += 1
    ethanol_fracs.append(eth_frac)
    non_ethanol_fracs.append(non_eth_total)
    if eth_frac >= 0.85:
        ethanol_85plus += 1
    if eth_frac >= 0.90:
        ethanol_90plus += 1

eth_arr = np.array(ethanol_fracs)
print(f"\n  Ethanol (64-17-5 or 25265-71-8) weight fraction per formula:")
print(f"    Formulas containing ethanol: {formulas_with_ethanol} / {N} ({formulas_with_ethanol/N*100:.1f}%)")
print(f"    Mean ethanol fraction: {eth_arr.mean():.4f}")
print(f"    Median ethanol fraction: {np.median(eth_arr):.4f}")
print(f"    Min ethanol fraction: {eth_arr.min():.4f}")
print(f"    Max ethanol fraction: {eth_arr.max():.4f}")
print(f"    Formulas with ethanol >= 85%: {ethanol_85plus} ({ethanol_85plus/N*100:.1f}%)")
print(f"    Formulas with ethanol >= 90%: {ethanol_90plus} ({ethanol_90plus/N*100:.1f}%)")

# Distribution buckets
buckets = [0, 0.01, 0.1, 0.3, 0.5, 0.7, 0.8, 0.85, 0.9, 0.95, 0.99, 1.01]
print(f"\n    Ethanol fraction buckets:")
for i in range(len(buckets) - 1):
    lo, hi = buckets[i], buckets[i+1]
    cnt = int(np.sum((eth_arr >= lo) & (eth_arr < hi)))
    if cnt > 0:
        print(f"      [{lo:.2f}, {hi:.2f}): {cnt} ({cnt/N*100:.1f}%)")

# Per-ingredient weight fraction stats (non-ethanol)
print(f"\n  Per-ingredient weight fraction (excluding ethanol):")
all_ingredient_fracs = []
for r in records:
    for ing in r["formula"]:
        if ing["cas"] not in ETHANOL_CAS:
            all_ingredient_fracs.append(ing["weight_fraction"])
if all_ingredient_fracs:
    ing_arr = np.array(all_ingredient_fracs)
    print(f"    n={len(ing_arr)} ingredient-slots")
    print(f"    mean={ing_arr.mean():.4f}, median={np.median(ing_arr):.4f}")
    print(f"    min={ing_arr.min():.6f}, max={ing_arr.max():.4f}")
    print(f"    percentiles: 10%={np.percentile(ing_arr,10):.4f}, 25%={np.percentile(ing_arr,25):.4f}, "
          f"75%={np.percentile(ing_arr,75):.4f}, 90%={np.percentile(ing_arr,90):.4f}")

# =====================================================================
# (3) PYRAMID TIER COVERAGE — formulas with at least one ingredient per tier
# =====================================================================
print(f"\n{'=' * 70}")
print("(3) PYRAMID TIER COVERAGE (top/mid/base)")
print(f"{'=' * 70}")
# pyramid_targets is 3 lists of 138 binary values
# Tier means at least one descriptor active in that tier

has_top = 0
has_mid = 0
has_base = 0
has_all_three = 0
has_none = 0
tier_sums = []

for r in records:
    pt = r.get("pyramid_targets")
    if pt is None:
        continue
    # Each tier is a list of 138 values
    top_active = sum(1 for x in pt[0] if x > 0)
    mid_active = sum(1 for x in pt[1] if x > 0)
    base_active = sum(1 for x in pt[2] if x > 0)
    tier_sums.append((top_active, mid_active, base_active))
    
    t_has = top_active > 0
    m_has = mid_active > 0
    b_has = base_active > 0
    if t_has: has_top += 1
    if m_has: has_mid += 1
    if b_has: has_base += 1
    if t_has and m_has and b_has: has_all_three += 1
    if not (t_has or m_has or b_has): has_none += 1

total_with_pyramid = len(tier_sums)
print(f"\n  Formulas with pyramid_targets: {total_with_pyramid}")
print(f"  Has top notes active: {has_top} ({has_top/total_with_pyramid*100:.1f}%)")
print(f"  Has mid notes active: {has_mid} ({has_mid/total_with_pyramid*100:.1f}%)")
print(f"  Has base notes active: {has_base} ({has_base/total_with_pyramid*100:.1f}%)")
print(f"  Has ALL THREE tiers active: {has_all_three} ({has_all_three/total_with_pyramid*100:.1f}%)")
print(f"  Has NO tiers active (all-zero): {has_none} ({has_none/total_with_pyramid*100:.1f}%)")

tier_arr = np.array(tier_sums)
print(f"\n  Avg active descriptors per tier:")
print(f"    Top:  {tier_arr[:,0].mean():.2f} (min={tier_arr[:,0].min()}, max={tier_arr[:,0].max()})")
print(f"    Mid:  {tier_arr[:,1].mean():.2f} (min={tier_arr[:,1].min()}, max={tier_arr[:,1].max()})")
print(f"    Base: {tier_arr[:,2].mean():.2f} (min={tier_arr[:,2].min()}, max={tier_arr[:,2].max()})")

# Also check only non-controls
has_all_nc = 0
for r in non_controls:
    pt = r.get("pyramid_targets")
    if pt and sum(1 for x in pt[0] if x>0) > 0 and sum(1 for x in pt[1] if x>0) > 0 and sum(1 for x in pt[2] if x>0) > 0:
        has_all_nc += 1
print(f"\n  [Non-controls only] Has ALL THREE tiers: {has_all_nc} / {len(non_controls)} ({has_all_nc/len(non_controls)*100:.1f}%)")

# =====================================================================
# (4) CHEMICAL FREQUENCY — over-represented chemicals
# =====================================================================
print(f"\n{'=' * 70}")
print("(4) CHEMICAL FREQUENCY / OVER-REPRESENTATION")
print(f"{'=' * 70}")

cas_counter = Counter()
cas_name = {}
for r in records:
    for ing in r["formula"]:
        cas = ing["cas"]
        cas_counter[cas] += 1
        if cas not in cas_name and ing.get("name"):
            cas_name[cas] = ing["name"]

print(f"\n  Unique chemicals (by CAS/SMILES key): {len(cas_counter)}")
print(f"\n  Top 30 most frequent chemicals:")
print(f"  {'Rank':>4} {'CAS':<20} {'Name':<30} {'Count':>6} {'% of formulas':>14}")
for i, (cas, cnt) in enumerate(cas_counter.most_common(30)):
    name = cas_name.get(cas, "(no name)")
    # Check if it's a SMILES key
    display_cas = cas[:18] + ".." if len(cas) > 20 else cas
    display_name = name[:28] + ".." if len(name) > 30 else name
    print(f"  {i+1:>4} {display_cas:<20} {display_name:<30} {cnt:>6} {cnt/N*100:>13.1f}%")

print(f"\n  Chemicals appearing in >50% of formulas: {sum(1 for _,c in cas_counter.items() if c/N > 0.5)}")
print(f"  Chemicals appearing in >25% of formulas: {sum(1 for _,c in cas_counter.items() if c/N > 0.25)}")
print(f"  Chemicals appearing in >10% of formulas: {sum(1 for _,c in cas_counter.items() if c/N > 0.10)}")
print(f"  Chemicals appearing in >5% of formulas:  {sum(1 for _,c in cas_counter.items() if c/N > 0.05)}")
print(f"  Chemicals appearing in only 1 formula:   {sum(1 for _,c in cas_counter.items() if c == 1)}")
print(f"  Chemicals appearing in <=2 formulas:     {sum(1 for _,c in cas_counter.items() if c <= 2)}")

# Also check non-controls only
cas_counter_nc = Counter()
for r in non_controls:
    for ing in r["formula"]:
        cas_counter_nc[ing["cas"]] += 1
print(f"\n  [Non-controls only] Top 15:")
for i, (cas, cnt) in enumerate(cas_counter_nc.most_common(15)):
    name = cas_name.get(cas, "")
    display_cas = cas[:18] + ".." if len(cas) > 20 else cas
    print(f"  {i+1:>4} {display_cas:<20} {cnt:>5} / {len(non_controls)} ({cnt/len(non_controls)*100:.1f}%)  {name[:40]}")

# =====================================================================
# (5) DESCRIPTOR CO-OCCURRENCE — redundancy analysis
# =====================================================================
print(f"\n{'=' * 70}")
print("(5) DESCRIPTOR CO-OCCURRENCE / REDUNDANCY")
print(f"{'=' * 70}")
# Use pyramid_targets flattened across all 3 tiers to find descriptor activity
# Actually, use objective_targets which is per-timestep (49 × 138)
# But for descriptor co-occurrence, aggregate: a descriptor is "active" in a formula
# if it's non-zero in any tier of pyramid_targets OR in objective_targets

# Use pyramid_targets: flatten 3×138 → 138 (OR across tiers)
descriptor_active = []  # per formula: set of active descriptor indices
for r in records:
    pt = r.get("pyramid_targets")
    if pt is None:
        descriptor_active.append(set())
        continue
    active = set()
    for tier in pt:
        for idx, val in enumerate(tier):
            if val > 0:
                active.add(idx)
    descriptor_active.append(active)

# Co-occurrence: count pairs
pair_counts = Counter()
single_counts = Counter()
for active_set in descriptor_active:
    active_sorted = sorted(active_set)
    for idx in active_sorted:
        single_counts[idx] += 1
    for i in range(len(active_sorted)):
        for j in range(i+1, len(active_sorted)):
            pair_counts[(active_sorted[i], active_sorted[j])] += 1

# Find always-together pairs: P(B|A) and P(A|B) both near 1.0
print(f"\n  Descriptor pairs with near-perfect co-occurrence (both directions >=0.95):")
perfect_pairs = []
for (a, b), co_occur in pair_counts.most_common():
    ca = single_counts[a]
    cb = single_counts[b]
    if ca == 0 or cb == 0:
        continue
    p_b_given_a = co_occur / ca
    p_a_given_b = co_occur / cb
    if p_b_given_a >= 0.95 and p_a_given_b >= 0.95 and co_occur >= 5:
        perfect_pairs.append((a, b, co_occur, p_b_given_a, p_a_given_b))
        
if perfect_pairs:
    print(f"    Found {len(perfect_pairs)} pairs:")
    for a, b, co, pba, pab in perfect_pairs[:30]:
        print(f"      [{DESCRIPTOR_NAMES[a]:>15}] ↔ [{DESCRIPTOR_NAMES[b]:>15}]  co={co:>5}  P(B|A)={pba:.3f}  P(A|B)={pab:.3f}")
else:
    print("    None found.")

# High co-occurrence pairs (>= 0.90)
print(f"\n  Descriptor pairs with high co-occurrence (both >=0.90, co-occur >=10):")
high_pairs = []
for (a, b), co_occur in pair_counts.most_common():
    ca = single_counts[a]
    cb = single_counts[b]
    if ca == 0 or cb == 0:
        continue
    p_b_given_a = co_occur / ca
    p_a_given_b = co_occur / cb
    if p_b_given_a >= 0.90 and p_a_given_b >= 0.90 and co_occur >= 10:
        high_pairs.append((a, b, co_occur, p_b_given_a, p_a_given_b))
print(f"    Found {len(high_pairs)} pairs:")
for a, b, co, pba, pab in high_pairs[:30]:
    print(f"      [{DESCRIPTOR_NAMES[a]:>15}] ↔ [{DESCRIPTOR_NAMES[b]:>15}]  co={co:>5}  P(B|A)={pba:.3f}  P(A|B)={pab:.3f}")

# Top 20 most co-occurring pairs by raw count
print(f"\n  Top 20 descriptor pairs by raw co-occurrence count:")
for (a, b), co in pair_counts.most_common(20):
    ca, cb = single_counts[a], single_counts[b]
    pba = co / ca if ca else 0
    pab = co / cb if cb else 0
    print(f"      [{DESCRIPTOR_NAMES[a]:>15}] + [{DESCRIPTOR_NAMES[b]:>15}]  co={co:>5}  P(B|A)={pba:.3f}  P(A|B)={pab:.3f}")

# =====================================================================
# (6) PER-DESCRIPTOR POSITIVE COUNT
# =====================================================================
print(f"\n{'=' * 70}")
print("(6) PER-DESCRIPTOR POSITIVE COUNT")
print(f"{'=' * 70}")

# Count: how many formulas have each descriptor active (in any tier)
desc_pos_counts = np.zeros(N_DESCRIPTORS, dtype=int)
for active_set in descriptor_active:
    for idx in active_set:
        desc_pos_counts[idx] += 1

print(f"\n  Total descriptors: {N_DESCRIPTORS}")
print(f"  Descriptors with 0 positives: {int(np.sum(desc_pos_counts == 0))}")
print(f"  Descriptors with <10 positives: {int(np.sum(desc_pos_counts < 10))}")
print(f"  Descriptors with <50 positives: {int(np.sum(desc_pos_counts < 50))}")
print(f"  Descriptors with <100 positives: {int(np.sum(desc_pos_counts < 100))}")
print(f"  Descriptors with >=100 positives: {int(np.sum(desc_pos_counts >= 100))}")
print(f"  Descriptors with >=500 positives: {int(np.sum(desc_pos_counts >= 500))}")

print(f"\n  All {N_DESCRIPTORS} descriptors sorted by positive count:")
print(f"  {'Idx':>4} {'Descriptor':<18} {'Positives':>9} {'% of N':>8}")
sorted_descs = sorted(range(N_DESCRIPTORS), key=lambda i: desc_pos_counts[i])
for i in sorted_descs:
    cnt = desc_pos_counts[i]
    name = DESCRIPTOR_NAMES[i] if i < len(DESCRIPTOR_NAMES) else f"idx_{i}"
    flag = " *** ZERO" if cnt == 0 else (" ** <10" if cnt < 10 else (" * <50" if cnt < 50 else ""))
    print(f"  {i:>4} {name:<18} {cnt:>9} {cnt/N*100:>7.1f}%{flag}")

# Per-tier descriptor counts
print(f"\n  Per-tier descriptor positive counts (pyramid_targets):")
for tier_idx, tier_name in enumerate(["TOP", "MID", "BASE"]):
    tier_counts = np.zeros(N_DESCRIPTORS, dtype=int)
    for r in records:
        pt = r.get("pyramid_targets")
        if pt and len(pt) > tier_idx:
            for d_idx, val in enumerate(pt[tier_idx]):
                if val > 0:
                    tier_counts[d_idx] += 1
    print(f"\n    [{tier_name}] Descriptors with 0 positives in this tier: {int(np.sum(tier_counts == 0))}")
    print(f"    [{tier_name}] Descriptors with <10 positives: {int(np.sum(tier_counts < 10))}")
    print(f"    [{tier_name}] Top 10: ", end="")
    top10 = sorted(range(N_DESCRIPTORS), key=lambda i: tier_counts[i], reverse=True)[:10]
    print(", ".join(f"{DESCRIPTOR_NAMES[i]}({tier_counts[i]})" for i in top10))

# =====================================================================
# (7) PSYCHOMETRIC TARGET DISTRIBUTION
# =====================================================================
print(f"\n{'=' * 70}")
print("(7) PSYCHOMETRIC TARGET DISTRIBUTION")
print(f"{'=' * 70}")

psy_records = [r for r in records if r.get("psychometric_targets") is not None]
print(f"\n  Records with psychometric_targets: {len(psy_records)}")

if psy_records:
    sample = psy_records[0]["psychometric_targets"]
    print(f"  Dimensionality: {len(sample)}")
    
    psy_arr = np.array([r["psychometric_targets"] for r in psy_records])
    
    # Per-dimension stats
    print(f"\n  Per-dimension statistics:")
    dim_names = ["season_0 (winter)", "season_1 (spring)", "season_2 (summer)", 
                 "season_3 (autumn)", "gender", "wear_day", "wear_night"]
    for d in range(min(len(sample), 7)):
        name = dim_names[d] if d < len(dim_names) else f"dim_{d}"
        col = psy_arr[:, d]
        print(f"    {name}: mean={col.mean():.4f}, std={col.std():.4f}, "
              f"min={col.min():.4f}, max={col.max():.4f}")
        
        # Check unique values
        unique_vals, counts = np.unique(col, return_counts=True)
        if len(unique_vals) <= 20:
            val_str = ", ".join(f"{v:.2f}({c})" for v, c in zip(unique_vals, counts))
            print(f"      Unique values: {val_str}")
        else:
            # Histogram
            percentiles = [0, 10, 25, 50, 75, 90, 100]
            print(f"      Percentiles: {', '.join(f'{p}%={np.percentile(col,p):.3f}' for p in percentiles)}")
    
    # Check if any dimension is constant
    print(f"\n  Constant dimensions (std < 0.01):")
    for d in range(min(len(sample), 7)):
        name = dim_names[d] if d < len(dim_names) else f"dim_{d}"
        col = psy_arr[:, d]
        if col.std() < 0.01:
            print(f"    {name}: CONSTANT at {col.mean():.4f}")
    
    # Seasonality: which season is highest per formula
    if len(sample) >= 4:
        seasons = psy_arr[:, :4]
        dominant_season = np.argmax(seasons, axis=1)
        season_names = ["winter", "spring", "summer", "autumn"]
        print(f"\n  Dominant season distribution:")
        for s in range(4):
            cnt = int(np.sum(dominant_season == s))
            print(f"    {season_names[s]}: {cnt} ({cnt/len(psy_arr)*100:.1f}%)")
    
    # Gender distribution
    if len(sample) >= 5:
        gender = psy_arr[:, 4]
        print(f"\n  Gender (dim 4):")
        unique_vals, counts = np.unique(np.round(gender, 2), return_counts=True)
        for v, c in zip(unique_vals, counts):
            print(f"    {v:.2f}: {c} ({c/len(gender)*100:.1f}%)")
    
    # Wear day/night
    if len(sample) >= 7:
        wear_day = psy_arr[:, 5]
        wear_night = psy_arr[:, 6]
        print(f"\n  Wear day (dim 5):")
        unique_vals, counts = np.unique(np.round(wear_day, 2), return_counts=True)
        for v, c in zip(unique_vals, counts):
            print(f"    {v:.2f}: {c} ({c/len(wear_day)*100:.1f}%)")
        print(f"\n  Wear night (dim 6):")
        unique_vals, counts = np.unique(np.round(wear_night, 2), return_counts=True)
        for v, c in zip(unique_vals, counts):
            print(f"    {v:.2f}: {c} ({c/len(wear_night)*100:.1f}%)")
        
        # Day vs night
        day_gt_night = int(np.sum(wear_day > wear_night))
        night_gt_day = int(np.sum(wear_night > wear_day))
        equal = int(np.sum(np.abs(wear_day - wear_night) < 0.01))
        print(f"\n  Day > Night: {day_gt_night}, Night > Day: {night_gt_day}, Equal: {equal}")

# Psychometric by genre
print(f"\n  Psychometric means by genre:")
for genre in sorted(set(r.get("genre") for r in records)):
    g_records = [r for r in psy_records if r.get("genre") == genre]
    if g_records:
        g_arr = np.array([r["psychometric_targets"] for r in g_records])
        means = g_arr.mean(axis=0)
        dim_strs = []
        for d in range(min(7, len(means))):
            name = dim_names[d] if d < len(dim_names) else f"d{d}"
            dim_strs.append(f"{name}={means[d]:.3f}")
        print(f"    [{genre}] n={len(g_records)}: {', '.join(dim_strs)}")

print(f"\n{'=' * 70}")
print("AUDIT COMPLETE")
print(f"{'=' * 70}")