src/pino/upload_data.py

from __future__ import annotations

import argparse
from collections import Counter, defaultdict
import hashlib
import json
import logging
import math
import os
import random
from pathlib import Path
from typing import Any

from datasets import ClassLabel, Dataset, DatasetDict, load_dataset

logger = logging.getLogger("pino.upload_data")

DEFAULT_TRAIN_RATIO = 0.85
DEFAULT_VAL_RATIO = 0.15


def _validate_ratio(name: str, value: float) -> None:
    if not 0.0 < value < 1.0:
        raise ValueError(f"{name} must be between 0 and 1, got {value!r}")


def _components(record: dict[str, Any]) -> list[dict[str, Any]]:
    return record.get("metadata", {}).get("initial_components") or record.get("formula", [])


def _formula_id(record: dict[str, Any], idx: int | None = None) -> str:
    formula_id = record.get("formula_id") or record.get("metadata", {}).get("formula_id")
    if formula_id:
        return str(formula_id)
    return f"record_{idx}" if idx is not None else "record"


def _is_ethanol_control(record: dict[str, Any], solvent_cas: str = "64-17-5") -> bool:
    if _formula_id(record).startswith(f"CONTROL_{solvent_cas}"):
        return True
    comps = _components(record)
    return bool(record.get("is_control")) and len(comps) == 1 and comps[0].get("cas") == solvent_cas


def _inchi_key_from_smiles(smiles: str | None) -> str | None:
    if not smiles:
        return None
    if smiles.startswith("NATURAL:"):
        return None
    try:
        from rdkit import Chem
        from rdkit.Chem import inchi

        mol = Chem.MolFromSmiles(smiles.removeprefix("SMILES:"))
        if mol is None:
            return None
        key = inchi.MolToInchiKey(mol)
    except Exception:
        return None
    return f"InChIKey:{key}" if key else None


def _canonical_molecule_key(comp: dict[str, Any]) -> str | None:
    """Canonical molecule identity used for leakage accounting.

    Prefer an RDKit-derived InChIKey from structure. When no parseable structure
    is available, fall back to the best identifier present so natural oils and
    unresolved legacy rows remain represented instead of silently disappearing.
    """
    cas = str(comp.get("cas") or "").strip()
    smiles = str(comp.get("smiles") or "").strip()
    if not smiles and cas.startswith("SMILES:"):
        smiles = cas.removeprefix("SMILES:")
    inchi_key = _inchi_key_from_smiles(smiles)
    if inchi_key:
        return inchi_key
    if cas:
        return cas
    name = str(comp.get("name") or "").strip()
    return f"NAME:{name.casefold()}" if name else None


def _canonical_alias(comp: dict[str, Any]) -> str:
    cas = str(comp.get("cas") or "").strip()
    name = str(comp.get("name") or "").strip()
    smiles = str(comp.get("smiles") or "").strip()
    if name and cas:
        return f"{name} [{cas}]"
    return name or cas or smiles or "unknown"


def _active_molecule_set(
    record: dict[str, Any],
    solvent_cas: str | None = None,
    *,
    exclude_controls: bool = True,
) -> set[str]:
    """Return active canonical molecule identities for a formula record.

    The directive format stores the canonical ingredient list in
    ``metadata.initial_components``. Older records fall back to the top-level
    ``formula`` field. Ethanol-only controls are excluded from data splits
    explicitly rather than treated as an empty formula.
    """
    solvent = solvent_cas or record.get("metadata", {}).get("solvent_cas", "64-17-5")
    if exclude_controls and _is_ethanol_control(record, solvent):
        return set()
    return {
        key
        for comp in _components(record)
        if comp.get("cas") != solvent
        for key in [_canonical_molecule_key(comp)]
        if key
    }


def _active_cas_set(record: dict[str, Any], solvent_cas: str | None = None) -> set[str]:
    """Backward-compatible alias for the canonical active molecule set."""
    return _active_molecule_set(record, solvent_cas)


