"""
Tempo drift mitigation for AI-generated tracks (Suno/Udio).

Detects fractional BPM and per-beat timing drift, then applies:
1. Per-beat warp map (timemap_stretch) to lock every beat to the grid
2. Final global stretch to ensure exact whole-number BPM
3. Sample-count trim for mathematically perfect BPM

Generates before/after beat-grid visual diagnostics.
"""

import numpy as np
from pathlib import Path
from typing import Optional, Tuple
from dataclasses import dataclass

import librosa
import soundfile as sf
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


COMMON_BPMS = sorted(set([
    70, 72, 75, 78, 80, 82, 85, 86, 88, 90, 92, 95, 98,
    100, 102, 104, 105, 108, 110, 112, 115, 118,
    120, 122, 124, 125, 126, 128, 130, 132, 133, 134, 135, 136, 138,
    140, 142, 144, 145, 148, 150, 152, 155, 158,
    160, 162, 165, 168, 170, 172, 174, 175, 178, 180,
]))


@dataclass
class TempoFixResult:
    original_bpm: float
    target_bpm: int
    beats_detected: int
    before_mean_dev_ms: float
    after_mean_dev_ms: float
    before_image: Optional[str]
    after_image: Optional[str]
    output_path: str


def _find_intended_bpm(detected: float) -> int:
    closest = min(COMMON_BPMS, key=lambda x: abs(x - detected))
    return closest if abs(closest - detected) < 3.0 else round(detected)


def _detect_bpm(mono: np.ndarray, sr: int) -> float:
    tempos = []
    chunk = sr * 15
    for s in range(0, len(mono) - chunk, chunk):
        t = librosa.beat.beat_track(y=mono[s:s+chunk], sr=sr)[0]
        tempos.append(float(t[0]) if hasattr(t, '__len__') else float(t))
    if not tempos:
        t = librosa.beat.beat_track(y=mono, sr=sr)[0]
        tempos = [float(t[0]) if hasattr(t, '__len__') else float(t)]
    return float(np.median(tempos))


def _detect_beats(mono: np.ndarray, sr: int, bpm: float) -> np.ndarray:
    hop = 64
    onset_env = librosa.onset.onset_strength(y=mono, sr=sr, hop_length=hop)
    _, frames = librosa.beat.beat_track(
        onset_envelope=onset_env, sr=sr, hop_length=hop,
        bpm=bpm, tightness=400, units='frames'
    )
    return librosa.frames_to_time(frames, sr=sr, hop_length=hop)


def _refine_to_onset(bt: float, mono: np.ndarray, sr: int, window_ms: float = 20.0) -> float:
    w = int(sr * window_ms / 1000)
    c = int(bt * sr)
    s, e = max(0, c - w), min(len(mono), c + w)
    if e - s < 10:
        return bt
    diff = np.diff(np.abs(mono[s:e]))
    return (s + np.argmax(diff)) / sr


def _find_phase(onset_env: np.ndarray, sr: int, hop: int, bpm: float) -> float:
    interval = 60.0 / bpm
    hop_t = hop / sr
    n = len(onset_env)
    click = np.zeros(n)
    for i in range(int(n * hop_t / interval)):
        f = int(i * interval / hop_t)
        if f < n:
            click[f] = 1.0
    max_off = int(interval / hop_t)
    corrs = np.array([np.sum(onset_env * np.roll(click, o)) for o in range(max_off)])
    best = np.argmax(corrs)
    if 0 < best < len(corrs) - 1:
        a, b, g = corrs[best-1], corrs[best], corrs[best+1]
        denom = a - 2*b + g
        if denom != 0:
            best += 0.5 * (a - g) / denom
    return float(best * hop_t)


def _compute_devs(beats: np.ndarray, bpm: float, phase: float):
    interval = 60.0 / bpm
    grid, devs = [], []
    for bt in beats:
        idx = round((bt - phase) / interval)
        gt = phase + idx * interval
        grid.append(gt)
        devs.append((bt - gt) * 1000)
    return np.array(grid), np.array(devs)


