#!/usr/bin/env python3
"""
KISUM - musik al reves.

Escucha un tema de audio (o el audio de un video) y lo devuelve como codigo
Strudel: tempo, bateria por bandas (bombo/caja/hats) y, opcionalmente, una
linea de bajo aproximada.

Acepta cualquier archivo con audio: mp3, wav, flac, ogg, m4a... y tambien
video (mp4, mkv, webm, mov...). Si el formato no es audio puro, extrae la
pista con ffmpeg automaticamente.

Uso:
    python3 kisum.py tema.mp3 [--bars 4] [--bass] [-o salida.strudel]
    python3 kisum.py clip.mp4                 # tambien vale un video
"""

import argparse
import os
import shutil
import subprocess
import sys
import tempfile
import warnings

import numpy as np

try:
    import librosa
    import soundfile as sf
except ImportError:
    sys.exit("Faltan dependencias. Instalalas:\n"
             "    python3 -m venv .venv && "
             ".venv/bin/pip install -r requirements.txt")

STEPS_PER_BEAT = 4          # semicorcheas
BEATS_PER_BAR = 4           # asumimos 4/4
STEPS_PER_BAR = STEPS_PER_BEAT * BEATS_PER_BAR

# Bandas de frecuencia -> sample de Strudel
BANDS = [
    ("bd", 20, 150),        # bombo
    ("sd", 150, 2500),      # caja / palmas
    ("hh", 5000, 11000),    # hats / brillos
]

NOTE_NAMES = ["c", "cs", "d", "eb", "e", "f",
              "fs", "g", "gs", "a", "bb", "b"]


def load_audio(path):
    """Carga el audio de `path` como (senal mono float32, sr).

    Primero intenta leerlo directamente con libsndfile (wav, flac, ogg,
    mp3...). Si no es un formato de audio que entienda (un video mp4/mkv/
    webm/mov, u otro contenedor), extrae la pista de audio con ffmpeg a un
    wav temporal. Asi KISUM acepta cualquier archivo con audio sin depender
    del backend audioread (deprecado en librosa)."""
    if not os.path.exists(path):
        sys.exit(f"No existe el archivo: {path}")

    # 1) intento directo: wav/flac/ogg/mp3 via libsndfile
    try:
        y, sr = sf.read(path, dtype="float32", always_2d=True)
        return y.mean(axis=1), sr           # mezcla a mono
    except Exception:
        pass                                # no es audio legible -> ffmpeg

    # 2) fallback: extraer audio con ffmpeg (videos y formatos raros)
    if not shutil.which("ffmpeg"):
        sys.exit(f"No puedo leer '{path}' directamente y no encuentro ffmpeg.\n"
                 "Instala ffmpeg (sudo apt install ffmpeg) o convierte el "
                 "archivo a wav/mp3.")

    tmp_wav = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    try:
        res = subprocess.run(
            ["ffmpeg", "-y", "-loglevel", "error", "-i", path,
             "-vn", "-ac", "1", "-ar", "44100", tmp_wav],
            capture_output=True, text=True)
        if res.returncode != 0 or not os.path.getsize(tmp_wav):
            sys.exit(f"ffmpeg no pudo extraer audio de '{path}'.\n"
                     f"{res.stderr.strip()}")
        y, sr = sf.read(tmp_wav, dtype="float32", always_2d=True)
        return y.mean(axis=1), sr
    finally:
        if os.path.exists(tmp_wav):
            os.remove(tmp_wav)


def band_onsets(y, sr, fmin, fmax):
    """Tiempos (s) de los golpes dentro de una banda de frecuencias."""
    with warnings.catch_warnings():
        # las bandas son estrechas: sobran filtros mel y librosa avisa.
        # Es inofensivo (esos canales quedan vacios), lo silenciamos.
        warnings.filterwarnings("ignore", message="Empty filters detected")
        S = librosa.feature.melspectrogram(
            y=y, sr=sr, n_mels=32, fmin=fmin, fmax=fmax)
    env = librosa.onset.onset_strength(S=librosa.power_to_db(S), sr=sr)
    frames = librosa.onset.onset_detect(
        onset_envelope=env, sr=sr, backtrack=False)
    return librosa.frames_to_time(frames, sr=sr)