def _canonical_aliases(records: list[dict[str, Any]], solvent_cas: str = "64-17-5") -> dict[str, list[str]]:
    aliases: dict[str, set[str]] = defaultdict(set)
    for record in records:
        if _is_ethanol_control(record, solvent_cas):
            continue
        for comp in _components(record):
            if comp.get("cas") == solvent_cas:
                continue
            key = _canonical_molecule_key(comp)
            if key:
                aliases[key].add(_canonical_alias(comp))
    return {key: sorted(vals) for key, vals in aliases.items()}


def _formula_fingerprint(record: dict[str, Any]) -> str:
    """Deterministic structural fingerprint from a recipe record."""
    components = sorted(_components(record), key=lambda x: _canonical_molecule_key(x) or x.get("cas", ""))
    return "".join(
        f"{_canonical_molecule_key(c) or c.get('cas', '')}_{c.get('weight_fraction', 0):.4f}" for c in components
    )


def _percentile(values: list[float], p: float) -> float | None:
    vals = sorted(v for v in values if math.isfinite(v))
    if not vals:
        return None
    idx = min(len(vals) - 1, max(0, int(round((len(vals) - 1) * p))))
    return float(vals[idx])


def _summary(values: list[float]) -> dict[str, float | None]:
    vals = [v for v in values if math.isfinite(v)]
    if not vals:
        return {"mean": None, "p50": None, "p90": None}
    return {
        "mean": float(sum(vals) / len(vals)),
        "p50": _percentile(vals, 0.5),
        "p90": _percentile(vals, 0.9),
    }


def _connected_components(nodes: set[str], edges: dict[str, set[str]]) -> list[set[str]]:
    seen: set[str] = set()
    components: list[set[str]] = []
    for node in sorted(nodes):
        if node in seen:
            continue
        stack = [node]
        seen.add(node)
        comp: set[str] = set()
        while stack:
            cur = stack.pop()
            comp.add(cur)
            for nxt in edges.get(cur, set()):
                if nxt not in seen:
                    seen.add(nxt)
                    stack.append(nxt)
        components.append(comp)
    components.sort(key=len, reverse=True)
    return components


def _component_stats(record_sets: list[set[str]], ignored_keys: set[str]) -> dict[str, Any]:
    nodes = {key for keys in record_sets for key in keys if key not in ignored_keys}
    edges: dict[str, set[str]] = defaultdict(set)
    for keys in record_sets:
        meaningful = sorted(key for key in keys if key not in ignored_keys)
        for key in meaningful:
            edges.setdefault(key, set())
        for i, left in enumerate(meaningful):
            for right in meaningful[i + 1:]:
                edges[left].add(right)
                edges[right].add(left)

    components = _connected_components(nodes, edges)
    node_to_component = {
        node: comp_idx
        for comp_idx, comp in enumerate(components)
        for node in comp
    }
    component_records: Counter[int] = Counter()
    covered_records = 0
    hub_only_records = 0
    for keys in record_sets:
        comp_ids = {node_to_component[key] for key in keys if key in node_to_component}
        if comp_ids:
            covered_records += 1
            for comp_id in comp_ids:
                component_records[comp_id] += 1
        elif keys:
            hub_only_records += 1

    record_counts = sorted(component_records.values(), reverse=True)
    total_records = len(record_sets)
    largest_records = record_counts[0] if record_counts else 0
    return {
        "n_components": len(components),
        "largest_component_molecules": len(components[0]) if components else 0,
        "largest_component_records": int(largest_records),
        "covered_records": int(covered_records),
        "hub_only_records": int(hub_only_records),
        "max_validation_records_if_largest_component_trains": int(max(0, total_records - largest_records)),
        "top_component_record_counts": [int(v) for v in record_counts[:10]],
        "node_to_component": node_to_component,
        "component_record_counts": component_records,
    }


