"""SoftChart — Gradio demo for Hugging Face Spaces.

Generate a Taiko no Tatsujin chart from any audio file, using the full system:
  - SoftChartGenerator  (main model, plan-conditioned)
  - SoftChartPlanner     (auto song-level planning)
  - SoftChartBeat        (beat/downbeat for barline anchoring)
All models load from the Hub via from_pretrained. MIT licensed.
"""

import os
import tempfile

import gradio as gr
import numpy as np
import torch

from softchart.generate import generate_song, load_hf
from softchart.hf import SoftChartPlanner
from softchart.rhythm import snap_chart
from softchart.vocab import FPS, HOP, N_FFT, N_MELS, SR, WINDOW

GEN_REPO = os.environ.get("SC_GEN", "JacobLinCool/softchart-generator")
BEAT_REPO = os.environ.get("SC_BEAT", "JacobLinCool/softchart-beat")
PLAN_REPO = os.environ.get("SC_PLAN", "JacobLinCool/softchart-planner")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

COURSE_DENS = {"easy": 1, "normal": 2, "hard": 4, "oni": 7}
DEFAULT_LEVEL = {"easy": 3, "normal": 5, "hard": 7, "oni": 9}
CHAR = {"don": "1", "ka": "2", "don_big": "3", "ka_big": "4",
        "roll": "5", "roll_big": "6", "balloon": "7"}
STYLE = {"don": ("#e8453c", 55), "ka": ("#4aa8c0", 55),
         "don_big": ("#e8453c", 150), "ka_big": ("#4aa8c0", 150),
         "roll": ("#e8b23a", 80), "roll_big": ("#e8b23a", 150), "balloon": ("#d47fda", 100)}
SUB = 96

_MODELS = {}


def get_models():
    if not _MODELS:
        _MODELS["gen"] = load_hf(GEN_REPO, device=DEVICE)
        try:
            _MODELS["beat"] = load_hf(BEAT_REPO, device=DEVICE)
        except Exception:
            _MODELS["beat"] = None
        try:
            _MODELS["plan"] = SoftChartPlanner.from_pretrained(PLAN_REPO).to(DEVICE).eval()
        except Exception:
            _MODELS["plan"] = None
    return _MODELS


def load_logmel(path):
    import librosa

    wav, _ = librosa.load(path, sr=SR, mono=True)
    fb = librosa.filters.mel(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=20.0, fmax=SR / 2)
    spec = torch.stft(torch.from_numpy(wav), N_FFT, hop_length=HOP,
                      window=torch.hann_window(N_FFT), center=True, return_complex=True)
    mel = np.log(fb @ spec.abs().pow(2).numpy() + 1e-5).astype(np.float32)
    return mel, wav


