"""Analysis, separation and mixing helpers for the accompaniment pipeline."""
import logging

import librosa
import numpy as np
import soundfile as sf
import torch

SR = 48000
# Krumhansl-Schmuckler profiles: correlating them against the averaged chroma
# is enough to name a key, which is all the style prompt needs.
MAJOR = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88])
MINOR = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17])
NOTES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]


def load(path, sr=SR, mono=False):
    audio, rate = librosa.load(path, sr=sr, mono=mono)
    if audio.ndim == 1:
        audio = audio[None, :]
    return audio, rate


def analyse(path):
    """Return tempo, key and the first beat of a recording."""
    mono, _ = librosa.load(path, sr=22050, mono=True)
    duration = float(len(mono)) / 22050
    _, beats = librosa.beat.beat_track(y=mono, sr=22050, units="time")
    # A confident tempo needs actual onsets; humming with no transients gives
    # too few beats, and claiming a tempo there would detune the whole backing.
    if len(beats) < 4:
        tempo, first = 0.0, 0.0
    else:
        # librosa's own tempo estimate is quantised to tempogram bins and can
        # land several BPM off (117 for a true 120), which drifts audibly over
        # a song. The beat times are far more accurate, so read it from those.
        spacing = float(np.median(np.diff(beats)))
        tempo = 60.0 / spacing if spacing > 0 else 0.0
        first = float(beats[0])
    if not np.isfinite(tempo) or not 30 <= tempo <= 300:
        tempo, first = 0.0, 0.0
    return round(tempo), key_of(mono, 22050), duration, first


def key_of(mono, sr):
    """Name the key of a recording, or return "" when it has no pitch content."""
    chroma = librosa.feature.chroma_cqt(y=librosa.effects.harmonic(mono), sr=sr)
    profile = chroma.mean(axis=1)
    if profile.sum() <= 0:
        return ""
    scores = []
    for index in range(12):
        rotated = np.roll(profile, -index)
        scores.append((np.corrcoef(rotated, MAJOR)[0, 1], f"{NOTES[index]} major"))
        scores.append((np.corrcoef(rotated, MINOR)[0, 1], f"{NOTES[index]} minor"))
    best = max(scores, key=lambda item: (item[0] if np.isfinite(item[0]) else -2))
    return best[1]


def separate(path, model, want="vocals"):
    """Split a mix with Demucs and return the requested stem plus the rest."""
    stems, names = _demucs(path, model)
    index = names.index(want) if want in names else names.index("vocals")
    target = stems[index].cpu().numpy()
    rest = sum(stems[i] for i in range(len(names)) if i != index).cpu().numpy()
    return target, rest, model.samplerate


def split_all(path, model):
    """Every Demucs stem by name: drums, bass, other and vocals."""
    stems, names = _demucs(path, model)
    return {name: stems[i].cpu().numpy() for i, name in enumerate(names)}, model.samplerate


def _demucs(path, model):
    from demucs.apply import apply_model

    audio, _ = load(path, sr=model.samplerate)
    if audio.shape[0] == 1:
        audio = np.repeat(audio, model.audio_channels, axis=0)
    tensor = torch.from_numpy(audio[: model.audio_channels]).float()
    # Demucs expects the loudness normalisation it was trained with, then the
    # original scale restored so the stems mix back to the source level.
    ref = tensor.mean(0)
    scale = ref.std() or torch.tensor(1.0)
    tensor = (tensor - ref.mean()) / scale
    with torch.inference_mode():
        stems = apply_model(model, tensor[None], device="cuda", progress=False)[0]
    return stems * scale + ref.mean(), list(model.sources)


def to_stereo(audio, sr, target_sr=SR):
    if sr != target_sr:
        audio = librosa.resample(audio, orig_sr=sr, target_sr=target_sr)
    if audio.ndim == 1:
        audio = audio[None, :]
    if audio.shape[0] == 1:
        audio = np.repeat(audio, 2, axis=0)
    return audio[:2]