def analyze_molecule_graph(
    records: list[dict[str, Any]],
    *,
    solvent_cas: str = "64-17-5",
    thresholds: tuple[float, ...] = (0.2, 0.1, 0.05, 0.02, 0.01),
) -> dict[str, Any]:
    """Diagnose formula co-occurrence graph fragmentation under hub exclusion."""
    data_records = [r for r in records if not _is_ethanol_control(r, solvent_cas)]
    record_sets = [_active_molecule_set(r, solvent_cas) for r in data_records]
    nonempty_sets = [s for s in record_sets if s]
    frequency = Counter(key for keys in nonempty_sets for key in keys)
    aliases = _canonical_aliases(data_records, solvent_cas)
    total = len(nonempty_sets)

    threshold_reports = []
    for threshold in thresholds:
        hubs = {key for key, count in frequency.items() if count / total > threshold}
        stats = _component_stats(nonempty_sets, hubs)
        stats.pop("node_to_component", None)
        stats.pop("component_record_counts", None)
        threshold_reports.append({
            "threshold": threshold,
            "hub_count": len(hubs),
            "hubs": [
                {
                    "key": key,
                    "formula_count": int(frequency[key]),
                    "formula_fraction": float(frequency[key] / total),
                    "aliases": aliases.get(key, [])[:5],
                }
                for key in sorted(hubs, key=lambda k: (-frequency[k], k))[:25]
            ],
            **stats,
        })

    duplicate_identities = [
        {"key": key, "aliases": vals}
        for key, vals in sorted(aliases.items())
        if len(vals) > 1
    ]
    return {
        "n_records": len(records),
        "n_data_records": len(data_records),
        "n_nonempty_records": total,
        "n_molecules": len(frequency),
        "top_molecules": [
            {
                "key": key,
                "formula_count": int(count),
                "formula_fraction": float(count / total),
                "aliases": aliases.get(key, [])[:5],
            }
            for key, count in frequency.most_common(25)
        ],
        "duplicate_identities_after_canonicalization": duplicate_identities[:100],
        "thresholds": threshold_reports,
    }


