src/pino/dataset_builder.py

from __future__ import annotations

import argparse
import json
import logging
import signal
import threading
import uuid
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
from pathlib import Path
from typing import Any

import numpy as np

from pino.draft_engine import FormulaGenerator
from pino.verifier import FragrancePipelineVerifier

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

GENRES = ["citrus_cologne", "fougere", "floral_woody", "amber_oriental", "wildcard"]

# Worker-thread local storage for generator and verifier reuse.
_worker_local = threading.local()


def _get_worker_generator(
    genre: str,
    rules_path: Path,
    registry_path: Path,
    literature_path: Path,
    base_seed: int,
) -> FormulaGenerator:
    """Return a thread-local FormulaGenerator for this genre."""
    current_genre = getattr(_worker_local, "genre", None)
    if current_genre != genre or not hasattr(_worker_local, "generator"):
        thread_id = threading.current_thread().ident or 0
        _worker_local.generator = FormulaGenerator(
            genre=genre,
            rules_path=rules_path,
            registry_path=registry_path,
            literature_path=literature_path,
            seed=base_seed + hash(genre) + thread_id,
            min_k=5,
            max_k=20,
        )
        _worker_local.genre = genre
    return _worker_local.generator


def _get_worker_verifier() -> FragrancePipelineVerifier:
    """Return a thread-local FragrancePipelineVerifier."""
    if not hasattr(_worker_local, "verifier"):
        _worker_local.verifier = FragrancePipelineVerifier()
    return _worker_local.verifier


