import os
import pickle
import numpy as np
import pandas as pd
import librosa
import gradio as gr
from scipy.io import wavfile
import soundfile as sf
import tempfile
import time

# Load the model
def load_model(model_path='cough_classification_model.pkl'):
    with open(model_path, 'rb') as f:
        components = pickle.load(f)
    
    return components

# Extract features from audio
def extract_all_features(audio_path, sample_rate=None):
    """Extract comprehensive set of audio features"""
    # Load audio file
    y, sr = librosa.load(audio_path, sr=sample_rate)
    
    # Basic features
    features = {}
    
    # Duration
    features['duration'] = librosa.get_duration(y=y, sr=sr)
    
    # RMS Energy
    features['rms_mean'] = np.mean(librosa.feature.rms(y=y)[0])
    features['rms_std'] = np.std(librosa.feature.rms(y=y)[0])
    
    # Zero Crossing Rate
    zcr = librosa.feature.zero_crossing_rate(y)[0]
    features['zcr_mean'] = np.mean(zcr)
    features['zcr_std'] = np.std(zcr)
    
    # Spectral Features
    # Spectral Centroid
    centroid = librosa.feature.spectral_centroid(y=y, sr=sr)[0]
    features['spectral_centroid_mean'] = np.mean(centroid)
    features['spectral_centroid_std'] = np.std(centroid)
    
    # Spectral Bandwidth
    bandwidth = librosa.feature.spectral_bandwidth(y=y, sr=sr)[0]
    features['spectral_bandwidth_mean'] = np.mean(bandwidth)
    features['spectral_bandwidth_std'] = np.std(bandwidth)
    
    # Spectral Contrast
    contrast = librosa.feature.spectral_contrast(y=y, sr=sr)
    features['spectral_contrast_mean'] = np.mean(contrast)
    features['spectral_contrast_std'] = np.std(contrast)
    
    # Spectral Rolloff
    rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr)[0]
    features['rolloff_mean'] = np.mean(rolloff)
    features['rolloff_std'] = np.std(rolloff)
    
    # MFCCs
    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    for i in range(13):
        features[f'mfcc{i+1}_mean'] = np.mean(mfccs[i])
        features[f'mfcc{i+1}_std'] = np.std(mfccs[i])
    
    # Chroma Features
    chroma = librosa.feature.chroma_stft(y=y, sr=sr)
    features['chroma_mean'] = np.mean(chroma)
    features['chroma_std'] = np.std(chroma)
    
    return features

# Function to segment cough from audio
def segment_cough(x, fs, cough_padding=0.2, min_cough_len=0.2, th_l_multiplier=0.1, th_h_multiplier=2):
    """Segment coughs from audio using a hysteresis comparator on the signal power"""
    cough_mask = np.array([False]*len(x))

    # Define hysteresis thresholds
    rms = np.sqrt(np.mean(np.square(x)))
    seg_th_l = th_l_multiplier * rms
    seg_th_h = th_h_multiplier * rms

    # Segment coughs
    coughSegments = []
    padding = round(fs*cough_padding)
    min_cough_samples = round(fs*min_cough_len)
    cough_start = 0
    cough_end = 0
    cough_in_progress = False
    tolerance = round(0.01*fs)
    below_th_counter = 0

    for i, sample in enumerate(x**2):
        if cough_in_progress:
            if sample < seg_th_l:
                below_th_counter += 1
                if below_th_counter > tolerance:
                    cough_end = i+padding if (i+padding < len(x)) else len(x)-1
                    cough_in_progress = False
                    if (cough_end+1-cough_start-2*padding > min_cough_samples):
                        coughSegments.append(x[cough_start:cough_end+1])
                        cough_mask[cough_start:cough_end+1] = True
            elif i == (len(x)-1):
                cough_end = i
                cough_in_progress = False
                if (cough_end+1-cough_start-2*padding > min_cough_samples):
                    coughSegments.append(x[cough_start:cough_end+1])
            else:
                below_th_counter = 0
        else:
            if sample > seg_th_h:
                cough_start = i-padding if (i-padding >= 0) else 0
                cough_in_progress = True

    return coughSegments, cough_mask

# Prediction function for uploaded audio
def predict_cough_health(audio_path, model_components=None):
    """Predict cough health status from audio file"""
    if model_components is None:
        model_components = load_model()
    
    model = model_components['model']
    scaler = model_components['scaler']
    label_encoder = model_components['label_encoder']
    feature_names = model_components['feature_names']
    
    # Extract features
    features = extract_all_features(audio_path)
    
    # Convert to DataFrame with correct feature order
    features_df = pd.DataFrame([features])
    features_df = features_df[feature_names]
    
    # Scale features
    features_scaled = scaler.transform(features_df)
    
    # Predict
    prediction_idx = model.predict(features_scaled)[0]
    prediction = label_encoder.inverse_transform([prediction_idx])[0]
    
    # Get probabilities
    probs = model.predict_proba(features_scaled)[0]
    class_probs = {label_encoder.inverse_transform([i])[0]: float(prob) for i, prob in enumerate(probs)}
    
    return prediction, class_probs

