src/pino/substitution_v11.py

from __future__ import annotations

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

import numpy as np

PIMT_V11 = "v11"
ANSWER_SHAPES = {"single_material", "any_of", "accord"}
# Minimum mutually-observed scoring components for a row to count as a
# trustworthy retrieval. Below this the rank is driven mostly by priors /
# file order, so it is reported as low-evidence, not a model miss or hit.
MIN_OBSERVED_COMPONENTS = 2


def load_jsonl(path: str | Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in Path(path).open(encoding="utf-8") if line.strip()]


def write_jsonl(path: str | Path, rows: Iterable[dict[str, Any]]) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        "".join(json.dumps(row, sort_keys=True, ensure_ascii=False) + "\n" for row in rows),
        encoding="utf-8",
    )


def wilson_interval(successes: int, total: int, z: float = 1.959963984540054) -> dict[str, Any]:
    """Return count-first 95% Wilson score reporting for a binomial proportion."""
    if total < 0 or successes < 0 or successes > total:
        raise ValueError("Wilson counts must satisfy 0 <= successes <= total")
    if total == 0:
        return {"successes": successes, "total": total, "wilson_95": None}
    proportion = successes / total
    z2 = z * z
    denominator = 1.0 + z2 / total
    center = (proportion + z2 / (2.0 * total)) / denominator
    half_width = z * math.sqrt(
        (proportion * (1.0 - proportion) + z2 / (4.0 * total)) / total
    ) / denominator
    return {
        "successes": successes,
        "total": total,
        "wilson_95": [max(0.0, center - half_width), min(1.0, center + half_width)],
    }


def explicit_eval_row(row: dict[str, Any]) -> dict[str, Any]:
    """Upgrade a legacy one-substitute row to the explicit v11.1 answer schema."""
    if "answer" in row:
        return row
    return {
        **row,
        "answer": {
            "shape": "single_material",
            "materials": [{"name": row.get("substitute"), "cas": row.get("substitute_cas")}],
        },
    }


def _answer(row: dict[str, Any]) -> tuple[str, list[dict[str, Any]]]:
    answer = explicit_eval_row(row)["answer"]
    shape = answer.get("shape")
    materials = answer.get("materials") or []
    if shape not in ANSWER_SHAPES:
        raise ValueError(f"unsupported answer shape for {row.get('pair_id')}: {shape!r}")
    if not materials or any(not (item.get("cas") or item.get("name")) for item in materials):
        raise ValueError(f"answer materials missing for {row.get('pair_id')}")
    if shape == "single_material" and len(materials) != 1:
        raise ValueError(f"single_material answer must contain exactly one material: {row.get('pair_id')}")
    return str(shape), materials


def normalize_name(value: str | None) -> str:
    if not value:
        return ""
    keep = []
    for ch in value.lower():
        keep.append(ch if ch.isalnum() else " ")
    return " ".join("".join(keep).split())


def profile_key(profile: dict[str, Any]) -> str:
    return str(profile.get("cas") or profile.get("profile_id") or normalize_name(profile.get("name")))


def descriptor_distribution_from_codes(codes: Iterable[str]) -> dict[str, float]:
    counts = Counter(code for code in codes if code)
    total = sum(counts.values())
    if not total:
        return {}
    return {code: count / total for code, count in sorted(counts.items())}


def cosine_similarity(left: dict[str, float], right: dict[str, float]) -> float:
    keys = sorted(set(left) | set(right))
    if not keys:
        return 0.0
    a = np.asarray([left.get(key, 0.0) for key in keys], dtype=float)
    b = np.asarray([right.get(key, 0.0) for key in keys], dtype=float)
    denom = float(np.linalg.norm(a) * np.linalg.norm(b))
    if denom <= 0:
        return 0.0
    return float(np.dot(a, b) / denom)


def numeric_proximity(a: Any, b: Any, scale: float) -> float | None:
    try:
        fa = float(a)
        fb = float(b)
    except (TypeError, ValueError):
        return None
    if not (math.isfinite(fa) and math.isfinite(fb)):
        return None
    return float(math.exp(-abs(fa - fb) / scale))


