import os
import shutil
import numpy as np
import librosa
import noisereduce as nr
import scipy.signal as signal
import torch
import torch.nn as nn
import gc # Added for memory management
from torchvision import models
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from scipy.ndimage import gaussian_filter1d
from scipy.signal import find_peaks

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- LOAD PYTORCH MODEL (CPU OPTIMIZED) ---
MODEL_PATH = "best_efficientnet_snoring.pth"
# Force CPU for Zero GPU environments
device = torch.device("cpu")

def load_model():
    model = models.efficientnet_b0(weights=None)
    num_ftrs = model.classifier[1].in_features
    model.classifier[1] = nn.Linear(num_ftrs, 1)
    
    try:
        # map_location="cpu" is mandatory for CPU-only servers
        state_dict = torch.load(MODEL_PATH, map_location=device)
        model.load_state_dict(state_dict, strict=False)
        print("✅ PyTorch EfficientNet Loaded on CPU")
    except Exception as e:
        print(f"❌ Load Error: {e}")
        
    model.to(device)
    model.eval()
    return model

# Global model instance
model = load_model()

# --- LOGIC FUNCTIONS ---

def clean_audio_stream(y, sr=16000):
    # Noise reduction can be heavy on CPU; reduced footprint
    y_denoised = nr.reduce_noise(y=y, sr=sr)
    b, a = signal.butter(4, [200/(sr/2), 2000/(sr/2)], btype='band')
    y_filtered = signal.filtfilt(b, a, y_denoised)
    return y_filtered

def is_snoring_sound_pytorch(y_segment, sr):
    try:
        rms = np.sqrt(np.mean(y_segment**2))
        if rms < 0.002: return False, 0.0

        # Ensure fixed length for model input
        y_fixed = librosa.util.fix_length(y_segment, size=16000)
        S = librosa.feature.melspectrogram(y=y_fixed, sr=16000, n_mels=128)
        S_db = librosa.power_to_db(S, ref=np.max)
        
        S_norm = (S_db - S_db.min()) / (S_db.max() - S_db.min() + 1e-6)
        input_tensor = torch.tensor(S_norm).float().unsqueeze(0).unsqueeze(0).repeat(1, 3, 1, 1)
        
        with torch.no_grad():
            output = model(input_tensor)
            confidence = torch.sigmoid(output).item()

        return confidence > 0.5, round(confidence, 2)
    except Exception:
        return False, 0.0

# --- UPDATED SNORE DETECTION FOR CPU ---

def detect_snores_accurate(y_clean, sr):
    # Optimized overlap for CPU: reduced from 0.67 to 0.5 to lower inference load
    segment_samples = int(1.5 * sr)
    hop_samples = int(segment_samples * 0.5) 
    
    snore_events = []
    
    # Process segments
    for start in range(0, len(y_clean) - segment_samples + 1, hop_samples):
        segment = y_clean[start : start + segment_samples]
        is_snore, conf = is_snoring_sound_pytorch(segment, sr)
        
        if is_snore:
            start_time = start / sr
            snore_events.append({
                'start_time': start_time,
                'end_time': start_time + 1.5,
                'duration': 1.5,
                'confidence': conf
            })
    
    return snore_events

# --- API ENDPOINTS ---

@app.post("/analyze")
async def analyze_audio(file: UploadFile = File(...)):
    temp_path = f"temp_{file.filename}"
    with open(temp_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)
            
    try:
        # Load with standardized SR to save RAM immediately
        y_orig, sr = librosa.load(temp_path, sr=16000)
        
        if len(y_orig) / sr < 20:
            return {"valid_recording": False, "reason": "Audio too short"}

        y_clean = clean_audio_stream(y_orig, sr)
        
        # Snore detection
        snore_events = detect_snores_accurate(y_clean, sr)
        
        # Apnea detection via gaps
        intervals = librosa.effects.split(y_clean, top_db=25)
        annotations = []
        prev_end = 0
        apnea_count = 0
                
        for start, end in intervals:
            gap_dur = (start - prev_end) / sr
            if 10.0 <= gap_dur <= 120.0:
                apnea_count += 1
                risk = "LOW" if gap_dur < 15.0 else ("MEDIUM" if gap_dur < 20.0 else "HIGH")
                annotations.append({
                    "label": "APNEA",
                    "start_sec": round(prev_end/sr, 2),
                    "end_sec": round(start/sr, 2),
                    "duration": round(gap_dur, 2),
                    "risk_level": risk
                })
            prev_end = end
        
        # Add snores to annotations
        for snore in snore_events:
            annotations.append({
                "label": "SNORING",
                "start_sec": round(snore['start_time'], 2),
                "end_sec": round(snore['end_time'], 2),
                "duration": round(snore['duration'], 2),
                "confidence": round(snore['confidence'], 2)
            })

        duration_hours = (len(y_orig) / sr) / 3600
        ahi = apnea_count / duration_hours if duration_hours > 0 else 0
        
        overall_risk = ""
        if ahi >= 20: overall_risk = "HIGH"
        elif ahi >= 15: overall_risk = "MEDIUM"
        elif ahi >= 10: overall_risk = "LOW"

        return {
            "valid_recording": True,
            "snore_count": len(snore_events),
            "apnea_count": apnea_count,
            "risk_level": overall_risk,
            "ahi_score": round(ahi, 1),
            "events": annotations
        }

    except Exception as e:
        print(f"SERVER ERROR: {e}") # Log error for debugging
        raise HTTPException(status_code=500, detail="Processing error")
    finally:
        if os.path.exists(temp_path):
            os.remove(temp_path)
        # FORCE CLEANUP for Zero GPU RAM
        if 'y_orig' in locals(): del y_orig
        if 'y_clean' in locals(): del y_clean
        gc.collect() 

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