def rms(audio):
    value = float(np.sqrt(np.mean(np.square(audio))))
    return value if value > 1e-6 else 1e-6


# About -34 dBFS: quieter than any usable take, so a failed separation or a
# whisper cannot drag the band down to silence with it.
QUIETEST = 0.02


def loudness(audio, sr=SR):
    """RMS over the parts that are actually sounding.

    A take is full of gaps between phrases; averaging them in makes a singer
    measure quiet and the band gets turned down to match. Only 50 ms frames
    within 30 dB of the loudest one count.
    """
    mono = np.asarray(audio).reshape(-1, np.asarray(audio).shape[-1]).mean(axis=0)
    frame = max(1, int(sr * 0.05))
    count = len(mono) // frame
    if count < 1:
        return rms(mono)
    levels = np.sqrt(np.mean(np.square(mono[: count * frame].reshape(count, frame)), axis=1))
    active = levels[levels > levels.max() * 10 ** (-30 / 20)]
    return float(np.sqrt(np.mean(np.square(active)))) if len(active) else rms(mono)


def mix(vocal, instrumental, path, vocal_db=0.0, instrumental_db=-3.0, offset=0.0, sr=SR, report=None):
    """Lay a vocal over a backing track and write a headroom-safe master."""
    vocal, instrumental = to_stereo(vocal, sr), to_stereo(instrumental, sr)
    # Nudging the backing track to the singer's first beat is the only timing
    # control available: the model matches tempo, not the exact performance.
    shift = int(round(offset * sr))
    if shift > 0:
        instrumental = np.pad(instrumental, ((0, 0), (shift, 0)))
    elif shift < 0:
        instrumental = instrumental[:, -shift:]
    length = max(vocal.shape[1], instrumental.shape[1])
    vocal = np.pad(vocal, ((0, 0), (0, length - vocal.shape[1])))
    instrumental = np.pad(instrumental, ((0, 0), (0, length - instrumental.shape[1])))
    # Match the backing track to the vocal before applying the requested
    # balance, so the blend does not depend on how loudly the user recorded.
    target = max(loudness(vocal, sr), QUIETEST)
    vocal_gain = 10 ** (vocal_db / 20)
    backing_gain = target / max(loudness(instrumental, sr), 1e-6) * (10 ** (instrumental_db / 20))
    blend = vocal * vocal_gain + instrumental * backing_gain
    peak = float(np.max(np.abs(blend)))
    if peak > 0.99:
        blend = blend * (0.99 / peak)
        vocal_gain *= 0.99 / peak
        backing_gain *= 0.99 / peak
    if report is not None:
        report["mix"] = {"vocal": float(vocal_gain), "backing": float(backing_gain)}
    sf.write(path, blend.T, sr, subtype="PCM_16")
    return path


def put_back(vocal, backing, reference, path, sr=SR):
    """Lay the untouched vocal over a regenerated backing, at the original balance.

    The new backing is brought to the loudness of the backing it replaces, so a
    restyle changes the sound of the music and never how loud the singer sits.
    """
    vocal, backing = to_stereo(vocal, sr), to_stereo(backing, sr)
    length = vocal.shape[1]
    backing = shift(backing, 0.0, sr, length=length)
    backing = backing * (loudness(reference, sr) / max(loudness(backing, sr), 1e-6))
    blend = vocal + backing
    peak = float(np.max(np.abs(blend)))
    if peak > 0.99:
        blend = blend * (0.99 / peak)
    sf.write(path, blend.T, sr, subtype="PCM_16")
    return path


def write(audio, path, sr=SR):
    audio = to_stereo(audio, sr)
    sf.write(path, audio.T, SR, subtype="PCM_16")
    return path


# ------------------------------------------------------------ band following
#
# ACE-Step reads "120 bpm, key of C major" in a prompt as a loose hint, so a
# band written from text alone starts on its own downbeat, wanders in tempo and
# plays its own chords. What it follows closely is audio: audio2audio keeps the
# timing and harmony of a reference while the prompt restyles it. So the band
# is written over a guide track — a plain synth arrangement built on the
# singer's own beat grid and chords — and the model turns that into the style.