def predict_downbeats(beat_model, mel):
    from scipy.signal import find_peaks
    from softchart.generate import _autocast

    T = mel.shape[1]
    L = WINDOW // 4
    acc = np.zeros(T // 4 + L)
    cnt = np.zeros(T // 4 + L)
    for st in range(0, max(T - 1, 1), WINDOW):
        w = torch.from_numpy(mel[:, st:st + WINDOW].astype(np.float32))
        if w.shape[1] < WINDOW:
            w = torch.nn.functional.pad(w, (0, WINDOW - w.shape[1]), value=float(np.log(1e-5)))
        with torch.no_grad(), _autocast(DEVICE):
            mem = beat_model.encode(w[None].to(DEVICE))
            pr = torch.sigmoid(beat_model.beat(mem[:, -L:]).float())[0, :, 1].cpu().numpy()
        acc[st // 4: st // 4 + L] += pr
        cnt[st // 4: st // 4 + L] += 1
    acc /= np.maximum(cnt, 1)
    return find_peaks(acc, height=0.4, distance=int(0.8 * FPS / 4))[0] * 4 / FPS


def auto_plan(mel, bpm, downbeats=None):
    T = mel.shape[1]
    dur = T / FPS
    flux = np.concatenate([[0], np.maximum(0, np.diff(mel, axis=1)).sum(0)])
    beat = 60.0 / bpm
    edges = (list(downbeats[::4]) + [dur]) if (downbeats is not None and len(downbeats) >= 2) \
        else list(np.arange(0, dur, 4 * beat)) + [dur]
    vals = [float(flux[int(a * FPS):int(b * FPS)].mean()) if int(b * FPS) > int(a * FPS) else 0.0
            for a, b in zip(edges, edges[1:])]
    if not vals:
        return None
    vals = np.array(vals)
    lo, hi = np.percentile(vals, 15), np.percentile(vals, 92)
    peak = int(np.argmax(vals))
    plan = []
    for i, (a, b) in enumerate(zip(edges, edges[1:])):
        frac = (vals[i] - lo) / max(hi - lo, 1e-6)
        d8 = int(np.clip(round(frac * 7), 0, 7))
        fl = 1 if (vals[i] <= lo and 0 < i < len(vals) - 1) else (2 if i == peak and vals[i] >= hi else 0)
        plan.append([round(a, 3), round(b, 3), d8, fl])
    return plan


def learned_plan(planner, mel, course, bpm, downbeats=None):
    dur = mel.shape[1] / FPS
    beat = 60.0 / bpm
    edges = (list(downbeats[::4]) + [dur]) if (downbeats is not None and len(downbeats) >= 2) \
        else list(np.arange(0, dur, 4 * beat)) + [dur]
    feats, spans = [], []
    for a, b in zip(edges, edges[1:]):
        seg = mel[:, int(a * FPS):int(b * FPS)]
        if seg.shape[1] < 2:
            continue
        fx = np.maximum(0, np.diff(seg, axis=1)).sum(0)
        feats.append(np.concatenate([seg.mean(1), seg.std(1), [fx.mean(), fx.std(), fx.max()]]))
        spans.append((round(float(a), 3), round(float(b), 3)))
    if not feats:
        return None
    cid = {"easy": 0, "normal": 1, "hard": 2, "oni": 3}[course]
    x = torch.tensor(np.array(feats), dtype=torch.float32)[None].to(DEVICE)
    with torch.no_grad():
        pd, pf = planner(x, torch.tensor([cid], device=DEVICE))
    d8 = pd[0].argmax(-1).cpu().numpy()
    fl = pf[0].argmax(-1).cpu().numpy()
    return [[a, b, int(d), int(f)] for (a, b), d, f in zip(spans, d8, fl)]


def group_quantize(times, phase, grid, min_run=3):
    n = len(times)
    slots = [0] * n
    i = 0
    while i < n:
        j = i
        while j + 1 < n:
            ioi = times[j + 1] - times[j]
            ref = (times[j] - times[i]) / (j - i) if j > i else ioi
            if 0.02 < ioi < 1.2 and abs(ioi - ref) < 0.22 * max(ref, 1e-6):
                j += 1
            else:
                break
        if j - i + 1 >= min_run:
            k = max(1, int(round((times[j] - times[i]) / (j - i) / grid)))
            anchor = int(round((times[i] - phase) / grid))
            for m in range(j - i + 1):
                slots[i + m] = anchor + m * k
        else:
            for m in range(i, j + 1):
                slots[m] = int(round((times[m] - phase) / grid))
        i = j + 1
    return slots


def write_tja(gen, bpm, title, course, level, downbeats=None):
    hits = sorted((h["t"], CHAR[h["type"]]) for h in gen["hits"])
    times = np.array([t for t, _ in hits]) if hits else np.array([0.0])
    beat = 60.0 / bpm
    grid = beat / (SUB / 4)
    cands = np.arange(0, beat, grid / 4)
    phase = float(cands[int(np.argmin([np.mean(np.abs(((times - o) / grid) - np.round((times - o) / grid))) for o in cands]))])
    slots = {}
    for idx, (t, ch) in zip(group_quantize([t for t, _ in hits], phase, grid), hits):
        if idx >= 0 and idx not in slots:
            slots[idx] = ch
    for sp in gen["spans"]:
        i0 = int(round((sp["t0"] - phase) / grid))
        i1 = int(round((sp["t1"] - phase) / grid))
        while i0 in slots:
            i0 += 1
        while i1 in slots or i1 <= i0:
            i1 += 1
        if i0 >= 0:
            slots[i0] = CHAR[sp["type"]]
            slots[i1] = "8"
    if slots:
        first_t = min(slots) * grid + phase
        anchor_t = None
        if downbeats is not None and len(downbeats):
            near = downbeats[downbeats <= first_t + 0.12]
            if len(near) and first_t - near[-1] < 4 * beat:
                anchor_t = near[-1]
        shift = int(round((anchor_t - phase) / grid)) if anchor_t is not None else min(slots)
        if shift:
            slots = {k - shift: v for k, v in slots.items()}
            phase += shift * grid
    n_meas = (max(slots) // SUB + 1) if slots else 1
    lines = ["".join(slots.get(m * SUB + k, "0") for k in range(SUB)) + "," for m in range(n_meas)]
    balloons = [10] * sum(1 for s in gen["spans"] if s["type"] == "balloon")
    return "\n".join([
        f"TITLE:{title} (SoftChart)", f"BPM:{bpm:g}", "WAVE:song.ogg",
        f"OFFSET:{-phase:.3f}", f"COURSE:{'Oni' if course == 'oni' else course.capitalize()}",
        f"LEVEL:{level}", f"BALLOON:{','.join(map(str, balloons))}" if balloons else "BALLOON:",
        "", "#START", *lines, "#END"]) + "\n"


def render(mel, gen, title, course):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    dur = mel.shape[1] / FPS
    t0, t1 = dur * 0.33, min(dur * 0.33 + 26, dur)
    fig = plt.figure(figsize=(13, 4))
    gs = fig.add_gridspec(2, 1, height_ratios=[2.3, 1], hspace=0.1)
    ax0 = fig.add_subplot(gs[0])
    ax0.imshow(mel[:, int(t0 * FPS):int(t1 * FPS)], aspect="auto", origin="lower",
               cmap="magma", extent=[t0, t1, 0, N_MELS])
    ax0.set_xticks([])
    ax0.set_ylabel("mel")
    ax0.set_title(f"{title} — {course}  |  red=don blue=ka big=large yellow=roll purple=balloon")
    ax1 = fig.add_subplot(gs[1])
    for h in gen["hits"]:
        if t0 <= h["t"] <= t1:
            c, sz = STYLE.get(h["type"], ("#999", 40))
            ax1.scatter(h["t"], 0, c=c, s=sz, edgecolors="k",
                        linewidths=1.2 if "big" in h["type"] else 0.5, zorder=3)
    for s in gen["spans"]:
        if s["t1"] >= t0 and s["t0"] <= t1:
            ax1.plot([max(s["t0"], t0), min(s["t1"], t1)], [0, 0],
                     c=STYLE.get(s["type"], ("#e8b23a",))[0], lw=6, alpha=0.45, zorder=2)
    ax1.set_xlim(t0, t1)
    ax1.set_ylim(-1, 1)
    ax1.set_yticks([])
    ax1.set_xlabel("time (s)")
    ax1.grid(axis="x", alpha=0.25)
    out = tempfile.mktemp(suffix=".png")
    fig.savefig(out, dpi=130, bbox_inches="tight")
    plt.close(fig)
    return out


def generate(audio, course, level, bpm_override, auto_plan_on, use_beat, use_planner, sampling):
    if audio is None:
        raise gr.Error("Please upload an audio file.")
    M = get_models()
    mel, wav = load_logmel(audio)
    dbs = None
    if use_beat and M["beat"] is not None:
        dbs = predict_downbeats(M["beat"], mel)
    if bpm_override and bpm_override > 0:
        bpm = float(bpm_override)
    elif dbs is not None and len(dbs) > 4:
        bpm = 60.0 / float(np.median(np.diff(dbs)))
    else:
        import librosa
        bpm = float(np.atleast_1d(librosa.beat.beat_track(y=wav, sr=SR)[0])[0])
    if abs(bpm - round(bpm)) < 0.06:
        bpm = float(round(bpm))

    plan = None
    if getattr(M["gen"], "_has_plan", False):
        if use_planner and M["plan"] is not None:
            plan = learned_plan(M["plan"], mel, course, bpm, dbs)
        elif auto_plan_on:
            plan = auto_plan(mel, bpm, dbs)

    g = generate_song(M["gen"], mel, course, level=int(level),
                      density_bucket=COURSE_DENS[course], greedy=not sampling,
                      seed=0, device=DEVICE, plan=plan)
    g = snap_chart(g, bpm)
    title = os.path.splitext(os.path.basename(audio))[0]
    tja = write_tja(g, bpm, title, course, int(level), dbs)
    tja_path = tempfile.mktemp(suffix=f"_{course}.tja")
    with open(tja_path, "w") as f:
        f.write(tja)
    img = render(mel, g, title, course)
    info = (f"BPM {bpm:.1f} · {len(g['hits'])} notes · {len(g['spans'])} spans"
            + (f" · plan: {sum(1 for b in plan if b[3]==1)} gaps, {sum(1 for b in plan if b[3]==2)} climax" if plan else ""))
    return img, tja_path, info


with gr.Blocks(title="SoftChart — AI Taiko chart generator") as demo:
    gr.Markdown(
        "# 🥁 SoftChart\n"
        "Generate a **Taiko no Tatsujin** chart from any song. 8M-param plan-conditioned "
        "model + learned planner + beat-anchoring — all MIT-licensed, loaded from the Hub.\n\n"
        "Outputs a preview and a playable `.tja` (for TJAPlayer3 / OpenTaiko). "
        "*Models trained on charts derived from a commercial game; for research/personal use.*"
    )
    with gr.Row():
        with gr.Column():
            audio = gr.Audio(type="filepath", label="Song (any format)")
            course = gr.Dropdown(["easy", "normal", "hard", "oni"], value="oni", label="Difficulty (course)")
            level = gr.Slider(1, 10, value=9, step=1, label="Level (1–10 stars)")
            bpm = gr.Number(label="BPM (0 = auto-detect)", value=0)
            with gr.Row():
                auto_plan_on = gr.Checkbox(label="Auto-plan (heuristic structure)", value=True)
                use_planner = gr.Checkbox(label="Learned planner", value=True)
            with gr.Row():
                use_beat = gr.Checkbox(label="Beat-anchor barlines", value=True)
                sampling = gr.Checkbox(label="Sampling (diverse) vs greedy (best)", value=False)
            btn = gr.Button("Generate chart", variant="primary")
        with gr.Column():
            out_img = gr.Image(label="Preview (mel + chart)")
            out_info = gr.Textbox(label="Result", interactive=False)
            out_tja = gr.File(label="Download .tja")
    btn.click(generate, [audio, course, level, bpm, auto_plan_on, use_beat, use_planner, sampling],
              [out_img, out_tja, out_info])

if __name__ == "__main__":
    demo.launch()