class DatasetBuilder:
    """
    Streaming, multi-threaded dataset builder for PINO synthetic trajectories.

    Runs one stratified sweep per genre, writes verified records append-only to
    a JSON-Lines file, and emits a final analytics card.
    """

    def __init__(
        self,
        output_path: str | Path,
        rules_path: str | Path,
        registry_path: str | Path,
        literature_path: str | Path,
        per_genre: int = 1000,
        max_workers: int = 16,
        seed: int = 2026,
    ) -> None:
        self.output_path = Path(output_path)
        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        self.rules_path = Path(rules_path)
        self.registry_path = Path(registry_path)
        self.literature_path = Path(literature_path)
        self.per_genre = per_genre
        self.max_workers = max_workers
        self.seed = seed
        self._shutdown = False

    def _signal_handler(self, signum: int, frame: Any) -> None:
        logger.warning("Shutdown signal received; finishing in-flight work ...")
        self._shutdown = True

    def _task(self, idx: int, genre: str) -> dict[str, Any] | None:
        """Generate and verify a single formula in a worker thread."""
        generator = _get_worker_generator(
            genre, self.rules_path, self.registry_path, self.literature_path, self.seed
        )
        formula, formula_id = generator.generate(idx=idx)

        if not generator.light_ifra_check(formula)["passed"]:
            return None

        try:
            verifier = _get_worker_verifier()
            result = verifier.run_sim(formula, duration_seconds=28800, interval_seconds=600)
        except Exception as exc:
            logger.debug("Verifier rejected %s: %s", formula_id, exc)
            return None

        if result.get("status") not in ("passed", "depleted"):
            return None

        # Build OAV-weighted descriptor targets using the new semantics module.
        from pino import semantics
        composition = [
            {"cas": c.get("cas"), "smiles": c.get("smiles")}
            for c in formula
        ]
        C_gas = np.array([step["C_gas_mg_m3"] for step in result["trajectory"]])
        # Align concentration matrix by canonical CAS ordering.
        cas_order = [c["cas"] for c in composition]
        concentrations = np.zeros((C_gas.shape[0], len(cas_order)), dtype=np.float32)
        for t_idx, step in enumerate(result["trajectory"]):
            for c_idx, cas in enumerate(cas_order):
                concentrations[t_idx, c_idx] = step["C_gas_mg_m3"].get(cas, 0.0)
        objective_targets = semantics.compute_oav_targets(composition, concentrations)
        psychometric_targets = semantics.get_psychometric_targets(genre)

        return {
            "status": result.get("status"),
            "formula_id": formula_id,
            "genre": genre,
            "metadata": {
                "generation_strategy": genre,
                "active_components_count": len(formula) - 1,  # exclude solvent
                "formula_id": formula_id,
            },
            "formula": formula,
            "depletion_rates": result.get("depletion_rates", {}),
            "trajectory": result.get("trajectory", []),
            "objective_targets": objective_targets.tolist(),
            "psychometric_targets": psychometric_targets.tolist(),
        }

    def _append_record(self, record: dict[str, Any]) -> None:
        with self.output_path.open("a") as f:
            f.write(json.dumps(record) + "\n")

    def _count_existing(self, genre: str) -> int:
        if not self.output_path.exists():
            return 0
        count = 0
        with self.output_path.open("r") as f:
            for line in f:
                try:
                    rec = json.loads(line)
                    if rec.get("metadata", {}).get("generation_strategy") == genre:
                        count += 1
                except Exception:
                    continue
        return count

    def _run_genre(self, genre: str, global_start_idx: int) -> dict[str, Any]:
        """Generate and verify `per_genre` records for a single genre."""
        existing = self._count_existing(genre)
        target = self.per_genre - existing
        if target <= 0:
            logger.info("Genre %s already has %d records; skipping", genre, existing)
            return {"genre": genre, "verified": existing, "drafted": 0, "rejected": 0}

        logger.info("Starting genre %s | need %d more records", genre, target)

        verified = existing
        drafted = 0
        rejected = 0
        next_idx = global_start_idx

        with ThreadPoolExecutor(max_workers=self.max_workers, thread_name_prefix=f"pino-{genre}") as executor:
            futures: set[Any] = set()

            while verified < self.per_genre and not self._shutdown:
                # Keep the worker pool saturated.
                while len(futures) < self.max_workers * 4 and not self._shutdown:
                    futures.add(executor.submit(self._task, next_idx, genre))
                    next_idx += 1
                    drafted += 1
                    if drafted >= target * 5:
                        # Safety valve: stop drafting if rejection rate is too high.
                        break

                if not futures:
                    break

                done, futures = wait(futures, return_when=FIRST_COMPLETED)
                for future in done:
                    result = future.result()
                    if result is None:
                        rejected += 1
                        continue
                    self._append_record(result)
                    verified += 1

                    if verified % 100 == 0 and verified > existing:
                        logger.info(
                            "Genre %s: verified %d/%d | drafted %d | rejected %d",
                            genre,
                            verified,
                            self.per_genre,
                            drafted,
                            rejected,
                        )
                    if verified >= self.per_genre:
                        break

                # Safety valve: if rejection rate is high, move on anyway.
                if drafted > max(100, verified * 3) and verified < self.per_genre:
                    logger.warning(
                        "Genre %s has high rejection rate (verified %d, drafted %d); continuing",
                        genre, verified, drafted
                    )

        logger.info(
            "Genre %s complete: verified %d/%d | drafted %d | rejected %d",
            genre,
            verified,
            self.per_genre,
            drafted,
            rejected,
        )
        return {"genre": genre, "verified": verified, "drafted": drafted, "rejected": rejected}

    def build(self) -> dict[str, Any]:
        """Run the full stratified sweep and return the analytics card."""
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)

        # Wipe only if empty; otherwise resume.
        if not self.output_path.exists() or self.output_path.stat().st_size == 0:
            self.output_path.write_text("")

        card = {"genres": [], "total_verified": 0, "total_drafted": 0, "total_rejected": 0}
        global_idx = 0
        for genre in GENRES:
            summary = self._run_genre(genre, global_idx)
            card["genres"].append(summary)
            card["total_verified"] += summary["verified"]
            card["total_drafted"] += summary["drafted"]
            card["total_rejected"] += summary["rejected"]
            global_idx += self.per_genre

        logger.info(
            "Dataset complete: %d verified | %d drafted | %d rejected",
            card["total_verified"],
            card["total_drafted"],
            card["total_rejected"],
        )
        return card


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Build PINO synthetic dataset v2")
    parser.add_argument("--output", default="data/synthetic_dataset_v2.jsonl")
    parser.add_argument("--rules", default="data/genre_rules.json")
    parser.add_argument("--registry", default="src/pino/registry.db")
    parser.add_argument("--literature", default="data/literature_formulas.json")
    parser.add_argument("--per-genre", type=int, default=1000)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--seed", type=int, default=2026)
    parser.add_argument("--log-level", default="INFO")
    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level.upper()),
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    builder = DatasetBuilder(
        output_path=args.output,
        rules_path=args.rules,
        registry_path=args.registry,
        literature_path=args.literature,
        per_genre=args.per_genre,
        max_workers=args.workers,
        seed=args.seed,
    )
    card = builder.build()
    print(json.dumps(card, indent=2))