"""
Core music analysis engine (audio-only, no lyrics).

Pure-DSP analysis with librosa/scikit-learn: tempo, key, time signature,
structure, energy and timbre. Fast on CPU, no model downloads needed.

Deep-learning tagging (instruments, genre, mood) and stem separation live in
tagging.py and stems.py and are imported lazily by the app, so this core
always works even if those heavier models are absent.
"""
from __future__ import annotations
import numpy as np
import librosa
import scipy.ndimage
import scipy.linalg
import scipy.sparse.csgraph
import sklearn.cluster

NOTES = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B']

# Krumhansl-Schmuckler key profiles
KS_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])
KS_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])

HOP = 512


def _scalar(x) -> float:
    return float(np.atleast_1d(x).ravel()[0])


def load_audio(path, sr=22050, mono=True, max_seconds=480):
    """Load audio, trimming very long files so analysis stays responsive."""
    y, sr = librosa.load(path, sr=sr, mono=mono)
    duration = len(y) / sr
    if max_seconds and duration > max_seconds:
        y = y[: int(max_seconds * sr)]
    return y, sr, duration


def detect_tempo(y, sr, onset_env=None):
    if onset_env is None:
        onset_env = librosa.onset.onset_strength(y=y, sr=sr, hop_length=HOP)
    tempo, beats = librosa.beat.beat_track(onset_envelope=onset_env, sr=sr, hop_length=HOP)
    tempo = _scalar(tempo)
    beat_times = librosa.frames_to_time(beats, sr=sr, hop_length=HOP)
    if len(beat_times) > 3:
        ibi = np.diff(beat_times)
        regularity = float(1.0 - min(1.0, np.std(ibi) / (np.mean(ibi) + 1e-9)))
    else:
        regularity = 0.0
    return {
        "bpm": round(tempo, 1),
        "beats": beats,
        "beat_times": beat_times,
        "beat_count": int(len(beat_times)),
        "regularity": round(regularity, 3),
    }


def detect_key(y, sr, chroma=None):
    if chroma is None:
        chroma = librosa.feature.chroma_cqt(y=y, sr=sr, hop_length=HOP)
    cm = chroma.mean(axis=1)
    cmc = cm - cm.mean()
    results = []
    for i in range(12):
        maj = np.corrcoef(np.roll(KS_MAJOR - KS_MAJOR.mean(), i), cmc)[0, 1]
        minr = np.corrcoef(np.roll(KS_MINOR - KS_MINOR.mean(), i), cmc)[0, 1]
        results.append((NOTES[i], "majeur", maj))
        results.append((NOTES[i], "mineur", minr))
    results.sort(key=lambda r: r[2], reverse=True)
    best, second = results[0], results[1]
    return {
        "key": best[0],
        "mode": best[1],
        "key_name": f"{best[0]} {best[1]}",
        "correlation": round(float(best[2]), 3),
        "confidence": round(float(max(0.0, best[2] - second[2])), 3),
        "chroma_mean": cm.tolist(),
        "alternatives": [f"{r[0]} {r[1]}" for r in results[1:4]],
    }


def estimate_time_signature(onset_env, sr, beat_times):
    """Best-effort beats-per-bar guess. Meter detection is inherently noisy, so
    when the cue is weak we fall back to the 4/4 prior and flag low confidence."""
    if len(beat_times) < 8:
        return {"beats_per_bar": None, "label": "onbekend", "confidence": 0.0}
    beat_frames = librosa.time_to_frames(beat_times, sr=sr, hop_length=HOP)
    beat_frames = beat_frames[(beat_frames >= 0) & (beat_frames < len(onset_env))]
    bs = onset_env[beat_frames]
    bs = bs - bs.mean()
    ac = librosa.autocorrelate(bs)
    ac = ac / (ac[0] + 1e-9)
    cand = {b: (ac[b] if b < len(ac) else -1) for b in (2, 3, 4, 6)}
    best = max(cand, key=cand.get)
    if best == 6:
        best = 3
    others = [v for k, v in cand.items() if k != max(cand, key=cand.get)]
    confidence = float(max(0.0, min(1.0, cand[max(cand, key=cand.get)] - (max(others) if others else 0))))
    if confidence < 0.05:
        return {"beats_per_bar": 4, "label": "4/4 (waarschijnlijk)", "confidence": round(confidence, 3)}
    return {"beats_per_bar": int(best), "label": f"{best}/4", "confidence": round(confidence, 3)}


