render_engine.py

"""Shared preview/full-set rendering entrypoints.

This module exists so preview exports and full-set exports cannot drift into
separate tempo, pitch, and headroom policies again. Both paths call the same
AutomationIR renderer and therefore share:

* source-tempo default rendering;
* pitch-preserving stretch only when explicitly enabled in plan metadata;
* component-lane handoff back to the untouched full source clip;
* attenuation-only output protection.
"""

from __future__ import annotations

from typing import Any, Callable, Mapping

from automation_ir import AutomationIR, render_automation_ir, build_transition_ir
from automation_set_renderer import build_set_automation_ir

AudioLoader = Callable[[str, float, float, int], Any]
TimeStretcher = Callable[[Any, float], Any]
StemResolver = Callable[[Any, str, Any, int], Any | None]


def render_ir(
    ir: AutomationIR,
    *,
    load_audio_segment: AudioLoader,
    time_stretch_audio: TimeStretcher,
    stem_resolver: StemResolver | None = None,
) -> Any:
    return render_automation_ir(
        ir,
        load_audio_segment=load_audio_segment,
        time_stretch_audio=time_stretch_audio,
        stem_resolver=stem_resolver,
    )


def render_transition_preview(
    plan: Any,
    track_a: Any,
    track_b: Any,
    *,
    candidate_rank: int = 0,
    load_audio_segment: AudioLoader,
    time_stretch_audio: TimeStretcher,
    stem_resolver: StemResolver | None = None,
    sr: int = 44100,
) -> tuple[Any, AutomationIR, Mapping[str, Any] | None]:
    candidate = None
    alternatives = list(getattr(plan, "alternatives", []) or [])
    if candidate_rank > 0 and candidate_rank <= len(alternatives):
        candidate = alternatives[candidate_rank - 1]
    ir = build_transition_ir(plan, track_a, track_b, candidate=candidate, sr=sr)
    audio = render_ir(
        ir,
        load_audio_segment=load_audio_segment,
        time_stretch_audio=time_stretch_audio,
        stem_resolver=stem_resolver,
    )
    try:
        from transition_diagnostics import diagnose_transition_audio
        diag = diagnose_transition_audio(audio, sr=sr, anchor_seconds=ir.anchor_seconds)
        ir.metadata["rendered_diagnostics"] = diag
        metrics = dict(diag.get("metrics", {}) or {})
        rendered_scores = {
            "kick_overlap_score": round(float(metrics.get("onset_regularity", 0.0) or 0.0), 3),
            "bass_overlap_score": round(max(0.0, min(1.0, 1.0 / max(float(metrics.get("anchor_low_balance_ratio", 1.0) or 1.0), 1.0))), 3),
            "energy_dip_jump_score": round(max(0.0, min(1.0, float(metrics.get("anchor_min_rms_ratio", 1.0) or 1.0) / max(float(metrics.get("anchor_rms_balance_ratio", 1.0) or 1.0), 1.0))), 3),
            "spectral_harshness_score": round(max(0.0, min(1.0, 1.0 / max(float(metrics.get("high_band_spike_ratio", 1.0) or 1.0) / 2.0, 1.0))), 3),
            "phrase_arrival_confidence": float((ir.metadata.get("candidate", {}) or {}).get("score_breakdown", {}).get("arrival_quality", 0.0) or 0.0),
        }
        ir.metadata["rendered_candidate_scores"] = rendered_scores
        if isinstance(candidate, dict):
            candidate["rendered_diagnostics"] = diag
            candidate["rendered_candidate_scores"] = rendered_scores
    except Exception:
        pass
    return audio, ir, candidate


def render_full_set(
    tracks: list[Any],
    order: list[int],
    transitions: list[Any],
    *,
    load_audio_segment: AudioLoader,
    time_stretch_audio: TimeStretcher,
    stem_resolver: StemResolver | None = None,
    sr: int = 44100,
) -> tuple[Any, dict[str, Any], AutomationIR]:
    ir = build_set_automation_ir(tracks, order, transitions, sr=sr)
    audio = render_ir(
        ir,
        load_audio_segment=load_audio_segment,
        time_stretch_audio=time_stretch_audio,
        stem_resolver=stem_resolver,
    )
    info = {
        "tracks": ir.metadata.get("tracks", []),
        "transitions": ir.metadata.get("transitions", []),
        "total_duration": audio.shape[-1] / sr,
        "automation_ir": {
            "clips": len(ir.clips),
            "lanes": len(ir.lanes),
            "duration_seconds": ir.duration_seconds,
            "component_lane_method": ir.metadata.get("component_lane_method"),
            "render_protection": ir.metadata.get("render_protection", {}),
        },
    }
    return audio, info, ir