grid_fixer_v2.py

import librosa
import numpy as np
import subprocess
import os
import sys
import argparse
import json
import soundfile as sf
import tempfile
from pathlib import Path

def analyze_audio(input_file, verbose=False):
    """Análise robusta: BPM, beats, consistency (IBI), key/tonalidade, fase."""
    if verbose: print(f"[*] A analisar: {input_file}")

    y, sr = librosa.load(input_file, sr=None, mono=False)
    y_mono = librosa.to_mono(y) if y.ndim > 1 else y

    # --- Deteção de BPM e Beats ---
    if verbose: print("[*] A detetar transientes e pulso rítmico...")
    
    # Usar onset strength para o beat_track
    onset_env = librosa.onset.onset_strength(y=y_mono, sr=sr)
    
    # Tentar detetar o BPM com maior precisão
    tempo, beat_frames = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr)
    
    if isinstance(tempo, np.ndarray):
        tempo = float(tempo[0])
    else:
        tempo = float(tempo)

    beat_times = librosa.frames_to_time(beat_frames, sr=sr)
    
    # --- Análise de Consistência (IBI - Inter-Beat Interval) ---
    if len(beat_times) > 1:
        intervals = np.diff(beat_times)
        ibi_mean = np.mean(intervals)
        ibi_std = np.std(intervals)
        ibi_cv = ibi_std / ibi_mean if ibi_mean > 0 else 1.0
        bpm_from_ibi = 60.0 / ibi_mean if ibi_mean > 0 else tempo
    else:
        ibi_cv = 0.0
        bpm_from_ibi = tempo

    is_metronomic = ibi_cv < 0.015 
    
    if verbose:
        print(f"    BPM Detetado (Librosa): {tempo:.2f}")
        print(f"    BPM Médio (IBI): {bpm_from_ibi:.2f}")
        print(f"    Variação Rítmica (CV): {ibi_cv:.4f} ({'Metronómico' if is_metronomic else 'Variável'})")

    # --- Estimativa de Tonalidade (Key Detection) Melhorada ---
    if verbose: print("[*] A analisar tonalidade (Chroma)...")
    try:
        chroma = librosa.feature.chroma_cqt(y=y_mono, sr=sr)
        chroma_avg = np.mean(chroma, axis=1)
        keys = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']
        key_index = np.argmax(chroma_avg)
        estimated_key = keys[key_index]
        minor_index = (key_index + 3) % 12
        major_index = (key_index + 4) % 12
        if chroma_avg[minor_index] > chroma_avg[major_index]:
            estimated_key += " Minor"
        else:
            estimated_key += " Major"
    except Exception:
        estimated_key = "N/A"

    # --- Deteção de Fase ---
    if len(beat_times) > 0:
        first_beat_time = float(beat_times[0])
        beat_period = 60.0 / bpm_from_ibi
        phase_offset = (first_beat_time % beat_period) / beat_period * 360.0
    else:
        first_beat_time = 0.0
        phase_offset = 0.0

    analysis = {
        "y": y,
        "sr": sr,
        "y_mono": y_mono,
        "tempo": bpm_from_ibi,
        "librosa_tempo": tempo,
        "beat_frames": beat_frames,
        "beat_times": beat_times,
        "ibi_cv": ibi_cv,
        "is_metronomic": is_metronomic,
        "estimated_key": estimated_key,
        "first_beat_time": first_beat_time,
        "phase_offset_degrees": phase_offset,
        "beat_period": 60.0 / bpm_from_ibi
    }

    return analysis


def create_smart_timemap(beat_times, target_bpm, sr, first_beat_time, 
                          is_metronomic, phase_lock=True, verbose=False):
    """Cria timemap inteligente: linear para metronómico, warp suave para variável."""
    t_ideal_interval = 60.0 / target_bpm
    
    if phase_lock:
        first_target_time = 0.0
    else:
        first_target_time = first_beat_time

    timemap_entries = []
    
    if is_metronomic:
        if verbose: print("[*] Modo Metronómico: Usando stretch linear")
        source_f0 = int(beat_times[0] * sr)
        target_f0 = int(first_target_time * sr)
        timemap_entries.append((source_f0, target_f0))
        
        delta_beats = round((beat_times[-1] - beat_times[0]) / (60.0 / target_bpm))
        t_target_last = first_target_time + (delta_beats * t_ideal_interval)
        
        source_f1 = int(beat_times[-1] * sr)
        target_f1 = int(t_target_last * sr)
        timemap_entries.append((source_f1, target_f1))
    else:
        if verbose: print("[*] Modo Variável: Usando warping")
        last_target_frame = -1
        for i, t_real in enumerate(beat_times):
            delta_src = t_real - beat_times[0]
            beats_elapsed = round(delta_src / t_ideal_interval)
            t_target = first_target_time + (beats_elapsed * t_ideal_interval)

            source_frame = int(t_real * sr)
            target_frame = int(t_target * sr)

            if target_frame > last_target_frame:
                timemap_entries.append((source_frame, target_frame))
                last_target_frame = target_frame

    return timemap_entries


