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
import noisereduce as nr

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)
    
    # 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:
        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:]
                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:
            new_counting_model = models.Sequential([
                layers.Input(shape=(128, None, 1)),
                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_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',
        '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


# ========================================
# NOISE REDUCTION AND AUDIO PREPROCESSING
# ========================================
def clean_and_prepare_audio(y, sr):
    """
    Step 1: Apply noise reduction to clean the audio
    Step 2: Format audio to proper shape for models
    Ensures consistent format regardless of input source
    """
    print("   🧹 Applying noise reduction...")
    try:
        # Ensure audio is at the correct sample rate
        if sr != SAMPLE_RATE:
            print(f"   ⚙️ Resampling from {sr} Hz to {SAMPLE_RATE} Hz...")
            y = librosa.resample(y, orig_sr=sr, target_sr=SAMPLE_RATE)
            sr = SAMPLE_RATE
        
        # Apply noise reduction using spectral gating method
        y_clean = nr.reduce_noise(y=y, sr=sr, stationary=True, prop_decrease=0.8)
        print("   ✅ Noise reduction completed")
    except Exception as e:
        print(f"   ⚠️ Noise reduction failed: {e}, using original audio")
        y_clean = y
    
    # Normalize audio to proper amplitude range
    print("   📐 Formatting audio shape for models...")
    y_clean = y_clean.astype(np.float32)
    
    # Remove any DC offset
    y_clean = y_clean - np.mean(y_clean)
    
    # Normalize to [-1, 1] range
    max_val = np.max(np.abs(y_clean))
    if max_val > 0:
        y_clean = y_clean / max_val
    
    # Ensure no NaN or Inf values
    y_clean = np.nan_to_num(y_clean, nan=0.0, posinf=1.0, neginf=-1.0)
    
    print("   ✅ Audio prepared and shaped correctly")
    return y_clean, SAMPLE_RATE


# ========================================
# 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)
    
    if isinstance(prediction, np.ndarray):
        if prediction.ndim > 1:
            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)"""
    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
    
    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]
    
    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)
    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:
        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
        
        window_size = 2 * sr
        step_size = int(0.5 * sr)
        
        for start in range(0, len(y_segment) - window_size, step_size):
            chunk = y_segment[start : start + window_size]
            
            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)
            
            prediction = model.predict(log_spec[np.newaxis, ..., np.newaxis], verbose=0)
            
            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
            elif pred < threshold:
                is_snoring = False
        
        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:
        if sr != 16000:
            y_segment = librosa.resample(y_segment, orig_sr=sr, target_sr=16000)
            sr = 16000
        
        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')
        
        y_segment = y_segment / (np.max(np.abs(y_segment)) + 1e-10)
        
        scores, embeddings, spectrogram = yamnet_model(y_segment)
        mean_scores = np.mean(scores, axis=0)
        
        snore_keywords = [ 'Snoring', '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
        
        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:
        if len(y_segment) / sr < 2:
            return True, "Segment too short for validation"
        
        rms = librosa.feature.rms(y=y_segment)[0]
        rms_mean = np.mean(rms)
        
        if rms_mean < 0.005:
            return False, "Too quiet (likely silence)"
        if rms_mean > 0.5:
            return False, "Too loud (unlikely to be snoring)"
        
        zcr = librosa.feature.zero_crossing_rate(y_segment)[0]
        zcr_mean = np.mean(zcr)
        
        if zcr_mean > 0.2:
            return False, f"High zero-crossing rate ({zcr_mean:.3f}) - speech-like"
        
        stft = np.abs(librosa.stft(y_segment))
        freq_bins = librosa.fft_frequencies(sr=sr)
        
        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
        
        if low_ratio < 0.2 and high_ratio > 0.5:
            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"


# ========================================
# 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
    """
    segment_22k = segment
    if sr != SAMPLE_RATE:
        segment_22k = librosa.resample(segment, orig_sr=sr, target_sr=SAMPLE_RATE)
    
    finetuned_prob = predict_snoring_segment(segment_22k, SAMPLE_RATE, finetuned_model)
    
    finetuned_threshold = 0.55
    finetuned_detects = finetuned_prob > finetuned_threshold
    
    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
        yamnet_detects = yamnet_conf > yamnet_threshold
    
    has_snoring_characteristics, validation_reason = validate_snoring_characteristics(segment, sr)
    
    if finetuned_detects and yamnet_detects:
        final_decision = True
        confidence_level = "HIGH"
        combined_prob = (finetuned_prob * 0.6) + (yamnet_conf * 0.4)
        
    elif finetuned_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:
        if finetuned_prob > 0.85 and yamnet_detects:
            final_decision = True
            confidence_level = "MEDIUM"
            combined_prob = (finetuned_prob * 0.5) + (yamnet_conf * 0.5)
        else:
            final_decision = False
            confidence_level = "REJECTED"
            combined_prob = finetuned_prob
            
    elif yamnet_detects and not finetuned_detects:
        final_decision = False
        confidence_level = "REJECTED"
        combined_prob = 0.0
        
    else:
        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')
        
        result = analyze_segment_with_dual_models(segment, sr, model, use_yamnet)
        
        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
        }
        
        segment_results.append(segment_result)
        
        if result['is_snoring']:
            if result['confidence_level'] == "HIGH":
                emoji = "✅✅"
                detail = f"High confidence"
            else:
                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