# Function to process uploaded audio file
def process_audio_file(audio_file):
    """Process uploaded audio file and return prediction"""
    model_components = load_model()
    prediction, probabilities = predict_cough_health(audio_file, model_components)
    
    # Format the output
    result = f"Prediction: {prediction}\n\nProbabilities:\n"
    for cls, prob in probabilities.items():
        result += f"{cls}: {prob:.4f}\n"
    
    return result

# Function to process recorded audio
def process_recorded_audio(audio_array, sample_rate):
    """Process recorded audio and return prediction"""
    # Save the recorded audio to a temporary file
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_audio:
        temp_filename = temp_audio.name
        sf.write(temp_filename, audio_array, sample_rate)
    
    model_components = load_model()
    prediction, probabilities = predict_cough_health(temp_filename, model_components)
    
    # Clean up the temporary file
    os.unlink(temp_filename)
    
    # Format the output
    result = f"Prediction: {prediction}\n\nProbabilities:\n"
    for cls, prob in probabilities.items():
        result += f"{cls}: {prob:.4f}\n"
    
    return result

# Function to detect coughs in live audio stream
def detect_coughs_in_stream(audio_array, sample_rate):
    """Detect coughs in live audio stream and analyze them"""
    # Save the audio chunk to a temporary file
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_audio:
        temp_filename = temp_audio.name
        sf.write(temp_filename, audio_array, sample_rate)
    
    # Load the audio for processing
    audio, sr = librosa.load(temp_filename, sr=None)
    
    # Segment coughs
    cough_segments, cough_mask = segment_cough(audio, sr)
    
    # Clean up the temporary file
    os.unlink(temp_filename)
    
    if len(cough_segments) > 0:
        # Save the first detected cough segment to a temporary file
        with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as cough_file:
            cough_filename = cough_file.name
            sf.write(cough_filename, cough_segments[0], sr)
        
        # Analyze the cough
        model_components = load_model()
        prediction, probabilities = predict_cough_health(cough_filename, model_components)
        
        # Clean up the temporary file
        os.unlink(cough_filename)
        
        # Format the output
        result = f"Cough detected! Prediction: {prediction}\n\nProbabilities:\n"
        for cls, prob in probabilities.items():
            result += f"{cls}: {prob:.4f}\n"
        
        return result
    else:
        return "No cough detected in this audio segment."

# Create Gradio interface
def create_interface():
    model_components = load_model()
    
    with gr.Blocks(title="Cough Health Analyzer") as demo:
        gr.Markdown("# Cough Health Analyzer")
        gr.Markdown("Upload an audio file containing a cough or record a cough to analyze its health status.")
        
        with gr.Tab("Upload Audio"):
            with gr.Row():
                audio_input = gr.Audio(type="filepath", label="Upload Audio File")
                output_text = gr.Textbox(label="Prediction Result")
            
            analyze_btn = gr.Button("Analyze Cough")
            analyze_btn.click(fn=process_audio_file, inputs=audio_input, outputs=output_text)
        
        with gr.Tab("Record Audio"):
            with gr.Row():
                audio_recorder = gr.Audio(type="numpy", source="microphone", label="Record Audio")
                record_output = gr.Textbox(label="Prediction Result")
            
            record_btn = gr.Button("Analyze Recording")
            record_btn.click(fn=process_recorded_audio, inputs=[audio_recorder], outputs=record_output)
        
        with gr.Tab("Live Cough Detection"):
            gr.Markdown("This tab continuously monitors audio for coughs and analyzes them when detected.")
            
            with gr.Row():
                live_audio = gr.Audio(type="numpy", source="microphone", streaming=True, label="Live Audio Stream")
                live_output = gr.Textbox(label="Detection Result")
            
            live_audio.stream(fn=detect_coughs_in_stream, outputs=live_output)
        
        gr.Markdown("## API Usage")
        gr.Markdown("""
        This Gradio app can be accessed via API calls. Example Python code:
        ```python
        import requests
        import json
        
        # For file upload
        files = {'audio': open('your_audio_file.wav', 'rb')}
        response = requests.post('YOUR_GRADIO_URL/api/predict', files=files)
        result = json.loads(response.content)
        print(result)
        ```
        """)
    
    return demo

# Main function
if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=True)