#!/usr/bin/env python3
"""Download and process GoodScents/Pyrfume substantivity data.
Substantivity = how long a material's odor lasts on a substrate (hours).
This is the physical basis for top/middle/base classification.
"""
import pandas as pd
import json
import numpy as np
from collections import defaultdict

BASE = "https://raw.githubusercontent.com/pyrfume/pyrfume-data/main/goodscents"

# Download all relevant files
print("Downloading GoodScents data from Pyrfume...")
opl = pd.read_csv(f"{BASE}/opl.csv")
data_rw_opl = pd.read_csv(f"{BASE}/data_rw_opl.csv")
data_rw_odor = pd.read_csv(f"{BASE}/data_rw_odor.csv")
stimuli = pd.read_csv(f"{BASE}/stimuli.csv")
molecules = pd.read_csv(f"{BASE}/molecules.csv")

print(f"  opl: {opl.shape}")
print(f"  data_rw_opl: {data_rw_opl.shape}")
print(f"  data_rw_odor: {data_rw_odor.shape}")
print(f"  stimuli: {stimuli.shape}")
print(f"  molecules: {molecules.shape}")

# --- Step 1: Build TGSC ID -> CAS mapping ---
# data_rw_opl links TGSC ID -> CAS Number via TGSC OPL ID
# Some entries have multiple OPL IDs per TGSC ID
tgsc_to_cas = {}
for _, row in data_rw_opl.iterrows():
    tgsc_id = row["TGSC ID"]
    cas = str(row.get("CAS Number", "")).strip()
    if cas and cas != "nan" and cas not in tgsc_to_cas:
        tgsc_to_cas[tgsc_id] = cas

print(f"\nTGSC ID -> CAS mapping: {len(tgsc_to_cas)} entries")

# --- Step 2: Build TGSC ID -> Substantivity mapping ---
# data_rw_odor has the substantivity columns
tgsc_to_subst = {}
for _, row in data_rw_odor.iterrows():
    tgsc_id = row["TGSC ID"]
    subst = row.get("Substantivity (Hours)")
    subst_min = row.get("Substantivity Min (Hours)")
    subst_max = row.get("Substantivity Max (Hours)")

    if pd.notna(subst) and subst > 0:
        tgsc_to_subst[tgsc_id] = float(subst)
    elif pd.notna(subst_min) and pd.notna(subst_max):
        tgsc_to_subst[tgsc_id] = (float(subst_min) + float(subst_max)) / 2

print(f"TGSC ID -> Substantivity: {len(tgsc_to_subst)} entries")

# Substantivity distribution
if tgsc_to_subst:
    vals = np.array(list(tgsc_to_subst.values()))
    print(f"  Range: {vals.min():.1f} - {vals.max():.1f} hours")
    print(f"  Median: {np.median(vals):.1f}h, Mean: {vals.mean():.1f}h")
    print(f"  Quartiles: {np.percentile(vals, 25):.1f} / {np.percentile(vals, 50):.1f} / {np.percentile(vals, 75):.1f}h")

# --- Step 3: Build CAS -> Substantivity mapping ---
cas_to_subst = {}
for tgsc_id, subst in tgsc_to_subst.items():
    cas = tgsc_to_cas.get(tgsc_id)
    if cas:
        # If multiple entries, keep the one with higher substantivity (pure material)
        if cas not in cas_to_subst or subst > cas_to_subst[cas]:
            cas_to_subst[cas] = subst

print(f"\nCAS -> Substantivity: {len(cas_to_subst)} entries")

# --- Step 4: Derive tier from substantivity ---
# Poucher thresholds in hours (approximate):
# Top notes: < 2 hours (evaporate quickly)
# Middle notes: 2-8 hours
# Base notes: > 8 hours (long-lasting)
# These are based on perfumery practice: materials are smelled on blotter at intervals
cas_to_tier_gs = {}
for cas, subst in cas_to_subst.items():
    if subst < 2:
        cas_to_tier_gs[cas] = "top"
    elif subst < 8:
        cas_to_tier_gs[cas] = "mid"
    else:
        cas_to_tier_gs[cas] = "base"

tier_counts = defaultdict(int)
for tier in cas_to_tier_gs.values():
    tier_counts[tier] += 1
print(f"CAS -> Tier (from substantivity):")
for t in ["top", "mid", "base"]:
    print(f"  {t}: {tier_counts[t]}")