# ========================================
# HYBRID APNEA DETECTION - ACCURATE
# ========================================
def count_precise_apneas(y, sr, segment_results, min_apnea=1, max_apnea=60):
    """
    ACCURATE APNEA DETECTION
    RULES:
    1. Apnea is only valid when a confirmed snoring segment (HIGH or MEDIUM)
       exists IMMEDIATELY before AND immediately after the silence gap.
       (Consecutive snoring periods in the snoring_periods list.)
    2. Gap duration must be between 1 and 60 seconds. Longer gaps are skipped.
    3. DUAL VALIDATION - both checks must pass:
       a) RMS energy check: gap audio must be significantly quieter than the
          surrounding snoring segments (rms_ratio < 0.5).
          This avoids the padding problem — we measure actual audio energy
          directly, not a padded model score.
       b) ML score check: gap audio scored by the fine-tuned model must be
          < 0.40 (model must agree it is NOT snoring).
    4. If total snoring count = 0, apnea is forced to 0.
    5. If total individual snores = 0, apnea is forced to 0.
    Risk Levels:
    - 10-15 seconds: LOW
    - 15-20 seconds: MEDIUM
    - 20+ seconds: HIGH
    Returns: (apnea_count, apnea_events_for_annotations)
    """
    print("\n🎯 APNEA DETECTION (Accurate - Immediate snoring before & after required)")
    print("   Stage 1: Identifying confirmed snoring periods...")

    # ── GUARD: Zero snoring → zero apnea ─────────────────────────────────
    snore_count_total = sum(1 for seg in segment_results if seg['is_snoring'])
    if snore_count_total == 0:
        print("      ⚠️ No snoring detected at all - apnea count forced to 0")
        return 0, []
    
    # ── GUARD: Zero individual snores → zero apnea ───────────────────────
    total_individual_snores = sum(seg.get('individual_snore_count', 0) for seg in segment_results)
    if total_individual_snores == 0:
        print("      ⚠️ No individual snores detected - apnea count forced to 0")
        return 0, []

    # STAGE 1: Collect confirmed snoring periods in time order
    all_segments_ordered = sorted(segment_results, key=lambda s: s['start_time'])
    snoring_periods = []
    for seg in all_segments_ordered:
        if seg['is_snoring']:
            snoring_periods.append({
                'start': seg['start_time'],
                'end': seg['end_time'],
                'confidence': seg['confidence_level']
            })

    print(f"      Found {len(snoring_periods)} snoring periods")

    if len(snoring_periods) < 2:
        print("      ⚠️ Need at least 2 snoring periods to detect apnea")
        return 0, []

    # STAGE 2: Find gaps between CONSECUTIVE snoring periods only (i → i+1).
    # This guarantees snoring is immediately before AND immediately after.
    print("   Stage 2: Finding gaps between consecutive snoring periods (1-60s only)...")

    apnea_candidates = []

    for i in range(len(snoring_periods) - 1):
        current_snore = snoring_periods[i]
        next_snore = snoring_periods[i + 1]

        gap_start = current_snore['end']
        gap_end = next_snore['start']
        gap_duration = gap_end - gap_start

        # Must be a real positive gap
        if gap_duration <= 0:
            continue

        # Only gaps in 1–60 second range
        if min_apnea <= gap_duration <= max_apnea:
            apnea_candidates.append({
                'start': gap_start,
                'end': gap_end,
                'duration': gap_duration,
                'before_snore': current_snore,
                'after_snore': next_snore
            })
            print(f"      Found gap: {gap_start:.1f}s → {gap_end:.1f}s (duration: {gap_duration:.1f}s)")
        else:
            print(f"      Skipping gap: {gap_start:.1f}s → {gap_end:.1f}s "
                  f"(duration: {gap_duration:.1f}s) - outside 1-60s range")

    print(f"      Found {len(apnea_candidates)} potential apnea candidates")

    if len(apnea_candidates) == 0:
        print("   ✅ Final: 0 validated apnea events\n")
        return 0, []

    # STAGE 3: DUAL VALIDATION
    # Check A: RMS energy ratio (gap must be quieter than surrounding snoring)
    # Check B: ML model score (gap must score < 0.40)
    # BOTH must pass. This fixes the padding problem and reduces false positives.
    print("   Stage 3: DUAL validation (RMS energy + ML score)...")
    print("      Rules: gap RMS ratio < 0.55 AND ML score < 0.40")

    validated_apneas = []

    for idx, candidate in enumerate(apnea_candidates):
        gap_start_sample = int(candidate['start'] * sr)
        gap_end_sample = int(min(candidate['end'] * sr, len(y)))
        gap_audio = y[gap_start_sample:gap_end_sample]

        if len(gap_audio) < 100:  # Too short to analyze
            print(f"      ⚠️ Skipping candidate {idx+1}: gap audio too short")
            continue

        # ── CHECK A: RMS energy ratio ─────────────────────────────────────
        # Compare gap RMS to the average RMS of the snoring segment before
        # and the snoring segment after. Real apnea = silence = very low RMS.
        try:
            before_start_s = int(candidate['before_snore']['start'] * sr)
            before_end_s   = int(min(candidate['before_snore']['end'] * sr, len(y)))
            after_start_s  = int(candidate['after_snore']['start'] * sr)
            after_end_s    = int(min(candidate['after_snore']['end'] * sr, len(y)))

            before_audio = y[before_start_s:before_end_s]
            after_audio  = y[after_start_s:after_end_s]

            gap_rms    = float(np.sqrt(np.mean(gap_audio ** 2)))
            before_rms = float(np.sqrt(np.mean(before_audio ** 2))) if len(before_audio) > 0 else 0.001
            after_rms  = float(np.sqrt(np.mean(after_audio ** 2)))  if len(after_audio) > 0  else 0.001
            avg_snore_rms = (before_rms + after_rms) / 2.0

            # Avoid division by zero
            rms_ratio = gap_rms / (avg_snore_rms + 1e-10)

            rms_ok = rms_ratio < 0.65  # Relaxed from 0.5 to 0.55 to account for noise reduction

        except Exception as e:
            print(f"      ⚠️ RMS check error for candidate {idx+1}: {e}")
            rms_ok = False
            rms_ratio = 999.0
            gap_rms = 0.0
            avg_snore_rms = 0.0

        # ── CHECK B: ML model score ───────────────────────────────────────
        # Run the gap audio through the fine-tuned model WITHOUT padding
        # by using the actual gap audio directly (clipped to DURATION if needed).
        # Score < 0.40 means model says: not snoring.
        ml_score = 1.0  # Default to "snoring" if check fails
        ml_ok = False

        if finetuned_snore_model is not None:
            try:
                # Resample to 22050 for the model
                gap_audio_22k = gap_audio
                if sr != SAMPLE_RATE:
                    gap_audio_22k = librosa.resample(gap_audio, orig_sr=sr, target_sr=SAMPLE_RATE)

                # Only run ML if gap is at least 1 second of real audio
                if len(gap_audio_22k) >= SAMPLE_RATE:
                    ml_score = predict_snoring_segment(gap_audio_22k, SAMPLE_RATE, finetuned_snore_model)
                    ml_ok = ml_score < 0.60
                else:
                    # Very short gap — trust RMS only, set ml_ok = rms_ok
                    ml_score = 0.0
                    ml_ok = rms_ok

            except Exception as e:
                print(f"      ⚠️ ML check error for candidate {idx+1}: {e}")
                ml_ok = False
                ml_score = 1.0

        # ── DECISION ─────────────────────────────────────────────────────
        is_valid = rms_ok and ml_ok

        print(f"      Candidate {idx+1}: gap={candidate['start']:.1f}s→{candidate['end']:.1f}s "
              f"| RMS ratio={rms_ratio:.3f} ({'✓' if rms_ok else '✗'}) "
              f"| ML score={ml_score:.3f} ({'✓' if ml_ok else '✗'}) "
              f"| {'✅ VALID' if is_valid else '❌ REJECTED'}")

        if is_valid:
            # Confidence based on how silent the gap is
            if rms_ratio < 0.2 and ml_score < 0.20:
                validation_confidence = "HIGH"
            elif rms_ratio < 0.35 and ml_score < 0.30:
                validation_confidence = "MEDIUM"
            else:
                validation_confidence = "LOW"

            # Risk based on duration
            duration = candidate['duration']
            if duration < 15:
                risk_level = "LOW"
            elif duration < 20:
                risk_level = "MEDIUM"
            else:
                risk_level = "HIGH"

            validated_apneas.append({
                'label': 'APNEA',
                'start_sec': round(candidate['start'], 2),
                'end_sec': round(candidate['end'], 2),
                'duration': round(duration, 2),
                'risk_level': risk_level,
                'confidence': {
                    'before_snoring': candidate['before_snore']['confidence'],
                    'gap_rms_ratio': round(rms_ratio, 3),
                    'gap_ml_score': round(ml_score, 3),
                    'gap_silence': round(1 - ml_score, 3),
                    'after_snoring': candidate['after_snore']['confidence'],
                    'validation_confidence': validation_confidence
                }
            })
            print(f"         → VALIDATED ({validation_confidence}) - Risk: {risk_level}")
        else:
            reasons = []
            if not rms_ok:
                reasons.append(f"gap not quiet enough (RMS ratio {rms_ratio:.3f} >= 0.55)")
            if not ml_ok:
                reasons.append(f"ML says snoring in gap (score {ml_score:.3f} >= 0.40)")
            print(f"         → REJECTED: {', '.join(reasons)}")

    print(f"   ✅ Final: {len(validated_apneas)} validated apnea events\n")
    return len(validated_apneas), validated_apneas


