src/pino/ingest_formulas.py

from __future__ import annotations

"""
Literature fragrance formula ingestion and normalization.

Reads raw recipe manifests (e.g., from Poucher, Arctander, Calkin & Jellinek),
recursively unpacks proprietary sub-bases into pure CAS constituents, and scales
all masses to strict unit weight fractions (Σ w_i = 1.0).

The output is directly consumable by the PINO optimizer and verifier.
"""

import json
import logging
from collections import defaultdict
from pathlib import Path
from typing import Any

from .registry import AromaRegistry
from .thermo.unifac import validate_subgroups

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


ETHANOL_CAS = "64-17-5"


def _recursive_unpack(
    component: dict[str, Any],
    accumulator: dict[str, float],
    parent_multiplier: float = 1.0,
) -> None:
    """
    Recursively unpack a component into pure CAS mass contributions.

    Parameters
    ----------
    component: dict
        A recipe component with keys 'type', 'amount', 'cas'/'name'/'composition'.
    accumulator: dict
        Running map of CAS -> absolute mass.
    parent_multiplier: float
        Scale factor inherited from parent bases (product of ancestor amounts).
    """
    comp_type = component.get("type", "pure")
    amount = float(component.get("amount", 0.0)) * parent_multiplier

    if comp_type == "pure":
        cas = component.get("cas")
        if not cas:
            raise ValueError(f"Pure component missing CAS: {component}")
        accumulator[cas] += amount

    elif comp_type == "sub_base":
        base_composition = component.get("composition") or []
        if not base_composition:
            logger.warning("Sub-base '%s' has no composition; skipping", component.get("name"))
            return

        # Determine how sub-base internal fractions are expressed.
        sub_total = sum(float(sub.get("amount", 0.0)) for sub in base_composition)
        if sub_total <= 0:
            raise ValueError(f"Sub-base '{component.get('name')}' composition sums to <= 0")

        for sub in base_composition:
            sub_component = dict(sub)
            sub_component["amount"] = float(sub.get("amount", 0.0)) / sub_total * amount
            _recursive_unpack(sub_component, accumulator, parent_multiplier=1.0)

    else:
        raise ValueError(f"Unknown component type: {comp_type}")


def normalize_and_unpack_recipe(raw_recipe_path: str | Path) -> dict[str, Any]:
    """
    Load a raw literature recipe and return a normalized pure-CAS formula.

    The returned dict has:
      - formula_id, name, source
      - components: list of {"cas", "weight_fraction"} with Σ = 1.0
    """
    path = Path(raw_recipe_path)
    data = json.loads(path.read_text(encoding="utf-8"))

    flat_composition: dict[str, float] = defaultdict(float)

    for comp in data.get("components", []):
        _recursive_unpack(comp, flat_composition)

    absolute_sum = sum(flat_composition.values())
    if absolute_sum <= 0:
        raise ValueError(f"Recipe {data.get('formula_id')} has zero total mass")

    normalized_components = []
    for cas in sorted(flat_composition.keys()):
        weight_fraction = round(flat_composition[cas] / absolute_sum, 6)
        normalized_components.append({"cas": cas, "weight_fraction": weight_fraction})

    # Micro-adjust to ensure exact 1.0 after rounding.
    total_fraction = sum(c["weight_fraction"] for c in normalized_components)
    if total_fraction != 1.0 and normalized_components:
        normalized_components[0]["weight_fraction"] = round(
            normalized_components[0]["weight_fraction"] + (1.0 - total_fraction), 6
        )

    return {
        "formula_id": data.get("formula_id"),
        "name": data.get("name"),
        "source": data.get("source"),
        "components": normalized_components,
    }


def load_literature_manifest(manifest_path: str | Path) -> list[dict[str, Any]]:
    """Load a JSON file that may be either a single recipe or a list of recipes."""
    path = Path(manifest_path)
    data = json.loads(path.read_text(encoding="utf-8"))
    if isinstance(data, list):
        return data
    return [data]


def normalize_and_unpack_recipe_from_dict(data: dict[str, Any]) -> dict[str, Any]:
    """In-memory version of normalize_and_unpack_recipe for a single recipe dict."""
    flat_composition: dict[str, float] = defaultdict(float)

    for comp in data.get("components", []):
        _recursive_unpack(comp, flat_composition)

    absolute_sum = sum(flat_composition.values())
    if absolute_sum <= 0:
        raise ValueError(f"Recipe {data.get('formula_id')} has zero total mass")

    normalized_components = []
    for cas in sorted(flat_composition.keys()):
        weight_fraction = round(flat_composition[cas] / absolute_sum, 6)
        normalized_components.append({"cas": cas, "weight_fraction": weight_fraction})

    total_fraction = sum(c["weight_fraction"] for c in normalized_components)
    if total_fraction != 1.0 and normalized_components:
        normalized_components[0]["weight_fraction"] = round(
            normalized_components[0]["weight_fraction"] + (1.0 - total_fraction), 6
        )

    return {
        "formula_id": data.get("formula_id"),
        "name": data.get("name"),
        "source": data.get("source"),
        "components": normalized_components,
    }