# --- Step 5: Compare with our dataset ---
with open("data/empirical_dataset_v8.jsonl") as f:
    records = [json.loads(l) for l in f]

dataset_cas = set()
for r in records:
    if r.get("is_control"):
        continue
    for comp in r.get("formula", []):
        cas = comp.get("cas", "")
        if cas:
            dataset_cas.add(cas)

matched = dataset_cas.intersection(set(cas_to_subst.keys()))
print(f"\n=== Dataset Coverage ===")
print(f"Dataset CAS: {len(dataset_cas)}")
print(f"Matched with GoodScents substantivity: {len(matched)} ({100*len(matched)/len(dataset_cas):.1f}%)")

# Also check with Poucher tiers
with open("data/poucher_tier_lookup_expanded.json") as f:
    poucher_data = json.load(f)
poucher_tiers = poucher_data["cas_to_tier"]
poucher_matched = dataset_cas.intersection(set(poucher_tiers.keys()))
combined_matched = dataset_cas.intersection(set(poucher_tiers.keys()) | set(cas_to_tier_gs.keys()))
print(f"Poucher tier matched: {len(poucher_matched)} ({100*len(poucher_matched)/len(dataset_cas):.1f}%)")
print(f"Combined (Poucher + GoodScents): {len(combined_matched)} ({100*len(combined_matched)/len(dataset_cas):.1f}%)")

# --- Step 6: Validate against Poucher ---
# Where we have both Poucher tier and GoodScents substantivity, check agreement
overlap = set(poucher_tiers.keys()).intersection(set(cas_to_tier_gs.keys()))
print(f"\n=== Validation: Poucher tier vs GoodScents tier (n={len(overlap)}) ===")
if overlap:
    from sklearn.metrics import cohen_kappa_score
    poucher_labels = [poucher_tiers[c] for c in overlap]
    gs_labels = [cas_to_tier_gs[c] for c in overlap]
    kappa = cohen_kappa_score(poucher_labels, gs_labels)
    print(f"Cohen's kappa: {kappa:.3f}")

    # Confusion matrix
    tier_to_idx = {"top": 0, "mid": 1, "base": 2}
    matrix = np.zeros((3, 3), dtype=int)
    for c in overlap:
        matrix[tier_to_idx[poucher_tiers[c]], tier_to_idx[gs_labels[list(overlap).index(c)]]] += 1

    print(f"\nConfusion (rows=Poucher, cols=GoodScents):")
    print(f"  {'':12s} {'GS-top':>10s} {'GS-mid':>10s} {'GS-base':>10s} {'Total':>10s}")
    for i, tier in enumerate(["top", "mid", "base"]):
        row = matrix[i]
        print(f"  Poucher-{tier:5s} {row[0]:>10d} {row[1]:>10d} {row[2]:>10d} {sum(row):>10d}")

# --- Step 7: Save everything ---
output = {
    "cas_to_substantivity_hours": {k: round(v, 2) for k, v in sorted(cas_to_subst.items())},
    "cas_to_tier": {k: v for k, v in sorted(cas_to_tier_gs.items())},
    "tier_thresholds": {"top": "<2h", "mid": "2-8h", "base": ">8h"},
    "total_cas": len(cas_to_subst),
}
with open("data/goodscents_substantivity.json", "w") as f:
    json.dump(output, f, indent=2)
print(f"\nSaved to data/goodscents_substantivity.json")

# Merge with Poucher tiers (Poucher takes priority)
merged_tiers = dict(cas_to_tier_gs)
for cas, tier in poucher_tiers.items():
    if cas not in merged_tiers:
        merged_tiers[cas] = tier
    # Poucher is ground truth - override
    else:
        merged_tiers[cas] = tier

with open("data/perfumer_tier_lookup.json", "w") as f:
    json.dump({k: v for k, v in sorted(merged_tiers.items())}, f, indent=2)
print(f"Merged tier lookup: {len(merged_tiers)} CAS")

merged_dist = defaultdict(int)
for t in merged_tiers.values():
    merged_dist[t] += 1
for t in ["top", "mid", "base"]:
    print(f"  {t}: {merged_dist[t]}")

# Coverage with merged
merged_matched = dataset_cas.intersection(set(merged_tiers.keys()))
print(f"Dataset coverage with merged: {len(merged_matched)}/{len(dataset_cas)} ({100*len(merged_matched)/len(dataset_cas):.1f}%)")