import librosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np

def generate_visualizations(audio_path):
    """
    Takes an audio file path, processes it, and returns two matplotlib figures:
    1. A Time-Domain PCG Waveform
    2. A Frequency-Domain Mel-Spectrogram tuned for heart sounds.
    """
    if not audio_path:
        return None, None

    try:
        # Load the audio file
        # sr=None preserves the original sample rate. 
        y, sr = librosa.load(audio_path, sr=None)
        
        # --- 1. Generate PCG Waveform (Time Domain) ---
        # Close any existing plots to prevent memory leaks in Gradio
        plt.close('all') 
        
        fig_pcg, ax_pcg = plt.subplots(figsize=(10, 4))
        librosa.display.waveshow(y, sr=sr, ax=ax_pcg, color="#1f77b4", alpha=0.8)
        
        ax_pcg.set_title("Phonocardiogram (PCG) Waveform", fontsize=14, fontweight="bold")
        ax_pcg.set_xlabel("Time (seconds)")
        ax_pcg.set_ylabel("Amplitude")
        ax_pcg.grid(True, linestyle="--", alpha=0.6)
        fig_pcg.tight_layout()

        # --- 2. Generate Mel-Spectrogram (Frequency Domain) ---
        fig_spec, ax_spec = plt.subplots(figsize=(10, 4))
        
        # Compute the Mel-spectrogram
        # We set fmax=2000 because heart sounds and murmurs rarely exceed 2000 Hz.
        # This focuses the graph on the relevant medical data.
        S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128, fmax=2000)
        
        # Convert power to decibels (log scale) for better visualization
        S_dB = librosa.power_to_db(S, ref=np.max)
        
        # Display the spectrogram
        img = librosa.display.specshow(
            S_dB, 
            x_axis='time', 
            y_axis='mel', 
            sr=sr, 
            fmax=2000, 
            ax=ax_spec, 
            cmap='magma' # 'magma' or 'viridis' looks very professional for medical imaging
        )
        
        ax_spec.set_title("Mel-Spectrogram (Low-Frequency Focus)", fontsize=14, fontweight="bold")
        ax_spec.set_xlabel("Time (seconds)")
        ax_spec.set_ylabel("Frequency (Hz)")
        fig_spec.colorbar(img, ax=ax_spec, format='%+2.0f dB', label="Intensity (dB)")
        fig_spec.tight_layout()

        return fig_pcg, fig_spec

    except Exception as e:
        print(f"Error generating visualizations: {e}")
        # Return empty figures in case of an error so the UI doesn't crash
        fig, ax = plt.subplots()
        ax.text(0.5, 0.5, "Error generating graph", ha="center")
        return fig, fig