supervised_export.py

#!/usr/bin/env python3
"""Render supervised semantic state into edited sample-pack artifacts.

The batch manifest remains immutable. This module takes the mutable
``supervision_state.json`` layer, excludes suppressed hits, honors explicit
representatives/favorites, and writes a separate ``supervised/`` export tree.
"""

from __future__ import annotations

import json
import os
import re
import shutil
import time
from dataclasses import asdict
from pathlib import Path
from typing import Any

import numpy as np
import soundfile as sf

from sample_extractor import (
    Cluster,
    Hit,
    build_archive,
    export_midi,
    render_midi_with_samples,
    sample_quality_score,
    select_best,
    synthesize_from_cluster,
)
from supervised_state import load_manifest, load_or_create_state, now, recompute_scores


def _safe_label(value: Any, fallback: str) -> str:
    text = str(value or fallback).strip() or fallback
    text = re.sub(r"[^A-Za-z0-9._-]+", "_", text)
    text = re.sub(r"_+", "_", text).strip("._-")
    return text or fallback


def _read_hit_audio(output_dir: Path, hit: dict[str, Any]) -> tuple[np.ndarray, int]:
    rel = hit.get("file")
    if not rel:
        raise FileNotFoundError(f"Hit {hit.get('id')} does not have a file path")
    path = (output_dir / rel).resolve()
    path.relative_to(output_dir.resolve())
    if not path.exists():
        raise FileNotFoundError(f"Hit audio missing for {hit.get('id')}: {rel}")
    audio, sr = sf.read(path, dtype="float32", always_2d=False)
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    return np.asarray(audio, dtype=np.float32), int(sr)


def _state_to_clusters(output_dir: Path, state: dict[str, Any]) -> list[Cluster]:
    hits_by_id = state.get("hits", {})
    clusters: list[Cluster] = []
    used_labels: set[str] = set()

    for ordinal, raw_cluster in enumerate(state.get("clusters", {}).values()):
        active_hit_ids = [
            hid
            for hid in raw_cluster.get("hit_ids", [])
            if hid in hits_by_id and not hits_by_id[hid].get("suppressed")
        ]
        if not active_hit_ids:
            continue

        label = _safe_label(raw_cluster.get("label"), f"cluster_{ordinal}")
        base = label
        suffix = 1
        while label in used_labels:
            suffix += 1
            label = f"{base}_{suffix}"
        used_labels.add(label)

        converted_hits: list[Hit] = []
        for hid in active_hit_ids:
            raw_hit = hits_by_id[hid]
            audio, sr = _read_hit_audio(output_dir, raw_hit)
            converted_hits.append(
                Hit(
                    audio=audio,
                    sr=sr,
                    onset_time=float(raw_hit.get("onset_sec") or 0.0),
                    duration=float(raw_hit.get("duration_ms") or 0.0) / 1000.0 or (len(audio) / max(sr, 1)),
                    index=int(raw_hit.get("index") or 0),
                    rms_energy=float(raw_hit.get("rms_energy") or 0.0),
                    spectral_centroid=float(raw_hit.get("spectral_centroid_hz") or 0.0),
                    label=str(raw_hit.get("label") or raw_cluster.get("classification") or "other"),
                    cluster_id=ordinal,
                )
            )

        cluster = Cluster(cluster_id=ordinal, label=label, hits=converted_hits)
        representative = raw_cluster.get("representative_hit_id")
        pinned = False
        if representative in active_hit_ids:
            cluster.best_hit_idx = active_hit_ids.index(representative)
            pinned = True
        else:
            favorite_idx = next((i for i, hid in enumerate(active_hit_ids) if hits_by_id[hid].get("favorite")), None)
            if favorite_idx is not None:
                cluster.best_hit_idx = favorite_idx
                pinned = True
        setattr(cluster, "_pinned_by_state", pinned)
        clusters.append(cluster)

    # Score representatives only where the supervision state did not pin one.
    for cluster in clusters:
        if cluster.count <= 1:
            cluster.best_hit_idx = 0
    unpinned = [cluster for cluster in clusters if cluster.count > 1 and not getattr(cluster, "_pinned_by_state", False)]
    if unpinned:
        select_best(unpinned)
    return clusters


def _write_audio(path: Path, audio: np.ndarray, sr: int, subtype: str = "PCM_16") -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    sf.write(path, audio, sr, subtype=subtype)


def _mono(audio: np.ndarray) -> np.ndarray:
    audio = np.asarray(audio, dtype=np.float32)
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    return audio.astype(np.float32)


def _pad_or_trim(audio: np.ndarray, length: int) -> np.ndarray:
    audio = _mono(audio)
    if len(audio) == length:
        return audio
    if len(audio) > length:
        return audio[:length]
    return np.pad(audio, (0, max(0, length - len(audio)))).astype(np.float32)