def _dedup_beats(beats, grid, devs):
    seen = {}
    for i, gt in enumerate(grid):
        k = round(gt * 1000)
        if k not in seen or abs(devs[i]) < abs(devs[seen[k]]):
            seen[k] = i
    idx = sorted(seen.values())
    return beats[idx], grid[idx], devs[idx]


def _gen_visual(mono, sr, beats, grid, devs, bpm, path, title):
    dur = len(mono) / sr
    hop = 512
    env = librosa.feature.rms(y=mono, hop_length=hop)[0]
    env_t = librosa.frames_to_time(np.arange(len(env)), sr=sr, hop_length=hop)
    interval = 60.0 / bpm

    fig, (a1, a2) = plt.subplots(2, 1, figsize=(24, 8), gridspec_kw={'height_ratios': [2, 1]})
    fig.suptitle(f'{title} — {bpm} BPM — {len(beats)} beats', fontsize=14)
    a1.fill_between(env_t, -env, env, color='#ccc', alpha=0.7)

    gt = 0
    while gt < dur:
        a1.axvline(gt, color='#0a0', alpha=0.12, linewidth=0.5)
        gt += interval

    for bt, d in zip(beats, devs):
        ad = abs(d)
        c = '#00f' if ad < 1 else '#0c0' if ad < 5 else '#fa0' if ad < 15 else '#f00'
        a1.axvline(bt, color=c, alpha=0.6, linewidth=1)
    a1.set_ylabel('Amplitude')
    a1.set_xlim(0, dur)

    colors = ['#00f' if abs(d)<1 else '#0c0' if abs(d)<5 else '#fa0' if abs(d)<15 else '#f00' for d in devs]
    a2.bar(beats, devs, width=interval*0.6, color=colors, alpha=0.8)
    a2.axhline(0, color='black', linewidth=1)
    for y in [1, -1, 5, -5]:
        a2.axhline(y, color='blue' if abs(y)==1 else 'orange', linewidth=0.5, linestyle=':', alpha=0.3)
    a2.set_ylabel('Deviation (ms)')
    a2.set_xlabel('Time (s)')
    a2.set_xlim(0, dur)

    w1 = np.sum(np.abs(devs) < 1) / max(len(devs), 1) * 100
    w5 = np.sum(np.abs(devs) < 5) / max(len(devs), 1) * 100
    a2.set_title(f"Max: ±{np.max(np.abs(devs)):.1f}ms | Mean: ±{np.mean(np.abs(devs)):.1f}ms | <1ms: {w1:.0f}% | <5ms: {w5:.0f}%")
    plt.tight_layout()
    Path(path).parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(path, dpi=150, bbox_inches='tight')
    plt.close(fig)