# ========================================
# VALIDATION
# ========================================
def validate_sleep_recording(y, sr):
    """Validate that this is a sleep recording with detailed diagnostics"""
    duration = len(y) / sr
    
    if duration < 20:
        return False, 0.0, "Audio too short (< 20 seconds)"
    
    rms_energy = np.sqrt(np.mean(y**2))
    
    # Calculate zero-crossing rate
    zcr = librosa.feature.zero_crossing_rate(y)[0]
    zcr_mean = np.mean(zcr)
    
    # Calculate spectral characteristics
    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
    
    # DIAGNOSTIC LOGGING
    print(f"   📊 AUDIO QUALITY DIAGNOSTICS:")
    print(f"      RMS Energy: {rms_energy:.4f} (normal: 0.001-0.6)")
    print(f"      Zero-Crossing Rate: {zcr_mean:.4f} (speech-like if > 0.2)")
    print(f"      Low Freq (<2kHz): {low_freq_ratio:.3f} (should be > 0.4)")
    print(f"      High Freq (>4kHz): {high_freq_ratio:.3f} (should be < 0.3)")
    
    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)"
    
    if low_freq_ratio < 0.4:
        return False, 0.0, f"Frequency distribution doesn't match sleep audio (low_freq={low_freq_ratio:.3f}, needs >0.4) - FRONTEND AUDIO QUALITY ISSUE"
    
    if high_freq_ratio > 0.3:
        return False, 0.0, f"Too much high-frequency content (high_freq={high_freq_ratio:.3f}, needs <0.3) - FRONTEND RECORDING HAS NOISE/ARTIFACTS"
    
    if zcr_mean > 0.2:
        return False, 0.0, f"Zero-crossing rate too high ({zcr_mean:.3f}) - Audio sounds like speech/noise, not snoring - FRONTEND RECORDING ISSUE"
    
    return True, 1.0, "Valid sleep recording detected"


