# %% [markdown]
# ## Working from the Consolidated Data

# %%
from suno_utils.utils.text import read_jsonl
import pandas as pd
import re

# %%
data = read_jsonl('/home/sara/sfx/extreme_stems_consolidated.jsonl')
data_df = pd.DataFrame(data)
data_df.head()

# %%
print(len(data_df))

# %%
def filter_by_stem_name(stem_name):
    if "Stem" in str(stem_name):
        return False
    return True

df_test = data_df[data_df["stem_name"].apply(filter_by_stem_name)]
print(len(df_test))

# %%
def strip_to_alphanumeric(text):
    stripped = re.sub(r'[^a-zA-Z0-9\s]', '', text).lower()
    return stripped.replace("gtr", "guitar")

df_test['stem_name'] = df_test['stem_name'].apply(strip_to_alphanumeric)

# %%
df_test.head()

# %%
df_test = df_test[df_test['stem_name'].str.len() < 20]
print(len(df_test))

# %%
df_unique = df_test.drop_duplicates(subset=['id', 'stem_name'], keep=False)
print(len(df_unique))

# %% [markdown]
# ## Filter for Silence

# %%
import librosa
import numpy as np
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
import json

# %%
def analyze_audio_volume(filepath, rms_threshold=0.01, silence_ratio_threshold=0.8):
    """
    Analyze if an audio file is mostly silent or very quiet.
    
    Args:
        filepath: Path to audio file
        rms_threshold: RMS threshold below which audio is considered "quiet"
        silence_ratio_threshold: Fraction of file that must be quiet to be considered "mostly silent"
    
    Returns:
        dict: Analysis results with is_mostly_silent boolean and stats
    """
    try:
        # Load audio file
        y, sr = librosa.load(filepath, sr=None)
        
        # Calculate RMS energy in overlapping frames
        rms = librosa.feature.rms(y=y, frame_length=2048, hop_length=512)[0]
        
        # Calculate percentage of frames below threshold
        quiet_frames = np.sum(rms < rms_threshold)
        total_frames = len(rms)
        silence_ratio = quiet_frames / total_frames if total_frames > 0 else 1.0
        
        # Overall volume statistics
        max_rms = np.max(rms) if len(rms) > 0 else 0
        mean_rms = np.mean(rms) if len(rms) > 0 else 0
        
        is_mostly_silent = silence_ratio >= silence_ratio_threshold
        
        return {
            'filepath': filepath,
            'is_mostly_silent': bool(is_mostly_silent),
            'silence_ratio': float(silence_ratio),
            'max_rms': float(max_rms),
            'mean_rms': float(mean_rms),
            'duration': float(len(y) / sr),
            'status': 'success'
        }
        
    except Exception as e:
        return {
            'filepath': filepath,
            'is_mostly_silent': True,  # Assume problematic files are "bad"
            'error': str(e),
            'status': 'error'
        }

def filter_audio_files(filepaths, max_workers=8, rms_threshold=0.01, silence_ratio_threshold=0.8):
    """
    Filter out mostly silent audio files from a list of filepaths.
    
    Args:
        filepaths: List of audio file paths
        max_workers: Number of parallel workers
        rms_threshold: RMS threshold for quiet detection
        silence_ratio_threshold: Minimum ratio of quiet frames to be considered mostly silent
    
    Returns:
        tuple: (good_files, silent_files, analysis_results)
    """
    results = []
    good_files = []
    silent_files = []
    
    print(f"Analyzing {len(filepaths)} audio files...")
    
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit all tasks
        future_to_filepath = {
            executor.submit(analyze_audio_volume, fp, rms_threshold, silence_ratio_threshold): fp 
            for fp in filepaths
        }
        
        # Process completed tasks with progress bar
        for future in tqdm(as_completed(future_to_filepath), total=len(filepaths)):
            result = future.result()
            results.append(result)
            
            if result['status'] == 'success' and not result['is_mostly_silent']:
                good_files.append(result['filepath'])
            else:
                silent_files.append(result['filepath'])
    
    return good_files, silent_files, results

# %%
filepaths = list(df_unique['local_path'])
print(len(filepaths))
filepaths = filepaths

# %%
good_files, silent_files, analysis_results = filter_audio_files(
    filepaths,
    max_workers=100,          # Adjust based on your system
    rms_threshold=0.01,     # Adjust sensitivity
    silence_ratio_threshold=0.8  # 80% of file must be quiet
)

print("\nResults:")
print(f"Total files analyzed: {len(filepaths)}")
print(f"Good files (not mostly silent): {len(good_files)}")
print(f"Silent/quiet files: {len(silent_files)}")

with open('audio_analysis.json', 'w') as f:
    json.dump(analysis_results, f, indent=2)