def _build_registry_cas_sets(registry: AromaRegistry) -> tuple[set[str], set[str]]:
    """Return (known_cas, unifac_valid_cas) from the registry."""
    rows = registry._conn.execute(
        "SELECT cas, name, unifac_groups FROM aroma_chemicals"
    ).fetchall()
    known_cas = {row[0] for row in rows}

    unifac_valid_cas: set[str] = set()
    for cas, name, groups_json in rows:
        try:
            groups = json.loads(groups_json or "{}")
        except Exception:
            groups = {}
        valid, _ = validate_subgroups(groups)
        if valid:
            unifac_valid_cas.add(cas)

    return known_cas, unifac_valid_cas


def _filter_and_rescale_components(
    components: list[dict[str, Any]],
    keep_cas: set[str],
    formula_id: str | None,
) -> list[dict[str, Any]]:
    """Drop filtered components and rescale the remaining weight fractions to 1.0."""
    if len(components) == len(keep_cas):
        return components

    filtered = [c for c in components if c["cas"] in keep_cas]
    total = sum(c["weight_fraction"] for c in filtered)
    if total <= 0:
        logger.warning("Recipe %s has no valid components after filtering", formula_id)
        return []

    for c in filtered:
        c["weight_fraction"] = round(c["weight_fraction"] / total, 6)

    adj = sum(c["weight_fraction"] for c in filtered)
    if filtered and adj != 1.0:
        filtered[0]["weight_fraction"] = round(
            filtered[0]["weight_fraction"] + (1.0 - adj), 6
        )

    return filtered


def normalize_manifest(
    manifest_path: str | Path,
    registry_path: str | Path | None = None,
    *,
    strict_unifac: bool = True,
) -> list[dict[str, Any]]:
    """
    Normalize every recipe in a manifest and validate each CAS against the registry.

    Recipes that contain unknown CAS numbers are flagged but still returned so
    the caller can decide whether to drop, map, or back-fill them.

    If strict_unifac is True, components whose stored UNIFAC subgroups are not
    available in DOUFSG are rejected (ethanol, the canonical solvent, is always
    allowed). Weight fractions are re-scaled after any filtering.
    """
    registry = AromaRegistry(registry_path)
    known_cas, unifac_valid_cas = _build_registry_cas_sets(registry)

    results = []
    for raw_recipe in load_literature_manifest(manifest_path):
        tmp = dict(raw_recipe)
        normalized = normalize_and_unpack_recipe_from_dict(tmp)

        unknown = [c["cas"] for c in normalized["components"] if c["cas"] not in known_cas]
        if unknown:
            logger.warning(
                "Recipe %s contains unknown CAS entries: %s",
                normalized["formula_id"],
                unknown,
            )
            normalized["registry_status"] = "unknown_cas"
            normalized["unknown_cas"] = unknown
        else:
            normalized["registry_status"] = "ok"

        if strict_unifac:
            unifac_missing = [
                c["cas"] for c in normalized["components"]
                if c["cas"] not in unifac_valid_cas and c["cas"] != ETHANOL_CAS
            ]
            if unifac_missing:
                logger.warning(
                    "Recipe %s contains components without valid DOUFSG subgroups: %s",
                    normalized["formula_id"],
                    unifac_missing,
                )
                normalized["unifac_status"] = "invalid_subgroups"
                normalized["unifac_invalid_cas"] = unifac_missing
            else:
                normalized["unifac_status"] = "ok"
        else:
            normalized["unifac_status"] = "skipped"

        # Build the set of components to keep after registry + UNIFAC filtering.
        keep_cas = {
            c["cas"] for c in normalized["components"]
            if c["cas"] in known_cas
            and (not strict_unifac or c["cas"] in unifac_valid_cas or c["cas"] == ETHANOL_CAS)
        }

        normalized["components"] = _filter_and_rescale_components(
            normalized["components"],
            keep_cas,
            normalized["formula_id"],
        )
        if not normalized["components"]:
            normalized["registry_status"] = "no_valid_components"

        results.append(normalized)
    return results


def recipe_to_formula_records(recipe: dict[str, Any]) -> list[dict[str, Any]]:
    """Convert a normalized recipe into the verifier-friendly formula format."""
    return [
        {"cas": c["cas"], "weight_fraction": c["weight_fraction"]}
        for c in recipe.get("components", [])
    ]


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Normalize literature fragrance recipes")
    parser.add_argument("manifest", help="Path to raw recipe JSON or JSON list")
    parser.add_argument("--registry", default="src/pino/registry.db", help="SQLite registry path")
    parser.add_argument("--output", help="Optional output JSON path")
    parser.add_argument(
        "--no-strict-unifac",
        action="store_true",
        help="Allow components with missing DOUFSG subgroups to pass validation",
    )
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

    results = normalize_manifest(
        args.manifest,
        registry_path=args.registry,
        strict_unifac=not args.no_strict_unifac,
    )
    out_text = json.dumps(results, indent=2, ensure_ascii=False)
    if args.output:
        Path(args.output).write_text(out_text, encoding="utf-8")
    else:
        print(out_text)