# ========================================
# APNEA RISK CALCULATION - MODIFIED
# ========================================
def calculate_apnea_risk(apnea_events: List[Dict[str, Any]]) -> str:
    """
    Calculate apnea risk based on the longest apnea duration
    Rules:
    - NO: No apnea detected
    - LOW: Apnea duration 10-15 seconds
    - MEDIUM: Apnea duration 15-20 seconds
    - HIGH: Apnea duration above 20 seconds
    """
    if not apnea_events or len(apnea_events) == 0:
        return "NO"
    
    max_duration = max(event.get('duration', 0) for event in apnea_events)
    
    if max_duration < 10:
        return "NO"
    elif 10 <= max_duration < 15:
        return "LOW"
    elif 15 <= max_duration < 20:
        return "MEDIUM"
    else:
        return "HIGH"


# --- 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",
        "apnea_detection": "HYBRID (Fast screening + ML validation)",
        "model_accuracy": "95.33%",
        "threshold_note": "Fine-tuned threshold: 0.55 (NO LOW confidence detection)",
        "noise_reduction": "Enabled - audio cleaned before analysis"
    }


@app.post("/analyze")
async def analyze_audio(file: UploadFile = File(...)):
    """Analyze audio using STRICT dual model validation (HIGH and MEDIUM only) + HYBRID apnea detection"""
    
    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 - FORCE CONSISTENT SAMPLE RATE
        print(f"\n🔊 Loading audio: {file.filename}")
        # Always load at 22050 Hz to ensure consistency
        y, sr = librosa.load(temp_path, sr=SAMPLE_RATE, dtype=np.float32)
        print(f"   Duration: {len(y)/sr:.1f}s, Sample rate: {sr} Hz (forced to {SAMPLE_RATE})")
        
        # 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. CLEAN AND PREPARE AUDIO (NEW STEP)
        print("\n🧹 Cleaning and preparing audio...")
        y_clean, sr = clean_and_prepare_audio(y, sr)
        
        # 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
        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')
        
        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
        annotations = []
        
        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']
                }
                
                if seg.get('individual_snore_count', 0) > 0:
                    annotation['individual_snore_count'] = seg['individual_snore_count']
                
                annotations.append(annotation)
        
        # ACCURATE APNEA DETECTION
        apnea_count, apnea_events = count_precise_apneas(
            y_clean,
            sr,
            segment_results,
            min_apnea=1,
            max_apnea=60
        )
        
        annotations.extend(apnea_events)
        
        # 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
        
        risk_summary = calculate_apnea_risk(apnea_events)
        
        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,
            "total_individual_snores": total_individual_snores,
            "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",
            "apnea_detection_method": "Hybrid (Fast screening + ML validation)",
            "noise_reduction_applied": True,
            "events": annotations
        }
        
        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)