#!/usr/bin/env python3
"""Build and freeze the v2 genre challenge set without inspecting embeddings."""

from __future__ import annotations

import argparse
from collections import Counter, defaultdict
import hashlib
import json
from pathlib import Path
from typing import Any, Iterable

import numpy as np
from rdkit import Chem, RDLogger
from rdkit.Chem import Descriptors

from pino.genre_benchmark import SOLVENTS
from pino.registry import AromaRegistry
from pino.thermo.naturals import NATURAL_PROFILES


SUBSTANTIVE_GENRES = ("amber_oriental", "citrus_cologne", "floral_woody", "fougere")
ELEMENTS = ("C", "H", "N", "O", "F", "P", "S", "Cl", "Br", "I", "B", "Si")
RDLogger.DisableLog("rdApp.error")


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def identity(component: dict[str, Any]) -> str:
    return str(component.get("cas") or component.get("smiles") or component.get("name") or "").strip().lower()


def genre(row: dict[str, Any]) -> str:
    return str(row.get("genre") or row.get("metadata", {}).get("generation_strategy") or "wildcard")


def positive_active_components(row: dict[str, Any]) -> list[dict[str, Any]]:
    return [
        component for component in row.get("formula", [])
        if identity(component) not in SOLVENTS and float(component.get("weight_fraction", 0.0)) > 0
    ]


def _natural_profile(identifier: str) -> dict[str, Any] | None:
    key = identifier.strip()
    return NATURAL_PROFILES.get(key) or NATURAL_PROFILES.get(f"NATURAL:{key}")


def _resolve_smiles(component: dict[str, Any], registry: AromaRegistry | None) -> str | None:
    """Resolve the structure carried inline, in a SMILES: key, or by registry id."""
    inline = str(component.get("smiles") or "").strip()
    identifier = str(component.get("cas") or component.get("name") or "").strip()
    candidates = [inline]
    if identifier.upper().startswith("SMILES:"):
        candidates.append(identifier.split(":", 1)[1])
    if registry is not None and identifier:
        record = registry.get(identifier)
        if record:
            candidates.append(str(record.get("smiles") or "").strip())
    for candidate in candidates:
        if not candidate:
            continue
        if candidate.upper().startswith("SMILES:"):
            candidate = candidate.split(":", 1)[1]
        mol = Chem.MolFromSmiles(candidate)
        if mol is not None and len(Chem.GetMolFrags(mol)) == 1:
            return Chem.MolToSmiles(mol, canonical=True)
    return None


def resolved_active_components(
    row: dict[str, Any], registry: AromaRegistry | None = None
) -> tuple[list[dict[str, Any]], list[str]]:
    """Expand complex materials and aggregate their resolved pure constituents."""
    aggregated: dict[str, dict[str, Any]] = {}
    unresolved: list[str] = []

    def add(component: dict[str, Any], weight: float, lineage: tuple[str, ...] = ()) -> None:
        identifier = str(component.get("cas") or component.get("name") or "").strip()
        profile = _natural_profile(identifier)
        if profile:
            profile_key = identifier.removeprefix("NATURAL:")
            if profile_key in lineage:
                unresolved.append(identifier)
                return
            for constituent, fraction in profile["constituents"].items():
                add({"cas": constituent}, weight * float(fraction), lineage + (profile_key,))
            return
        smiles = _resolve_smiles(component, registry)
        if not smiles:
            unresolved.append(identifier or str(component.get("smiles") or "<missing identifier>"))
            return
        entry = aggregated.setdefault(smiles, {"cas": identifier, "smiles": smiles, "weight_fraction": 0.0})
        entry["weight_fraction"] += weight

    for component in positive_active_components(row):
        add(component, float(component["weight_fraction"]))
    return list(aggregated.values()), unresolved