def analyze_structure(y, sr, beats=None, max_k=6, min_section_sec=6.0):
    """Laplacian structural segmentation (McFee & Ellis): cluster beat-synchronous
    CQT into repeated sections, then label them A/B/C by order of appearance."""
    dur = len(y) / sr
    if beats is None:
        _, beats = librosa.beat.beat_track(y=y, sr=sr, hop_length=HOP, trim=False)
    beat_times = librosa.frames_to_time(beats, sr=sr, hop_length=HOP)
    if len(beats) < 12:
        return {"sections": [{"label": "A", "start": 0.0, "end": round(dur, 2),
                              "duration": round(dur, 2)}], "n_unique": 1}

    C = librosa.amplitude_to_db(np.abs(librosa.cqt(y=y, sr=sr, hop_length=HOP,
                                bins_per_octave=36, n_bins=36 * 7)), ref=np.max)
    Csync = librosa.util.sync(C, beats, aggregate=np.median)
    n = Csync.shape[1]

    R = librosa.segment.recurrence_matrix(Csync, width=3, mode='affinity', sym=True)
    df = librosa.segment.timelag_filter(scipy.ndimage.median_filter)
    Rf = df(R, size=(1, 7))

    mfcc = librosa.feature.mfcc(y=y, sr=sr, hop_length=HOP)
    Msync = librosa.util.sync(mfcc, beats)
    path_dist = np.sum(np.diff(Msync, axis=1) ** 2, axis=0)
    sigma = np.median(path_dist) + 1e-9
    path_sim = np.exp(-path_dist / sigma)
    R_path = np.diag(path_sim, 1) + np.diag(path_sim, -1)

    deg_path = np.sum(R_path, axis=1)
    deg_rec = np.sum(Rf, axis=1)
    mu = deg_path.dot(deg_path + deg_rec) / (np.sum((deg_path + deg_rec) ** 2) + 1e-9)
    A = mu * Rf + (1 - mu) * R_path

    L = scipy.sparse.csgraph.laplacian(A, normed=True)
    _, evecs = scipy.linalg.eigh(L)
    evecs = scipy.ndimage.median_filter(evecs, size=(9, 1))

    k = int(np.clip(round(dur / 25) + 1, 2, max_k))
    k = min(k, n)
    Cn = np.cumsum(evecs ** 2, axis=1) ** 0.5
    X = evecs[:, :k] / (Cn[:, k - 1:k] + 1e-9)
    seg_ids = sklearn.cluster.KMeans(n_clusters=k, n_init=10, random_state=0).fit_predict(X)

    def btime(j):
        if j <= 0:
            return 0.0
        if j >= n:
            return float(dur)
        return float(beat_times[min(j - 1, len(beat_times) - 1)])

    # contiguous runs of identical cluster id -> raw sections
    bounds = np.flatnonzero(np.diff(seg_ids)) + 1
    starts = np.concatenate([[0], bounds])
    ends = np.concatenate([bounds, [n]])
    raw = [[int(seg_ids[s]), btime(s), btime(e)] for s, e in zip(starts, ends)]

    # merge sections shorter than the minimum into the previous one
    merged = []
    for cid, s, e in raw:
        if merged and (e - s) < min_section_sec:
            merged[-1][2] = e
        else:
            merged.append([cid, s, e])
    if len(merged) > 1 and (merged[0][2] - merged[0][1]) < min_section_sec:
        merged[1][1] = merged[0][1]
        merged.pop(0)

    # relabel cluster ids -> A/B/C by first appearance
    mapping, order = {}, 0
    sections = []
    for cid, s, e in merged:
        if cid not in mapping:
            mapping[cid] = chr(ord('A') + order)
            order += 1
        sections.append({"label": mapping[cid], "start": round(s, 2),
                         "end": round(e, 2), "duration": round(e - s, 2)})
    return {"sections": sections, "n_unique": len(mapping)}