def create_hub_excluded_component_split_from_records(
    records: list[dict[str, Any]],
    train_ratio: float = DEFAULT_TRAIN_RATIO,
    seed: int = 42,
    solvent_cas: str = "64-17-5",
    hub_threshold: float = 0.1,
) -> dict[str, Any]:
    """Split formulas by non-hub co-occurrence components.

    Molecules above ``hub_threshold`` are ignored only for leakage grouping. They
    remain in the returned records and may appear on both sides.
    """
    _validate_ratio("train_ratio", train_ratio)

    record_sets = [_active_molecule_set(r, solvent_cas) for r in records]
    data_indices = [
        idx for idx, record in enumerate(records)
        if not _is_ethanol_control(record, solvent_cas) and record_sets[idx]
    ]
    data_sets = [record_sets[idx] for idx in data_indices]
    frequency = Counter(key for keys in data_sets for key in keys)
    hubs = {key for key, count in frequency.items() if count / len(data_sets) > hub_threshold}
    stats = _component_stats(data_sets, hubs)
    node_to_component = stats["node_to_component"]

    component_to_indices: dict[int, list[int]] = defaultdict(list)
    hub_only_indices: list[int] = []
    excluded_indices = [
        idx for idx, record in enumerate(records)
        if _is_ethanol_control(record, solvent_cas) or not record_sets[idx]
    ]
    for local_idx, record_idx in enumerate(data_indices):
        comp_ids = {node_to_component[key] for key in data_sets[local_idx] if key in node_to_component}
        if not comp_ids:
            hub_only_indices.append(record_idx)
            continue
        if len(comp_ids) > 1:
            raise RuntimeError("Record spans multiple non-hub components; graph construction is inconsistent")
        component_to_indices[next(iter(comp_ids))].append(record_idx)

    target_val = max(1, int(round(len(data_indices) * (1.0 - train_ratio))))
    rng = random.Random(seed)
    # Deterministic greedy assignment: walk components smallest-first (ties broken by
    # the seeded RNG) and add a component to validation only while doing so does not
    # overshoot the target. This avoids forcing an oversized first component into val
    # and keeps the split reproducible for a given seed.
    component_items = [
        (rng.random(), comp_id, idxs)
        for comp_id, idxs in component_to_indices.items()
    ]
    component_items.sort(key=lambda item: (len(item[2]), item[0]))
    val_components: set[int] = set()
    val_count = 0
    for _tie, comp_id, idxs in component_items:
        if val_count < target_val and val_count + len(idxs) <= target_val:
            val_components.add(comp_id)
            val_count += len(idxs)
    # Fallback: if nothing fit under the target (tiny/fragmented graph), take the
    # single smallest component so validation is non-empty.
    if not val_components and component_items:
        _tie, comp_id, idxs = component_items[0]
        val_components.add(comp_id)
        val_count += len(idxs)

    validation_indices = sorted(
        idx for comp_id in val_components for idx in component_to_indices[comp_id]
    )
    train_indices = sorted(
        idx
        for comp_id, idxs in component_to_indices.items()
        if comp_id not in val_components
        for idx in idxs
    ) + sorted(hub_only_indices)

    train_meaningful = {
        key
        for idx in train_indices
        for key in record_sets[idx]
        if key not in hubs
    }
    val_meaningful = {
        key
        for idx in validation_indices
        for key in record_sets[idx]
        if key not in hubs
    }
    overlap = train_meaningful & val_meaningful
    if overlap:
        raise RuntimeError(f"Meaningful molecule leakage detected: {len(overlap)} compounds appear in both splits")

    return {
        "train": [records[idx] for idx in train_indices],
        "validation": [records[idx] for idx in validation_indices],
        "train_indices": train_indices,
        "validation_indices": validation_indices,
        "excluded_indices": excluded_indices,
        "hub_only_indices": sorted(hub_only_indices),
        "hub_compounds": sorted(hubs),
        "train_compounds": sorted(train_meaningful),
        "validation_compounds": sorted(val_meaningful),
        "diagnostics": {
            "hub_threshold": hub_threshold,
            "target_validation_records": target_val,
            "component_count": len(component_to_indices),
            "top_component_record_counts": sorted((len(v) for v in component_to_indices.values()), reverse=True)[:10],
        },
    }


def leakage_spectrum(
    train_records: list[dict[str, Any]],
    validation_records: list[dict[str, Any]],
    *,
    solvent_cas: str = "64-17-5",
    ignored_molecules: set[str] | None = None,
) -> dict[str, Any]:
    ignored = ignored_molecules or set()
    train_sets = [
        _active_molecule_set(r, solvent_cas) - ignored
        for r in train_records
        if _active_molecule_set(r, solvent_cas) - ignored
    ]
    nearest: list[float] = []
    for record in validation_records:
        val_set = _active_molecule_set(record, solvent_cas) - ignored
        if not val_set or not train_sets:
            continue
        nearest.append(max(len(val_set & train_set) / len(val_set | train_set) for train_set in train_sets))
    return {
        "n_validation_with_fingerprint": len(nearest),
        "nearest_train_jaccard": {
            "mean": _summary(nearest)["mean"],
            "p50": _percentile(nearest, 0.5),
            "p90": _percentile(nearest, 0.9),
            "p95": _percentile(nearest, 0.95),
            "max": max(nearest) if nearest else None,
        },
    }


def selection_bias_report(
    records: list[dict[str, Any]],
    included_indices: list[int],
    excluded_indices: list[int],
    *,
    solvent_cas: str = "64-17-5",
) -> dict[str, Any]:
    record_sets = [_active_molecule_set(r, solvent_cas) for r in records]
    frequency = Counter(key for keys in record_sets for key in keys)

    def group(indices: list[int]) -> dict[str, Any]:
        formula_sizes = [float(len(record_sets[idx])) for idx in indices]
        mean_freqs = [
            sum(frequency[key] for key in record_sets[idx]) / len(record_sets[idx])
            for idx in indices
            if record_sets[idx]
        ]
        objective_lengths = [
            float(len(records[idx].get("objective_targets") or []))
            for idx in indices
            if records[idx].get("objective_targets") is not None
        ]
        psychometric_means = [
            float(sum(vals) / len(vals))
            for idx in indices
            for vals in [records[idx].get("psychometric_targets") or []]
            if vals
        ]
        return {
            "n": len(indices),
            "formula_size": _summary(formula_sizes),
            "mean_molecule_formula_frequency": _summary(mean_freqs),
            "objective_target_length": _summary(objective_lengths),
            "psychometric_target_mean": _summary(psychometric_means),
        }

    return {
        "included": group(included_indices),
        "excluded": group(excluded_indices),
    }