def characterize(
    row: dict[str, Any], source_index: int, registry: AromaRegistry | None = None
) -> dict[str, Any]:
    active, unresolved = resolved_active_components(row, registry)
    weights = np.asarray([float(c["weight_fraction"]) for c in active], dtype=float)
    element_mass = Counter()
    valid = not unresolved
    for component, weight in zip(active, weights):
        smiles = str(component["smiles"])
        mol = Chem.MolFromSmiles(smiles)
        if mol is None or len(Chem.GetMolFrags(mol)) != 1:
            valid = False
            continue
        molecular_weight = Descriptors.MolWt(mol)
        if molecular_weight <= 0:
            valid = False
            continue
        for atom in mol.GetAtoms():
            element_mass[atom.GetSymbol()] += weight * atom.GetMass() / molecular_weight
    total = float(weights.sum())
    proportions = weights / total if total else weights
    hhi = float(np.square(proportions).sum()) if len(proportions) else 0.0
    return {
        "source_index": source_index,
        "genre": genre(row),
        "active_count": len(active),
        "structurally_valid": bool(active) and valid,
        "unresolved_identifiers": sorted(set(unresolved)),
        "element_vector": [float(element_mass[e]) for e in ELEMENTS],
        "element_other": float(sum(v for e, v in element_mass.items() if e not in ELEMENTS)),
        "hhi": hhi,
        "top_weight_share": float(proportions.max()) if len(proportions) else 0.0,
        "ingredients": sorted({str(c["smiles"]) for c in active}),
        "natural_weight_share": natural_weight_share(row),
    }


def natural_weight_share(row: dict[str, Any]) -> float:
    """Fraction of the active formula entered as a natural/complex material."""
    active = positive_active_components(row)
    total = sum(float(c["weight_fraction"]) for c in active)
    if not total:
        return 0.0
    natural = sum(
        float(c["weight_fraction"]) for c in active
        if _natural_profile(str(c.get("cas") or c.get("name") or "")) is not None
    )
    return natural / total


def size_bin(active_count: int) -> str:
    if active_count == 2:
        return "2"
    if active_count <= 5:
        return "3-5"
    if active_count <= 10:
        return "6-10"
    return "11+"


def composition_match(a: dict[str, Any], b: dict[str, Any]) -> bool:
    """Predeclared hard-negative caliper using label-free formula properties."""
    if size_bin(a["active_count"]) != size_bin(b["active_count"]):
        return False
    va, vb = np.asarray(a["element_vector"]), np.asarray(b["element_vector"])
    denom = float(np.linalg.norm(va) * np.linalg.norm(vb))
    cosine = float(np.dot(va, vb) / denom) if denom else 0.0
    return (
        cosine >= 0.98
        and abs(a["hhi"] - b["hhi"]) <= 0.10
        and abs(a["top_weight_share"] - b["top_weight_share"]) <= 0.10
        and not (set(a["ingredients"]) & set(b["ingredients"]))
    )


def exact_formula_key(row: dict[str, Any]) -> tuple[tuple[str, int], ...]:
    """Order-independent resolved formula key, with weights rounded to 1e-6."""
    return tuple(sorted(
        (str(c["smiles"]), round(float(c["weight_fraction"]) * 1_000_000))
        for c in row["resolved_components"]
    ))


def near_duplicate(a: dict[str, Any], b: dict[str, Any]) -> bool:
    """Conservative near-duplicate rule: same ingredients and nearly same weights."""
    aw = {c["smiles"]: float(c["weight_fraction"]) for c in a["resolved_components"]}
    bw = {c["smiles"]: float(c["weight_fraction"]) for c in b["resolved_components"]}
    if aw.keys() != bw.keys():
        return False
    at, bt = sum(aw.values()), sum(bw.values())
    return max(abs(aw[k] / at - bw[k] / bt) for k in aw) <= 0.02


