#!/usr/bin/env python3
"""Task 0.5.1 — Discrimination probe.

Loads current tier labels, filters the descriptor vocabulary to a
perfumery-relevant subset, and recomputes cross-genre Jaccard to determine
whether the vocabulary carries recoverable signal (Option A) or must be
replaced entirely (Option B).

Decision rule:
  - Filtered Jaccard ≤ 0.70 → Path A (remapping is viable)
  - Filtered Jaccard > 0.85  → Path B (vocabulary is non-discriminative)
  - In between              → Partial; escalate
"""
from __future__ import annotations

import json
import numpy as np
from pathlib import Path
from collections import Counter

DATA = Path("data")
ARTIFACTS = Path("artifacts")
ARTIFACTS.mkdir(exist_ok=True)

# ─── 1. Descriptor keep/drop mapping ─────────────────────────────────────────

# Perfumery-relevant descriptors — terms used in actual fragrance pyramids.
# Includes: floral families, citrus, woods, resins, spices, gourmand, musk/amber,
# green/herbal, aldehydic, and common qualitative notes (clean, powdery, sweet).
KEEP = {
    # Floral
    "floral", "rose", "jasmin", "jasmine", "muguet", "hyacinth", "lily",
    "violet", "geranium", "orris",
    # Citrus
    "citrus", "lemon", "orange", "grapefruit", "bergamot",
    # Woody / resin
    "woody", "cedar", "sandalwood", "pine", "cypress",
    # Amber / balsamic / resinous
    "amber", "balsamic", "vanilla", "coumarinic", "tonka",
    # Musk / animalic (perfumery-relevant)
    "musk", "leathery", "animal",
    # Spicy
    "spicy", "cinnamon", "clove", "pepper",
    # Green / herbal
    "green", "herbal", "grassy", "leafy", "lavender", "tea", "tobacco",
    "hay", "aromatic",
    # Aldehydic / clean
    "aldehydic", "clean", "fresh", "soapy",
    # Fruity (perfumery subset)
    "fruity", "apple", "peach", "pear", "berry", "plum",
    # Gourmand / sweet
    "sweet", "honey", "caramellic", "chocolate", "cocoa", "coconut",
    "creamy", "milky",
    # Powder / soft
    "powdery", "warm",
    # Earthy / mossy (perfumery-relevant)
    "earthy", "mossy", "oakmoss",
    # Marine / ozonic
    "ozone",
    # Miscellaneous perfumery
    "waxy", "dry",
}

# Off-odor / raw-material GC-O descriptors that have no place in perfume pyramids
DROP_EXPLICIT = {
    # Food/spoilage off-odors
    "sulfurous", "fishy", "cheesy", "burnt", "popcorn", "fatty", "garlic",
    "onion", "cabbage", "radish", "horseradish", "potato", "tomato",
    "vegetable", "brothy", "beefy", "meaty", "chicken", "mushroom",
    "fermented", "rummy", "brandy", "cognac", "winey", "rum",
    # Chemical / industrial
    "acidic", "alliaceous", "camphoreous", "ketonic", "lactonic", "estery",
    "ethereal", "solvent", "gasoline", "metallic", "medicinal", "phenolic",
    "musty", "moldy",
    # Redundant / too generic to be discriminative
    "bland", "mild", "odorless", "cooked", "roasted", "ripe", "juicy",
    "sour", "sharp", "terpenic", "weedy", "celery",
    # Food-specific (non-perfumery)
    "banana", "cherry", "grape", "melon", "pineapple", "raspberry",
    "strawberry", "tropical", "cucumber", "cortex", "almond", "hazelnut",
    "nutty", "malty", "bready", "brown", "buttery", "dairy",
    # Duplicate / variant
    "minty", "mint", "mentholic",  # keep "cooling" or "fresh" instead
    "caramellic",  # keep "sweet" instead
    "chamomile",  # rarely used in perfume pyramids
    "tropical",
}

# Any descriptor not in KEEP or DROP_EXPLICIT gets flagged for manual review
NEEDS_REVIEW: set[str] = set()


