#!/usr/bin/env python3
"""
GN‑v7 – Fast Raw Mode Extractor
================================
Extract raw Koopman (PCA) modes from any audio file.
No ICA, no plots – just the raw modes as WAV files.
"""

import os
import shutil
import tempfile
import zipfile
from typing import List

import gradio as gr
import numpy as np
from numpy.lib.stride_tricks import sliding_window_view
import scipy.signal as signal
from sklearn.utils.extmath import randomized_svd
import soundfile as sf
import spaces

# ----------------------------------------------------------------------
# Fast embedding
# ----------------------------------------------------------------------
def embed(x: np.ndarray, dim: int, tau: int) -> np.ndarray:
    L = (dim - 1) * tau + 1
    win = sliding_window_view(x, L)[:, ::tau]
    return win[:, ::-1].copy()

# ----------------------------------------------------------------------
# Fast raw mode extraction using randomized SVD
# ----------------------------------------------------------------------
def extract_raw_modes(x: np.ndarray, dim: int = 32, tau: int = 3, k: int = 8) -> List[np.ndarray]:
    """
    Returns list of raw PCA mode signals (time domain) of length same as input.
    Each mode is a reconstructed time series from one singular vector.
    """
    X = embed(x, dim, tau)
    a = np.linalg.norm(X, axis=1) + 1e-9
    Xv = X / a[:, None]

    U, S, Vt = randomized_svd(Xv, n_components=k, n_iter=5, random_state=0)
    W = Vt  # shape (k, dim)

    Y = Xv @ W.T  # activations (T x k)

    raw_modes = []
    for i in range(k):
        mode_embedded = Y[:, i:i+1] @ W[i:i+1, :]  # (T, dim)
        mode_signal = np.real(mode_embedded[:, 0]) * a
        raw_modes.append(mode_signal)

    return raw_modes

# ----------------------------------------------------------------------
# Gradio interface (ZeroGPU compatible)
# ----------------------------------------------------------------------
@spaces.GPU
def process_to_raw_modes(file, dim: int, tau: int, k: int):
    if file is None:
        raise gr.Error("Please upload an audio file.")

    try:
        sig, fs = sf.read(file)
    except Exception as e:
        raise gr.Error(f"Failed to read: {e}")

    if sig.ndim > 1:
        sig = np.mean(sig, axis=1)

    target_sr = 16000
    if fs != target_sr:
        sig = signal.resample_poly(sig, target_sr, fs)
        fs = target_sr

    sig = sig / (np.max(np.abs(sig)) + 1e-9)

    duration = len(sig) / fs
    gr.Info(f"Processing {duration:.1f} sec song with dim={dim}, tau={tau}, k={k} ...")

    try:
        raw_modes = extract_raw_modes(sig, dim=dim, tau=tau, k=k)
    except Exception as e:
        raise gr.Error(f"Extraction failed: {e}")

    temp_dir = tempfile.mkdtemp()
    try:
        for i, mode in enumerate(raw_modes):
            sf.write(os.path.join(temp_dir, f"raw_mode_{i:02d}.wav"), mode, fs)

        zip_path = os.path.join(tempfile.gettempdir(), "raw_modes.zip")
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zf:
            for fname in os.listdir(temp_dir):
                zf.write(os.path.join(temp_dir, fname), arcname=fname)
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)

    preview_audio = (fs, raw_modes[0])
    return preview_audio, zip_path

# ----------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------
with gr.Blocks(title="Fast Raw Mode Extractor", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🔧 Fast Raw Mode Extractor (GN‑v7, PCA only)
    Upload your song → get raw Koopman/PCA modes as WAV files.
    **No ICA, no plots – just the raw modes.**  
    These are the `raw_mode_XX.wav` files from the original GN‑v7.
    """)
    with gr.Row():
        with gr.Column():
            audio_in = gr.Audio(type="filepath", label="Your song (any length)")
            dim_slider = gr.Slider(16, 64, value=32, step=2, label="Embedding dim (lower = faster)")
            tau_slider = gr.Slider(1, 8, value=3, step=1, label="Delay tau")
            k_slider = gr.Slider(2, 16, value=8, step=1, label="Number of raw modes (k)")
            run_btn = gr.Button("Extract Raw Modes", variant="primary")
        with gr.Column():
            preview = gr.Audio(label="Preview of first raw mode", type="numpy")
            zip_output = gr.File(label="📦 Download all raw modes (ZIP)")

    run_btn.click(
        fn=process_to_raw_modes,
        inputs=[audio_in, dim_slider, tau_slider, k_slider],
        outputs=[preview, zip_output]
    )

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