#!/usr/bin/env python3
"""
Audio Stem Cropping Tool
Extracts clean audio samples from music stems with smooth fade-in/out transitions.
"""

import os
import librosa
import soundfile as sf
import numpy as np
from pathlib import Path
import argparse


def find_diverse_regions(audio, sr, sample_duration, num_samples, hop_length=512, frame_length=2048):
    """Find diverse audio regions using multiple sampling strategies tailored to each track."""
    import random

    # Calculate RMS energy
    rms = librosa.feature.rms(y=audio, hop_length=hop_length, frame_length=frame_length)[0]
    times = librosa.frames_to_time(np.arange(len(rms)), sr=sr, hop_length=hop_length)

    audio_duration = len(audio) / sr
    sample_frames = int(sample_duration * sr / hop_length)

    # Create a seed based on audio characteristics to ensure reproducible but unique results
    audio_hash = hash(tuple(rms[::100])) % 10000  # Sample RMS every 100 frames for seed
    random.seed(audio_hash)
    np.random.seed(audio_hash)

    # Strategy 1: Content-aware sampling based on energy distribution
    rms_mean = np.mean(rms)
    rms_std = np.std(rms)

    # Find different energy zones
    high_energy = rms > (rms_mean + 0.5 * rms_std)
    medium_energy = (rms > (rms_mean - 0.3 * rms_std)) & (rms < (rms_mean + 0.5 * rms_std))
    low_energy = rms < (rms_mean - 0.3 * rms_std)

    energy_based_regions = []

    # Sample from each energy zone
    for energy_mask, zone_name in [
        (high_energy, "high"),
        (medium_energy, "medium"),
        (low_energy, "low"),
    ]:
        zone_regions = []
        i = 0
        while i < len(energy_mask):
            if energy_mask[i]:
                start = i
                while i < len(energy_mask) and energy_mask[i]:
                    i += 1
                end = i
                if end - start >= sample_frames:
                    # Add some randomness within the zone
                    possible_starts = list(range(start, max(start + 1, end - sample_frames + 1)))
                    if possible_starts:
                        random_start = random.choice(possible_starts)
                        zone_regions.append(times[random_start])
            else:
                i += 1

        # Randomly select from this zone
        if zone_regions:
            selected = random.sample(zone_regions, min(len(zone_regions), num_samples))
            energy_based_regions.extend(selected)

    # Strategy 2: Dynamic range sampling (find regions with good variation)
    dynamic_regions = []
    window_size = sample_frames
    for i in range(len(rms) - window_size + 1):
        window_rms = rms[i : i + window_size]
        dynamic_range = np.max(window_rms) - np.min(window_rms)
        dynamic_regions.append((times[i], dynamic_range))

    # Sort by dynamic range and add randomness
    dynamic_regions.sort(key=lambda x: x[1], reverse=True)
    # Take top candidates but with some randomness
    top_dynamic_count = min(len(dynamic_regions), num_samples * 3)
    top_candidates = dynamic_regions[:top_dynamic_count]
    dynamic_times = [
        region[0] for region in random.sample(top_candidates, min(len(top_candidates), num_samples * 2))
    ]

    # Strategy 3: Temporal diversity - avoid clustering by dividing track into sections
    num_sections = max(3, num_samples + 1)
    section_duration = audio_duration / num_sections
    section_based_regions = []

    for section in range(num_sections):
        section_start = section * section_duration
        section_end = min((section + 1) * section_duration, audio_duration - sample_duration)

        if section_end > section_start:
            # Find best region within this section
            start_idx = int(section_start * sr / hop_length)
            end_idx = int(section_end * sr / hop_length)

            if start_idx < end_idx and end_idx <= len(rms):
                section_rms = rms[start_idx:end_idx]
                if len(section_rms) >= sample_frames:
                    # Find peak energy within section
                    windowed_section_rms = []
                    for i in range(len(section_rms) - sample_frames + 1):
                        window_energy = np.mean(section_rms[i : i + sample_frames])
                        windowed_section_rms.append((start_idx + i, window_energy))

                    if windowed_section_rms:
                        # Add randomness - pick from top 50% of energy windows in this section
                        windowed_section_rms.sort(key=lambda x: x[1], reverse=True)
                        top_half = windowed_section_rms[: max(1, len(windowed_section_rms) // 2)]
                        best_idx, _ = random.choice(top_half)
                        section_based_regions.append(times[best_idx])

    # Combine all strategies
    all_candidates = energy_based_regions + dynamic_times + section_based_regions

    # Filter valid times and remove overlaps
    selected_regions = []
    used_intervals = []

    for candidate_time in all_candidates:
        if candidate_time + sample_duration > audio_duration:
            continue

        # Check for overlap with already selected regions
        overlap = any(
            not (candidate_time + sample_duration <= used_start or candidate_time >= used_end)
            for used_start, used_end in used_intervals
        )

        if not overlap:
            selected_regions.append(candidate_time)
            used_intervals.append((candidate_time, candidate_time + sample_duration))

            if len(selected_regions) >= num_samples:
                break

    return np.array(selected_regions)


def find_content_regions(
    audio, sr, sample_duration, num_samples, min_energy=0.005, hop_length=512, frame_length=2048
):
    """Find regions with actual audio content for vocal tracks."""
    # Calculate RMS energy
    rms = librosa.feature.rms(y=audio, hop_length=hop_length, frame_length=frame_length)[0]
    times = librosa.frames_to_time(np.arange(len(rms)), sr=sr, hop_length=hop_length)

    # Find regions above minimum energy threshold
    content_mask = rms > min_energy
    min_content_frames = int(sample_duration * sr / hop_length)  # Frames needed for full sample

    content_regions = []
    i = 0
    while i < len(content_mask):
        if content_mask[i]:
            start = i
            while i < len(content_mask) and content_mask[i]:
                i += 1
            end = i
            if end - start >= min_content_frames:
                start_time = times[start]
                end_time = times[end - 1] if end - 1 < len(times) else times[-1]
                # Use start of content region as sample point
                content_regions.append(start_time)
        else:
            i += 1

    # Select up to num_samples regions, spread across the track
    if len(content_regions) >= num_samples:
        indices = np.linspace(0, len(content_regions) - 1, num_samples, dtype=int)
        selected_times = np.array([content_regions[i] for i in indices])
    else:
        selected_times = np.array(content_regions)

    return selected_times


def apply_fade(audio, sr, fade_duration=0.1):
    """Apply fade-in and fade-out to audio segment."""
    fade_samples = int(fade_duration * sr)

    if len(audio.shape) == 1:  # Mono
        # Fade in
        audio[:fade_samples] *= np.linspace(0, 1, fade_samples)
        # Fade out
        audio[-fade_samples:] *= np.linspace(1, 0, fade_samples)
    else:  # Stereo (channels x samples)
        fade_in = np.linspace(0, 1, fade_samples)
        fade_out = np.linspace(1, 0, fade_samples)

        # Apply fade to each channel
        for channel in range(audio.shape[0]):
            audio[channel, :fade_samples] *= fade_in
            audio[channel, -fade_samples:] *= fade_out

    return audio


def validate_sample_volume(audio_segment, sr, min_rms_threshold=0.008):
    """Check if audio segment meets minimum volume threshold."""
    # Calculate RMS for each channel
    if len(audio_segment.shape) == 1:
        rms = np.sqrt(np.mean(audio_segment**2))
    else:
        rms_values = []
        for ch in range(audio_segment.shape[0]):
            rms = np.sqrt(np.mean(audio_segment[ch] ** 2))
            rms_values.append(rms)
        rms = np.mean(rms_values)

    return rms >= min_rms_threshold, rms


def get_sample_times_only(audio_file, sample_duration=7.0, num_samples=2):
    """Get sample start times without extracting audio - for visualization use."""
    # Load audio
    audio, sr = librosa.load(audio_file, sr=None, mono=False)

    # Ensure stereo format for consistency
    if len(audio.shape) == 1:
        audio = np.stack([audio, audio])  # Convert mono to stereo

    start_times = []

    # Apply content detection to all tracks (not just vocals)
    candidate_times = find_content_regions(audio[0], sr, sample_duration, num_samples)
    start_times = list(candidate_times)

    # Apply fallback if needed
    if len(start_times) < num_samples:
        loudest_times = find_loudest_regions(
            audio[0], sr, sample_duration, num_samples * 2
        )  # Get extra candidates

        # Add loudest regions that don't overlap with existing ones
        used_regions = [(t, t + sample_duration) for t in start_times]

        for loud_time in loudest_times:
            if len(start_times) >= num_samples:
                break

            # Check for overlap
            overlap = any(
                not (loud_time + sample_duration <= used_start or loud_time >= used_end)
                for used_start, used_end in used_regions
            )

            if not overlap:
                start_times.append(loud_time)
                used_regions.append((loud_time, loud_time + sample_duration))

    # Ensure we have exactly num_samples
    start_times = np.array(start_times[:num_samples])
    return start_times


def find_loudest_regions(audio, sr, sample_duration, num_samples, hop_length=512, frame_length=2048):
    """Find the loudest regions in the audio for sampling."""
    # Calculate RMS energy
    rms = librosa.feature.rms(y=audio, hop_length=hop_length, frame_length=frame_length)[0]
    times = librosa.frames_to_time(np.arange(len(rms)), sr=sr, hop_length=hop_length)

    # Calculate windowed RMS for potential sample regions
    sample_frames = int(sample_duration * sr / hop_length)
    windowed_rms = []
    valid_start_times = []

    for i in range(len(rms) - sample_frames + 1):
        window_rms = np.mean(rms[i : i + sample_frames])
        windowed_rms.append(window_rms)
        valid_start_times.append(times[i])

    if not windowed_rms:
        return np.array([])

    # Sort by RMS energy (loudest first)
    sorted_indices = np.argsort(windowed_rms)[::-1]

    # Select non-overlapping regions
    selected_times = []
    used_regions = []

    for idx in sorted_indices:
        start_time = valid_start_times[idx]
        end_time = start_time + sample_duration

        # Check if this region overlaps with already selected regions
        overlap = any(
            not (end_time <= used_start or start_time >= used_end)
            for used_start, used_end in used_regions
        )

        if not overlap:
            selected_times.append(start_time)
            used_regions.append((start_time, end_time))

            if len(selected_times) >= num_samples:
                break

    return np.array(selected_times)


def extract_samples(
    audio_file,
    output_dir,
    sample_duration=7.0,
    num_samples=2,
    fade_duration=0.1,
    min_rms_threshold=0.008,
):
    """Extract clean audio samples from a stem file with volume validation."""
    print(f"Processing: {os.path.basename(audio_file)}")

    # Load audio
    audio, sr = librosa.load(audio_file, sr=None, mono=False)

    # Ensure stereo format for consistency
    if len(audio.shape) == 1:
        audio = np.stack([audio, audio])  # Convert mono to stereo

    audio_duration = audio.shape[1] / sr

    if audio_duration < sample_duration:
        print(
            f"Warning: Audio file is shorter ({audio_duration:.1f}s) than requested sample duration ({sample_duration}s)"
        )
        return

    # Strategy 1: Try vocal content detection for vocal tracks
    stem_name = Path(audio_file).stem.lower()
    is_vocal_track = "vocal" in stem_name or "voice" in stem_name

    start_times = []

    # Apply content detection to all tracks (not just vocals)
    candidate_times = find_content_regions(audio[0], sr, sample_duration, num_samples)
    start_times = list(candidate_times)

    if len(start_times) < num_samples:
        print(f"  Warning: Only {len(start_times)} content regions found.")

    # Strategy 2: If we don't have enough samples or they might be too quiet,
    # find the loudest regions as backup
    if len(start_times) < num_samples:
        print(f"  Using loudest regions to fill remaining samples.")
        loudest_times = find_loudest_regions(
            audio[0], sr, sample_duration, num_samples * 2
        )  # Get extra candidates

        # Add loudest regions that don't overlap with existing ones
        used_regions = [(t, t + sample_duration) for t in start_times]

        for loud_time in loudest_times:
            if len(start_times) >= num_samples:
                break

            # Check for overlap
            overlap = any(
                not (loud_time + sample_duration <= used_start or loud_time >= used_end)
                for used_start, used_end in used_regions
            )

            if not overlap:
                start_times.append(loud_time)
                used_regions.append((loud_time, loud_time + sample_duration))

    # Ensure we have exactly num_samples
    start_times = np.array(start_times[:num_samples])

    # Extract and save samples with volume validation
    stem_name = Path(audio_file).stem
    samples_extracted = 0

    # Get loudest regions as backup candidates
    all_loudest_times = find_loudest_regions(audio[0], sr, sample_duration, num_samples * 3)
    backup_idx = 0

    for i, start_time in enumerate(start_times):
        max_attempts = 5  # Try up to 5 different regions
        attempt = 0

        while attempt < max_attempts:
            current_start_time = (
                start_time
                if attempt == 0
                else (
                    all_loudest_times[backup_idx] if backup_idx < len(all_loudest_times) else start_time
                )
            )

            start_sample = int(current_start_time * sr)
            end_sample = int((current_start_time + sample_duration) * sr)

            # Extract audio segment
            audio_segment = audio[:, start_sample:end_sample].copy()

            # Validate volume before applying fade
            is_loud_enough, rms_value = validate_sample_volume(audio_segment, sr, min_rms_threshold)

            if is_loud_enough or attempt == max_attempts - 1:  # Accept if loud enough or last attempt
                # Apply fade-in/out for smooth transitions
                audio_segment = apply_fade(audio_segment, sr, fade_duration)

                # Generate filename with start and end times
                start_time_str = f"{current_start_time:.1f}s"
                end_time_str = f"{current_start_time + sample_duration:.1f}s"
                output_filename = f"{stem_name}_sample_{i+1:02d}_{start_time_str}-{end_time_str}.wav"
                output_path = output_dir / output_filename

                # Save audio (transpose for soundfile: samples x channels)
                sf.write(output_path, audio_segment.T, sr)

                samples_extracted += 1
                volume_status = "✓" if is_loud_enough else "⚠️ (quiet)"
                print(
                    f"  {volume_status} Extracted sample {i+1}: {current_start_time:.1f}s - {current_start_time + sample_duration:.1f}s (RMS: {rms_value:.6f})"
                )
                break
            else:
                print(
                    f"    Sample {i+1} attempt {attempt+1} too quiet (RMS: {rms_value:.6f}), trying next loudest region..."
                )
                backup_idx += 1
                if backup_idx >= len(all_loudest_times):
                    break  # No more backup regions

            attempt += 1

    return samples_extracted


def main():
    parser = argparse.ArgumentParser(description="Extract clean audio samples from music stems")
    parser.add_argument(
        "--stems-dir", default="now, here, nothing Stems", help="Directory containing stem files"
    )
    parser.add_argument(
        "--output-dir", default="extracted_samples", help="Output directory for extracted samples"
    )
    parser.add_argument(
        "--sample-duration",
        type=float,
        default=7.0,
        help="Duration of each sample in seconds (default: 7.0)",
    )
    parser.add_argument(
        "--num-samples", type=int, default=2, help="Number of samples to extract per stem (default: 2)"
    )
    parser.add_argument(
        "--fade-duration", type=float, default=0.1, help="Fade in/out duration in seconds (default: 0.1)"
    )

    args = parser.parse_args()

    # Setup paths
    script_dir = Path(__file__).parent
    stems_dir = script_dir / args.stems_dir
    output_base_dir = script_dir / args.output_dir

    if not stems_dir.exists():
        print(f"Error: Stems directory not found: {stems_dir}")
        return

    # Find all audio files
    audio_extensions = {".mp3", ".wav", ".flac", ".m4a", ".aac"}
    stem_files = [f for f in stems_dir.iterdir() if f.suffix.lower() in audio_extensions]

    if not stem_files:
        print(f"No audio files found in {stems_dir}")
        return

    print(f"Found {len(stem_files)} stem files")
    print(f"Sample duration: {args.sample_duration}s")
    print(f"Samples per stem: {args.num_samples}")
    print(f"Fade duration: {args.fade_duration}s")
    print(f"Output directory: {output_base_dir}")
    print("-" * 50)

    # Create output directory structure
    total_samples = 0

    for stem_file in sorted(stem_files):
        # Create subfolder for each stem
        stem_name = stem_file.stem
        stem_output_dir = output_base_dir / stem_name
        stem_output_dir.mkdir(parents=True, exist_ok=True)

        try:
            samples_extracted = extract_samples(
                stem_file,
                stem_output_dir,
                sample_duration=args.sample_duration,
                num_samples=args.num_samples,
                fade_duration=args.fade_duration,
            )
            total_samples += samples_extracted or 0

        except Exception as e:
            print(f"Error processing {stem_file.name}: {e}")

    print("-" * 50)
    print(f"✅ Extraction complete! Total samples created: {total_samples}")
    print(f"Output saved to: {output_base_dir}")


if __name__ == "__main__":
    main()