def build_descriptor_map(vocab: list[str]) -> dict[str, dict]:
    """Classify each vocabulary descriptor into keep/drop/review."""
    mapping = {}
    for idx, desc in enumerate(vocab):
        desc_lower = desc.lower()
        if desc_lower in KEEP:
            status = "keep"
        elif desc_lower in DROP_EXPLICIT:
            status = "drop"
        elif desc_lower in {d.lower() for d in KEEP}:  # case-insensitive catch
            status = "keep"
        elif desc_lower in {d.lower() for d in DROP_EXPLICIT}:
            status = "drop"
        else:
            status = "review"
            NEEDS_REVIEW.add(desc)
        mapping[desc] = {"index": idx, "status": status}
    return mapping


def main() -> None:
    # Load data
    vocab = json.loads((DATA / "pyrfume_vocabulary.json").read_text())["vocabulary"]
    assert len(vocab) == 138

    with open(DATA / "empirical_dataset_v8.jsonl") as f:
        records = [json.loads(l) for l in f]

    # Build descriptor mapping
    mapping = build_descriptor_map(vocab)
    keep_indices = [v["index"] for v in mapping.values() if v["status"] == "keep"]
    drop_indices = [v["index"] for v in mapping.values() if v["status"] == "drop"]
    review_indices = [v["index"] for v in mapping.values() if v["status"] == "review"]

    print(f"=== Descriptor Mapping ({len(vocab)} total) ===")
    print(f"  KEEP:   {len(keep_indices):3d} descriptors")
    print(f"  DROP:   {len(drop_indices):3d} descriptors")
    print(f"  REVIEW: {len(review_indices):3d} descriptors")
    if review_indices:
        print(f"  Review items: {[vocab[i] for i in review_indices]}")
    print()

    # Auto-resolve review items: default to DROP for anything not explicitly kept
    # (conservative — prefer signal over noise)
    for desc in NEEDS_REVIEW:
        mapping[desc]["status"] = "drop_auto"
    final_keep = [v["index"] for v in mapping.values() if v["status"] in ("keep",)]
    final_drop = [v["index"] for v in mapping.values() if v["status"] in ("drop", "drop_auto")]

    print(f"After auto-resolution (review → drop):")
    print(f"  Final KEEP: {len(final_keep)} descriptors")
    print(f"  Final DROP: {len(final_drop)} descriptors")
    print(f"  Kept: {[vocab[i] for i in sorted(final_keep)]}")
    print()

    # Persist mapping artifact
    artifact = {
        "version": "v1",
        "total_descriptors": len(vocab),
        "keep_count": len(final_keep),
        "drop_count": len(final_drop),
        "keep_indices": sorted(final_keep),
        "drop_indices": sorted(final_drop),
        "mapping": mapping,
    }
    (ARTIFACTS / "descriptor_map_v1.json").write_text(json.dumps(artifact, indent=2))
    print(f"Mapping persisted to artifacts/descriptor_map_v1.json")
    print()

    # ─── 2. Recompute cross-genre Jaccard on filtered subset ──────────────────

    # Build target matrix
    targets = []
    genres = []
    for r in records:
        if r.get("is_control"):
            continue
        targets.append(np.array(r["pyramid_targets"]))
        genres.append(r.get("genre", "unknown"))

    targets = np.stack(targets)  # (N, 3, 138)
    genres = np.array(genres)

    keep_arr = np.array(sorted(final_keep))

    # ─── ORIGINAL (unfiltered) Jaccard ────────────────────────────────────────
    # Jaccard is computed per-sample (binary vectors), then averaged within
    # each genre pair. Genre means are fractional (0-1 across samples), so
    # we must binarize at a lower threshold that reflects "this descriptor
    # appears in a meaningful fraction of formulas of this genre."
    BINARIZE_THRESHOLD = 0.05  # descriptor active in ≥5% of genre formulas

    print(f"=== ORIGINAL (unfiltered) Cross-Genre Jaccard — TOP tier ===")
    print(f"  (binarizing genre-mean vectors at {BINARIZE_THRESHOLD} frequency threshold)")
    genre_means = {}
    for genre in sorted(set(genres)):
        mask = genres == genre
        if mask.sum() < 10:
            continue
        genre_means[genre] = targets[mask, 0, :].mean(axis=0)

    genres_list = list(genre_means.keys())
    orig_jaccards_top = []
    for gi in range(len(genres_list)):
        for gj in range(gi + 1, len(genres_list)):
            v1, v2 = genre_means[genres_list[gi]], genre_means[genres_list[gj]]
            b1 = (v1 >= BINARIZE_THRESHOLD).astype(float)
            b2 = (v2 >= BINARIZE_THRESHOLD).astype(float)
            intersection = np.sum((b1 > 0) & (b2 > 0))
            union = np.sum((b1 > 0) | (b2 > 0))
            jac = intersection / max(1, union)
            orig_jaccards_top.append(jac)
            print(f"  {genres_list[gi]:16s} vs {genres_list[gj]:16s}: Jaccard={jac:.3f}  (|A∩B|={int(intersection)}, |A∪B|={int(union)})")
    print(f"  MEAN original Jaccard (top): {np.mean(orig_jaccards_top):.3f}")
    print()

    # ─── Filtered Jaccard — TOP ───────────────────────────────────────────────
    print("=== FILTERED Cross-Genre Jaccard — TOP tier ===")
    filt_jaccards_top = []
    for gi in range(len(genres_list)):
        for gj in range(gi + 1, len(genres_list)):
            v1 = genre_means[genres_list[gi]][keep_arr]
            v2 = genre_means[genres_list[gj]][keep_arr]
            b1 = (v1 >= BINARIZE_THRESHOLD).astype(float)
            b2 = (v2 >= BINARIZE_THRESHOLD).astype(float)
            intersection = np.sum((b1 > 0) & (b2 > 0))
            union = np.sum((b1 > 0) | (b2 > 0))
            jac = intersection / max(1, union)
            filt_jaccards_top.append(jac)
            print(f"  {genres_list[gi]:16s} vs {genres_list[gj]:16s}: Jaccard={jac:.3f}  (|A∩B|={int(intersection)}, |A∪B|={int(union)})")
    mean_filt_top = np.mean(filt_jaccards_top)
    print(f"  MEAN filtered Jaccard (top): {mean_filt_top:.3f}")
    print()

    # ─── Filtered Jaccard — MIDDLE ────────────────────────────────────────────
    genre_means_mid = {}
    for genre in sorted(set(genres)):
        mask = genres == genre
        if mask.sum() < 10:
            continue
        genre_means_mid[genre] = targets[mask, 1, :].mean(axis=0)

    print("=== FILTERED Cross-Genre Jaccard — MIDDLE tier ===")
    filt_jaccards_mid = []
    for gi in range(len(genres_list)):
        for gj in range(gi + 1, len(genres_list)):
            v1 = genre_means_mid[genres_list[gi]][keep_arr]
            v2 = genre_means_mid[genres_list[gj]][keep_arr]
            b1 = (v1 >= BINARIZE_THRESHOLD).astype(float)
            b2 = (v2 >= BINARIZE_THRESHOLD).astype(float)
            intersection = np.sum((b1 > 0) & (b2 > 0))
            union = np.sum((b1 > 0) | (b2 > 0))
            jac = intersection / max(1, union)
            filt_jaccards_mid.append(jac)
            print(f"  {genres_list[gi]:16s} vs {genres_list[gj]:16s}: Jaccard={jac:.3f}  (|A∩B|={int(intersection)}, |A∪B|={int(union)})")
    mean_filt_mid = np.mean(filt_jaccards_mid)
    print(f"  MEAN filtered Jaccard (mid): {mean_filt_mid:.3f}")
    print()

    # ─── 3. Decision ──────────────────────────────────────────────────────────
    print("=" * 60)
    print("=== DECISION ===")
    print(f"  Filtered top Jaccard:   {mean_filt_top:.3f}")
    print(f"  Filtered mid Jaccard:   {mean_filt_mid:.3f}")
    print()

    worst = max(mean_filt_top, mean_filt_mid)
    if worst <= 0.70:
        decision = "PATH A — Vocabulary carries recoverable signal. Remapping is viable."
    elif worst > 0.85:
        decision = "PATH B — Vocabulary is fundamentally non-discriminative. Label SOURCE must change."
    else:
        decision = "PARTIAL — Vocabulary has some signal but insufficient. Escalate to Matt."

    print(f"  Decision: {decision}")
    print("=" * 60)


if __name__ == "__main__":
    main()