def note_role(profile: dict[str, Any]) -> str:
    role = profile.get("poucher_tier")
    if role in {"top", "heart", "mid", "base"}:
        return "heart" if role == "mid" else str(role)
    log_vp = profile.get("log10_vapor_pressure_pa")
    if log_vp is None:
        return "unknown"
    try:
        value = float(log_vp)
    except (TypeError, ValueError):
        return "unknown"
    if value >= 2.0:
        return "top"
    if value <= 0.0:
        return "base"
    return "heart"


def role_score(left: dict[str, Any], right: dict[str, Any]) -> float | None:
    l_role = note_role(left)
    r_role = note_role(right)
    if "unknown" in {l_role, r_role}:
        return None
    return 1.0 if l_role == r_role else 0.25


def profile_lookup(profiles: Iterable[dict[str, Any]]) -> dict[str, dict[str, Any]]:
    out: dict[str, dict[str, Any]] = {}
    for profile in profiles:
        keys = {profile_key(profile), normalize_name(profile.get("name"))}
        cas = profile.get("cas")
        if cas:
            keys.add(str(cas))
        for alias in profile.get("aliases") or []:
            keys.add(normalize_name(alias))
        for key in keys:
            if key:
                out[key] = profile
    return out


MODE_WEIGHTS: dict[str, dict[str, float]] = {
    "odor_match": {"odor": 0.70, "substantivity": 0.10, "volatility": 0.10, "role": 0.10},
    "function_match": {"odor": 0.25, "substantivity": 0.25, "volatility": 0.25, "role": 0.25},
    "evaporation_match": {"odor": 0.10, "substantivity": 0.35, "volatility": 0.40, "role": 0.15},
    "regulatory_match": {"odor": 0.30, "substantivity": 0.20, "volatility": 0.20, "role": 0.15, "constraints": 0.15},
    "cost_match": {"odor": 0.35, "substantivity": 0.20, "volatility": 0.20, "role": 0.15, "constraints": 0.10},
    "accord_rebuild": {"odor": 0.35, "substantivity": 0.20, "volatility": 0.20, "role": 0.25},
}

MODE_REQUIRED_FIELDS = {
    "regulatory_match": ("regulatory_status", "ifra_restrictions", "allergens"),
    "cost_match": ("cost_tier", "price", "availability_tier"),
}


def mode_supported(mode: str, target: dict[str, Any], profiles: Iterable[dict[str, Any]]) -> bool:
    """Whether a mode has at least target and catalogue evidence for its advertised axis."""
    fields = MODE_REQUIRED_FIELDS.get(mode)
    if not fields:
        return True
    return any(target.get(field) is not None for field in fields) and any(
        any(profile.get(field) is not None for field in fields) for profile in profiles
    )


def mode_axis_similarity(mode: str, target: dict[str, Any], candidate: dict[str, Any]) -> float | None:
    """Compare observed mode-specific fields; never manufacture a neutral value."""
    if mode == "cost_match":
        numeric = numeric_proximity(target.get("price"), candidate.get("price"), 1.0)
        if numeric is not None:
            return numeric
        for field in ("cost_tier", "availability_tier"):
            if target.get(field) is not None and candidate.get(field) is not None:
                return 1.0 if target[field] == candidate[field] else 0.0
    if mode == "regulatory_match":
        scores = []
        if target.get("regulatory_status") is not None and candidate.get("regulatory_status") is not None:
            scores.append(1.0 if target["regulatory_status"] == candidate["regulatory_status"] else 0.0)
        for field in ("ifra_restrictions", "allergens"):
            left, right = set(target.get(field) or []), set(candidate.get(field) or [])
            if left or right:
                scores.append(len(left & right) / len(left | right))
        if scores:
            return sum(scores) / len(scores)
    return None