def apply_rubberband_processing(temp_wav, output_file, tempo, target_bpm, 
                                 pitch_semitones, timemap_entries, 
                                 verbose=False):
    """Aplica Rubber Band com flags otimizadas."""

    fd, timemap_path = tempfile.mkstemp(suffix=".txt")
    os.close(fd)

    with open(timemap_path, "w") as f:
        for src, tgt in timemap_entries:
            f.write(f"{src} {tgt}\n")

    # Ratio global (obrigatório com timemap)
    tempo_ratio = target_bpm / tempo

    # Construir comando Rubber Band
    cmd = ["rubberband", "-q", "--fine"]

    # Tempo ratio
    cmd.extend(["-T", str(tempo_ratio)])

    # Pitch shift
    if pitch_semitones != 0:
        cmd.extend(["-p", str(pitch_semitones)])
        cmd.append("--pitch-hq")

    # Timemap
    cmd.extend(["-M", timemap_path])
    
    # Flags de qualidade
    cmd.append("--formant") # Preservar timbre
    
    cmd.extend([temp_wav, output_file])

    if verbose: print(f"[*] A executar: {' '.join(cmd)}")

    result = subprocess.run(cmd, capture_output=True, text=True)

    if os.path.exists(timemap_path):
        os.remove(timemap_path)

    if result.returncode != 0:
        print(f"[!] Erro Rubber Band: {result.stderr}")
        return False

    return True


def fix_grid_advanced(input_file, output_file, target_bpm=None, pitch_semitones=0,
                       phase_lock=True, verbose=False):
    temp_wav = None
    try:
        analysis = analyze_audio(input_file, verbose=verbose)
        y = analysis["y"]
        sr = analysis["sr"]
        tempo = analysis["tempo"]
        beat_times = analysis["beat_times"]
        first_beat_time_orig = analysis["first_beat_time"]

        if len(beat_times) < 2:
            print("[!] Erro: Beats insuficientes para análise.")
            return None

        if target_bpm is None or target_bpm <= 0:
            target_bpm = round(tempo)
        
        fd, temp_wav = tempfile.mkstemp(suffix=".wav")
        os.close(fd)
        sf.write(temp_wav, y.T if y.ndim > 1 else y, sr)

        timemap_entries = create_smart_timemap(
            beat_times, target_bpm, sr, first_beat_time_orig,
            analysis["is_metronomic"], phase_lock=phase_lock, verbose=verbose
        )

        success = apply_rubberband_processing(
            temp_wav, output_file, tempo, target_bpm,
            pitch_semitones, timemap_entries, verbose
        )

        if not success:
            return None

        metadata = {
            "input": input_file,
            "original_bpm": round(float(tempo), 2),
            "target_bpm": float(target_bpm),
            "pitch_shift_semitons": float(pitch_semitones),
            "phase_lock": bool(phase_lock),
            "is_metronomic": bool(analysis["is_metronomic"]),
            "ibi_cv": round(float(analysis["ibi_cv"]), 5),
            "beats_count": int(len(beat_times)),
            "first_beat_time_original": round(float(first_beat_time_orig), 3),
            "estimated_key": str(analysis["estimated_key"]),
            "phase_offset_degrees": round(float(analysis["phase_offset_degrees"]), 1),
            "output": output_file
        }

        meta_path = str(Path(output_file).with_suffix(".json"))
        with open(meta_path, "w", encoding="utf-8") as f:
            json.dump(metadata, f, indent=4, ensure_ascii=False)

        return metadata

    except Exception as e:
        print(f"[!] Erro: {str(e)}")
        import traceback
        traceback.print_exc()
        return None
    finally:
        if temp_wav and os.path.exists(temp_wav):
            os.remove(temp_wav)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Forca-Grid v2 Enhanced: Beatgrid + Pitch + Phase Lock"
    )
    parser.add_argument("input", help="Ficheiro de áudio de entrada")
    parser.add_argument("output", help="Ficheiro de áudio de saída")
    parser.add_argument("--bpm", type=float, default=None, help="BPM alvo")
    parser.add_argument("--pitch", type=float, default=0, 
                        help="Pitch shift em semitons (+/-)")
    parser.add_argument("--no-phase-lock", action="store_true",
                        help="Desativar phase-locking")
    parser.add_argument("-v", "--verbose", action="store_true")

    args = parser.parse_args()

    fix_grid_advanced(
        args.input, args.output,
        target_bpm=args.bpm,
        pitch_semitones=args.pitch,
        phase_lock=not args.no_phase_lock,
        verbose=args.verbose
    )