import shlex
import subprocess
import spaces
import torch

# install packages for mamba
def install_mamba():
    #subprocess.run(shlex.split("pip install torch==2.2.2 torchvision==0.17.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu118"))
    #subprocess.run(shlex.split("pip install https://github.com/Dao-AILab/causal-conv1d/releases/download/v1.4.0/causal_conv1d-1.4.0+cu122torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"))
    subprocess.run(shlex.split("pip install https://github.com/state-spaces/mamba/releases/download/v2.2.2/mamba_ssm-2.2.2+cu122torch2.3cxx11abiFALSE-cp310-cp310-linux_x86_64.whl"))
    subprocess.run(shlex.split("pip install numpy==1.26.4"))

    subprocess.run(shlex.split("ls"))


install_mamba()


import gradio as gr
import torch
import yaml
import librosa
from huggingface_hub import hf_hub_download
from models.stfts    import mag_phase_stft, mag_phase_istft
from models.generator import SEMamba
from models.pcs400   import cal_pcs

# download model files from your HF repo
#ckpt  = hf_hub_download("rc19477/Speech_Enhancement_Mamba",
#                        "ckpts/SEMamba_advanced.pth")
#cfg_f = hf_hub_download("rc19477/Speech_Enhancement_Mamba",
#                        "recipes/SEMamba_advanced.yaml")

ckpt = "ckpts/SEMamba_advanced.pth"
cfg_f = "recipes/SEMamba_advanced.yaml"


# load config
with open(cfg_f) as f:
    cfg = yaml.safe_load(f)

stft_cfg    = cfg["stft_cfg"]
model_cfg   = cfg["model_cfg"]
sr          = stft_cfg["sampling_rate"]
n_fft       = stft_cfg["n_fft"]
hop_size    = stft_cfg["hop_size"]
win_size    = stft_cfg["win_size"]
compress_ff = model_cfg["compress_factor"]

# init model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = SEMamba(cfg).to(device)
sdict  = torch.load(ckpt, map_location=device)
model.load_state_dict(sdict["generator"])
model.eval()

def enhance(audio, do_pcs):
    orig_sr, wav_np = audio
    # 1) resample to 16 kHz if needed
    if orig_sr != sr:
        wav_np = librosa.resample(wav_np, orig_sr, sr)
    wav = torch.from_numpy(wav_np).float().to(device)

    # normalize
    norm = torch.sqrt(len(wav) / torch.sum(wav**2))
    wav  = (wav * norm).unsqueeze(0)

    # STFT → model → ISTFT
    amp, pha, _ = mag_phase_stft(wav, n_fft, hop_size, win_size, compress_ff)
    amp_g, pha_g = model(amp, pha)
    out = mag_phase_istft(amp_g, pha_g, n_fft, hop_size, win_size, compress_ff)
    out = (out / norm).squeeze().cpu().numpy()

    # optional PCS filter
    if do_pcs:
        out = cal_pcs(out)

    # 2) resample back to original rate
    if orig_sr != sr:
        out = librosa.resample(out, sr, orig_sr)

    return orig_sr, out

with gr.Blocks() as demo:
    gr.Markdown("## SEMamba Speech Enhancement demo")
    with gr.Row():
        upload = gr.Audio(label="Upload WAV", type="numpy")
        record = gr.Audio(label="Record via mic", type="numpy")
    pcs    = gr.Checkbox(label="Apply PCS post-processing", value=False)
    btn    = gr.Button("Enhance")
    out    = gr.Audio(label="Enhanced WAV", type="numpy")

    @spaces.GPU
    def runner(up, rec, do_pcs):
        return enhance(up if up else rec, do_pcs)

    btn.click(runner, [upload, record, pcs], out)


demo.launch()