def make_grid(beat_times, bars):
    """Rejilla de semicorcheas: subdivide cada pulso en 4. Devuelve los
    tiempos (s) de cada paso, recortados a `bars` compases completos."""
    steps = []
    for i in range(len(beat_times) - 1):
        t0, t1 = beat_times[i], beat_times[i + 1]
        for k in range(STEPS_PER_BEAT):
            steps.append(t0 + (t1 - t0) * k / STEPS_PER_BEAT)
    return steps[: bars * STEPS_PER_BAR]


def quantize(onsets, grid):
    """Marca en que paso de la rejilla cae cada golpe (al mas cercano,
    con tolerancia de medio paso)."""
    hits = [False] * len(grid)
    if not len(grid):
        return hits
    step = np.median(np.diff(grid)) if len(grid) > 1 else 0.1
    for t in onsets:
        i = int(np.argmin(np.abs(np.asarray(grid) - t)))
        if abs(grid[i] - t) <= step / 2:
            hits[i] = True
    return hits


def pattern_string(hits, token):
    """Convierte la lista de pasos en mini-notacion Strudel: un compas por
    [ ] y, si hay varios, alternando con < >."""
    bars = []
    for b in range(0, len(hits), STEPS_PER_BAR):
        bar = hits[b:b + STEPS_PER_BAR]
        if len(bar) < STEPS_PER_BAR:
            break
        bars.append("[" + " ".join(token if h else "~" for h in bar) + "]")
    if not bars:
        return None
    if len(bars) == 1:
        return bars[0][1:-1]                    # sin corchetes si es uno
    if all(b == bars[0] for b in bars):
        return bars[0][1:-1]                    # todos iguales: uno basta
    return "<" + " ".join(bars) + ">"


def bass_line(y, sr, beat_times, bars):
    """Nota dominante por pulso sobre el audio grave (experimental)."""
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr, fmin=32.7)  # desde C1
    beats = librosa.time_to_frames(beat_times, sr=sr)
    notes = []
    for i in range(min(len(beats) - 1, bars * BEATS_PER_BAR)):
        seg = chroma[:, beats[i]:beats[i + 1]]
        if seg.size == 0 or seg.max() < 0.5:
            notes.append("~")
        else:
            notes.append(NOTE_NAMES[int(seg.mean(axis=1).argmax())] + "2")
    if not notes or all(n == "~" for n in notes):
        return None
    return " ".join(notes)


def analyze(path, bars, want_bass):
    y, sr = load_audio(path)
    tempo, beat_times = librosa.beat.beat_track(y=y, sr=sr, units="time")
    tempo = float(np.atleast_1d(tempo)[0])
    if len(beat_times) < BEATS_PER_BAR + 1:
        sys.exit("No encuentro un pulso claro en el audio. "
                 "Prueba con un fragmento con percusion.")
    grid = make_grid(beat_times, bars)

    lines = []
    for token, fmin, fmax in BANDS:
        hits = quantize(band_onsets(y, sr, fmin, fmax), grid)
        pat = pattern_string(hits, token)
        if pat:
            gain = '.gain(0.7)' if token == "hh" else ""
            lines.append(f'  s("{pat}"){gain},')

    if want_bass:
        bl = bass_line(y, sr, beat_times, bars)
        if bl:
            lines.append(f'  note("{bl}").s("sawtooth").lpf(400)'
                         f'.gain(0.6),  // bajo experimental')

    bpm = round(tempo, 1)
    code = [
        f"// KISUM v0 :: {path}",
        f"// bpm detectado ~ {bpm} (si suena a mitad/doble, corrige aqui)",
        f"setcps({bpm}/60/4)",
        "stack(",
        *lines,
        ")",
    ]
    return "\n".join(code) + "\n"


def main():
    ap = argparse.ArgumentParser(
        description="KISUM: de un tema de audio a codigo Strudel.")
    ap.add_argument("audio",
                    help="archivo con audio: mp3, wav, flac, m4a, ogg... "
                         "o video (mp4, mkv, webm, mov...)")
    ap.add_argument("--bars", type=int, default=4,
                    help="compases a transcribir (defecto: 4)")
    ap.add_argument("--bass", action="store_true",
                    help="anade linea de bajo aproximada")
    ap.add_argument("-o", "--out", help="guardar tambien en un archivo")
    args = ap.parse_args()

    code = analyze(args.audio, max(1, args.bars), args.bass)
    print(code)
    if args.out:
        with open(args.out, "w", encoding="utf-8") as fh:
            fh.write(code)
        print(f"// guardado en {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()
