feature_prompt.py

"""Format features + rule-derived candidates into a prompt for Nemotron.

The model is asked to PICK among grounded candidates and EXPLAIN, not to invent
faults from scratch. This collapses its job to something a 4B model does well.
"""
from __future__ import annotations

from typing import List

from audio_analyzer import AudioFeatures
from fault_rules import Candidate

SYSTEM_PROMPT = (
    "You are a master appliance and machinery technician with 20 years of "
    "experience diagnosing faults by ear. You explain clearly to non-experts. "
    "You only state conclusions supported by the measured evidence given to you."
)


def _build_feature_analysis(features: AudioFeatures, appliance: str) -> str:
    """Build a feature-based analysis section when rules return Inconclusive."""
    lines = []

    # Spectral analysis
    if features.spectral_centroid_hz > 3000:
        lines.append("- HIGH spectral centroid (>{:.0f} Hz) suggests harsh/grinding/friction noise — "
                      "bearings, metal-on-metal contact, or broken fan blades".format(
                          features.spectral_centroid_hz))
    elif features.spectral_centroid_hz > 1500:
        lines.append("- Moderate spectral centroid ({:.0f} Hz) — could be belt noise, "
                      "pump impeller, or motor hum with some harmonics".format(
                          features.spectral_centroid_hz))
    elif features.spectral_centroid_hz > 500:
        lines.append("- Low spectral centroid ({:.0f} Hz) suggests rumbling — unbalanced load, "
                      "loose mounting, or worn bearing".format(
                          features.spectral_centroid_hz))

    # ZCR analysis
    if features.zero_crossing_rate > 0.15:
        lines.append("- HIGH zero-crossing rate ({:.4f}) indicates harsh/grinding sound — "
                      "metal contact, broken gear teeth, or abrasive wear".format(
                          features.zero_crossing_rate))

    # Pattern analysis
    if features.has_regular_pattern:
        iv = features.pattern_interval_ms
        if iv < 50:
            lines.append("- FAST regular pattern (every {:.0f} ms) suggests high-speed "
                          "rotational defect — bearing race or gear tooth".format(iv))
        elif iv < 200:
            lines.append("- Regular pattern (every {:.0f} ms) suggests mechanical periodicity — "
                          "bearing defect, belt joint, or pump impeller".format(iv))
        else:
            lines.append("- Slow regular pattern (every {:.0f} ms) suggests intermittent "
                          "fault — loose part or cyclical load shift".format(iv))

    # Harmonic analysis
    if features.harmonic_ratio > 0.7:
        lines.append("- Strong harmonic content (ratio {:.2f}) — tonal/motor hum dominant, "
                      "possibly pump or compressor".format(features.harmonic_ratio))
    elif features.harmonic_ratio < 0.3:
        lines.append("- Low harmonic content (ratio {:.2f}) — broadband noise dominant, "
                      "possibly air leak, turbulence, or abrasive wear".format(
                          features.harmonic_ratio))

    # Onset analysis
    if features.onset_rate_per_sec > 5:
        lines.append("- HIGH click/knock rate ({:.1f}/s) — rapid impacts suggest loose "
                      "components or high-speed mechanical contact".format(
                          features.onset_rate_per_sec))
    elif features.onset_rate_per_sec > 1:
        lines.append("- Moderate click rate ({:.1f}/s) — intermittent impacts suggest "
                      "loose part or cyclical mechanical contact".format(
                          features.onset_rate_per_sec))

    # Anomaly score
    if features.anomaly_score > 0.5:
        lines.append("- HIGH anomaly score ({:.2f}) — sound is significantly different "
                      "from typical machine audio".format(features.anomaly_score))
    elif features.anomaly_score > 0.3:
        lines.append("- Moderate anomaly score ({:.2f}) — sound shows some unusual "
                      "characteristics".format(features.anomaly_score))

    if not lines:
        return ""

    return (
        "\n## FEATURE-BASED ANALYSIS\n"
        "The rule engine found no exact match, but the acoustic features suggest:\n"
        + "\n".join(lines)
        + "\n\nBased on these features and the appliance type ({appliance}), "
        "suggest the most likely fault category. If the sound appears normal, "
        "say 'Inconclusive'."
    ).format(appliance=appliance)


def build_diagnosis_prompt(
    features: AudioFeatures,
    candidates: List[Candidate],
    appliance: str,
) -> str:
    pattern = (
        f"YES, every {round(features.pattern_interval_ms)} ms"
        if features.has_regular_pattern else "No regular pattern"
    )

    candidate_block = "\n".join(
        f"  {i+1}. {c.name} (urgency {c.urgency}, acoustic support "
        f"{c.weight:.0%}) — {c.evidence}"
        for i, c in enumerate(candidates)
    )

    # Check if rules returned Inconclusive
    is_inconclusive = (
        len(candidates) == 1 and candidates[0].name == "Inconclusive"
    )

    feature_section = ""
    if is_inconclusive:
        feature_section = _build_feature_analysis(features, appliance)

    return f"""A {appliance} is making an unusual sound. Deterministic audio analysis produced these MEASURED facts:

AUDIO CHARACTERISTICS:
- Duration: {features.duration_s:.1f} s
- Loudness: {features.rms_db:.1f} dB avg, {features.peak_db:.1f} dB peak
- Loudness variation: {features.rms_variance:.4f} (high = intermittent)
- Harshness (zero-crossing rate): {features.zero_crossing_rate:.4f} (>0.15 = grinding)
- Spectral centroid: {features.spectral_centroid_hz:.0f} Hz (high = harsh, low = rumble)
- Spectral bandwidth: {features.spectral_bandwidth_hz:.0f} Hz (wide = complex noise)
- Dominant frequency: {features.dominant_frequency_hz:.0f} Hz
- Harmonic ratio: {features.harmonic_ratio:.2f} (1.0 = pure tone, 0.0 = pure noise)
- Click/knock rate: {features.onset_rate_per_sec:.1f} per second
- Regular clicking pattern: {pattern}
- Anomaly severity: {features.anomaly_score:.2f} / 1.0

The deterministic rule engine ranked these CANDIDATE faults (strongest first):
{candidate_block}
{feature_section}

Choose the single best-supported candidate above (or say "Inconclusive" if none
is convincing). Do NOT invent a fault that is not justified by the measured facts.

Respond ONLY with JSON in exactly this shape:
{{"fault": "<short fault name>", "urgency": "CRITICAL|HIGH|MEDIUM|LOW", "checks": ["step 1", "step 2", "step 3"], "safety": "<one safety note or 'None'>", "confidence": <integer 0-100>}}"""