GUIDE_SR = 44100
ROOT_FREQ = 130.81  # C3
PARTS = ("drums", "bass", "chords")
QUALITIES = {"major": (0, 4, 7), "minor": (0, 3, 7)}
MAJOR_SCALE = (0, 2, 4, 5, 7, 9, 11)
MINOR_SCALE = (0, 2, 3, 5, 7, 8, 10)


def beat_grid(mono, sr, bpm=0.0):
    """Beat times across the whole recording, and the tempo they imply.

    A cappella singing has soft onsets, so tracked beats jitter. When they fit
    a straight line the take is steady and the fitted grid is used, which keeps
    the band from stumbling on every late syllable; a rubato take keeps its
    real beats. `bpm` overrides detection when the user knows the tempo.
    """
    duration = len(mono) / sr
    kwargs = {"start_bpm": float(bpm)} if bpm else {}
    _, beats = librosa.beat.beat_track(y=mono, sr=sr, units="time", **kwargs)
    if bpm:
        step = 60.0 / float(bpm)
        phase = float(beats[0]) % step if len(beats) else 0.0
        return np.arange(phase, duration, step), float(bpm)
    if len(beats) < 4:
        return np.array([]), 0.0
    index = np.arange(len(beats))
    slope, start = np.polyfit(index, beats, 1)
    if not 0.2 <= slope <= 2.0:
        return np.array([]), 0.0
    residual = float(np.median(np.abs(beats - (start + slope * index))))
    if residual < 0.05:
        grid = np.arange(start % slope, duration, slope)
    else:
        # Rubato: keep the tracked beats and extend them at the average spacing.
        head = np.arange(beats[0] - slope, 0, -slope)[::-1]
        tail = np.arange(beats[-1] + slope, duration, slope)
        grid = np.concatenate([head, beats, tail])
    grid = grid - _lag(mono, sr, grid, slope)
    return grid[grid >= 0], 60.0 / slope


def _lag(mono, sr, grid, spacing):
    """How far the tracker's beats trail the note attacks they stand for.

    Beat frames sit on onset-strength peaks, a frame or two after the attack
    a listener hears, and a band that late sounds like it is dragging.
    """
    onsets = librosa.onset.onset_detect(y=mono, sr=sr, units="time", backtrack=True)
    if len(onsets) < 4 or len(grid) < 2:
        return 0.0
    nearest = onsets[np.clip(np.searchsorted(onsets, grid), 1, len(onsets) - 1)]
    earlier = onsets[np.clip(np.searchsorted(onsets, grid) - 1, 0, len(onsets) - 1)]
    closest = np.where(np.abs(nearest - grid) < np.abs(earlier - grid), nearest, earlier)
    gap = grid - closest
    near = gap[np.abs(gap) < spacing / 4]
    return float(np.median(near)) if len(near) >= 4 else 0.0


def _beat_chroma(mono, sr, beats):
    chroma = librosa.feature.chroma_cqt(y=librosa.effects.harmonic(mono), sr=sr)
    frames = np.clip(librosa.time_to_frames(beats, sr=sr), 0, chroma.shape[1] - 1)
    edges = list(frames) + [chroma.shape[1]]
    return np.stack([chroma[:, a:max(b, a + 1)].mean(axis=1) for a, b in zip(edges[:-1], edges[1:])], axis=1)


def _templates(key):
    tonic, scale = 0, MAJOR_SCALE
    if key:
        name, mode = key.split()
        tonic = NOTES.index(name)
        scale = MAJOR_SCALE if mode == "major" else MINOR_SCALE
    diatonic = {(tonic + degree) % 12 for degree in scale}
    templates = []
    for root in range(12):
        for quality, shape in QUALITIES.items():
            template = np.zeros(12)
            for weight, offset in zip((1.0, 0.8, 0.9), shape):
                template[(root + offset) % 12] = weight
            in_key = all((root + offset) % 12 in diatonic for offset in shape)
            templates.append((root, quality, template / np.linalg.norm(template), in_key))
    return templates