def collapse_duplicates(rows: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Keep the lowest source index from each exact/near formula family."""
    families: dict[tuple[str, ...], list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        families[tuple(row["ingredients"])].append(row)
    kept, clusters = [], []
    for ingredient_key in sorted(families):
        representatives: list[dict[str, Any]] = []
        for row in sorted(families[ingredient_key], key=lambda r: r["source_index"]):
            match = next((r for r in representatives if near_duplicate(r, row)), None)
            if match is None:
                representatives.append(row)
                kept.append(row)
            else:
                clusters.append({"representative": match["source_index"], "collapsed": row["source_index"]})
    return sorted(kept, key=lambda r: r["source_index"]), {
        "input_records": len(rows), "representatives": len(kept),
        "collapsed_records": len(clusters), "clusters": clusters,
        "rule": "identical resolved ingredient set and maximum normalized weight difference <= 0.02",
    }


def selection_distance(a: dict[str, Any], b: dict[str, Any]) -> float:
    va, vb = np.asarray(a["element_vector"]), np.asarray(b["element_vector"])
    cosine = float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
    return ((1 - cosine) / 0.02 + abs(a["hhi"] - b["hhi"]) / .10
            + abs(a["top_weight_share"] - b["top_weight_share"]) / .10
            + abs(a["natural_weight_share"] - b["natural_weight_share"]))


def select_balanced_challenge(
    records: list[dict[str, Any]], registry: AromaRegistry | None, per_genre: int
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    """Deterministically select disjoint, pairwise-calipered four-genre blocks."""
    candidates = []
    for index, source in enumerate(records):
        row = characterize(source, index, registry)
        if row["genre"] not in SUBSTANTIVE_GENRES or row["active_count"] < 2 or not row["structurally_valid"]:
            continue
        row["formula_id"] = source.get("formula_id") or source.get("metadata", {}).get("formula_id")
        row["resolved_components"], _ = resolved_active_components(source, registry)
        candidates.append(row)
    candidates, duplicate_audit = collapse_duplicates(candidates)
    by_genre = {g: [r for r in candidates if r["genre"] == g] for g in SUBSTANTIVE_GENRES}
    used: set[int] = set()
    blocks: list[list[dict[str, Any]]] = []
    anchor_genre = min(SUBSTANTIVE_GENRES, key=lambda g: len(by_genre[g]))
    bins = ("3-5", "6-10", "11+")
    common = {b: min(sum(size_bin(r["active_count"]) == b for r in by_genre[g])
                         for g in SUBSTANTIVE_GENRES) for b in bins}
    denominator = sum(common.values())
    raw_quota = {b: per_genre * common[b] / denominator for b in bins}
    bin_quota = {b: int(raw_quota[b]) for b in bins}
    for b in sorted(bins, key=lambda x: (-(raw_quota[x] - bin_quota[x]), x))[:per_genre - sum(bin_quota.values())]:
        bin_quota[b] += 1
    anchors = sorted(by_genre[anchor_genre], key=lambda r: (size_bin(r["active_count"]), r["hhi"], r["source_index"]))
    selected_bins: Counter[str] = Counter()
    for anchor in anchors:
        if anchor["source_index"] in used:
            continue
        anchor_bin = size_bin(anchor["active_count"])
        if anchor_bin not in bin_quota or selected_bins[anchor_bin] >= bin_quota[anchor_bin]:
            continue
        block = [anchor]
        for name in SUBSTANTIVE_GENRES:
            if name == anchor["genre"]:
                continue
            feasible = [r for r in by_genre[name] if r["source_index"] not in used
                        and all(composition_match(r, chosen) for chosen in block)]
            if not feasible:
                break
            block.append(min(feasible, key=lambda r: (sum(selection_distance(r, x) for x in block), r["source_index"])))
        if len(block) == len(SUBSTANTIVE_GENRES):
            blocks.append(block)
            used.update(r["source_index"] for r in block)
            selected_bins[anchor_bin] += 1
            if len(blocks) == per_genre:
                break
    # If a scarce bin cannot realize its proportional quota as a complete
    # pairwise clique, fill from any remaining bin without relaxing calipers.
    if len(blocks) < per_genre:
        for anchor in anchors:
            if anchor["source_index"] in used:
                continue
            block = [anchor]
            for name in SUBSTANTIVE_GENRES:
                if name == anchor["genre"]:
                    continue
                feasible = [r for r in by_genre[name] if r["source_index"] not in used
                            and all(composition_match(r, chosen) for chosen in block)]
                if not feasible:
                    break
                block.append(min(feasible, key=lambda r: (
                    sum(selection_distance(r, x) for x in block), r["source_index"])))
            if len(block) == len(SUBSTANTIVE_GENRES):
                blocks.append(block)
                used.update(r["source_index"] for r in block)
                selected_bins[size_bin(anchor["active_count"])] += 1
                if len(blocks) == per_genre:
                    break
    if len(blocks) < per_genre:
        raise ValueError(f"only {len(blocks)} pairwise-matched blocks available; requested {per_genre}")
    selected = [r for block in blocks for r in sorted(block, key=lambda r: r["genre"])]
    lean = [{k: r[k] for k in ("source_index", "formula_id", "genre", "active_count", "element_vector",
                                "element_other", "hhi", "top_weight_share", "natural_weight_share", "ingredients")}
            for r in selected]
    audit = {
        "selection_method": "deterministic greedy pairwise-calipered four-genre blocks",
        "blocks": len(blocks), "records": len(lean), "per_genre": count_by_genre(lean),
        "size_bin_quota": bin_quota,
        "source_indices_unique": len({r["source_index"] for r in lean}) == len(lean),
        "duplicate_collapse": duplicate_audit,
        "size_bins_by_genre": {g: dict(sorted(Counter(size_bin(r["active_count"]) for r in lean if r["genre"] == g).items())) for g in SUBSTANTIVE_GENRES},
        "natural_weight_share_by_genre": {g: {"mean": float(np.mean([r["natural_weight_share"] for r in lean if r["genre"] == g])), "nonzero": sum(r["natural_weight_share"] > 0 for r in lean if r["genre"] == g)} for g in SUBSTANTIVE_GENRES},
    }
    return lean, audit


def count_by_genre(rows: Iterable[dict[str, Any]]) -> dict[str, int]:
    counts = Counter(row["genre"] for row in rows)
    return {name: counts.get(name, 0) for name in (*SUBSTANTIVE_GENRES, "wildcard")}


def census(records: list[dict[str, Any]], registry: AromaRegistry | None = None) -> dict[str, Any]:
    characterized = [characterize(row, i, registry) for i, row in enumerate(records)]
    multi = [r for r in characterized if r["active_count"] >= 2]
    valid = [r for r in multi if r["structurally_valid"]]
    substantive = [r for r in valid if r["genre"] in SUBSTANTIVE_GENRES]

    partners: dict[int, set[str]] = defaultdict(set)
    pair_counts: Counter[tuple[str, str]] = Counter()
    for i, left in enumerate(substantive):
        for right in substantive[i + 1:]:
            if left["genre"] == right["genre"] or not composition_match(left, right):
                continue
            pair = tuple(sorted((left["genre"], right["genre"])))
            pair_counts[pair] += 1
            partners[left["source_index"]].add(right["genre"])
            partners[right["source_index"]].add(left["genre"])

    matchable = [r for r in substantive if partners[r["source_index"]]]
    per_genre = {}
    for name in SUBSTANTIVE_GENRES:
        eligible = [r for r in substantive if r["genre"] == name]
        matched = [r for r in eligible if partners[r["source_index"]]]
        per_genre[name] = {
            "eligible": len(eligible),
            "composition_matchable": len(matched),
            "matchable_to_all_three_other_genres": sum(len(partners[r["source_index"]]) == 3 for r in eligible),
            "by_size_bin": dict(sorted(Counter(size_bin(r["active_count"]) for r in eligible).items())),
        }
    return {
        "records": len(records),
        "resolution": {
            "records_with_unresolved_active_identifiers": sum(bool(r["unresolved_identifiers"]) for r in characterized),
            "unresolved_identifier_counts": dict(sorted(Counter(
                identifier for r in characterized for identifier in r["unresolved_identifiers"]
            ).items())),
        },
        "attrition": {
            "all": count_by_genre(characterized),
            "multi_component": count_by_genre(multi),
            "multi_component_structurally_valid": count_by_genre(valid),
            "substantive_composition_matchable": count_by_genre(matchable),
        },
        "per_substantive_genre": per_genre,
        "cross_genre_candidate_pairs": {"__vs__".join(pair): count for pair, count in sorted(pair_counts.items())},
        "maximum_balanced_pool_from_individually_matchable_records": min(
            (per_genre[g]["composition_matchable"] for g in SUBSTANTIVE_GENRES), default=0
        ),
        "wildcard_open_set_pool": {
            "multi_component_structurally_valid": sum(r["genre"] == "wildcard" for r in valid),
            "singleton_sanity_check": sum(r["genre"] == "wildcard" and r["active_count"] == 1 for r in characterized),
        },
    }


def specification(dataset: Path, report: dict[str, Any], selected: bool = False) -> dict[str, Any]:
    return {
        "protocol_version": 2,
        "status": ("frozen_challenge_rows_selected; embeddings_and_outcomes_not_inspected" if selected else
                   "frozen_design_and_feasibility_census; challenge_rows_not_selected; embeddings_not_inspected"),
        "dataset": {"path": str(dataset), "sha256": sha256(dataset), "records": report["records"]},
        "primary_population": {
            "genres": list(SUBSTANTIVE_GENRES),
            "requirements": [
                "at least two positive-weight non-solvent resolved constituents",
                "natural/complex materials expanded recursively",
                "CAS and SMILES: identifiers resolve to one-fragment canonical SMILES",
                "duplicate resolved constituents aggregated before characterization",
            ],
            "excluded": "wildcard and singleton records",
        },
        "matching": {
            "size_bins": ["2", "3-5", "6-10", "11+"],
            "element_mass_cosine_minimum": 0.98,
            "concentration_hhi_absolute_tolerance": 0.10,
            "top_weight_share_absolute_tolerance": 0.10,
            "shared_active_ingredient_maximum": 0,
            "selection_rule": "balanced across four genres; optimize match coverage without using embeddings or outcomes",
        },
        "tasks": {
            "primary": [
                "real formula versus non-identity weight permutation",
                "real formula versus ingredient-matched decoy",
                "same-genre versus different-genre composition-matched retrieval",
            ],
            "residual_value": "nested cross-validation: composition alone versus composition plus learned embedding",
            "open_set": "wildcard is rejection-only and never a fifth primary class",
            "sanity_check": "report singleton wildcard classification separately",
        },
        "metrics": {
            "classification": ["macro_f1", "balanced_accuracy", "per_class_recall"],
            "contrastive": ["paired_accuracy", "roc_auc"],
            "retrieval": ["same_genre_enrichment", "mean_reciprocal_rank"],
            "uncertainty": "formula-level paired bootstrap 95% intervals",
            "raw_accuracy_is_primary": False,
        },
        "leakage_controls": [
            "all preprocessing and hyperparameter selection fit inside training folds",
            "exact formula and active ingredient identities disjoint across outer folds",
            "decoys and contrastive variants remain in the source formula's fold",
            "final row selection is frozen before learned embeddings are inspected",
        ],
        "go_no_go": {
            "minimum_individually_matchable_records_per_genre": 20,
            "minimum_candidate_pairs_per_genre_pair": 20,
            "pass": report["maximum_balanced_pool_from_individually_matchable_records"] >= 20
                and all(v >= 20 for v in report["cross_genre_candidate_pairs"].values())
                and len(report["cross_genre_candidate_pairs"]) == 6,
        },
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--registry", default=None, help="Aroma registry SQLite path")
    parser.add_argument("--per-genre", type=int, default=50,
                        help="Rows per substantive genre in the frozen challenge set")
    args = parser.parse_args()
    dataset, output = Path(args.dataset), Path(args.output)
    with dataset.open(encoding="utf-8") as handle:
        records = [json.loads(line) for line in handle if line.strip()]
    with AromaRegistry(args.registry) as registry:
        report = census(records, registry)
        selected, audit = select_balanced_challenge(records, registry, args.per_genre)
    result = {"specification": specification(dataset, report, selected=True),
              "feasibility_census": report, "selection_audit": audit,
              "challenge_records": selected}
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()