def fix_tempo(input_path: str, output_path: Optional[str] = None,
              target_bpm: Optional[int] = None,
              output_dir: Optional[str] = None) -> TempoFixResult:
    """Fix tempo drift in an audio file.

    Args:
        input_path: Input audio file
        output_path: Output file (default: input_fixed_tempo.wav)
        target_bpm: Force target BPM (auto-detect if None)
        output_dir: Directory for diagnostic images

    Returns:
        TempoFixResult with paths and stats
    """
    import pyrubberband as pyrb

    audio, sr = sf.read(input_path, dtype='float32')
    mono = np.mean(audio, axis=1) if audio.ndim > 1 else audio

    if output_path is None:
        p = Path(input_path)
        output_path = str(p.with_stem(p.stem + "_fixed_tempo"))
    if output_dir is None:
        output_dir = str(Path(output_path).parent)

    # Detect BPM
    detected = _detect_bpm(mono, sr)
    if target_bpm is None:
        target_bpm = _find_intended_bpm(detected)

    print(f"  Tempo fix: {detected:.2f} BPM → {target_bpm} BPM")

    # High-precision beat detection + onset refinement
    beats = _detect_beats(mono, sr, target_bpm)
    if len(beats) < 4:
        print(f"  Too few beats ({len(beats)}), skipping")
        sf.write(output_path, audio, sr)
        return TempoFixResult(detected, target_bpm, len(beats), 0, 0, None, None, output_path)

    refined = np.array([_refine_to_onset(bt, mono, sr) for bt in beats])
    clean = [refined[0]]
    for bt in refined[1:]:
        if bt > clean[-1]:
            clean.append(bt)
    refined = np.array(clean)

    # Grid phase + deviations
    hop = 64
    onset_env = librosa.onset.onset_strength(y=mono, sr=sr, hop_length=hop)
    phase = _find_phase(onset_env, sr, hop, target_bpm)
    grid, devs = _compute_devs(refined, target_bpm, phase)
    refined, grid, devs = _dedup_beats(refined, grid, devs)

    before_mean = float(np.mean(np.abs(devs)))
    print(f"  Before: mean ±{before_mean:.1f}ms, {len(refined)} beats")

    # Before visual
    before_img = str(Path(output_dir) / "tempo_grid_before.png")
    _gen_visual(mono, sr, refined, grid, devs, target_bpm, before_img, "BEFORE Fix")

    # Build warp map
    total = len(audio) if audio.ndim == 1 else audio.shape[0]
    timemap = [(0, 0)]
    for bt, gt, d in zip(refined, grid, devs):
        if abs(d) < 0.3:
            continue
        inp, out = int(bt * sr), int(gt * sr)
        prev_i, prev_o = timemap[-1]
        if inp > prev_i and out > prev_o:
            di, do = inp - prev_i, out - prev_o
            if di > 0 and do > 0 and abs(do/di - 1) < 0.15:
                timemap.append((inp, out))

    last_i, last_o = timemap[-1]
    timemap.append((total, last_o + (total - last_i)))

    print(f"  Warp map: {len(timemap)} keyframes")

    # Apply warp
    if len(timemap) > 2:
        corrected = pyrb.timemap_stretch(audio.astype(np.float32), sr, timemap)
    else:
        corrected = audio.copy()

    # Final global stretch to exact BPM
    corr_mono = np.mean(corrected, axis=1) if corrected.ndim > 1 else corrected
    current = _detect_bpm(corr_mono, sr)
    if abs(current - target_bpm) > 0.05:
        ratio = current / target_bpm
        corrected = pyrb.time_stretch(corrected.astype(np.float32), sr, ratio)

    # Trim to exact beat count
    corr_n = len(corrected) if corrected.ndim == 1 else corrected.shape[0]
    n_beats = round((corr_n / sr) * target_bpm / 60.0)
    exact_n = int(n_beats * 60.0 / target_bpm * sr)
    if corrected.ndim > 1:
        corrected = corrected[:exact_n] if corr_n > exact_n else np.pad(corrected, ((0, exact_n-corr_n), (0,0)))
    else:
        corrected = corrected[:exact_n] if corr_n > exact_n else np.pad(corrected, (0, exact_n-corr_n))

    final_bpm = n_beats * 60.0 / (exact_n / sr)
    print(f"  Final: {exact_n/sr:.3f}s, {n_beats} beats, BPM={final_bpm:.4f}")

    # Verify + after visual
    after_mono = np.mean(corrected, axis=1) if corrected.ndim > 1 else corrected
    after_beats = _detect_beats(after_mono, sr, target_bpm)
    after_mean = 0.0
    after_img = None
    if len(after_beats) >= 4:
        ar = np.array([_refine_to_onset(bt, after_mono, sr) for bt in after_beats])
        ac = [ar[0]]
        for bt in ar[1:]:
            if bt > ac[-1]:
                ac.append(bt)
        ar = np.array(ac)
        ae = librosa.onset.onset_strength(y=after_mono, sr=sr, hop_length=hop)
        ap = _find_phase(ae, sr, hop, target_bpm)
        ag, ad = _compute_devs(ar, target_bpm, ap)
        ar, ag, ad = _dedup_beats(ar, ag, ad)
        after_mean = float(np.mean(np.abs(ad)))
        print(f"  After: mean ±{after_mean:.1f}ms")
        after_img = str(Path(output_dir) / "tempo_grid_after.png")
        _gen_visual(after_mono, sr, ar, ag, ad, target_bpm, after_img, "AFTER Fix")

    sf.write(output_path, corrected.astype(np.float32), sr, subtype='PCM_24')
    print(f"  Saved: {output_path}")

    return TempoFixResult(
        detected, target_bpm, len(refined), before_mean, after_mean,
        before_img, after_img, output_path
    )