import gradio as gr
import librosa
import numpy as np
import tempfile
import zipfile
import os
import shutil

def key_bpm_embedding(audio):
    y, sr = librosa.load(audio, sr=None, mono=True)
    bpm, _ = librosa.beat.beat_track(y=y, sr=sr)
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr)
    chroma_mean = chroma.mean(axis=1)
    key_int = np.argmax(chroma_mean)
    key_list = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
    key = key_list[key_int]
    mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    embedding = np.mean(mfcc, axis=1).tolist()
    return {"bpm": float(bpm), "key": key, "embedding": embedding}

def order_tracks(zipfile_bytes):
    with tempfile.TemporaryDirectory() as td:
        zip_path = os.path.join(td, "tracks.zip")
        shutil.copy(zipfile_bytes.name, zip_path)
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(td)
            filelist = sorted(zip_ref.namelist())
        emb_map = {}
        for fname in filelist:
            audio_path = os.path.join(td, fname)
            try:
                y, sr = librosa.load(audio_path, sr=None, mono=True)
                mfcc = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
                emb_map[fname] = np.mean(mfcc, axis=1)
            except Exception:
                emb_map[fname] = None
        valid_files = [f for f in filelist if emb_map[f] is not None]
        if len(valid_files) == 0:
            return {"order": []}
        embs = np.stack([emb_map[f] for f in valid_files])
        sim = np.matmul(embs, embs.T)
        out_order = [0]
        used = {0}
        for _ in range(1, len(valid_files)):
            last = out_order[-1]
            next_idx = np.argmax([sim[last, j] if j not in used else -1e9 for j in range(len(valid_files))])
            out_order.append(next_idx)
            used.add(next_idx)
        ordered_files = [valid_files[i] for i in out_order]
        return {"order": ordered_files}

track_analysis = gr.Interface(
    fn=key_bpm_embedding,
    inputs=gr.Audio(type="filepath", label="Audio"),
    outputs="json"
)

track_orderer = gr.Interface(
    fn=order_tracks,
    inputs=gr.File(label="ZIP archive with tracks"),
    outputs="json"
)

demo = gr.TabbedInterface(
    [track_analysis, track_orderer],
    ["Track Analysis", "Order Tracks"]
)

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