#!/usr/bin/env python3
"""Enrich empirical_dataset_v1.jsonl with physics-derived time-varying targets.

For each record we compute how the OAV budget is distributed across top / heart /
base notes at every timestep. The objective target is then a time-varying blend
of the corresponding descriptor profiles (citrus/fresh/green for top notes,
floral/rose/jasmin for heart notes, woody/musk/amber for base notes). This
produces a target that is physically grounded in the diffusion trajectory and
varies over time because top notes fade and base notes emerge.

If Pyrfume annotations are available for an ingredient, the note class is
reconciled with the Pyrfume descriptor profile by using the ingredient's OAV
peak time as the primary signal.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import numpy as np

DESCRIPTOR_DIM = 138

# Descriptor profiles for the three fragrance notes in the Pyrfume vocabulary.
TOP_PROFILE = ["citrus", "fresh", "fruity", "green", "aldehydic", "lemon", "grapefruit", "orange"]
HEART_PROFILE = ["floral", "rose", "jasmin", "jasmine", "muguet", "lily", "neroli", "ylang", "hyacinth", "lavender"]
BASE_PROFILE = ["woody", "musk", "amber", "vanilla", "balsamic", "sweet", "cedar", "leathery", "tobacco"]


def build_profile_vectors(vocab: list[str]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    def _vec(words: list[str]) -> np.ndarray:
        arr = np.zeros(DESCRIPTOR_DIM, dtype=np.float32)
        for w in words:
            if w in vocab:
                arr[vocab.index(w)] = 1.0
        return arr / (arr.sum() + 1e-8)
    return _vec(TOP_PROFILE), _vec(HEART_PROFILE), _vec(BASE_PROFILE)


def classify_ingredients(oav: np.ndarray) -> list[str]:
    """Classify each ingredient as top/heart/base by its OAV peak location."""
    t_steps = oav.shape[0]
    notes: list[str] = []
    for s in range(oav.shape[1]):
        traj = oav[:, s]
        if traj.max() < 1e-12:
            notes.append("heart")
            continue
        peak_t = int(np.argmax(traj))
        frac = peak_t / (t_steps - 1) if t_steps > 1 else 0.5
        if frac < 0.33:
            notes.append("top")
        elif frac > 0.66:
            notes.append("base")
        else:
            notes.append("heart")
    return notes


def enrich_record(
    record: dict[str, Any],
    top_vec: np.ndarray,
    heart_vec: np.ndarray,
    base_vec: np.ndarray,
) -> dict[str, Any]:
    record = dict(record)
    formula = record.get("formula", [])
    trajectory = record.get("trajectory", [])
    if not formula or not trajectory:
        return record

    t_steps = len(trajectory)
    s_ing = len(formula)
    oav = np.zeros((t_steps, s_ing), dtype=np.float32)
    for t, step in enumerate(trajectory):
        oav_step = step.get("OAV", {})
        for s, comp in enumerate(formula):
            cas = comp.get("cas", "")
            lookup_cas = cas[len("NATURAL:"):] if cas.startswith("NATURAL:") else cas
            oav[t, s] = oav_step.get(lookup_cas, 0.0)

    notes = classify_ingredients(oav)

    # Build a time-varying descriptor matrix per ingredient.
    ingredient_matrix = np.zeros((t_steps, s_ing, DESCRIPTOR_DIM), dtype=np.float32)
    for s, note in enumerate(notes):
        if note == "top":
            profile = top_vec
        elif note == "base":
            profile = base_vec
        else:
            profile = heart_vec
        # The ingredient's descriptor contribution is modulated by its OAV at
        # each timestep so that it fades in/out of the mixture profile.
        traj = oav[:, s]
        max_oav = traj.max()
        if max_oav < 1e-12:
            weights = np.zeros(t_steps, dtype=np.float32)
        else:
            weights = traj / max_oav
        ingredient_matrix[:, s, :] = weights[:, None] * profile[None, :]

    # Mixture target per timestep is the sum over ingredients.
    targets = ingredient_matrix.sum(axis=1)  # (T, 138)
    targets = np.clip(targets, 0.0, 1.0)

    # Normalize each timestep to a unit-length descriptor direction while
    # preserving the relative descriptor weights within the timestep.
    norms = np.linalg.norm(targets, axis=1, keepdims=True)
    targets = np.divide(targets, norms, out=np.zeros_like(targets), where=norms > 1e-8)

    record["objective_targets"] = targets.astype(np.float32).tolist()
    return record


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", default="data/empirical_dataset_v1.jsonl")
    parser.add_argument("--output", default="data/empirical_dataset_v2.jsonl")
    parser.add_argument("--vocab", default="data/pyrfume_vocabulary.json")
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output)

    vocab_data = json.loads(Path(args.vocab).read_text())
    vocab = vocab_data["vocabulary"] if isinstance(vocab_data, dict) else vocab_data
    assert len(vocab) == DESCRIPTOR_DIM, f"Unexpected vocab size {len(vocab)}"
    print(f"Loaded vocab: {len(vocab)} descriptors")

    top_vec, heart_vec, base_vec = build_profile_vectors(vocab)

    records = [json.loads(line) for line in input_path.open()]
    print(f"Enriching {len(records)} records...")

    out_file = output_path.open("w")
    variance_list: list[float] = []
    note_counts = {"top": 0, "heart": 0, "base": 0}
    for rec in records:
        enriched = enrich_record(rec, top_vec, heart_vec, base_vec)
        obj = np.array(enriched["objective_targets"])
        variance_list.append(float(obj.var(axis=0).mean()))
        out_file.write(json.dumps(enriched, default=float) + "\n")

        # Count notes for diagnostics.
        t_steps = len(rec.get("trajectory", []))
        if t_steps:
            s_ing = len(rec.get("formula", []))
            oav = np.zeros((t_steps, s_ing), dtype=np.float32)
            for t, step in enumerate(rec["trajectory"]):
                oav_step = step.get("OAV", {})
                for s, comp in enumerate(rec["formula"]):
                    cas = comp.get("cas", "")
                    lookup_cas = cas[len("NATURAL:"):] if cas.startswith("NATURAL:") else cas
                    oav[t, s] = oav_step.get(lookup_cas, 0.0)
            for note in classify_ingredients(oav):
                note_counts[note] += 1
    out_file.close()

    variances = np.array(variance_list)
    print(f"Wrote {len(records)} records to {output_path}")
    print(f"Note distribution: {note_counts}")
    print(f"Mean objective variance: {variances.mean():.6f}")
    print(f"Median objective variance: {np.median(variances):.6f}")
    print(f"Max objective variance: {variances.max():.6f}")
    print(f"Fraction > 1e-3: {(variances > 1e-3).mean():.3f}")
    print(f"Fraction > 1e-2: {(variances > 1e-2).mean():.3f}")
    print(f"Fraction > 5e-2: {(variances > 5e-2).mean():.3f}")


if __name__ == "__main__":
    main()