def _label(per_beat, beats, bounds, templates):
    """Best triad for each span of beats, and how well the spans fit overall."""
    found, previous, fits = [], None, []
    for first, last in bounds:
        column = per_beat[:, first:last].mean(axis=1)
        norm = float(np.linalg.norm(column))
        time = 0.0 if first == 0 else float(beats[first])
        if norm < 1e-6:
            if previous:
                found.append({"time": time, **previous})
            continue
        best, choice, fit = -9.0, None, 0.0
        for root, quality, template, in_key in templates:
            match = float(column @ template) / norm
            score = match + (0.12 if in_key else 0.0)
            if previous and previous["root"] == root and previous["quality"] == quality:
                score += 0.06
            if score > best:
                best, choice, fit = score, {"root": root, "quality": quality}, match
        previous = choice
        fits.append(fit)
        found.append({"time": time, **choice})
    return found, (float(np.mean(fits)) if fits else 0.0)


def chords(mono, sr, beats, key, per_bar=4):
    """Name one triad per bar from the harmony the recording implies.

    Each bar's chroma is scored against every major and minor triad. Chords in
    the detected key get a bonus, as does repeating the previous chord, which
    keeps a vocal line — a few notes per bar — from yielding a new chord on
    every passing tone.

    Where the bar lines fall is not known either. Neighbouring chords share two
    notes, so harmonic change is a weak clue; instead each of the `per_bar`
    possible placements is labelled and the one whose bars fit triads best
    wins, since a bar line in the wrong place blends two chords into a muddle.
    """
    if len(beats) < per_bar + 1:
        return []
    per_beat = _beat_chroma(mono, sr, beats)
    templates = _templates(key)
    count = per_beat.shape[1]
    best, result = -1.0, []
    for phase in range(per_bar):
        lines = list(range(phase, count, per_bar))
        bounds = [(a, min(a + per_bar, count)) for a in lines]
        if phase:
            bounds.insert(0, (0, phase))
        found, fit = _label(per_beat, beats, bounds, templates)
        if fit > best + 1e-6:
            best, result = fit, found
    # A chord held across bars is one chord; the guide sustains it.
    merged = []
    for chord in result:
        if not merged or (merged[-1]["root"], merged[-1]["quality"]) != (chord["root"], chord["quality"]):
            merged.append(chord)
    return merged


def chord_name(chord):
    return NOTES[chord["root"]] + ("m" if chord["quality"] == "minor" else "")


