"""Analyse-Pipeline: RMVPE-F0-Kurve, MIDI aus F0, Whisper-Transkript mit
Wort-Timestamps, und ein Zeit-Mapping Wort <-> Note <-> F0-Segment.

Alle Outputs liegen auf derselben Zeitachse (RMVPE: 100 fps / 10 ms Grid,
identisch zur pretty_midi-Piano-Roll mit fs=100, die Diff-Pitcher liest).
"""

import json
import os
import tempfile

import numpy as np
import pretty_midi

F0_FPS = 100           # RMVPE native hop: 10 ms
CONF_THRESHOLD = 0.03  # RMVPE-Voicing-Schwelle
MIN_NOTE_FRAMES = 3    # <30 ms -> kein eigenstaendiger Note-Wert
MERGE_GAP_FRAMES = 3   # gleichhohe Segmente mit <=30 ms Luecke verschmelzen

NOTE_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]


def note_name(midi_note):
    m = int(round(midi_note))
    return f"{NOTE_NAMES[m % 12]}{m // 12 - 1}"


# ---------------------------------------------------------------- RMVPE

_RMVPE = None


def get_rmvpe():
    global _RMVPE
    if _RMVPE is None:
        from rmvpe_onnx import RMVPE
        _RMVPE = RMVPE()  # laedt die Gewichte beim ersten Aufruf automatisch
    return _RMVPE


def rmvpe_f0(audio, sr):
    """F0-Kurve (Hz) @100 fps. Unvoiced/low-confidence -> 0.0."""
    t, freq, conf, _act = get_rmvpe().predict(audio=audio, sr=sr)
    freq = np.asarray(freq, dtype=float).copy()
    conf = np.asarray(conf, dtype=float)
    freq[conf < CONF_THRESHOLD] = 0.0
    return t, freq


# ---------------------------------------------------------------- MIDI aus F0


def f0_to_notes(freq, shift=0):
    """Segmentiert die F0-Kurve zu [pitch, start_s, end_s]-Noten.

    Saemtliche Zeiten liegen auf dem 10-ms-Raster, damit die Piano-Roll
    (fs=100) sauber ist — genau den Onset-Contract, den Diff-Pitchers
    prepare_midi_wav erwartet (roll>0 -> 100, Onset-Frame genullt).
    """
    voiced = freq > 0
    semis = np.zeros(len(freq))
    semis[voiced] = np.round(69 + 12 * np.log2(freq[voiced] / 440.0))

    notes = []
    n = len(freq)
    i = 0
    while i < n:
        if not voiced[i]:
            i += 1
            continue
        j = i
        while j < n and voiced[j]:
            j += 1
        if j - i >= MIN_NOTE_FRAMES:
            pitch = int(round(float(np.median(semis[i:j])))) + shift
            notes.append([pitch, round(i / 100.0, 2), round(j / 100.0, 2)])
        i = j

    merged = []
    for pitch, start, end in notes:
        if merged and merged[-1][0] == pitch and start - merged[-1][2] <= 0.03:
            merged[-1][2] = end
        else:
            merged.append([pitch, start, end])
    return merged


def notes_to_midi(notes, path, tempo=120):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    pm = pretty_midi.PrettyMIDI(initial_tempo=tempo)
    inst = pretty_midi.Instrument(program=0, name="rmvpe-derived")
    for pitch, start, end in notes:
        inst.notes.append(pretty_midi.Note(
            velocity=100, pitch=int(pitch), start=float(start), end=float(end)))
    pm.instruments.append(inst)
    pm.write(path)
    return path


# ---------------------------------------------------------------- Whisper

_WHISPER = {}


def get_whisper(model_size="small", device="cpu", compute_type="int8"):
    from faster_whisper import WhisperModel
    key = (model_size, device, compute_type)
    if key not in _WHISPER:
        _WHISPER[key] = WhisperModel(model_size, device=device, compute_type=compute_type)
    return _WHISPER[key]


def whisper_words(audio_path, model_size="small", language=None, device="cpu"):
    """Transkript + Wort-Timestamps via faster-whisper."""
    compute_type = "float16" if device == "cuda" else "int8"
    model = get_whisper(model_size, device, compute_type)
    segments, info = model.transcribe(
        os.fspath(audio_path),
        language=language,
        word_timestamps=True,
        vad_filter=True,
    )
    words = []
    for seg in segments:
        for w in (seg.words or []):
            words.append({
                "word": w.word.strip(),
                "start": round(float(w.start), 3),
                "end": round(float(w.end), 3),
                "probability": round(float(w.probability), 3),
            })
    return words, {"language": info.language,
                   "language_probability": round(float(info.language_probability), 3)}


# ---------------------------------------------------------------- Mapping