def _rms(audio: np.ndarray) -> float:
    audio = _mono(audio)
    return float(np.sqrt(np.mean(np.square(audio, dtype=np.float64)))) if audio.size else 0.0


def _match_rms(rendered: np.ndarray, reference: np.ndarray, *, min_gain: float = 0.05, max_gain: float = 8.0) -> np.ndarray:
    rendered_rms = _rms(rendered)
    reference_rms = _rms(reference)
    if rendered_rms <= 1e-10 or reference_rms <= 1e-10:
        return _mono(rendered)
    return (_mono(rendered) * float(np.clip(reference_rms / rendered_rms, min_gain, max_gain))).astype(np.float32)


def _soft_limit(audio: np.ndarray, ceiling: float = 0.98) -> np.ndarray:
    audio = _mono(audio).astype(np.float32)
    peak = float(np.max(np.abs(audio))) if audio.size else 0.0
    if peak > ceiling > 0:
        audio = audio * (ceiling / peak)
    return audio.astype(np.float32)


def _read_optional_audio(path: Path) -> tuple[np.ndarray | None, int | None]:
    if not path.exists():
        return None, None
    audio, sr = sf.read(path, dtype="float32", always_2d=False)
    return _mono(audio), int(sr)


def _make_reproduction_mix(target_reconstruction: np.ndarray, context_bed: np.ndarray, length: int) -> np.ndarray:
    return _soft_limit(_pad_or_trim(context_bed, length) + _pad_or_trim(target_reconstruction, length))