def score_candidate(
    target: dict[str, Any],
    candidate: dict[str, Any],
    mode: str = "function_match",
    constraints: dict[str, Any] | None = None,
) -> dict[str, Any]:
    constraints = constraints or {}
    banned = {normalize_name(v) for v in constraints.get("banned_materials", [])}
    banned.update(str(v) for v in constraints.get("banned_cas", []))
    c_key = profile_key(candidate)
    if c_key == profile_key(target):
        return {"score": -1.0, "reasons": {"self": True}}
    if c_key in banned or normalize_name(candidate.get("name")) in banned:
        return {"score": -1.0, "reasons": {"banned": True}}

    target_odor = target.get("descriptor_distribution") or {}
    candidate_odor = candidate.get("descriptor_distribution") or {}
    components: dict[str, float | None] = {
        "odor": cosine_similarity(target_odor, candidate_odor) if target_odor and candidate_odor else None,
        "substantivity": numeric_proximity(
            target.get("substantivity_log10_predicted"),
            candidate.get("substantivity_log10_predicted"),
            0.45,
        ),
        "volatility": numeric_proximity(target.get("log10_vapor_pressure_pa"), candidate.get("log10_vapor_pressure_pa"), 1.0),
        "role": role_score(target, candidate),
        "constraints": mode_axis_similarity(mode, target, candidate),
    }
    weights = MODE_WEIGHTS.get(mode, MODE_WEIGHTS["function_match"])
    observed = {name: value for name, value in components.items() if value is not None and name in weights}
    observed_weight = sum(weights[name] for name in observed)
    score = sum(float(value) * weights[name] for name, value in observed.items()) / observed_weight if observed_weight else 0.0
    return {
        "score": float(score),
        "reasons": components,
        "evidence_coverage": observed_weight / sum(weights.values()),
        "observed_components": sorted(observed),
    }


def substitute(
    target: str,
    profiles: list[dict[str, Any]],
    constraints: dict[str, Any] | None = None,
    mode: str = "function_match",
    top_k: int = 5,
) -> list[dict[str, Any]]:
    lookup = profile_lookup(profiles)
    target_profile = lookup.get(str(target)) or lookup.get(normalize_name(target))
    if target_profile is None:
        raise KeyError(f"target material not found in v11 profile table: {target}")
    ranked = []
    for candidate in profiles:
        scored = score_candidate(target_profile, candidate, mode=mode, constraints=constraints)
        if scored["score"] < 0:
            continue
        ranked.append(
            {
                "candidate": candidate.get("name"),
                "candidate_cas": candidate.get("cas"),
                "score": scored["score"],
                "components": scored["reasons"],
                "evidence_coverage": scored["evidence_coverage"],
                "observed_components": scored["observed_components"],
                "note_role": note_role(candidate),
            }
        )
    ranked.sort(
        key=lambda row: (
            -row["score"],
            -row["evidence_coverage"],
            str(row.get("candidate_cas") or ""),
            normalize_name(row.get("candidate")),
        )
    )
    return ranked[:top_k]


def evaluate_material(material: str, profiles: list[dict[str, Any]]) -> dict[str, Any]:
    lookup = profile_lookup(profiles)
    profile = lookup.get(str(material)) or lookup.get(normalize_name(material))
    if profile is None:
        raise KeyError(f"material not found in v11 profile table: {material}")
    return {
        "pimt_version": PIMT_V11,
        "type": "material",
        "material": profile.get("name"),
        "cas": profile.get("cas"),
        "odor_family_fit": profile.get("descriptor_distribution") or {},
        "predicted_substantivity_log10": profile.get("substantivity_log10_predicted"),
        "volatility": {
            "vapor_pressure_pa": profile.get("vapor_pressure_pa"),
            "log10_vapor_pressure_pa": profile.get("log10_vapor_pressure_pa"),
            "boiling_point_k": profile.get("boiling_point_k"),
            "note_role": note_role(profile),
        },
        "substitution_risks": profile.get("substitution_risks") or [],
        "enrichment_pending": profile.get("enrichment_pending") or [],
    }