def build_mapping(freq, notes, words):
    """Wort <-> Note <-> F0-Segment auf der 100-fps-Zeitachse."""

    def f0_stats(a, b):
        i0 = max(0, int(a * F0_FPS))
        i1 = min(len(freq), int(np.ceil(b * F0_FPS)))
        v = freq[i0:i1]
        v = v[v > 0]
        if len(v) == 0:
            return {"voiced_frames": 0, "mean_hz": None, "stdev_semitones": None}
        return {
            "voiced_frames": int(len(v)),
            "mean_hz": round(float(np.mean(v)), 2),
            "stdev_semitones": round(float(np.std(12 * np.log2(v / v.mean()))), 3),
        }

    words_out = []
    for w in words:
        overlaps = [
            {"pitch": int(p), "name": note_name(p), "start": s, "end": e,
             "overlap_s": round(max(0.0, min(w["end"], e) - max(w["start"], s)), 3)}
            for p, s, e in notes
            if min(w["end"], e) - max(w["start"], s) > 0
        ]
        words_out.append({**w, "notes": overlaps, "f0": f0_stats(w["start"], w["end"])})

    notes_out = [
        {"pitch": int(p), "name": note_name(p), "start": s, "end": e, "f0": f0_stats(s, e)}
        for p, s, e in notes
    ]
    return {"fps": F0_FPS, "words": words_out, "notes": notes_out}


# ---------------------------------------------------------------- Plot


def render_plot(freq, notes, words, path):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, ax = plt.subplots(figsize=(14, 4))
    for p, s, e in notes:
        ax.add_patch(plt.Rectangle((s, p - 0.4), e - s, 0.8,
                                   facecolor="tab:blue", alpha=0.5, edgecolor="none"))
    t = np.arange(len(freq)) / F0_FPS
    voiced = freq > 0
    ax.plot(t[voiced], 69 + 12 * np.log2(freq[voiced] / 440.0),
            ".", ms=1.5, color="k", label="RMVPE F0")
    for w in words:
        ax.axvline(w["start"], color="red", lw=0.6, alpha=0.5)
        if w["word"]:
            ax.text(w["start"], 84, w["word"], rotation=90, fontsize=6, va="bottom")
    ax.set_xlabel("Zeit (s)")
    ax.set_ylabel("MIDI-Note")
    ax.set_title("Piano-Roll (MIDI aus F0) + RMVPE-F0 + Wortgrenzen")
    ax.legend(loc="lower right")
    fig.tight_layout()
    fig.savefig(path, dpi=110)
    plt.close(fig)
    return path


# ---------------------------------------------------------------- Top-Level


def analyze(audio_path, sr=24000, max_seconds=20, whisper_model="small",
            language=None, out_dir=None, device="cpu"):
    """Komplette Analyse eines Recordings.

    Liefert dict mit summary, transcript, mapping und Datei-Pfaden
    (midi/json/png/npy). Alle Outputs teilen sich die 10-ms-Zeitachse.
    """
    import librosa

    out_dir = out_dir or tempfile.mkdtemp(prefix="analysis_")
    os.makedirs(out_dir, exist_ok=True)
    wav, _ = librosa.load(audio_path, sr=sr, mono=True, duration=max_seconds)

    t, freq = rmvpe_f0(wav, sr)
    notes = f0_to_notes(freq)

    midi_path = notes_to_midi(notes, os.path.join(out_dir, "derived_from_f0.mid"))
    words, lang_info = whisper_words(audio_path, model_size=whisper_model,
                                     language=language, device=device)

    mapping = build_mapping(freq, notes, words)
    mapping["audio"] = {"sr": sr, "duration_s": round(len(wav) / sr, 3), **lang_info}
    mapping["transcript"] = " ".join(w["word"] for w in words)

    f0_path = os.path.join(out_dir, "f0_rmvpe_100fps.npy")
    np.save(f0_path, np.stack([t, freq]))
    json_path = os.path.join(out_dir, "mapping.json")
    with open(json_path, "w", encoding="utf-8") as fh:
        json.dump(mapping, fh, ensure_ascii=False, indent=2)
    plot_path = render_plot(freq, notes, words, os.path.join(out_dir, "overview.png"))

    summary = {
        "n_words": len(words),
        "n_notes": len(notes),
        "voiced_ratio": round(float((freq > 0).mean()), 3),
        "language": lang_info.get("language"),
        "whisper_model": whisper_model,
        "alignment": "faster-whisper word_timestamps (MFA ist conda-only und im Space nicht installierbar)",
    }
    return {
        "summary": summary,
        "transcript": " ".join(w["word"] for w in words),
        "mapping": mapping,
        "midi_path": midi_path,
        "json_path": json_path,
        "plot_path": plot_path,
        "f0_path": f0_path,
        "midi_notes": notes,
    }