def _tone(freq, length, sr, harmonics=4, decay=0.0):
    t = np.arange(length) / sr
    wave = sum(np.sin(2 * np.pi * freq * h * t) / h for h in range(1, harmonics + 1))
    envelope = np.exp(-decay * t) if decay else np.ones(length)
    fade = min(length // 2, int(sr * 0.01))
    if fade:
        envelope[:fade] *= np.linspace(0, 1, fade)
        envelope[-fade:] *= np.linspace(1, 0, fade)
    return wave * envelope


def _place(track, sound, at):
    start = int(at)
    if start < 0 or start >= len(track):
        return
    end = min(len(track), start + len(sound))
    track[start:end] += sound[: end - start]


def guide(beats, progression, duration, parts=PARTS, sr=GUIDE_SR):
    """Synthesise a plain arrangement on the given beat grid and chords.

    It is deliberately simple — a kick and snare backbeat, a root-note bass and
    a sustained triad — because it only has to tell the model where the beats
    and changes are; the prompt supplies the actual sound.
    """
    length = int(duration * sr)
    drums, bass, pad = np.zeros(length), np.zeros(length), np.zeros(length)
    rng = np.random.default_rng(7)
    spacing = float(np.median(np.diff(beats))) if len(beats) > 1 else 0.5
    t = np.arange(int(sr * 0.18)) / sr
    kick = np.sin(2 * np.pi * (50 + 90 * np.exp(-t * 30)) * t) * np.exp(-t * 18)
    snare = rng.standard_normal(int(sr * 0.14)) * np.exp(-np.arange(int(sr * 0.14)) / sr * 28)
    hat = np.diff(rng.standard_normal(int(sr * 0.04)), prepend=0) * np.exp(
        -np.arange(int(sr * 0.04)) / sr * 90)
    bar_lines = {round(chord["time"], 3) for chord in progression}
    counter = 0
    for beat in beats:
        # Count from each bar line so the kick lands on the downbeat.
        if round(float(beat), 3) in bar_lines:
            counter = 0
        _place(drums, kick if counter % 2 == 0 else 0.55 * snare, beat * sr)
        _place(drums, 0.18 * hat, beat * sr)
        _place(drums, 0.12 * hat, (beat + spacing / 2) * sr)
        counter += 1
    for index, chord in enumerate(progression):
        start = chord["time"]
        end = progression[index + 1]["time"] if index + 1 < len(progression) else duration
        span = int((end - start) * sr)
        if span <= 0:
            continue
        root = ROOT_FREQ * 2 ** (chord["root"] / 12)
        triad = sum(_tone(root * 2 ** (offset / 12), span, sr, harmonics=3)
                    for offset in QUALITIES[chord["quality"]])
        _place(pad, triad / 3, start * sr)
        for number, beat in enumerate(b for b in beats if start - 1e-3 <= b < end):
            if number % 2 == 0:
                note_len = int(min(end - beat, 2 * spacing) * sr)
                if note_len > 0:
                    _place(bass, _tone(root / 2, note_len, sr, harmonics=2, decay=3.0), beat * sr)
    layers = {"drums": drums, "bass": bass, "chords": pad}
    out = np.zeros(length)
    for name, level in (("drums", 0.5), ("bass", 0.6), ("chords", 0.35)):
        if name in parts and np.any(layers[name]):
            out += level * layers[name] / np.max(np.abs(layers[name]))
    peak = float(np.max(np.abs(out))) if length else 0.0
    return (out / peak * 0.8 if peak > 0 else out).astype(np.float32)


def has_backing(vocal, rest):
    """Whether Demucs found real music under the voice rather than room noise."""
    return rms(rest) / rms(vocal) > 0.2


def has_voice(vocal, rest, sr=SR):
    """Whether Demucs found a voice worth keeping at all.

    A take with no singing — a guitar riff, a synth line — leaves the vocal
    stem close to silent. Lifting that would swap the user's recording for
    nothing, so such a take is followed as it is instead.
    """
    voice = loudness(vocal, sr)
    return voice >= 0.005 and voice >= 0.15 * loudness(rest, sr)


def offset_between(reference, candidate, sr, window=0.35):
    """Seconds to delay `candidate` so its onsets line up with `reference`."""
    hop = 512
    a = librosa.onset.onset_strength(y=np.asarray(reference, dtype=np.float32), sr=sr, hop_length=hop)
    b = librosa.onset.onset_strength(y=np.asarray(candidate, dtype=np.float32), sr=sr, hop_length=hop)
    n = min(len(a), len(b))
    if n < 8:
        return 0.0
    a, b = a[:n] - a[:n].mean(), b[:n] - b[:n].mean()
    reach = int(window * sr / hop)
    lags = list(range(-reach, reach + 1))
    scores = [float(np.dot(a[max(0, lag):n + min(0, lag)], b[max(0, -lag):n - max(0, lag)]))
              for lag in lags]
    return lags[int(np.argmax(scores))] * hop / sr


def shift(audio, seconds, sr, length=None):
    """Delay (positive) or advance audio, keeping it at `length` samples."""
    audio = audio if audio.ndim == 2 else audio[None, :]
    samples = int(round(seconds * sr))
    if samples > 0:
        audio = np.pad(audio, ((0, 0), (samples, 0)))
    elif samples < 0:
        audio = audio[:, -samples:]
    length = length or audio.shape[1]
    if audio.shape[1] < length:
        audio = np.pad(audio, ((0, 0), (0, length - audio.shape[1])))
    return audio[:, :length]