#!/usr/bin/env python3
"""Design the prospective formula benchmark (in-silico, preregistered).

Generates 30-50 deliberately-designed formulas from a restricted palette of
common hobbyist-available aroma chemicals (option held open for a future
community physical-compounding arm). Each formula spans a genre and a target
family profile with a top/heart/base substantivity structure.

Blinding: formulas and their *intended* target profiles are stored separately.
- data/benchmarks/prospective_formulas/formulas.jsonl  : formula + structure (shareable)
- data/benchmarks/prospective_formulas/labels.sequestered.json : intended targets (EVAL ONLY)
The labels file carries access:"evaluation_only" and must not enter training or
model-selection inputs.

Deterministic: --seed fixes all sampling so the set is reproducible.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import random
from pathlib import Path

GENRES = {
    "citrus_cologne":  ["citrus_fresh", "aromatic_herbal", "green"],
    "floral_woody":    ["floral_sweet", "woody_amber", "musk_clean"],
    "amber_oriental":  ["gourmand", "woody_amber", "musk_clean"],
    "fougere":         ["aromatic_herbal", "green", "woody_amber"],
    "wildcard":        None,  # any combination
}

# Substantivity bands (log10 hours-ish predicted) used to slot top/heart/base.
def band(sub):
    if sub is None:
        return "heart"
    if sub < 1.0:
        return "top"
    if sub < 1.9:
        return "heart"
    return "base"


def load_palette(path: Path) -> dict[str, dict]:
    return {k: v for k, v in json.loads(path.read_text()).items()}


def pick(seq, rng, n):
    seq = list(seq)
    rng.shuffle(seq)
    return seq[:n]


def build_formula(fid, genre, palette, rng):
    fams = GENRES[genre]
    if fams is None:
        fams = rng.sample(
            ["citrus_fresh", "floral_sweet", "woody_amber", "gourmand", "aromatic_herbal", "musk_clean", "green"],
            k=rng.choice([2, 3]),
        )
    # Pool of materials across the chosen families, slotted by substantivity band.
    by_band = {"top": [], "heart": [], "base": []}
    for name, meta in palette.items():
        if meta["family"] in fams:
            by_band[band(meta.get("substantivity_log10"))].append((name, meta))
    n_ing = rng.randint(6, 10)
    n_top = max(1, round(n_ing * 0.3))
    n_base = max(1, round(n_ing * 0.3))
    n_heart = n_ing - n_top - n_base

    chosen = []
    chosen += pick(by_band["top"], rng, min(n_top, len(by_band["top"])))
    chosen += pick(by_band["heart"], rng, min(n_heart, len(by_band["heart"])))
    chosen += pick(by_band["base"], rng, min(n_base, len(by_band["base"])))
    # Top up if a band ran dry.
    if len(chosen) < n_ing:
        rest = [(n, m) for b in by_band.values() for (n, m) in b if n not in {c[0] for c in chosen}]
        chosen += pick(rest, rng, min(n_ing - len(chosen), len(rest)))
    chosen = chosen[:n_ing]

    # Dosing: base notes heavier, top notes lighter, with jitter; then normalise.
    raw = []
    for name, meta in chosen:
        b = band(meta.get("substantivity_log10"))
        w = {"top": rng.uniform(0.5, 2.0), "heart": rng.uniform(1.0, 4.0), "base": rng.uniform(2.0, 8.0)}[b]
        raw.append((name, meta, w))
    total = sum(w for _, _, w in raw)
    ingredients = [
        {
            "name": meta["name"],
            "cas": meta.get("cas"),
            "smiles": meta.get("smiles"),
            "family": meta["family"],
            "weight_fraction": round(w / total, 4),
        }
        for (name, meta, w) in raw
    ]

    # Intended target: weight-averaged family profile (the sequestered label).
    fam_target: dict[str, float] = {}
    for ing in ingredients:
        fam_target[ing["family"]] = fam_target.get(ing["family"], 0.0) + ing["weight_fraction"]
    fam_target = {k: round(v, 4) for k, v in sorted(fam_target.items())}

    formula = {
        "formula_id": fid,
        "genre": genre,
        "intended_families": fams,
        "ingredients": ingredients,
        "n_ingredients": len(ingredients),
        "design": "in_silico_preregistered",
    }
    label = {
        "formula_id": fid,
        "target_family_profile": fam_target,
        "access": "evaluation_only",
    }
    return formula, label


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--palette", default="artifacts/prospective/common_palette.json")
    ap.add_argument("--out-dir", default="data/benchmarks/prospective_formulas")
    ap.add_argument("--n", type=int, default=40)
    ap.add_argument("--seed", type=int, default=20260717)
    args = ap.parse_args()

    palette = load_palette(Path(args.palette))
    rng = random.Random(args.seed)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    genres = list(GENRES)
    formulas, labels = [], []
    for i in range(args.n):
        genre = genres[i % len(genres)]
        fid = f"PROSP_{args.seed}_{i:03d}"
        f, l = build_formula(fid, genre, palette, rng)
        formulas.append(f)
        labels.append(l)

    fpath = out_dir / "formulas.jsonl"
    with fpath.open("w") as fh:
        for f in formulas:
            fh.write(json.dumps(f) + "\n")
    lpath = out_dir / "labels.sequestered.json"
    lpath.write_text(json.dumps({
        "access": "evaluation_only",
        "note": "Intended target family profiles. Forbidden in training and model selection.",
        "seed": args.seed,
        "labels": labels,
    }, indent=1))

    manifest = {
        "benchmark": "prospective_formulas",
        "created": "2026-07-17",
        "seed": args.seed,
        "n_formulas": len(formulas),
        "genres": {g: sum(1 for f in formulas if f["genre"] == g) for g in genres},
        "palette_size": len(palette),
        "evaluation": "in_silico (model scoring/retrieval vs sequestered intended profiles)",
        "community_arm": "option held open; palette restricted to hobbyist-available materials",
        "label_access": "evaluation_only; forbidden in training and model selection",
        "formula_sha256": hashlib.sha256(fpath.read_bytes()).hexdigest(),
        "label_sha256": hashlib.sha256(lpath.read_bytes()).hexdigest(),
    }
    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=1))
    print(f"wrote {len(formulas)} formulas -> {fpath}")
    print(f"sequestered labels -> {lpath}")
    print("genre counts:", manifest["genres"])


if __name__ == "__main__":
    main()