def export_supervised_state(
    output_dir: str | os.PathLike[str],
    job_id: str,
    *,
    synthesize: bool = True,
    quantize: bool | None = None,
    subdivision: int | None = None,
    selected_labels: set[str] | list[str] | None = None,
    export_dir_name: str = "supervised",
    kind: str = "supervised-export",
) -> dict[str, Any]:
    """Create edited artifacts from ``supervision_state.json``.

    Returns the JSON manifest written to ``supervised/manifest.json``.
    """
    out = Path(output_dir)
    manifest = load_manifest(out)
    state = load_or_create_state(job_id, out)
    recompute_scores(state)

    safe_export_dir_name = "".join(ch if ch.isalnum() or ch in {"-", "_"} else "_" for ch in str(export_dir_name or "supervised")).strip("_") or "supervised"
    export_prefix = safe_export_dir_name
    selected_label_set = {str(label) for label in selected_labels} if selected_labels else None

    export_dir = out / safe_export_dir_name
    if export_dir.exists():
        shutil.rmtree(export_dir)
    samples_dir = export_dir / "samples"
    samples_dir.mkdir(parents=True, exist_ok=True)

    clusters = _state_to_clusters(out, state)
    if selected_label_set is not None:
        clusters = [cluster for cluster in clusters if cluster.label in selected_label_set]
        missing = sorted(selected_label_set - {cluster.label for cluster in clusters})
        if missing:
            raise ValueError(f"Selected sample label(s) not found in current state: {', '.join(missing)}")
    bpm = float(manifest.get("bpm") or 120.0)
    sr = int(manifest.get("sample_rate") or 44100)
    params = manifest.get("params") or {}
    if quantize is None:
        quantize = bool(params.get("quantize_midi", True))
    if subdivision is None:
        subdivision = int(params.get("subdivision", 16))

    started = time.perf_counter()
    files: dict[str, str] = {}
    samples: list[dict[str, Any]] = []

    context_bed, context_sr = _read_optional_audio(out / "context_bed.wav")
    source_audio, source_sr = _read_optional_audio(out / "source.wav")
    stem_audio, stem_file_sr = _read_optional_audio(out / "stem.wav")
    if context_sr:
        sr = int(context_sr)
    elif source_sr:
        sr = int(source_sr)
    elif stem_file_sr:
        sr = int(stem_file_sr)
    source_length = len(source_audio) if source_audio is not None else max((len(context_bed) if context_bed is not None else 0), sr)
    if context_bed is None:
        context_bed = np.zeros(source_length, dtype=np.float32)
    if stem_audio is None:
        stem_audio = np.zeros(source_length, dtype=np.float32)

    midi_path = export_dir / "reconstruction.mid"
    if clusters:
        export_midi(clusters, str(midi_path), bpm=bpm, quantize=quantize, subdivision=int(subdivision))
        target_rendered = render_midi_with_samples(clusters, sr=sr)
        target_rendered = _match_rms(target_rendered, stem_audio)
        if synthesize:
            for cluster in clusters:
                if cluster.count >= 2:
                    cluster.synthesized = synthesize_from_cluster(cluster)
    else:
        midi_path.write_bytes(b"")
        target_rendered = np.zeros(source_length, dtype=np.float32)

    rendered = _make_reproduction_mix(target_rendered, context_bed, max(source_length, len(target_rendered)))
    _write_audio(export_dir / "target_reconstruction.wav", _soft_limit(target_rendered), sr, subtype="PCM_16")
    _write_audio(export_dir / "reconstruction.wav", rendered, sr, subtype="PCM_16")
    files["midi"] = f"{export_prefix}/reconstruction.mid"
    files["target_reconstruction"] = f"{export_prefix}/target_reconstruction.wav"
    files["reconstruction"] = f"{export_prefix}/reconstruction.wav"

    for cluster in sorted(clusters, key=lambda item: item.count, reverse=True):
        best = cluster.best_hit
        sample_file = f"{export_prefix}/samples/{cluster.label}.wav"
        best.save(str(out / sample_file))
        quality = sample_quality_score(best.audio, best.sr, cluster.label.rsplit("_", 1)[0])
        samples.append(
            {
                "label": cluster.label,
                "classification": cluster.label.rsplit("_", 1)[0],
                "hits": int(cluster.count),
                "midi_note": int(cluster.midi_note),
                "score": round(float(quality["total"]), 2),
                "cleanness": round(float(quality["cleanness"]), 4),
                "completeness": round(float(quality["completeness"]), 4),
                "duration_ms": round(float(best.duration * 1000), 1),
                "first_onset_sec": round(float(min(hit.onset_time for hit in cluster.hits)), 4),
                "file": sample_file,
            }
        )
        if cluster.synthesized is not None:
            _write_audio(out / f"{export_prefix}/samples/{cluster.label}__synth.wav", cluster.synthesized, sr, subtype="PCM_24")

    archive_tmp = build_archive(
        clusters,
        bpm,
        sr,
        midi_path=str(midi_path),
        rendered_audio=rendered,
        target_rendered_audio=target_rendered,
    )
    archive_rel = f"{export_prefix}/sample-pack.zip"
    shutil.copyfile(archive_tmp, out / archive_rel)
    try:
        os.unlink(archive_tmp)
    except OSError:
        pass
    files["archive"] = archive_rel

    active_hits = [hit for hit in state.get("hits", {}).values() if not hit.get("suppressed")]
    export_manifest = {
        "kind": kind,
        "job_id": job_id,
        "created_at": now(),
        "duration_sec": round(time.perf_counter() - started, 6),
        "source_manifest_fingerprint": state.get("manifest_fingerprint"),
        "state_updated_at": state.get("updated_at"),
        "bpm": bpm,
        "sample_rate": sr,
        "hit_count": len(active_hits),
        "suppressed_hit_count": sum(1 for hit in state.get("hits", {}).values() if hit.get("suppressed")),
        "cluster_count": len(clusters),
        "selected_labels": sorted(selected_label_set) if selected_label_set is not None else None,
        "quantize_midi": bool(quantize),
        "subdivision": int(subdivision),
        "samples": samples,
        "files": files,
        "state_summary": {
            "constraint_count": len(state.get("constraints", [])),
            "event_count": len(state.get("events", [])),
            "open_suggestion_count": len([s for s in state.get("suggestions", []) if s.get("status") == "open"]),
        },
    }
    (export_dir / "manifest.json").write_text(json.dumps(export_manifest, indent=2, sort_keys=True), encoding="utf-8")

    state.setdefault("exports", []).append(
        {
            "created_at": export_manifest["created_at"],
            "path": f"{export_prefix}/manifest.json",
            "kind": kind,
            "hit_count": export_manifest["hit_count"],
            "cluster_count": export_manifest["cluster_count"],
            "suppressed_hit_count": export_manifest["suppressed_hit_count"],
        }
    )
    state["latest_export"] = state["exports"][-1]
    state.setdefault("events", []).append(
        {
            "id": f"event:export:{int(export_manifest['created_at'] * 1000)}",
            "type": "supervised.exported",
            "source": "system",
            "created_at": export_manifest["created_at"],
            "payload": {
                "hit_count": export_manifest["hit_count"],
                "cluster_count": export_manifest["cluster_count"],
                "archive": archive_rel,
            },
        }
    )
    state_path = out / "supervision_state.json"
    state["updated_at"] = now()
    state_path.write_text(json.dumps(state, indent=2, sort_keys=True), encoding="utf-8")

    return export_manifest


def export_selected_samples(
    output_dir: str | os.PathLike[str],
    job_id: str,
    *,
    selected_labels: list[str] | set[str],
    synthesize: bool = True,
    quantize: bool | None = None,
    subdivision: int | None = None,
) -> dict[str, Any]:
    if not selected_labels:
        raise ValueError("selected_labels must contain at least one sample label")
    return export_supervised_state(
        output_dir,
        job_id,
        synthesize=synthesize,
        quantize=quantize,
        subdivision=subdivision,
        selected_labels=set(map(str, selected_labels)),
        export_dir_name="selected",
        kind="selected-sample-export",
    )