{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import random\n",
    "from suno_utils.audio import Audio\n",
    "from suno_utils.audio.midi import Midi\n",
    "import shutil\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torchcrepe\n",
    "import tempfile\n",
    "import torch\n",
    "import pesto\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import librosa\n",
    "import copy\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trombone_champ_dir = \"/app2/suno/data/victor/trombone_champ_vocals\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MidiPair:\n",
    "    def __init__(self, midi_file: str, audio_file: str):\n",
    "        self.midi_file = midi_file\n",
    "        self.audio_file = audio_file\n",
    "        self.midi = None\n",
    "        self.audio = None\n",
    "\n",
    "    def load_midi(self):\n",
    "        if self.midi is not None:\n",
    "            return self.midi\n",
    "        self.midi = Midi.from_path(self.midi_file)\n",
    "        return self.midi\n",
    "\n",
    "    def load_audio(self):\n",
    "        if self.audio is not None:\n",
    "            return self.audio\n",
    "        self.audio = Audio.from_file(self.audio_file)\n",
    "        return self.audio\n",
    "\n",
    "    def load_all(self):\n",
    "        self.load_midi()\n",
    "        self.load_audio()\n",
    "\n",
    "    def __str__(self):\n",
    "        return f\"MidiPair(midi_file={self.midi_file}, audio_file={self.audio_file})\"\n",
    "\n",
    "    def __repr__(self):\n",
    "        return self.__str__()\n",
    "\n",
    "    def play(self):\n",
    "        self.load_all()\n",
    "        stereo_audio = self.midi.make_stereo_comparison(self.audio)\n",
    "        stereo_audio.play()\n",
    "\n",
    "\n",
    "def load_pairs(dir: str, audio_ext: str = \"mp3\", midi_ext: str = \"mid\"):\n",
    "    midi_files = glob.glob(os.path.join(dir, \"**\", f\"*.{midi_ext}\"), recursive=True)\n",
    "    audio_files = glob.glob(os.path.join(dir, \"**\", f\"*.{audio_ext}\"), recursive=True)\n",
    "\n",
    "    # Create a mapping of base filenames to audio files\n",
    "    audio_map = {}\n",
    "    for audio_file in audio_files:\n",
    "        base_name = os.path.splitext(os.path.basename(audio_file))[0]\n",
    "        audio_map[base_name] = audio_file\n",
    "\n",
    "    pairs = []\n",
    "    for midi_file in midi_files:\n",
    "        base_name = os.path.splitext(os.path.basename(midi_file))[0]\n",
    "        if base_name in audio_map:\n",
    "            pairs.append(MidiPair(midi_file, audio_map[base_name]))\n",
    "\n",
    "    return pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## trombone champ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trombone_champ_pairs = load_pairs(trombone_champ_dir, audio_ext=\"opus\")\n",
    "print(f\"Loaded {len(trombone_champ_pairs)} trombone champ pairs\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## fix octave shifts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_f0_comparison(results, stem, midi):\n",
    "    # Convert frequencies to MIDI note numbers (pitches)\n",
    "    frequencies = torchcrepe.filter.median(results, 30)[0]\n",
    "\n",
    "    # zero out sections that are silent\n",
    "    hz = int(len(frequencies) / stem.duration_s)\n",
    "    for i in range(int(stem.duration_s)):\n",
    "        if stem.get_slice(i, i + 1).loudness < -35:\n",
    "            frequencies[i * hz : (i + 1) * hz] = 20\n",
    "\n",
    "    pitches = librosa.hz_to_midi(frequencies)\n",
    "\n",
    "    # Create time axis\n",
    "    time_axis = [t / hz for t in range(len(frequencies))]\n",
    "\n",
    "    plt.figure(figsize=(12, 6))\n",
    "    plt.plot(time_axis, pitches, label=\"Vocal F0\")\n",
    "\n",
    "    # filter out notes that are silent\n",
    "    for i in reversed(range(len(midi.pmidi.instruments[0].notes))):\n",
    "        audio_slice = stem.get_slice(\n",
    "            midi.pmidi.instruments[0].notes[i].start,\n",
    "            midi.pmidi.instruments[0].notes[i].end + 1,\n",
    "        )\n",
    "        if np.mean(audio_slice.array_float**2) < 1e-6:\n",
    "            midi.pmidi.instruments[0].notes.pop(i)\n",
    "\n",
    "    # Overlay MIDI notes\n",
    "    for note in midi.pmidi.instruments[0].notes:\n",
    "        start_time = note.start\n",
    "        end_time = note.end\n",
    "        pitch = note.pitch\n",
    "        plt.hlines(\n",
    "            pitch,\n",
    "            start_time,\n",
    "            end_time,\n",
    "            colors=\"red\",\n",
    "            linewidth=2,\n",
    "            alpha=0.7,\n",
    "            label=\"MIDI Notes\" if note == midi.pmidi.instruments[0].notes[0] else \"\",\n",
    "        )\n",
    "\n",
    "    plt.xlabel(\"Time (s)\")\n",
    "    plt.ylabel(\"MIDI Note Number\")\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_f0_pesto(audio):\n",
    "    with tempfile.NamedTemporaryFile(suffix=\".wav\") as f:\n",
    "        audio.write_wav(f.name)\n",
    "\n",
    "        x, sr = torchaudio.load(f.name)\n",
    "        x = x.mean(dim=0)\n",
    "\n",
    "        _, pitch, confidence, _ = pesto.predict(x, sr)\n",
    "\n",
    "    return pitch.clone(), confidence.clone()\n",
    "\n",
    "\n",
    "def get_f0(audio):\n",
    "    # Process in 60s chunks to avoid OOM\n",
    "    chunk_duration = 60\n",
    "    total_duration = audio.duration_s  # min(240, audio.duration_s)\n",
    "    all_results = []\n",
    "    all_period_results = []\n",
    "\n",
    "    for start_time in range(0, int(total_duration), chunk_duration):\n",
    "        end_time = min(start_time + chunk_duration, total_duration)\n",
    "\n",
    "        with tempfile.NamedTemporaryFile(suffix=\".wav\") as f:\n",
    "            audio.get_slice(start_time, end_time).write_wav(f.name)\n",
    "            chunk_results, period_results = torchcrepe.predict_from_file(\n",
    "                f.name,\n",
    "                device=\"cuda:3\",\n",
    "                decoder=torchcrepe.decode.weighted_argmax,\n",
    "                pad=False,\n",
    "                return_periodicity=True,\n",
    "            )\n",
    "        all_results.append(chunk_results)\n",
    "        all_period_results.append(period_results)\n",
    "\n",
    "    # Concatenate results\n",
    "    results = torch.cat([r[0] for r in all_results], dim=0).unsqueeze(0)\n",
    "    period_final_results = torch.cat(\n",
    "        [r[0] for r in all_period_results], dim=0\n",
    "    ).unsqueeze(0)\n",
    "    return results, period_final_results\n",
    "\n",
    "\n",
    "def octave_dist(note1, note2):\n",
    "    diff = abs(note1 - note2)\n",
    "    return min(diff, 12 - diff)\n",
    "\n",
    "\n",
    "def fix_midi_octave_errors(\n",
    "    pair, verbose=False, pitch_algo=\"crepe\", λ=500, confidence_threshold=0.05\n",
    "):\n",
    "    midi = pair.load_midi()\n",
    "    stem = pair.load_audio().resample(22050)\n",
    "\n",
    "    # get vocal f0\n",
    "    if pitch_algo == \"pesto\":\n",
    "        frequencies, confidences = get_f0_pesto(stem)\n",
    "        results = frequencies.unsqueeze(0)\n",
    "    else:\n",
    "        results, confidences = get_f0(stem)\n",
    "        confidences = confidences[0]\n",
    "        frequencies = torchcrepe.filter.median(results, 30)[0]\n",
    "\n",
    "    # zero out sections that are silent OR low confidence\n",
    "    hz = int(len(frequencies) / stem.duration_s)\n",
    "    device = frequencies.device\n",
    "    for i in range(int(stem.duration_s)):\n",
    "        if stem.get_slice(i, i + 1).loudness < -35:\n",
    "            frequencies[i * hz : (i + 1) * hz] = 20\n",
    "            if confidences is not None:\n",
    "                confidences[i * hz : (i + 1) * hz] = 0\n",
    "\n",
    "    # viterbi octave correction\n",
    "    Ks = torch.arange(-3, 4, device=device)  # [-3, -2, -1, 0, 1, 2, 3]\n",
    "    N = len(midi.pmidi.instruments[0].notes)\n",
    "\n",
    "    dp = torch.full((N, len(Ks)), float(\"inf\"), device=device)\n",
    "    prev = torch.zeros((N, len(Ks)), dtype=torch.long, device=device)\n",
    "    note_estimates = []\n",
    "    note_confidences = []\n",
    "\n",
    "    for note in midi.pmidi.instruments[0].notes:\n",
    "        start_idx, end_idx = int(note.start * hz), int(note.end * hz)\n",
    "        note_pitches = frequencies[start_idx:end_idx]\n",
    "\n",
    "        if confidences is not None:\n",
    "            note_confs = confidences[start_idx:end_idx]\n",
    "            # filter by confidence and silence\n",
    "            valid_mask = (note_pitches > 20) & (note_confs > confidence_threshold)\n",
    "            valid_pitches = note_pitches[valid_mask]\n",
    "            valid_confs = note_confs[valid_mask]\n",
    "\n",
    "            if len(valid_pitches) > 0:\n",
    "                # weighted median\n",
    "                median_pitch = weighted_median_torch(valid_pitches, valid_confs)\n",
    "                # average confidence for this note\n",
    "                avg_confidence = torch.mean(valid_confs)\n",
    "            else:\n",
    "                median_pitch = torch.tensor(float(\"nan\"), device=device)\n",
    "                avg_confidence = torch.tensor(0.0, device=device)\n",
    "        else:\n",
    "            # fallback for CREPE\n",
    "            valid_pitches = note_pitches[note_pitches > 20]\n",
    "            if len(valid_pitches) > 0:\n",
    "                median_pitch = torch.median(valid_pitches)\n",
    "            else:\n",
    "                median_pitch = torch.tensor(float(\"nan\"), device=device)\n",
    "            avg_confidence = torch.tensor(1.0, device=device)\n",
    "\n",
    "        # Convert Hz to MIDI\n",
    "        if torch.isnan(median_pitch):\n",
    "            note_estimates.append(median_pitch)\n",
    "        else:\n",
    "            midi_pitch = 69 + 12 * torch.log2(median_pitch / 440)\n",
    "            note_estimates.append(midi_pitch)\n",
    "        note_confidences.append(avg_confidence)\n",
    "\n",
    "    notes = midi.pmidi.instruments[0].notes\n",
    "\n",
    "    # init with confidence weighting\n",
    "    for ik, k in enumerate(Ks):\n",
    "        if torch.isnan(note_estimates[0]):\n",
    "            dp[0, ik] = 0\n",
    "        else:\n",
    "            error = (note_estimates[0] - (notes[0].pitch + 12 * k)) ** 2\n",
    "            confidence_weight = 1.0 / torch.clamp(note_confidences[0], min=0.1)\n",
    "            dp[0, ik] = error * confidence_weight\n",
    "\n",
    "    # fill with confidence weighting\n",
    "    for i in range(1, N):\n",
    "        for ik, k in enumerate(Ks):\n",
    "            if torch.isnan(note_estimates[i]):\n",
    "                obs = torch.tensor(0.0, device=device)\n",
    "            else:\n",
    "                error = (note_estimates[i] - (notes[i].pitch + 12 * k)) ** 2\n",
    "                confidence_weight = 1.0 / torch.clamp(note_confidences[i], min=0.1)\n",
    "                obs = error * confidence_weight\n",
    "\n",
    "            # transition costs\n",
    "            transition_costs = λ * torch.abs(k - Ks)\n",
    "            costs = dp[i - 1, :] + transition_costs\n",
    "            dp[i, ik] = obs + torch.min(costs)\n",
    "            prev[i, ik] = torch.argmin(costs)\n",
    "\n",
    "    # backtrack\n",
    "    best_path = torch.zeros(N, dtype=torch.long, device=device)\n",
    "    best_path[-1] = torch.argmin(dp[-1, :])\n",
    "    for i in range(N - 2, -1, -1):\n",
    "        best_path[i] = prev[i + 1, best_path[i + 1]]\n",
    "\n",
    "    if verbose:\n",
    "        print(best_path.cpu().numpy())\n",
    "\n",
    "    cost = torch.tensor(0.0, device=device)\n",
    "    for i in range(N):\n",
    "        if torch.isnan(note_estimates[i]):\n",
    "            cost += 0\n",
    "        else:\n",
    "            diff = note_estimates[i] - (notes[i].pitch + 12 * Ks[best_path[i]])\n",
    "            cost += diff**2\n",
    "    normalized_cost = cost / N\n",
    "    if verbose:\n",
    "        print(f\"Cost: {cost.item()}, Normalized cost: {normalized_cost.item()}\")\n",
    "\n",
    "    new_midi = copy.deepcopy(midi)\n",
    "    # apply shifts - convert back to CPU for MIDI manipulation\n",
    "    best_path_cpu = best_path.cpu()\n",
    "    Ks_cpu = Ks.cpu()\n",
    "    for i, note in enumerate(notes):\n",
    "        note.pitch += 12 * Ks_cpu[best_path_cpu[i]].item()\n",
    "\n",
    "    return results, new_midi\n",
    "\n",
    "\n",
    "def weighted_median_torch(values, weights):\n",
    "    \"\"\"Compute weighted median using torch operations\"\"\"\n",
    "    if len(values) == 0:\n",
    "        return torch.tensor(float(\"nan\"), device=values.device)\n",
    "\n",
    "    sorted_indices = torch.argsort(values)\n",
    "    sorted_values = values[sorted_indices]\n",
    "    sorted_weights = weights[sorted_indices]\n",
    "\n",
    "    cumsum = torch.cumsum(sorted_weights, dim=0)\n",
    "    total_weight = cumsum[-1]\n",
    "\n",
    "    median_pos = total_weight / 2\n",
    "    median_idx = torch.searchsorted(cumsum, median_pos)\n",
    "\n",
    "    # Clamp to valid range\n",
    "    median_idx = torch.clamp(median_idx, 0, len(sorted_values) - 1)\n",
    "    return sorted_values[median_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# midi_pair = random.choice(trombone_champ_pairs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stem = midi_pair.load_audio().resample(22050)\n",
    "# midi = midi_pair.load_midi()\n",
    "# pitches, new_midi = fix_midi_octave_errors(midi_pair,λ=500)\n",
    "# plot_f0_comparison(pitches, midi_pair.load_audio(), new_midi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stem = midi_pair.load_audio().resample(22050)\n",
    "# midi = midi_pair.load_midi()\n",
    "# pitches, new_midi = fix_midi_octave_errors(midi_pair, pitch_algo=\"pesto\",λ=350)\n",
    "# new_midi.make_stereo_comparison(stem).play()\n",
    "# plot_f0_comparison(pitches, midi_pair.load_audio(), new_midi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pitch_algo = \"pesto\"\n",
    "λ = 250\n",
    "\n",
    "output_dir = f\"/app2/suno/data/sara/trombone_champ_vocals_octaved_{pitch_algo}_{λ}\"\n",
    "os.makedirs(output_dir, exist_ok=True)\n",
    "\n",
    "for midi_pair in tqdm(trombone_champ_pairs):\n",
    "    pitches, new_midi = fix_midi_octave_errors(midi_pair, pitch_algo=pitch_algo, λ=λ)\n",
    "    filename = midi_pair.audio_file.split(\"/\")[-1].split(\".\")[0]\n",
    "    midi_path = os.path.join(output_dir, filename + \".mid\")\n",
    "    audio_path = os.path.join(output_dir, filename + \".opus\")\n",
    "    print(midi_path)\n",
    "    print(audio_path)\n",
    "    new_midi.write(midi_path)\n",
    "    shutil.copy(midi_pair.audio_file, audio_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_clean",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}