from __future__ import annotations

import gradio as gr
try:
    import spaces
except ImportError:  # 'spaces' is only provided by Hugging Face Spaces
    import types as _types

    def _gpu(*args, **kwargs):
        if len(args) == 1 and callable(args[0]) and not kwargs:
            return args[0]

        def _decorator(func):
            return func

        return _decorator

    spaces = _types.SimpleNamespace(GPU=_gpu)


import tensorflow as tf
import gin
import ddsp.training.inference
import ddsp.core
import librosa
import soundfile
import numpy as np
import requests
import zipfile
import io
import tempfile
import os
import crepe

# Global variables for model and parameters
MODEL = None
SAMPLE_RATE = None
N_SAMPLES = None
N_FRAMES = None

def _download_and_extract_model(model_name):
    model_url = f'https://storage.googleapis.com/ddsp/models/{model_name}/{model_name}.zip'
    print(f"Downloading model from: {model_url}")
    response = requests.get(model_url)
    response.raise_for_status() # Raise an exception for HTTP errors

    # Create a temporary directory for the model
    temp_dir = tempfile.mkdtemp()
    zip_path = os.path.join(temp_dir, f'{model_name}.zip')

    with open(zip_path, 'wb') as f:
        f.write(response.content)

    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(temp_dir)

    # The actual model checkpoint is usually in a subdirectory
    # Find the directory that contains the 'operative_config.gin'
    model_ckpt_dir = None
    for root, dirs, files in os.walk(temp_dir):
        if 'operative_config.gin' in files:
            model_ckpt_dir = root
            break
    
    if not model_ckpt_dir:
        raise FileNotFoundError("Could not find operative_config.gin in the downloaded model.")

    print(f"Model extracted to: {model_ckpt_dir}")
    return model_ckpt_dir

# --- Model Loading ---
MODEL_NAME = 'violin_ddsp_2020_03_03' # Example model from DDSP demos
CKPT_DIR = _download_and_extract_model(MODEL_NAME)

# Initialize the model with fixed parameters for this deployment
LENGTH_SECONDS = 4 # Default length for output audio
REMOVE_REVERB = True # Whether to remove reverb from the output

# Need to unlock gin config before parsing
with gin.unlock_config():
    MODEL = ddsp.training.inference.AutoencoderInference(
        ckpt=CKPT_DIR,
        length_seconds=LENGTH_SECONDS,
        remove_reverb=REMOVE_REVERB
    )

# Extract model parameters after initialization
SAMPLE_RATE = MODEL.sample_rate
N_SAMPLES = MODEL.n_samples
N_FRAMES = MODEL.n_frames

# Build the network by running a fake batch to initialize weights
MODEL.build_network()

# --- Helper functions for audio processing ---
def _load_audio_for_inference(audio_path, sr, n_samples):
    audio, _ = librosa.load(audio_path, sr=sr, mono=True)
    # Pad or trim audio to n_samples
    if len(audio) < n_samples:
        audio = np.pad(audio, (0, n_samples - len(audio)), mode='constant')
    else:
        audio = audio[:n_samples]
    return audio

def _save_audio_for_inference(audio_tensor, sr):
    output_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
    audio_np = audio_tensor.numpy() if tf.is_tensor(audio_tensor) else audio_tensor
    soundfile.write(output_path, audio_np, sr)
    return output_path


@spaces.GPU
def predict(input_audio):
    # Load audio
    audio_np = _load_audio_for_inference(input_audio, SAMPLE_RATE, N_SAMPLES)

    # Extract pitch using crepe
    # crepe.predict expects a 1D numpy array
    # Calculate step_size in milliseconds for crepe based on model's frame rate
    step_size_ms = int(1000 * N_SAMPLES / (SAMPLE_RATE * N_FRAMES))
    _, f0_hz, _, _ = crepe.predict(
        audio_np, SAMPLE_RATE, viterbi=True, step_size=step_size_ms
    )

    # Resample f0_hz to N_FRAMES
    f0_hz_resampled = ddsp.core.resample(
        tf.convert_to_tensor(f0_hz, dtype=tf.float32), N_FRAMES
    )

    # Compute loudness
    loudness_db = ddsp.core.compute_loudness(
        tf.expand_dims(tf.convert_to_tensor(audio_np, dtype=tf.float32), axis=0),
        SAMPLE_RATE
    )

    # Prepare features dictionary
    audio_features = {
        'audio': tf.expand_dims(tf.convert_to_tensor(audio_np, dtype=tf.float32), axis=0),
        'f0_hz': tf.expand_dims(f0_hz_resampled, axis=0),
        'loudness_db': loudness_db,
    }

    # Preprocess features using the model's preprocessor
    processed_features = MODEL.preprocessor(audio_features)

    # Run inference
    outputs = MODEL(processed_features, training=False)

    # Get synthesized audio from outputs
    synthesized_audio = MODEL.get_audio_from_outputs(outputs)

    # Save and return the output audio file path
    output_audio_path = _save_audio_for_inference(synthesized_audio[0], SAMPLE_RATE)
    return output_audio_path


demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Audio(type="filepath", label="Input Audio"),
    ],
    outputs=[
        gr.Audio(type="filepath", label="Output Audio"),
    ],
    title="DDSP Autoencoder Inference",
    description="Resynthesize input audio using a pre-trained DDSP autoencoder model. This model takes an audio input, extracts its fundamental frequency (f0) and loudness, and then uses a DDSP autoencoder to synthesize a new audio output, effectively performing a timbre transfer based on the loaded model.",
)

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)