def diagnose_split_strategy(
    records: list[dict[str, Any]],
    *,
    train_ratio: float = DEFAULT_TRAIN_RATIO,
    seed: int = 42,
    solvent_cas: str = "64-17-5",
    hub_thresholds: tuple[float, ...] = (0.2, 0.1, 0.05, 0.02, 0.01),
    chosen_hub_threshold: float = 0.1,
) -> dict[str, Any]:
    strict = create_molecule_disjoint_split_from_records(records, train_ratio, seed, solvent_cas)
    hub_split = create_hub_excluded_component_split_from_records(
        records,
        train_ratio=train_ratio,
        seed=seed,
        solvent_cas=solvent_cas,
        hub_threshold=chosen_hub_threshold,
    )
    return {
        "canonicalization": {
            "identifier": "InChIKey from RDKit structure when available; fallback to CAS/name for unresolved rows",
            "ethanol_control": "CONTROL_64-17-5 excluded as a control",
        },
        "strict_random_molecule_split": {
            "train_records": len(strict["train_indices"]),
            "validation_records": len(strict["validation_indices"]),
            "excluded_records": len(strict["excluded_indices"]),
            "selection_bias": selection_bias_report(
                records,
                strict["train_indices"] + strict["validation_indices"],
                strict["excluded_indices"],
                solvent_cas=solvent_cas,
            ),
        },
        "graph": analyze_molecule_graph(records, solvent_cas=solvent_cas, thresholds=hub_thresholds),
        "chosen_strategy": {
            "name": "canonicalized component split with optional hub exclusion",
            "hub_threshold": chosen_hub_threshold,
            "justification": (
                "Canonicalization fixes identifier fragmentation and includes nearly all records. "
                "The requested 20/10/5% hub thresholds do not identify any hubs in v9; lower "
                "thresholds are reported as sensitivity analysis because aggressive hub removal "
                "can ignore too many aroma materials."
            ),
            "train_records": len(hub_split["train_indices"]),
            "validation_records": len(hub_split["validation_indices"]),
            "excluded_records": len(hub_split["excluded_indices"]),
            "hub_count": len(hub_split["hub_compounds"]),
            "meaningful_train_validation_overlap": 0,
            "diagnostics": hub_split["diagnostics"],
            "leakage_spectrum": leakage_spectrum(
                hub_split["train"],
                hub_split["validation"],
                solvent_cas=solvent_cas,
                ignored_molecules=set(hub_split["hub_compounds"]),
            ),
            "selection_bias": selection_bias_report(
                records,
                hub_split["train_indices"] + hub_split["validation_indices"],
                hub_split["excluded_indices"],
                solvent_cas=solvent_cas,
            ),
        },
    }


def create_formula_disjoint_split_from_dataset(
    dataset: Dataset,
    train_ratio: float = DEFAULT_TRAIN_RATIO,
) -> DatasetDict:
    """
    Partition a Hugging Face Dataset of formula records at the formula level.
    """
    records = [dict(r) for r in dataset]
    train_records, val_records = create_formula_disjoint_split_from_records(
        records, train_ratio=train_ratio
    )
    return DatasetDict({
        "train": Dataset.from_list(train_records),
        "validation": Dataset.from_list(val_records),
    })