def analyze_dynamics(y, sr):
    rms = librosa.feature.rms(y=y, hop_length=HOP)[0]
    rms_db = librosa.amplitude_to_db(rms + 1e-9)
    centroid = librosa.feature.spectral_centroid(y=y, sr=sr, hop_length=HOP)[0]
    zcr = librosa.feature.zero_crossing_rate(y, hop_length=HOP)[0]
    return {
        "energy": round(float(np.clip((np.mean(rms) / 0.15) * 100, 0, 100)), 1),
        "dynamic_range_db": round(float(np.percentile(rms_db, 95) - np.percentile(rms_db, 5)), 1),
        "brightness_hz": round(float(np.mean(centroid)), 0),
        "rms_curve": rms.tolist(),
        "zcr_mean": round(float(np.mean(zcr)), 4),
    }


def estimate_danceability(tempo_info, dynamics):
    reg = tempo_info["regularity"]
    bpm = tempo_info["bpm"]
    tempo_fit = float(np.exp(-((bpm - 120) ** 2) / (2 * 35 ** 2)))
    e = dynamics["energy"] / 100.0
    return round(float(np.clip(100 * (0.5 * reg + 0.3 * tempo_fit + 0.2 * e), 0, 100)), 1)


def estimate_valence(key_info, tempo_info, dynamics):
    mode_term = 0.6 if key_info["mode"] == "majeur" else 0.3
    tempo_term = float(np.clip((tempo_info["bpm"] - 60) / 120, 0, 1)) * 0.25
    bright_term = float(np.clip(dynamics["brightness_hz"] / 4000, 0, 1)) * 0.15
    return round(float(np.clip((mode_term + tempo_term + bright_term) * 100, 0, 100)), 1)


# ----------------------------------------------------------------- chords
FLATS = ['C', 'Db', 'D', 'Eb', 'E', 'F', 'Gb', 'G', 'Ab', 'A', 'Bb', 'B']

# Roman-cijfers per interval t.o.v. de grondtoon (diatonisch majeur/mineur, met
# nette terugval voor niet-diatonische akkoorden). Hoofd-/kleine letter wordt
# later op basis van het werkelijke akkoord (maj/min) gezet.
_ROMAN_MAJOR = {0: "I", 2: "II", 4: "III", 5: "IV", 7: "V", 9: "VI", 11: "VII"}
_ROMAN_MINOR = {0: "I", 2: "II", 3: "III", 5: "IV", 7: "V", 8: "VI", 10: "VII"}
_ROMAN_FALLBACK = {0: "I", 1: "bII", 2: "II", 3: "bIII", 4: "III", 5: "IV",
                   6: "bV", 7: "V", 8: "bVI", 9: "VI", 10: "bVII", 11: "VII"}


def _chord_templates():
    """24 binaire triad-templates (12 majeur + 12 mineur), genormaliseerd."""
    temps, labels = [], []
    for r in range(12):
        for quality, third in (("maj", 4), ("min", 3)):
            v = np.zeros(12)
            v[[r, (r + third) % 12, (r + 7) % 12]] = 1.0
            temps.append(v / np.linalg.norm(v))
            labels.append((r, quality))
    return np.array(temps), labels


def _chord_name(chord):
    r, q = chord
    return FLATS[r] + ("m" if q == "min" else "")


def _roman_numeral(chord, key_root, mode):
    r, q = chord
    interval = (r - key_root) % 12
    table = _ROMAN_MINOR if mode == "mineur" else _ROMAN_MAJOR
    base = table.get(interval) or _ROMAN_FALLBACK[interval]
    acc = ""
    if base and base[0] == "b":
        acc, base = "b", base[1:]
    return acc + (base.lower() if q == "min" else base.upper())


def _find_loop(bars, lengths=(4, 8, 2, 3, 6)):
    """Vind de meest voorkomende terugkerende reeks maten (de 'loop')."""
    from collections import Counter
    n = len(bars)
    if n == 0:
        return []
    if n < 4:
        out = [bars[0]]
        for b in bars[1:]:
            if b != out[-1]:
                out.append(b)
        return out[:4]
    best, best_score = None, -1
    for L in lengths:
        if n < L * 2:
            continue
        windows = [tuple(bars[i:i + L]) for i in range(0, n - L + 1)]
        win, freq = Counter(windows).most_common(1)[0]
        score = freq * L
        if freq >= 2 and score > best_score:
            best, best_score = list(win), score
    if best is None:
        out = [bars[0]]
        for b in bars[1:]:
            if b != out[-1]:
                out.append(b)
        best = out[:4]
    return best


