app.py

import os
import shutil
import numpy as np
import librosa
import gc
import tensorflow as tf
from tensorflow.keras import layers, models
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from typing import Optional, Dict, Any, List
import warnings
warnings.filterwarnings('ignore')

# --- OPTIMIZE FOR CPU ---
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

app = FastAPI(title="Real Snore & Apnea Detector")

# Enable CORS for frontend integration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- CONFIGURATION FROM TRAINING ---
SAMPLE_RATE = 22050
N_FFT = 2048
HOP_LENGTH = 512
N_MELS = 128
FMAX = 8000
DURATION = 10

# NEW MODEL CONFIG (for snore counting)
NEW_MODEL_SAMPLE_RATE = 16000
NEW_MODEL_DURATION = 10

# --- LOAD YOUR FINE-TUNED MODEL (OLD MODEL FOR DETECTION) ---
FINETUNED_MODEL_PATH = "snoring_detection_model fine tuned.h5"
finetuned_snore_model = None

# --- LOAD NEW MODEL (FOR COUNTING) ---
NEW_MODEL_PATH = "snore_detector_v2.h5"
new_counting_model = None

def build_finetuned_model_architecture():
    """Reconstruct the EXACT fine-tuned model architecture from training"""
    inputs = layers.Input(shape=(128, 431, 1))
    
    # Block 1
    x = layers.Conv2D(32, (3, 3), padding='same')(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.MaxPooling2D((2, 2))(x)
    
    # Block 2
    x = layers.Conv2D(64, (3, 3), padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.MaxPooling2D((2, 2))(x)
    
    # Block 3
    x = layers.Conv2D(128, (3, 3), padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.MaxPooling2D((2, 2))(x)
    
    # Block 4
    x = layers.Conv2D(256, (3, 3), padding='same')(x)
    x = layers.BatchNormalization()(x)
    x = layers.Activation('relu')(x)
    x = layers.GlobalAveragePooling2D()(x)  # FIXED: Add GlobalAveragePooling2D
    
    # Dense layers
    x = layers.Dense(256, activation='relu')(x)
    x = layers.Dropout(0.5)(x)
    x = layers.Dense(128, activation='relu')(x)
    x = layers.Dropout(0.5)(x)
    x = layers.Dense(64, activation='relu')(x)
    x = layers.Dropout(0.5)(x)
    
    outputs = layers.Dense(1, activation='sigmoid')(x)
    
    model = models.Model(inputs=inputs, outputs=outputs)
    return model

# LOAD OLD MODEL (FINE-TUNED)
try:
    print("🔄 Loading fine-tuned snoring detection model...")
    
    try:
        class CompatibleInputLayer(layers.InputLayer):
            def __init__(self, batch_shape=None, input_shape=None, **kwargs):
                if batch_shape is not None and input_shape is None:
                    input_shape = batch_shape[1:]
                super().__init__(input_shape=input_shape, **kwargs)
        
        custom_objects = {'InputLayer': CompatibleInputLayer}
        finetuned_snore_model = tf.keras.models.load_model(
            FINETUNED_MODEL_PATH,
            custom_objects=custom_objects,
            compile=False
        )
        print("✅ Fine-tuned Model Loaded Successfully (Method 1)!")
    
    except Exception as e1:
        print(f"   Method 1 failed: {str(e1)[:100]}")
        print("   Trying Method 2: Weight reconstruction...")
        
        try:
            finetuned_snore_model = build_finetuned_model_architecture()
            finetuned_snore_model.load_weights(FINETUNED_MODEL_PATH)
            print("✅ Fine-tuned Model Loaded Successfully (Method 2 - Weight Reconstruction)!")
        
        except Exception as e2:
            print(f"   Method 2 failed: {str(e2)[:100]}")
            print("   ⚠️ Fine-tuned model not available.")
            finetuned_snore_model = None
    
    if finetuned_snore_model is not None:
        print(f"   Model input shape: {finetuned_snore_model.input_shape}")
        print(f"   Model output shape: {finetuned_snore_model.output_shape}")

except Exception as e:
    print(f"❌ Unexpected error loading fine-tuned model: {e}")
    finetuned_snore_model = None

# LOAD NEW MODEL (FOR COUNTING)
try:
    print("\n🔄 Loading NEW snore counting model (snore_detector_v2.h5)...")
    
    try:
        # Try method 1: Direct loading with custom objects
        class CompatibleInputLayer(layers.InputLayer):
            def __init__(self, batch_shape=None, input_shape=None, **kwargs):
                # Handle batch_shape parameter
                if batch_shape is not None and input_shape is None:
                    input_shape = batch_shape[1:]
                # Remove batch_shape from kwargs before passing to parent
                kwargs.pop('batch_shape', None)
                super().__init__(input_shape=input_shape, **kwargs)
        
        custom_objects = {'InputLayer': CompatibleInputLayer}
        new_counting_model = tf.keras.models.load_model(
            NEW_MODEL_PATH,
            custom_objects=custom_objects,
            compile=False
        )
        print("✅ New Counting Model Loaded Successfully (Method 1)!")
        print(f"   Model input shape: {new_counting_model.input_shape}")
        print(f"   Model output shape: {new_counting_model.output_shape}")
    
    except Exception as e1:
        print(f"   Method 1 failed: {str(e1)[:150]}")
        print("   Trying Method 2: Load weights only...")
        
        try:
            # Build a compatible architecture based on your latest model code
            new_counting_model = models.Sequential([
                layers.Input(shape=(128, None, 1)),  # Variable time dimension
                layers.Conv2D(32, (3, 3), activation='relu', padding='same'),
                layers.BatchNormalization(),
                layers.MaxPooling2D((2, 2)),
                
                layers.Conv2D(64, (3, 3), activation='relu', padding='same'),
                layers.BatchNormalization(),
                layers.MaxPooling2D((2, 2)),
                
                layers.Conv2D(128, (3, 3), activation='relu', padding='same'),
                layers.BatchNormalization(),
                layers.GlobalAveragePooling2D(),
                
                layers.Dense(128, activation='relu'),
                layers.Dropout(0.4),
                layers.Dense(1, activation='sigmoid')
            ])
            
            new_counting_model.load_weights(NEW_MODEL_PATH)
            print("✅ New Counting Model Loaded Successfully (Method 2 - Weight Reconstruction)!")
            print(f"   Model input shape: {new_counting_model.input_shape}")
            print(f"   Model output shape: {new_counting_model.output_shape}")
        
        except Exception as e2:
            print(f"   Method 2 failed: {str(e2)[:150]}")
            print("   ⚠️ New counting model not available. Continuing without counting feature...")
            new_counting_model = None

except Exception as e:
    print(f"❌ Unexpected error loading new counting model: {e}")
    new_counting_model = None

# --- LOAD YAMNet MODEL ---
yamnet_model = None
yamnet_class_names = []

try:
    print("\n🔄 Loading YAMNet model for additional validation...")
    import tensorflow_hub as hub
    
    YAMNET_MODEL_URL = 'https://tfhub.dev/google/yamnet/1'
    yamnet_model = hub.load(YAMNET_MODEL_URL)
    
    # YAMNet has 521 classes, "Snoring" is class index 38
    yamnet_class_names = [
        'Speech', 'Child speech, kid speaking', 'Conversation', 'Narration, monologue',
        'Babbling', 'Speech synthesizer', 'Shout', 'Bellow', 'Whoop', 'Yell',
        'Children shouting', 'Screaming', 'Whispering', 'Laughter', 'Baby laughter',
        'Giggle', 'Snicker', 'Belly laugh', 'Chuckle, chortle', 'Crying, sobbing',
        'Baby cry, infant cry', 'Whimper', 'Wail, moan', 'Sigh', 'Singing',
        'Choir', 'Yodeling', 'Chant', 'Mantra', 'Child singing',
        'Synthetic singing', 'Rapping', 'Humming', 'Groan', 'Grunt',
        'Whistling', 'Breathing', 'Wheeze', 'Snoring', 'Gasp',  # Index 38 is Snoring
        'Pant', 'Snort', 'Cough', 'Throat clearing', 'Sneeze'
    ]
    
    print(f"✅ YAMNet model loaded successfully")
    
except Exception as e:
    print(f"⚠️ YAMNet loading error: {e}")
    print("   Continuing with fine-tuned model only")
    yamnet_model = None

# ========================================
# KAGGLE METHOD - FINE-TUNED MODEL
# ========================================

def create_mel_spectrogram(y, sr):
    """Create mel-spectrogram EXACTLY as in Kaggle cleaning pipeline"""
    mel_spec = librosa.feature.melspectrogram(
        y=y, sr=sr, n_fft=N_FFT, hop_length=HOP_LENGTH,
        n_mels=N_MELS, fmax=FMAX, power=1.0
    )
    mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)
    mel_spec_db = (mel_spec_db - mel_spec_db.min()) / (mel_spec_db.max() - mel_spec_db.min() + 1e-10)
    return mel_spec_db

def predict_snoring_segment(y_segment, sr, model):
    """Predict snoring for a segment - EXACT from Kaggle cleaning pipeline"""
    target_length = SAMPLE_RATE * DURATION
    
    if len(y_segment) < target_length:
        padding = target_length - len(y_segment)
        y_segment = np.pad(y_segment, (0, padding), mode='constant')
    elif len(y_segment) > target_length:
        y_segment = y_segment[:target_length]
    
    y_segment = y_segment / (np.max(np.abs(y_segment)) + 1e-10)
    mel_spec = create_mel_spectrogram(y_segment, sr)
    mel_spec_input = mel_spec[np.newaxis, ..., np.newaxis]
    
    prediction = model.predict(mel_spec_input, verbose=0)
    
    # Handle different output shapes
    if isinstance(prediction, np.ndarray):
        if prediction.ndim > 1:
            # Flatten and take mean or first value
            probability = float(np.mean(prediction))
        else:
            probability = float(prediction[0])
    else:
        probability = float(prediction)
    
    return probability

# ========================================
# NEW MODEL - SNORE COUNTING WITHIN DETECTED SEGMENTS
# ========================================

def extract_spectrogram_for_new_model(y_segment, sr):
    """Extract spectrogram for new model (from latest model code)"""
    # Resample to 16kHz if needed
    if sr != NEW_MODEL_SAMPLE_RATE:
        y_segment = librosa.resample(y_segment, orig_sr=sr, target_sr=NEW_MODEL_SAMPLE_RATE)
        sr = NEW_MODEL_SAMPLE_RATE
    
    # Ensure exact length
    target_len = NEW_MODEL_SAMPLE_RATE * NEW_MODEL_DURATION
    if len(y_segment) < target_len:
        y_segment = np.pad(y_segment, (0, target_len - len(y_segment)))
    else:
        y_segment = y_segment[:target_len]
    
    # Generate Mel Spectrogram
    spec = librosa.feature.melspectrogram(y=y_segment, sr=sr, n_mels=128, n_fft=2048, hop_length=512)
    log_spec = librosa.power_to_db(spec, ref=np.max)
    # Normalize
    log_spec = (log_spec - log_spec.min()) / (log_spec.max() - log_spec.min() + 1e-6)
    return log_spec[..., np.newaxis]

def count_snores_in_segment(y_segment, sr, model, threshold=0.7):
    """
    Count individual snores within a detected snoring segment
    Uses sliding 2-second window to detect multiple snore events
    """
    if model is None:
        return 0
    
    try:
        # Resample to 16kHz for new model
        if sr != NEW_MODEL_SAMPLE_RATE:
            y_segment = librosa.resample(y_segment, orig_sr=sr, target_sr=NEW_MODEL_SAMPLE_RATE)
            sr = NEW_MODEL_SAMPLE_RATE
        
        snore_count = 0
        is_snoring = False
        
        # Slide a 2-second window across the segment
        window_size = 2 * sr
        step_size = int(0.5 * sr)  # 0.5s steps for high resolution
        
        for start in range(0, len(y_segment) - window_size, step_size):
            chunk = y_segment[start : start + window_size]
            
            # Process chunk to spectrogram
            spec = librosa.feature.melspectrogram(y=chunk, sr=sr, n_mels=128, n_fft=2048, hop_length=512)
            log_spec = librosa.power_to_db(spec, ref=np.max)
            log_spec = (log_spec - log_spec.min()) / (log_spec.max() - log_spec.min() + 1e-6)
            
            # Predict
            prediction = model.predict(log_spec[np.newaxis, ..., np.newaxis], verbose=0)
            
            # Handle different output shapes safely
            if isinstance(prediction, np.ndarray):
                if prediction.ndim > 1:
                    pred = float(np.mean(prediction))
                else:
                    pred = float(prediction[0])
            else:
                pred = float(prediction)
            
            if pred > threshold and not is_snoring:
                snore_count += 1
                is_snoring = True  # Entered a snore event
            elif pred < threshold:
                is_snoring = False  # Snore ended
        
        return snore_count
    
    except Exception as e:
        print(f"   ⚠️ Error counting snores in segment: {e}")
        return 0

# ========================================
# YAMNet DETECTION
# ========================================

def detect_snoring_with_yamnet(y_segment, sr):
    """
    Use YAMNet to detect snoring
    Returns: (snoring_confidence, top_class_name)
    """
    if yamnet_model is None:
        return 0.0, "YAMNet unavailable"
    
    try:
        # YAMNet expects 16kHz audio
        if sr != 16000:
            y_segment = librosa.resample(y_segment, orig_sr=sr, target_sr=16000)
            sr = 16000
        
        # Ensure minimum length
        min_length = int(0.96 * sr)
        if len(y_segment) < min_length:
            padding = min_length - len(y_segment)
            y_segment = np.pad(y_segment, (0, padding), mode='constant')
        
        # Normalize
        y_segment = y_segment / (np.max(np.abs(y_segment)) + 1e-10)
        
        # Run YAMNet
        scores, embeddings, spectrogram = yamnet_model(y_segment)
        
        # Get mean scores
        mean_scores = np.mean(scores, axis=0)
        
        # Find snoring-related classes
        snore_keywords = ['snore', 'snoring', 'breathing', 'breath', 'wheeze', 'gasp', 'snort']
        snore_indices = []
        
        for idx, class_name in enumerate(yamnet_class_names):
            if idx >= len(mean_scores):
                break
            class_lower = class_name.lower()
            for keyword in snore_keywords:
                if keyword in class_lower:
                    snore_indices.append(idx)
                    break
        
        # Get max snoring confidence
        if snore_indices:
            snore_scores = mean_scores[snore_indices]
            snoring_confidence = float(np.max(snore_scores))
            best_idx = snore_indices[np.argmax(snore_scores)]
            top_class = yamnet_class_names[best_idx] if best_idx < len(yamnet_class_names) else "Unknown"
        else:
            snoring_confidence = 0.0
            top_class = "No snoring class"
        
        return snoring_confidence, top_class
    
    except Exception as e:
        print(f"   ⚠️ YAMNet error: {e}")
        return 0.0, "YAMNet error"

# ========================================
# SNORING CHARACTERISTICS VALIDATION - LESS STRICT
# ========================================

def validate_snoring_characteristics(y_segment, sr):
    """
    Validate if audio segment has characteristics of real snoring
    LESS STRICT VERSION - only filter obvious non-snoring
    """
    try:
        # Skip validation for very short segments
        if len(y_segment) / sr < 2:
            return True, "Segment too short for validation"
        
        # 1. Check RMS energy - only filter extremes
        rms = librosa.feature.rms(y=y_segment)[0]
        rms_mean = np.mean(rms)
        
        # Only filter obvious extremes
        if rms_mean < 0.005:  # Very quiet - likely silence
            return False, "Too quiet (likely silence)"
        if rms_mean > 0.5:    # Very loud - likely not snoring
            return False, "Too loud (unlikely to be snoring)"
        
        # 2. Check zero-crossing rate - be lenient
        zcr = librosa.feature.zero_crossing_rate(y_segment)[0]
        zcr_mean = np.mean(zcr)
        
        # Snoring generally has low ZCR, but be lenient
        if zcr_mean > 0.2:  # Very high ZCR indicates speech/clicking sounds
            return False, f"High zero-crossing rate ({zcr_mean:.3f}) - speech-like"
        
        # 3. Check frequency content - lenient version
        stft = np.abs(librosa.stft(y_segment))
        freq_bins = librosa.fft_frequencies(sr=sr)
        
        # Calculate energy in different bands
        low_band = np.sum(stft[freq_bins < 500, :])
        mid_band = np.sum(stft[(freq_bins >= 500) & (freq_bins < 2000), :])
        high_band = np.sum(stft[freq_bins >= 2000, :])
        total_energy = low_band + mid_band + high_band + 1e-10
        
        low_ratio = low_band / total_energy
        high_ratio = high_band / total_energy
        
        # Snoring is mostly low-mid frequency, but be lenient
        if low_ratio < 0.2 and high_ratio > 0.5:
            # Too much high frequency, not enough low frequency
            return False, f"Spectral balance off (low: {low_ratio:.2f}, high: {high_ratio:.2f})"
        
        return True, "Passed basic snoring characteristics"
    
    except Exception as e:
        print(f"   ⚠️ Snoring validation error: {e}")
        return True, "Validation skipped"  # Default to true if validation fails

# ========================================
# BALANCED DUAL MODEL VALIDATION - NO LOW CONFIDENCE
# ========================================

def analyze_segment_with_dual_models(segment, sr, finetuned_model, use_yamnet=True):
    """
    Analyze segment with BALANCED approach
    CHANGED: Removed LOW confidence detection, only HIGH and MEDIUM
    """
    # Resample to 22050 Hz for fine-tuned model
    segment_22k = segment
    if sr != SAMPLE_RATE:
        segment_22k = librosa.resample(segment, orig_sr=sr, target_sr=SAMPLE_RATE)
    
    # Fine-tuned model prediction with MODERATE threshold
    finetuned_prob = predict_snoring_segment(segment_22k, SAMPLE_RATE, finetuned_model)
    
    # Use original threshold but with additional validation
    finetuned_threshold = 0.55  # Slightly higher than 0.5
    finetuned_detects = finetuned_prob > finetuned_threshold
    
    # YAMNet prediction
    yamnet_conf = 0.0
    yamnet_class = "N/A"
    yamnet_detects = False
    
    if use_yamnet and yamnet_model is not None:
        yamnet_conf, yamnet_class = detect_snoring_with_yamnet(segment, sr)
        yamnet_threshold = 0.15  # Lower threshold for YAMNet
        yamnet_detects = yamnet_conf > yamnet_threshold
    
    # Snoring characteristics validation (LENIENT)
    has_snoring_characteristics, validation_reason = validate_snoring_characteristics(segment, sr)
    
    # MODIFIED DECISION LOGIC - NO LOW CONFIDENCE DETECTION
    
    if finetuned_detects and yamnet_detects:
        # BOTH MODELS AGREE - HIGH CONFIDENCE
        final_decision = True
        confidence_level = "HIGH"
        combined_prob = (finetuned_prob * 0.6) + (yamnet_conf * 0.4)
    
    elif finetuned_detects and has_snoring_characteristics:
        # Fine-tuned detects AND has snoring characteristics
        if finetuned_prob > 0.7:
            final_decision = True
            confidence_level = "HIGH"
            combined_prob = finetuned_prob * 0.9
        else:
            final_decision = True
            confidence_level = "MEDIUM"
            combined_prob = finetuned_prob * 0.8
    
    elif finetuned_detects and not has_snoring_characteristics:
        # Fine-tuned detects but characteristics don't match
        # REMOVED: No longer accepting as LOW confidence
        if finetuned_prob > 0.85 and yamnet_detects:  # Very high threshold if no characteristics
            final_decision = True
            confidence_level = "MEDIUM"
            combined_prob = (finetuned_prob * 0.5) + (yamnet_conf * 0.5)
        else:
            # Reject - no characteristics match
            final_decision = False
            confidence_level = "REJECTED"
            combined_prob = finetuned_prob
    
    elif yamnet_detects and not finetuned_detects:
        # Only YAMNet detects - REMOVED: No longer accepting as LOW confidence
        final_decision = False
        confidence_level = "REJECTED"
        combined_prob = 0.0
    
    else:
        # No detection
        final_decision = False
        confidence_level = "NONE"
        combined_prob = finetuned_prob
    
    return {
        'is_snoring': final_decision,
        'confidence_level': confidence_level,
        'finetuned_prob': finetuned_prob,
        'finetuned_threshold_used': finetuned_threshold,
        'yamnet_conf': yamnet_conf,
        'yamnet_class': yamnet_class,
        'has_snoring_characteristics': has_snoring_characteristics,
        'validation_reason': validation_reason,
        'combined_prob': combined_prob
    }

def analyze_audio_segments(y, sr, model, segment_duration=10):
    """Analyze audio in segments with BALANCED dual model validation - NO LOW CONFIDENCE"""
    segment_length = int(segment_duration * sr)
    total_duration = len(y) / sr
    num_segments = int(np.ceil(total_duration / segment_duration))
    
    segment_results = []
    use_yamnet = yamnet_model is not None
    
    print(f"\n📊 Analyzing {num_segments} segments ({segment_duration}s each)...")
    print(f"   🔍 Using STRICT validation - ONLY HIGH and MEDIUM confidence snores")
    print(f"   ⚖️  Fine-tuned threshold: 0.55 (no LOW confidence detection)")
    if use_yamnet:
        print(f"   🔍 Using DUAL validation: Fine-tuned Model + YAMNet")
    else:
        print(f"   🔍 Using SINGLE validation: Fine-tuned Model only")
    
    if new_counting_model is not None:
        print(f"   🔢 NEW: Counting individual snores within detected segments")
    
    for i in range(num_segments):
        start_sample = i * segment_length
        end_sample = min(start_sample + segment_length, len(y))
        
        segment = y[start_sample:end_sample]
        
        if len(segment) < segment_length:
            padding = segment_length - len(segment)
            segment = np.pad(segment, (0, padding), mode='constant')
        
        # Balanced dual model analysis
        result = analyze_segment_with_dual_models(segment, sr, model, use_yamnet)
        
        # NEW: Count snores within this segment if detected as snoring
        snore_count_in_segment = 0
        if result['is_snoring'] and new_counting_model is not None:
            snore_count_in_segment = count_snores_in_segment(segment, sr, new_counting_model, threshold=0.7)
        
        segment_result = {
            'segment': i + 1,
            'start_time': i * segment_duration,
            'end_time': min((i + 1) * segment_duration, total_duration),
            'is_snoring': result['is_snoring'],
            'confidence_level': result['confidence_level'],
            'finetuned_prob': result['finetuned_prob'],
            'finetuned_threshold': result['finetuned_threshold_used'],
            'yamnet_conf': result['yamnet_conf'],
            'yamnet_class': result['yamnet_class'],
            'has_snoring_characteristics': result['has_snoring_characteristics'],
            'validation_reason': result['validation_reason'],
            'probability': result['combined_prob'],
            'confidence': result['combined_prob'] if result['is_snoring'] else 1 - result['combined_prob'],
            'individual_snore_count': snore_count_in_segment  # NEW FIELD
        }
        
        segment_results.append(segment_result)
        
        # Logging
        if result['is_snoring']:
            if result['confidence_level'] == "HIGH":
                emoji = "✅✅"
                detail = f"High confidence"
            else:  # MEDIUM
                emoji = "✅"
                detail = f"Medium confidence"
            
            count_msg = f" [{snore_count_in_segment} snores]" if snore_count_in_segment > 0 else ""
            print(f"   {emoji} Segment {i+1}/{num_segments}: SNORE [{result['confidence_level']}] (Fine-tuned: {result['finetuned_prob']:.3f}, YAMNet: {result['yamnet_conf']:.3f}) - {detail}{count_msg}")
        else:
            if result['confidence_level'] == "REJECTED":
                reason = ""
                if result['finetuned_prob'] > 0.55 and not result['has_snoring_characteristics']:
                    reason = f" - lacks snoring characteristics: {result['validation_reason']}"
                elif result['yamnet_conf'] > 0.15 and result['finetuned_prob'] < 0.55:
                    reason = f" - only YAMNet detected ({result['yamnet_conf']:.3f})"
                
                print(f"   🚫 Segment {i+1}/{num_segments}: REJECTED{reason}")
            else:
                print(f"   ⬜ Segment {i+1}/{num_segments}: No snore (Fine-tuned: {result['finetuned_prob']:.3f})")
    
    return segment_results

# ========================================
# VALIDATION
# ========================================

def validate_sleep_recording(y, sr):
    """Validate that this is a sleep recording"""
    duration = len(y) / sr
    
    if duration < 20:
        return False, 0.0, "Audio too short (< 20 seconds)"
    
    rms_energy = np.sqrt(np.mean(y**2))
    if rms_energy < 0.001:
        return False, 0.0, "Audio is blank/silent"
    
    if rms_energy > 0.6:
        return False, 0.0, "Audio too loud (not a sleep recording)"
    
    stft = np.abs(librosa.stft(y))
    freq_bins = librosa.fft_frequencies(sr=sr)
    
    low_freq_energy = np.sum(stft[freq_bins < 2000, :])
    high_freq_energy = np.sum(stft[freq_bins > 4000, :])
    total_energy = np.sum(stft) + 1e-10
    
    low_freq_ratio = low_freq_energy / total_energy
    high_freq_ratio = high_freq_energy / total_energy
    
    if low_freq_ratio < 0.4:
        return False, 0.0, "Frequency distribution doesn't match sleep audio"
    
    if high_freq_ratio > 0.3:
        return False, 0.0, "Too much high-frequency content (speech/TV detected)"
    
    return True, 1.0, "Valid sleep recording detected"

# --- API ENDPOINTS ---

@app.get("/")
async def root():
    """Health check endpoint"""
    return {
        "status": "online", 
        "finetuned_model_loaded": finetuned_snore_model is not None,
        "yamnet_loaded": yamnet_model is not None,
        "new_counting_model_loaded": new_counting_model is not None,
        "detection_mode": "STRICT DUAL validation (HIGH and MEDIUM confidence only)" if yamnet_model else "STRICT Single validation",
        "counting_feature": "Enabled - counts individual snores within detected segments" if new_counting_model else "Disabled",
        "model_accuracy": "95.33%",
        "threshold_note": "Fine-tuned threshold: 0.55 (NO LOW confidence detection)"
    }

@app.post("/analyze")
async def analyze_audio(file: UploadFile = File(...)):
    """Analyze audio using STRICT dual model validation (HIGH and MEDIUM only) + NEW snore counting"""
    
    if finetuned_snore_model is None:
        raise HTTPException(
            status_code=503,
            detail="Fine-tuned model not loaded. Please ensure model file is available."
        )
    
    temp_path = f"temp_{os.getpid()}_{file.filename}"
    
    with open(temp_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
    
    try:
        # 1. LOAD
        print(f"\n🔊 Loading audio: {file.filename}")
        y, sr = librosa.load(temp_path, sr=None, dtype=np.float32)
        print(f"   Duration: {len(y)/sr:.1f}s, Sample rate: {sr} Hz")
        
        # 2. VALIDATE
        print("\n🔍 Validating sleep recording...")
        is_valid, confidence, validation_reason = validate_sleep_recording(y, sr)
        
        print(f"   {'✅ PASS' if is_valid else '⚠️ WARNING'}: {validation_reason}")
        
        validation_warning = None if is_valid else validation_reason
        
        # 3. SKIP NOISE REDUCTION (REMOVED)
        y_clean = y  # Use original audio directly
        print("⏭️  Noise reduction skipped (removed from pipeline)")
        
        # 4. ANALYZE with STRICT DUAL MODELS + NEW COUNTING
        print("\n🛡️  Using STRICT detection (HIGH and MEDIUM confidence only)...")
        segment_results = analyze_audio_segments(y_clean, sr, finetuned_snore_model, segment_duration=10)
        
        # 5. COUNT RESULTS - NO LOW CONFIDENCE + NEW INDIVIDUAL SNORE COUNTS
        snore_count = sum(1 for seg in segment_results if seg['is_snoring'])
        high_conf_count = sum(1 for seg in segment_results if seg.get('confidence_level') == 'HIGH')
        medium_conf_count = sum(1 for seg in segment_results if seg.get('confidence_level') == 'MEDIUM')
        rejected_count = sum(1 for seg in segment_results if seg.get('confidence_level') == 'REJECTED')
        
        # NEW: Total individual snore count
        total_individual_snores = sum(seg.get('individual_snore_count', 0) for seg in segment_results)
        
        total_segments = len(segment_results)
        
        print(f"\n📊 STRICT Analysis Results (HIGH & MEDIUM only):")
        print(f"   Total segments: {total_segments}")
        print(f"   Snoring segments: {snore_count} (High: {high_conf_count}, Medium: {medium_conf_count})")
        print(f"   Rejected segments: {rejected_count}")
        print(f"   Snoring percentage: {(snore_count/total_segments)*100:.1f}%")
        if new_counting_model is not None:
            print(f"   🔢 Total individual snores detected: {total_individual_snores}")
        
        # 6. CREATE ANNOTATIONS - ONLY HIGH AND MEDIUM + NEW SNORE COUNTS
        annotations = []
        apnea_count = 0
        
        for seg in segment_results:
            if seg['is_snoring']:
                annotation = {
                    "label": "SNORING",
                    "start_sec": round(seg['start_time'], 2),
                    "end_sec": round(seg['end_time'], 2),
                    "duration": round(seg['end_time'] - seg['start_time'], 2),
                    "probability": round(seg['probability'], 4),
                    "confidence_level": seg['confidence_level'],
                    "finetuned_prob": round(seg['finetuned_prob'], 4),
                    "finetuned_threshold": seg['finetuned_threshold'],
                    "yamnet_conf": round(seg['yamnet_conf'], 4),
                    "has_snoring_characteristics": seg['has_snoring_characteristics'],
                    "validation_reason": seg['validation_reason']
                }
                
                # NEW: Add individual snore count if available
                if seg.get('individual_snore_count', 0) > 0:
                    annotation['individual_snore_count'] = seg['individual_snore_count']
                
                annotations.append(annotation)
        
        # Apnea detection - only with HIGH and MEDIUM confidence snoring
        for i in range(len(segment_results) - 1):
            current = segment_results[i]
            if current['is_snoring'] and current['confidence_level'] in ['HIGH', 'MEDIUM']:
                for j in range(i + 1, len(segment_results)):
                    if segment_results[j]['is_snoring'] and segment_results[j]['confidence_level'] in ['HIGH', 'MEDIUM']:
                        gap_duration = segment_results[j]['start_time'] - current['end_time']
                        
                        # Apnea detection with confidence requirement
                        if 10.0 <= gap_duration <= 90.0:
                            # Require higher confidence for longer gaps
                            if gap_duration > 30.0:
                                # For long gaps, require at least one HIGH confidence
                                if current['confidence_level'] == 'HIGH' or segment_results[j]['confidence_level'] == 'HIGH':
                                    apnea_count += 1
                                    risk = "MEDIUM" if gap_duration < 40.0 else "HIGH"
                                    annotations.append({
                                        "label": "APNEA",
                                        "start_sec": round(current['end_time'], 2),
                                        "end_sec": round(segment_results[j]['start_time'], 2),
                                        "duration": round(gap_duration, 2),
                                        "risk_level": risk,
                                        "note": "Detected between high/medium confidence snoring segments"
                                    })
                                    print(f"   ⚠️  Apnea: {gap_duration:.1f}s gap (between confident snoring segments)")
                            else:
                                # Shorter gaps - any confidence level (but already filtered for HIGH/MEDIUM)
                                apnea_count += 1
                                risk = "LOW" if gap_duration < 20.0 else "MEDIUM"
                                annotations.append({
                                    "label": "APNEA",
                                    "start_sec": round(current['end_time'], 2),
                                    "end_sec": round(segment_results[j]['start_time'], 2),
                                    "duration": round(gap_duration, 2),
                                    "risk_level": risk
                                })
                                print(f"   ⚠️  Apnea: {gap_duration:.1f}s gap")
                        break
        
        # 7. STATISTICS
        total_duration = segment_results[-1]['end_time']
        duration_hours = total_duration / 3600
        ahi = apnea_count / duration_hours if duration_hours > 0 else 0
        
        snoring_percentage = (snore_count / total_segments) * 100
        
        # Strict risk assessment
        if snoring_percentage > 40:
            risk_summary = "HIGH"
        elif snoring_percentage > 20:
            risk_summary = "MEDIUM"
        elif snoring_percentage > 5:
            risk_summary = "LOW"
        else:
            risk_summary = "NORMAL"
        
        # Add note about detection mode
        detection_note = f"Detected {snore_count} snoring segments (HIGH: {high_conf_count}, MEDIUM: {medium_conf_count})"
        if new_counting_model is not None and total_individual_snores > 0:
            detection_note += f" with {total_individual_snores} individual snores counted"
        if snore_count == 0:
            detection_note = f"No confident snoring detected (rejected {rejected_count} low-confidence segments)"
        
        # Cleanup
        del y, y_clean
        gc.collect()
        
        response_data = {
            "valid_recording": True,
            "validation_warning": validation_warning,
            "snore_count": snore_count,
            "total_snore_counts": total_individual_snores,  # NEW: Total individual snores counted by new model
            "total_individual_snores": total_individual_snores,  # ADDED: Same as total_snore_counts for clarity
            "high_confidence_snores": high_conf_count,
            "medium_confidence_snores": medium_conf_count,
            "rejected_potential_snores": rejected_count,
            "total_segments": total_segments,
            "snoring_percentage": round(snoring_percentage, 2),
            "apnea_count": apnea_count,
            "apnea_risk": risk_summary,
            "ahi_score": round(ahi, 1),
            "detection_mode": "STRICT: Only HIGH and MEDIUM confidence snoring detection",
            "detection_note": detection_note,
            "fine_tuned_threshold": 0.55,
            "validation_approach": "Strict - only detecting HIGH and MEDIUM confidence snores",
            "events": annotations
        }
        
        
        # NEW: Add individual snore counting data if available
        if new_counting_model is not None:
            response_data["individual_snore_count"] = total_individual_snores
            response_data["counting_model_enabled"] = True
        else:
            response_data["counting_model_enabled"] = False
        
        return response_data
        
    except Exception as e:
        print(f"🔥 ERROR: {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Processing Error: {str(e)}")
    
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)
        gc.collect()

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)