def create_formula_disjoint_split_from_records(
    records: list[dict[str, Any]],
    train_ratio: float = DEFAULT_TRAIN_RATIO,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """
    Partition a list of formula records at the formula level.

    Avoids row-by-row iteration over a Hugging Face Arrow Dataset and avoids the
    expensive Python-list <-> Arrow round-trip. Returns plain Python lists ready for
    FragranceTrajectoryDataset(records=...).
    """
    _validate_ratio("train_ratio", train_ratio)
    train_records = []
    val_records = []

    for r in records:
        formula_string = _formula_fingerprint(r)
        hash_val = int(hashlib.sha256(formula_string.encode("utf-8")).hexdigest(), 16)
        if (hash_val % 100) < (train_ratio * 100):
            train_records.append(r)
        else:
            val_records.append(r)

    print("✅ Formula-Disjoint Split Complete.")
    print(f"   - Training Recipes: {len(train_records)}")
    print(f"   - Validation Recipes: {len(val_records)}")

    return train_records, val_records


def create_formula_disjoint_split(
    jsonl_path: str | Path,
    train_ratio: float = DEFAULT_TRAIN_RATIO,
) -> DatasetDict:
    """
    Partition a JSON-L formulation corpus at the formula level.

    Thin wrapper around :func:`create_formula_disjoint_split_from_dataset` that
    loads the JSON-L file first.
    """
    jsonl_path = Path(jsonl_path)
    logger.info("Loading raw corpus from %s", jsonl_path)
    raw_dataset = load_dataset("json", data_files=str(jsonl_path), split="train")
    return create_formula_disjoint_split_from_dataset(raw_dataset, train_ratio=train_ratio)


def create_molecule_disjoint_split(
    jsonl_path: str | Path,
    train_ratio: float = DEFAULT_TRAIN_RATIO,
    seed: int = 42,
    solvent_cas: str = "64-17-5",
) -> DatasetDict:
    """
    Partition a JSON-L formulation corpus into molecule-disjoint train/validation sets.

    A recipe containing even one molecule assigned to the validation tier is pushed
    to the holdout validation set. The solvent (default ethanol) is ignored because
    it appears in every formula. Returns a Hugging Face DatasetDict.
    """
    jsonl_path = Path(jsonl_path)
    logger.info("Loading raw corpus from %s", jsonl_path)
    raw_dataset = load_dataset("json", data_files=str(jsonl_path), split="train")

    all_records = list(raw_dataset)
    split = create_molecule_disjoint_split_from_records(
        all_records,
        train_ratio=train_ratio,
        seed=seed,
        solvent_cas=solvent_cas,
    )

    logger.info(
        "Disjoint Split Complete. Train Formulas: %d | Val Formulas: %d | Excluded: %d",
        len(split["train_indices"]),
        len(split["validation_indices"]),
        len(split["excluded_indices"]),
    )

    return DatasetDict({
        "train": raw_dataset.select(split["train_indices"]),
        "validation": raw_dataset.select(split["validation_indices"]),
    })


def create_molecule_disjoint_split_from_records(
    records: list[dict[str, Any]],
    train_ratio: float = DEFAULT_TRAIN_RATIO,
    seed: int = 42,
    solvent_cas: str = "64-17-5",
) -> dict[str, Any]:
    """
    Partition records into train/validation sets with zero active-CAS overlap.

    Molecules are assigned deterministically from ``seed``. A record is included
    in training only when all active molecules are train molecules, included in
    validation only when all active molecules are held out, and excluded when it
    mixes molecules from both sides. The solvent CAS is ignored.
    """
    _validate_ratio("train_ratio", train_ratio)

    record_cas_sets = [_active_cas_set(rec, solvent_cas) for rec in records]
    all_molecules = sorted({cas for cas_set in record_cas_sets for cas in cas_set})
    if not all_molecules:
        raise ValueError("No active molecules found in records")

    rng = random.Random(seed)
    rng.shuffle(all_molecules)
    split_idx = int(len(all_molecules) * train_ratio)
    split_idx = min(max(split_idx, 1), len(all_molecules) - 1)
    train_molecules = set(all_molecules[:split_idx])
    held_out = set(all_molecules[split_idx:])

    train_indices: list[int] = []
    validation_indices: list[int] = []
    excluded_indices: list[int] = []
    for idx, cas_set in enumerate(record_cas_sets):
        if not cas_set:
            excluded_indices.append(idx)
        elif cas_set.issubset(train_molecules):
            train_indices.append(idx)
        elif cas_set.issubset(held_out):
            validation_indices.append(idx)
        else:
            excluded_indices.append(idx)

    train_cas = {cas for idx in train_indices for cas in record_cas_sets[idx]}
    validation_cas = {cas for idx in validation_indices for cas in record_cas_sets[idx]}
    overlap = train_cas & validation_cas
    if overlap:
        raise RuntimeError(f"Molecule leakage detected: {len(overlap)} compounds appear in both splits")

    return {
        "train": [records[idx] for idx in train_indices],
        "validation": [records[idx] for idx in validation_indices],
        "train_indices": train_indices,
        "validation_indices": validation_indices,
        "excluded_indices": excluded_indices,
        "train_compounds": sorted(train_cas),
        "validation_compounds": sorted(validation_cas),
        "held_out_compounds": sorted(held_out),
    }


def molecule_disjoint_split(
    jsonl_path: str | Path,
    val_ratio: float = DEFAULT_VAL_RATIO,
    seed: int = 42,
    solvent_cas: str = "64-17-5",
) -> dict[str, Any]:
    """
    Partition a JSON-L formulation corpus into molecule-disjoint train/validation sets.

    A held-out subset of individual active compounds is selected from the registry.
    Validation records are required to contain *only* held-out compounds; training
    records must contain *no* held-out compounds. The solvent (e.g. ethanol) is
    ignored because it is present in every formula.
    """
    jsonl_path = Path(jsonl_path)
    logger.info("Loading raw corpus from %s", jsonl_path)
    raw_dataset = load_dataset("json", data_files=str(jsonl_path), split="train")

    # Build per-record active compound sets and collect the universe of actives.
    all_records = list(raw_dataset)
    record_cas_sets = [_active_cas_set(rec, solvent_cas) for rec in all_records]
    active_universe = sorted({cas for cas_set in record_cas_sets for cas in cas_set})
    logger.info("Found %d unique active compounds across %d records", len(active_universe), len(all_records))

    # Deterministic hold-out of individual molecules.
    rng = sorted(active_universe)
    import random

    random.seed(seed)
    random.shuffle(rng)
    n_holdout = max(1, int(len(active_universe) * val_ratio))
    held_out = set(rng[:n_holdout])
    train_allowed = set(rng[n_holdout:])
    logger.info(
        "Held out %d/%d compounds for validation (%.1f%%)",
        len(held_out),
        len(active_universe),
        100 * len(held_out) / len(active_universe),
    )

    train_indices = []
    val_indices = []
    for idx, cas_set in enumerate(record_cas_sets):
        if not cas_set:
            continue  # skip degenerate records
        if cas_set.issubset(held_out):
            val_indices.append(idx)
        elif cas_set.isdisjoint(held_out):
            train_indices.append(idx)
        else:
            # Record mixes held-out and seen compounds; exclude to keep split strict.
            continue

    logger.info(
        "Split complete: %d train records | %d validation records | %d excluded",
        len(train_indices),
        len(val_indices),
        len(all_records) - len(train_indices) - len(val_indices),
    )

    # Verify disjointness.
    train_cas = {cas for idx in train_indices for cas in record_cas_sets[idx]}
    val_cas = {cas for idx in val_indices for cas in record_cas_sets[idx]}
    overlap = train_cas & val_cas
    if overlap:
        raise RuntimeError(f"Molecule leakage detected: {len(overlap)} compounds appear in both splits")
    if val_cas - held_out:
        raise RuntimeError("Validation set contains compounds not in the held-out set")
    if held_out - (train_cas | val_cas):
        logger.warning(
            "%d held-out compounds never appear in any record; split is still valid but coverage is incomplete",
            len(held_out - (train_cas | val_cas)),
        )

    logger.info("Validation split uses %d unique compounds, all held out from training", len(val_cas))
    logger.info("Active compound overlap between train and validation: 0%%")

    return {
        "train": raw_dataset.select(train_indices),
        "validation": raw_dataset.select(val_indices),
        "held_out_compounds": sorted(held_out),
        "train_compounds": sorted(train_cas),
        "validation_compounds": sorted(val_cas),
    }


def split_and_upload_dataset(
    jsonl_path: str,
    repo_id: str,
    val_ratio: float = DEFAULT_VAL_RATIO,
    seed: int = 42,
    solvent_cas: str = "64-17-5",
) -> None:
    """
    Load a production JSON-L corpus, enforce a molecule-disjoint train/validation
    split, and push the resulting DatasetDict to the Hugging Face Hub.
    """
    split = molecule_disjoint_split(jsonl_path, val_ratio=val_ratio, seed=seed, solvent_cas=solvent_cas)

    # Cast genre to ClassLabel for downstream stratification / metrics.
    genres = sorted({rec["genre"] for rec in split["train"] + split["validation"]})
    for key in ("train", "validation"):
        split[key] = split[key].cast_column("genre", ClassLabel(names=genres))

    dataset_dict = DatasetDict({"train": split["train"], "validation": split["validation"]})

    print(f"Train rows: {len(dataset_dict['train'])} | Validation rows: {len(dataset_dict['validation'])}")
    print(f"Held-out compounds: {len(split['held_out_compounds'])}")
    print("Active compound overlap between train and validation: 0%")

    token = os.environ.get("HF_TOKEN")
    print(f"Uploading molecule-disjoint DatasetDict to HF Hub: {repo_id} ...")
    dataset_dict.push_to_hub(repo_id, private=True, token=token)
    print("Dataset sync complete. Training partitions are fully live and molecule-disjoint.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Molecule-disjoint dataset split and upload for PINO")
    parser.add_argument("--jsonl", default="data/synthetic_dataset_v2_text.jsonl", help="Input JSON-L corpus")
    parser.add_argument("--repo-id", default="mattbitzesty/pino-synthetic-dataset", help="HF Hub dataset repo")
    parser.add_argument("--val-ratio", type=float, default=DEFAULT_VAL_RATIO, help="Fraction of compounds to hold out")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for hold-out selection")
    parser.add_argument("--solvent-cas", default="64-17-5", help="CAS number of the universal solvent")
    parser.add_argument("--diagnose-only", action="store_true", help="Write split diagnostics without uploading")
    parser.add_argument("--diagnostics-output", default="artifacts/data_split_diagnostics.json", help="Diagnostics JSON path")
    parser.add_argument("--hub-threshold", type=float, default=0.1, help="Formula-frequency threshold for hub-excluded component diagnostics")
    parser.add_argument("--log-level", default="INFO", help="Logging level")
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper(), logging.INFO),
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )

    if args.diagnose_only:
        raw_dataset = load_dataset("json", data_files=str(args.jsonl), split="train")
        report = diagnose_split_strategy(
            list(raw_dataset),
            train_ratio=1.0 - args.val_ratio,
            seed=args.seed,
            solvent_cas=args.solvent_cas,
            chosen_hub_threshold=args.hub_threshold,
        )
        output = Path(args.diagnostics_output)
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
        print(json.dumps({
            "output": str(output),
            "strict_random_molecule_split": report["strict_random_molecule_split"],
            "chosen_strategy": {
                k: v
                for k, v in report["chosen_strategy"].items()
                if k not in {"selection_bias", "leakage_spectrum", "diagnostics"}
            },
        }, indent=2))
    else:
        split_and_upload_dataset(
            jsonl_path=args.jsonl,
            repo_id=args.repo_id,
            val_ratio=args.val_ratio,
            seed=args.seed,
            solvent_cas=args.solvent_cas,
        )