def evaluate_formula(formula: list[dict[str, Any]], profiles: list[dict[str, Any]]) -> dict[str, Any]:
    lookup = profile_lookup(profiles)
    materials = []
    missing = []
    total_weight = 0.0
    role_weights = Counter()
    for item in formula:
        name = item.get("cas") or item.get("name")
        profile = lookup.get(str(name)) or lookup.get(normalize_name(str(name)))
        if profile is None:
            missing.append(name)
            continue
        weight = float(item.get("weight_fraction", item.get("percent", 0.0)) or 0.0)
        total_weight += weight
        role_weights[note_role(profile)] += weight
        materials.append(evaluate_material(str(name), profiles) | {"input_weight": weight})
    return {
        "pimt_version": PIMT_V11,
        "type": "formula",
        "n_materials": len(materials),
        "missing_materials": missing,
        "total_weight": total_weight,
        "top_heart_base_balance": dict(role_weights),
        "materials": materials,
        "substitution_risks": ["safety/use-level flags enrichment-pending; no safety claims emitted"],
    }


def evaluate_retriever(
    eval_rows: list[dict[str, Any]],
    profiles: list[dict[str, Any]],
    mode: str | None = None,
    top_k: int = 5,
    ranking_strategy: str = "model",
) -> dict[str, Any]:
    if ranking_strategy not in {"model", "random"}:
        raise ValueError(f"unknown ranking strategy: {ranking_strategy}")
    lookup = profile_lookup(profiles)
    evaluated = []
    excluded = []
    mode_counts = Counter()
    for legacy_row in eval_rows:
        row = explicit_eval_row(legacy_row)
        shape, answer_materials = _answer(row)
        target_key = row.get("target_cas") or row.get("target")
        if not (lookup.get(str(target_key)) or lookup.get(normalize_name(str(target_key)))):
            excluded.append({"pair_id": row["pair_id"], "reason": "target_missing_profile"})
            continue
        resolved_answers = []
        missing_answers = []
        for item in answer_materials:
            keys = [item.get("cas"), item.get("name")]
            profile = next(
                (lookup.get(str(key)) or lookup.get(normalize_name(str(key))) for key in keys if key),
                None,
            )
            if profile is None:
                missing_answers.append(item)
            else:
                resolved_answers.append(profile)
        insufficient = not resolved_answers if shape == "any_of" else bool(missing_answers)
        if insufficient:
            excluded.append(
                {
                    "pair_id": row["pair_id"],
                    "reason": "answer_material_missing_profile",
                    "missing_answer_materials": missing_answers,
                }
            )
            continue
        row_mode = mode or row.get("mode")
        if row_mode not in MODE_WEIGHTS:
            raise ValueError(f"unsupported retrieval mode for {row['pair_id']}: {row_mode!r}")
        target_profile = lookup.get(str(target_key)) or lookup.get(normalize_name(str(target_key)))
        target_profile_key = profile_key(target_profile)
        answer_keys = {profile_key(profile) for profile in resolved_answers}
        if target_profile_key in answer_keys:
            excluded.append({"pair_id": row["pair_id"], "reason": "answer_resolves_to_target"})
            continue
        if not mode_supported(row_mode, target_profile, profiles):
            excluded.append({"pair_id": row["pair_id"], "reason": "mode_features_unsupported", "mode": row_mode})
            continue
        mode_counts[row_mode] += 1
        if ranking_strategy == "model":
            ranked = substitute(str(target_key), profiles, mode=row_mode, top_k=top_k)
        else:
            candidates = [profile for profile in profiles if profile_key(profile) != target_profile_key]
            candidates.sort(
                key=lambda profile: hashlib.sha256(
                    f"{row['pair_id']}|{profile_key(profile)}".encode("utf-8")
                ).digest()
            )
            ranked = [
                {"candidate": profile.get("name"), "candidate_cas": profile.get("cas")}
                for profile in candidates[:top_k]
            ]
        ranked_keys = {
            profile_key(profile)
            for result in ranked
            for profile in [
                lookup.get(str(result.get("candidate_cas")))
                or lookup.get(normalize_name(result.get("candidate")))
            ]
            if profile
        }
        matched_keys = ranked_keys & answer_keys
        hit = answer_keys <= ranked_keys if shape == "accord" else bool(matched_keys)
        substitute_profile = resolved_answers[0]
        # Evidence floor: count mutually-observed scoring components between the
        # target and the documented answer. Rows under the floor are reported as
        # low-evidence so their rank is not read as a trustworthy model outcome.
        answer_scored = score_candidate(target_profile, substitute_profile, mode=row_mode) if target_profile is not None else {}
        answer_observed_components = len(answer_scored.get("observed_components") or [])
        low_evidence = answer_observed_components < MIN_OBSERVED_COMPONENTS
        evaluated.append(
            {
                "pair_id": row["pair_id"],
                "grade": row["grade"],
                "mode": row_mode,
                "answer_shape": shape,
                "hit_top_k": bool(hit),
                "top_k": ranked,
                "documented_answer": row["answer"],
                "answer_components_retrieved": len(matched_keys),
                "answer_components_total": len(answer_keys),
                "answer_observed_components": answer_observed_components,
                "answer_evidence_coverage": answer_scored.get("evidence_coverage"),
                "low_evidence": low_evidence,
                "unavailable_optional_answers": missing_answers if shape == "any_of" else [],
                "substantivity_delta_abs": (
                    abs(
                        float(target_profile.get("substantivity_log10_predicted"))
                        - float(substitute_profile.get("substantivity_log10_predicted"))
                    )
                    if target_profile
                    and substitute_profile
                    and target_profile.get("substantivity_log10_predicted") is not None
                    and substitute_profile.get("substantivity_log10_predicted") is not None
                    else None
                ),
            }
        )
    n = len(evaluated)
    n_hit = sum(1 for row in evaluated if row["hit_top_k"])
    n_missed = n - n_hit
    by_mode_top_k_hit = {}
    for row_mode in sorted(mode_counts):
        mode_rows = [row for row in evaluated if row["mode"] == row_mode]
        by_mode_top_k_hit[row_mode] = wilson_interval(
            sum(1 for row in mode_rows if row["hit_top_k"]),
            len(mode_rows),
        )
    subst_deltas = [row["substantivity_delta_abs"] for row in evaluated if row["substantivity_delta_abs"] is not None]
    n_low_evidence = sum(1 for row in evaluated if row.get("low_evidence"))
    # Trustworthy subset: rows meeting the evidence floor. Hits are also reported
    # restricted to this subset so the headline is not carried by prior-only ranks.
    trustworthy = [row for row in evaluated if not row.get("low_evidence")]
    n_trustworthy = len(trustworthy)
    n_trustworthy_hit = sum(1 for row in trustworthy if row["hit_top_k"])
    return {
        "pimt_version": PIMT_V11,
        "mode": mode or "per-row",
        "mode_counts": dict(sorted(mode_counts.items())),
        "ranking_strategy": ranking_strategy,
        "top_k": top_k,
        "n_eval_rows": len(eval_rows),
        "n_evaluable": n,
        "n_excluded": len(excluded),
        "n_hit": n_hit,
        "n_missed": n_missed,
        "n_low_evidence": n_low_evidence,
        "n_trustworthy": n_trustworthy,
        "n_trustworthy_hit": n_trustworthy_hit,
        "top_k_hit": wilson_interval(n_hit, n),
        "top_k_hit_trustworthy_only": wilson_interval(n_trustworthy_hit, n_trustworthy) if n_trustworthy else None,
        "by_mode_top_k_hit": by_mode_top_k_hit,
        "mean_substantivity_delta_abs": float(np.mean(subst_deltas)) if subst_deltas else None,
        "evaluated": evaluated,
        "excluded": excluded,
    }