def detect_chords(chroma, beats, beats_per_bar=4, key="C", mode="majeur"):
    """Schat akkoorden per maat en leid een progressie/harmonisch palet af.
    Akkoordherkenning uit audio is benaderend, zeker bij dichte mixen."""
    if chroma is None or beats is None or len(beats) < 4:
        return {"available": False}
    from collections import Counter
    bpb = max(1, int(beats_per_bar or 4))
    Csync = librosa.util.sync(chroma, beats, aggregate=np.median)   # [12, n_beats]
    nbeats = Csync.shape[1]
    temps, labels = _chord_templates()

    # akkoord per maat: gemiddelde chroma over de beats in de maat (ruisarm)
    bar_vecs = []
    for i in range(0, nbeats, bpb):
        seg = Csync[:, i:i + bpb]
        if seg.shape[1] == 0:
            break
        bar_vecs.append(seg.mean(axis=1))
    if not bar_vecs:
        return {"available": False}
    B = np.array(bar_vecs).T
    B = B / (np.linalg.norm(B, axis=0, keepdims=True) + 1e-9)
    sims = temps @ B
    best = np.argmax(sims, axis=0)
    conf = float(np.mean(np.max(sims, axis=0)))
    bars = [labels[int(b)] for b in best]
    key_root = NOTES.index(key) if key in NOTES else 0

    # harmonisch palet: meest voorkomende akkoorden, tonica eerst
    counts = Counter(bars)
    tonic = (key_root, "min" if mode == "mineur" else "maj")
    palette = ([tonic] if tonic in counts else []) + \
              [c for c, _ in counts.most_common() if c != tonic]
    palette = palette[:4]

    # terugkerende loop (alleen als zinnig: minstens 2 verschillende akkoorden)
    loop = _find_loop(bars)
    loop_ok = bool(loop) and len(set(loop)) >= 2
    prog = loop if loop_ok else palette

    return {
        "available": True,
        "is_loop": loop_ok,
        "progression": [_chord_name(c) for c in prog],
        "progression_abs": " – ".join(_chord_name(c) for c in prog),
        "progression_roman": " – ".join(_roman_numeral(c, key_root, mode) for c in prog),
        "palette": [_chord_name(c) for c in palette],
        "bars": [_chord_name(c) for c in bars[:32]],
        "confidence": round(conf, 3),
    }


def analyze_core(path):
    """Run the full DSP analysis and return one JSON-friendly dict."""
    y, sr, duration = load_audio(path)
    onset_env = librosa.onset.onset_strength(y=y, sr=sr, hop_length=HOP)
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr, hop_length=HOP)

    tempo_info = detect_tempo(y, sr, onset_env=onset_env)
    key_info = detect_key(y, sr, chroma=chroma)
    ts = estimate_time_signature(onset_env, sr, tempo_info["beat_times"])
    structure = analyze_structure(y, sr, beats=tempo_info["beats"])
    dynamics = analyze_dynamics(y, sr)
    danceability = estimate_danceability(tempo_info, dynamics)
    valence = estimate_valence(key_info, tempo_info, dynamics)
    chords = detect_chords(chroma, tempo_info["beats"],
                           ts.get("beats_per_bar") or 4, key_info["key"], key_info["mode"])

    return {
        "duration_sec": round(duration, 1),
        "sample_rate": sr,
        "tempo": {"bpm": tempo_info["bpm"], "beat_count": tempo_info["beat_count"],
                  "regularity": tempo_info["regularity"],
                  "beat_times": [round(float(t), 3) for t in tempo_info["beat_times"]]},
        "key": key_info,
        "time_signature": ts,
        "structure": structure,
        "dynamics": dynamics,
        "danceability": danceability,
        "valence": valence,
        "chords": chords,
    }


if __name__ == "__main__":
    import json, sys
    res = analyze_core(sys.argv[1] if len(sys.argv) > 1 else "test_track.wav")
    res["dynamics"].pop("rms_curve", None)
    res["tempo"].pop("beat_times", None)
    res["key"].pop("chroma_mean", None)
    print(json.dumps(res, indent=2, ensure_ascii=False))