import os
import sys
import numpy as np
import librosa
import soundfile as sf
import pandas as pd
from tqdm import tqdm
import gc
import tensorflow as tf
import time

# Add project root to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.hear_extractor import HeARExtractor

# --- Config ---
OUTPUT_DIR = r"c:\Users\ASUS\lung_ai_project\data\hear_embeddings_optimized"
TEMP_AUDIO_DIR = r"c:\Users\ASUS\lung_ai_project\data\temp_aug_audio"
ORIG_HEAR_DIR = r"c:\Users\ASUS\lung_ai_project\data\hear_embeddings"
RESP_BASE = r"c:\Users\ASUS\lung_ai_project\data\extracted_cough\Respiratory_Sound_Dataset-main"
COS_BASE = r"c:\Users\ASUS\lung_ai_project\data\coswara"
CHECKPOINT_INTERVAL = 50

# --- Augmentations ---
def add_noise(data, noise_factor=0.005):
    noise = np.random.randn(len(data))
    augmented_data = data + noise_factor * noise
    return augmented_data

def speed_change(data, speed_factor=0.9):
    # Resample is much faster than time_stretch/pitch_shift (FFT based)
    # This changes both pitch and speed, which is a valid augmentation
    new_len = int(len(data) / speed_factor)
    return librosa.resample(data, orig_sr=16000, target_sr=int(16000*speed_factor))

# --- Data Collection ---
def get_sick_files():
    files = []
    # Coswara
    csv_dir = os.path.join(COS_BASE, "csvs")
    data_dir = os.path.join(COS_BASE, "coswara_data", "kaggle_data")
    status_map = {}
    if os.path.exists(csv_dir):
        for csv_file in os.listdir(csv_dir):
            if csv_file.endswith(".csv"):
                df = pd.read_csv(os.path.join(csv_dir, csv_file))
                if 'id' in df.columns and 'covid_status' in df.columns:
                    for _, row in df.iterrows():
                        status_map[row['id']] = row['covid_status']

    if os.path.exists(data_dir):
        for pid in os.listdir(data_dir):
            status = status_map.get(pid)
            if status and status.lower() != "healthy":
                pid_dir = os.path.join(data_dir, pid)
                for af in ["cough.wav", "cough-heavy.wav", "cough-shallow.wav"]:
                    path = os.path.join(pid_dir, af)
                    if os.path.exists(path):
                        files.append(path)
                        break 
    # Respiratory
    resp_audio = os.path.join(RESP_BASE, "audio_and_txt_files")
    resp_csv = os.path.join(RESP_BASE, "patient_diagnosis.csv")
    if os.path.exists(resp_csv):
        df = pd.read_csv(resp_csv)
        diag_map = dict(zip(df['Patient_ID'], df['DIAGNOSIS']))
        for f in os.listdir(resp_audio):
            if f.endswith(".wav"):
                try:
                    pid = int(f.split('_')[0])
                    diag = diag_map.get(pid)
                    if diag and diag.lower() != "healthy":
                        files.append(os.path.join(resp_audio, f))
                except: continue
    return files

def main():
    # Setup directories
    if not os.path.exists(OUTPUT_DIR):
        os.makedirs(OUTPUT_DIR)
    if not os.path.exists(TEMP_AUDIO_DIR):
        os.makedirs(TEMP_AUDIO_DIR)
        
    print("Identifying Sick Files...")
    sick_files = get_sick_files()
    print(f"Found {len(sick_files)} sick files.")
    
    # Load feature lists
    features = []
    labels = []
    
    # Check for existing checkpoint
    checkpoint_path = os.path.join(OUTPUT_DIR, "checkpoint_indices.npy")
    start_idx = 0
    if os.path.exists(checkpoint_path):
        start_idx = np.load(checkpoint_path).item()
        print(f"Resuming from index {start_idx}")
        features = list(np.load(os.path.join(OUTPUT_DIR, "X_checkpoint.npy")))
        labels = list(np.load(os.path.join(OUTPUT_DIR, "y_checkpoint.npy")))
    
    print("Loading HeAR Extractor...")
    extractor = HeARExtractor()
    
    # Processing Loop
    processed_count = 0
    
    for i in tqdm(range(start_idx, len(sick_files))):
        file_path = sick_files[i]
        
        try:
            # Memory Cleanup
            if processed_count % CHECKPOINT_INTERVAL == 0 and processed_count > 0:
                gc.collect()
                tf.keras.backend.clear_session()
                # Save Checkpoint
                np.save(os.path.join(OUTPUT_DIR, "X_checkpoint.npy"), np.array(features))
                np.save(os.path.join(OUTPUT_DIR, "y_checkpoint.npy"), np.array(labels))
                np.save(checkpoint_path, i)
            
            # Load Audio (limit duration to 5s to save speed)
            y, sr = librosa.load(file_path, sr=16000, duration=5.0) 
            
            if len(y) < 2000: # Skip empty/too short
                continue

            # Aug 1: Noise (Fast)
            y_noise = add_noise(y)
            temp_path_1 = os.path.join(TEMP_AUDIO_DIR, "temp_noise.wav")
            sf.write(temp_path_1, y_noise, 16000)
            
            emb1 = extractor.extract(temp_path_1)
            if emb1 is not None:
                features.append(emb1)
                labels.append("sick")
                
            # Aug 2: Speed/Pitch Change (Resampling - Fast)
            y_speed = speed_change(y, speed_factor=0.9) # Slightly slower/deeper
            temp_path_2 = os.path.join(TEMP_AUDIO_DIR, "temp_speed.wav")
            sf.write(temp_path_2, y_speed, 16000)
            
            emb2 = extractor.extract(temp_path_2)
            if emb2 is not None:
                features.append(emb2)
                labels.append("sick")
                
        except Exception as e:
            print(f"Error on {file_path}: {e}")
            continue
            
        processed_count += 1

    # Merge
    print("Merging with Original Data...")
    if os.path.exists(os.path.join(ORIG_HEAR_DIR, "X_hear.npy")):
        X_orig = np.load(os.path.join(ORIG_HEAR_DIR, "X_hear.npy"))
        y_orig = np.load(os.path.join(ORIG_HEAR_DIR, "y_hear.npy"))
        
        X_final = np.concatenate([X_orig, np.array(features)])
        y_final = np.concatenate([y_orig, np.array(labels)])
    else:
        X_final = np.array(features)
        y_final = np.array(labels)
        
    np.save(os.path.join(OUTPUT_DIR, "X_hear_opt_merged.npy"), X_final)
    np.save(os.path.join(OUTPUT_DIR, "y_hear_opt_merged.npy"), y_final)
    
    print(f"DONE. Saved {len(X_final)} total samples to {OUTPUT_DIR}")
    
    # Cleanup Temp
    try:
        import shutil
        shutil.rmtree(TEMP_AUDIO_DIR)
    except: pass

if __name__ == "__main__":
    main()