"""
Smart Sample Selector Engine

Intelligently selects samples based on REAL compatibility analysis.
Uses coherence validator to ensure professional-grade sample compatibility.
"""

import os
import glob
from typing import List, Dict, Any, Optional, Tuple
import statistics


class SmartSampleSelector:
    """
    Intelligent sample selection engine that analyzes samples for compatibility.
    
    Features:
    - Scans folders for audio samples
    - Analyzes samples for spectral features
    - Finds best anchor sample (median energy/brightness)
    - Builds coherent kits with similarity matching
    - Validates coherence against threshold
    - Distributes samples across song sections
    """
    
    def __init__(self, validator):
        """
        Initialize SmartSampleSelector with coherence validator.
        
        Args:
            validator: CoherenceValidator instance for sample analysis
        """
        self.validator = validator
        self._analysis_cache = {}
    
    def scan_folder(self, folder_path: str) -> List[str]:
        """
        Scan folder and return all WAV files.
        
        Args:
            folder_path: Path to folder containing samples
            
        Returns:
            List of WAV file paths
        """
        if not os.path.exists(folder_path):
            return []
        
        # Search for WAV files (case-insensitive)
        pattern = os.path.join(folder_path, "**", "*.wav")
        wav_files = glob.glob(pattern, recursive=True)
        
        # Also check for .WAV uppercase
        pattern_upper = os.path.join(folder_path, "**", "*.WAV")
        wav_files_upper = glob.glob(pattern_upper, recursive=True)
        
        # Combine and remove duplicates
        all_files = list(set(wav_files + wav_files_upper))
        
        # Sort for consistent ordering
        all_files.sort()
        
        return all_files
    
    def _analyze_sample(self, sample_path: str) -> Optional[Dict[str, Any]]:
        """
        Analyze a single sample using validator.
        
        Args:
            sample_path: Path to audio file
            
        Returns:
            Analysis dict with features or None if analysis fails
        """
        if sample_path in self._analysis_cache:
            return self._analysis_cache[sample_path]
        
        try:
            features = self.validator.analyze_sample(sample_path)
            if features:
                self._analysis_cache[sample_path] = features
            return features
        except Exception:
            return None
    
    def _find_best_anchor(self, samples_data: List[Tuple[str, Dict[str, Any]]]) -> Tuple[str, Dict[str, Any]]:
        """
        Find the best anchor sample based on median energy and brightness.
        
        Args:
            samples_data: List of (path, features) tuples
            
        Returns:
            (anchor_path, anchor_features) tuple
        """
        if not samples_data:
            return None, None
        
        # Extract RMS energy values
        energies = []
        for path, features in samples_data:
            energy = features.get("rms_energy", 0.5)
            energies.append((path, features, energy))
        
        # Find median energy
        energy_values = [e[2] for e in energies]
        median_energy = statistics.median(energy_values)
        
        # Find sample closest to median energy
        closest_to_median = min(energies, key=lambda x: abs(x[2] - median_energy))
        
        return closest_to_median[0], closest_to_median[1]
    
    def _calculate_similarity(self, features1: Dict[str, Any], features2: Dict[str, Any]) -> float:
        """
        Calculate similarity score between two samples.
        
        Args:
            features1: Features dict for first sample
            features2: Features dict for second sample
            
        Returns:
            Similarity score (0.0 - 1.0)
        """
        # Use validator's coherence calculation
        return self.validator.calculate_coherence(features1, features2)
    
    def select_coherent_kit(
        self, 
        folder_path: str, 
        target_count: int = 12, 
        threshold: float = 0.90
    ) -> Dict[str, Any]:
        """
        Main method: Select a coherent kit of samples.
        
        Process:
        1. Scan folder for samples
        2. Analyze all samples
        3. Find best anchor (median energy/brightness)
        4. Find N-1 samples most similar to anchor
        5. Validate coherence >= threshold
        6. Return kit + metadata
        
        Args:
            folder_path: Path to folder containing samples
            target_count: Number of samples to select (default: 12)
            threshold: Minimum coherence threshold (default: 0.90)
            
        Returns:
            Dict with kit info:
            - samples: List of selected sample paths
            - anchor: Anchor sample path
            - coherence_score: Overall coherence score
            - threshold_met: Boolean if threshold was met
            - rejected_count: Number of samples rejected
            - section_distribution: Distribution across sections
        """
        # Step 1: Scan folder
        all_samples = self.scan_folder(folder_path)
        
        if len(all_samples) < target_count:
            return {
                "success": False,
                "error": f"Not enough samples. Found {len(all_samples)}, need {target_count}",
                "samples": [],
                "anchor": None,
                "coherence_score": 0.0,
                "threshold_met": False
            }
        
        # Step 2: Analyze all samples
        analyzed_samples = []
        failed_analyses = []
        
        for sample_path in all_samples:
            features = self._analyze_sample(sample_path)
            if features:
                analyzed_samples.append((sample_path, features))
            else:
                failed_analyses.append(sample_path)
        
        if len(analyzed_samples) < target_count:
            return {
                "success": False,
                "error": f"Only {len(analyzed_samples)} samples analyzable ({len(failed_analyses)} failed)",
                "samples": [],
                "anchor": None,
                "coherence_score": 0.0,
                "threshold_met": False
            }
        
        # Step 3: Find best anchor (median energy/brightness)
        anchor_path, anchor_features = self._find_best_anchor(analyzed_samples)
        
        if not anchor_path:
            return {
                "success": False,
                "error": "Could not determine anchor sample",
                "samples": [],
                "anchor": None,
                "coherence_score": 0.0,
                "threshold_met": False
            }
        
        # Step 4: Find N-1 samples most similar to anchor
        similarities = []
        for sample_path, features in analyzed_samples:
            if sample_path == anchor_path:
                continue
            
            similarity = self._calculate_similarity(anchor_features, features)
            similarities.append((sample_path, features, similarity))
        
        # Sort by similarity (highest first)
        similarities.sort(key=lambda x: x[2], reverse=True)
        
        # Select top N-1 samples
        selected_count = min(target_count - 1, len(similarities))
        selected_similar = similarities[:selected_count]
        
        # Build final kit: anchor + selected samples
        kit_samples = [anchor_path] + [s[0] for s in selected_similar]
        kit_features = [anchor_features] + [s[1] for s in selected_similar]
        
        # Calculate individual similarity scores to anchor
        individual_scores = [1.0] + [s[2] for s in selected_similar]
        
        # Step 5: Validate coherence
        # Calculate pairwise coherence across all kit samples
        pairwise_scores = []
        for i in range(len(kit_features)):
            for j in range(i + 1, len(kit_features)):
                score = self._calculate_similarity(kit_features[i], kit_features[j])
                pairwise_scores.append(score)
        
        # Overall coherence is average of all pairwise scores
        if pairwise_scores:
            coherence_score = sum(pairwise_scores) / len(pairwise_scores)
        else:
            coherence_score = 1.0  # Single sample
        
        threshold_met = coherence_score >= threshold
        
        # Step 6: Return kit + metadata
        rejected_count = len(all_samples) - len(kit_samples)
        
        result = {
            "success": True,
            "samples": kit_samples,
            "anchor": anchor_path,
            "coherence_score": round(coherence_score, 3),
            "threshold": threshold,
            "threshold_met": threshold_met,
            "rejected_count": rejected_count,
            "total_analyzed": len(analyzed_samples),
            "failed_analyses": len(failed_analyses),
            "individual_scores": {
                os.path.basename(path): round(score, 3) 
                for path, score in zip(kit_samples, individual_scores)
            },
            "section_distribution": self.distribute_for_sections(kit_samples, section_count=8)
        }
        
        return result
    
    def distribute_for_sections(
        self, 
        kit_samples: List[str], 
        section_count: int = 8
    ) -> Dict[str, List[str]]:
        """
        Distribute samples across song sections.
        
        Strategy:
        - Intro: 1-2 samples (sparse)
        - Verse: 2-3 samples (building)
        - Chorus: 3-4 samples (full energy)
        - Bridge: 2-3 samples (variation)
        - Outro: 1-2 samples (winding down)
        
        Args:
            kit_samples: List of selected sample paths
            section_count: Number of sections to distribute to (default: 8)
            
        Returns:
            Dict mapping section names to lists of sample paths
        """
        if not kit_samples:
            return {}
        
        total_samples = len(kit_samples)
        
        # Standard 5-section structure: Intro, Verse, Chorus, Bridge, Outro
        # With variations: Intro, Verse1, Chorus1, Verse2, Chorus2, Bridge, Chorus3, Outro
        section_names = [
            "intro",
            "verse_1", 
            "chorus_1",
            "verse_2",
            "chorus_2", 
            "bridge",
            "chorus_3",
            "outro"
        ]
        
        # Use requested section count
        sections = section_names[:section_count]
        
        distribution = {}
        sample_idx = 0
        
        # Calculate samples per section with weighting
        # Chorus gets more samples, intro/outro get fewer
        section_weights = {
            "intro": 0.5,
            "verse_1": 0.8,
            "verse_2": 0.8,
            "chorus_1": 1.2,
            "chorus_2": 1.2,
            "chorus_3": 1.2,
            "bridge": 1.0,
            "outro": 0.5
        }
        
        # Calculate total weight
        total_weight = sum(section_weights.get(s, 1.0) for s in sections)
        
        # Distribute samples
        for section in sections:
            weight = section_weights.get(section, 1.0)
            # Calculate count based on weight
            base_count = total_samples / total_weight * weight
            
            # Ensure minimum 1 sample for non-empty sections
            count = max(1, int(round(base_count)))
            
            # Don't exceed remaining samples
            remaining = total_samples - sample_idx
            count = min(count, remaining)
            
            # For last section, use all remaining
            if section == sections[-1]:
                count = remaining
            
            if count > 0:
                # Select samples (with overlap between sections for continuity)
                if sample_idx == 0:
                    # First section: start from beginning
                    section_samples = kit_samples[:count]
                else:
                    # Other sections: overlap with previous section
                    overlap = min(1, len(distribution.get(sections[sections.index(section)-1], [])))
                    start_idx = max(0, sample_idx - overlap)
                    section_samples = kit_samples[start_idx:start_idx + count]
                
                distribution[section] = section_samples
                sample_idx += max(0, count - overlap) if sample_idx > 0 else count
            else:
                distribution[section] = []
        
        # Ensure all sections have at least the anchor if empty
        if kit_samples:
            anchor = kit_samples[0]
            for section in sections:
                if not distribution.get(section):
                    distribution[section] = [anchor]
        
        return distribution
    
    def clear_cache(self):
        """Clear analysis cache."""
        self._analysis_cache.clear()
    
    def get_cache_stats(self) -> Dict[str, int]:
        """
        Get cache statistics.
        
        Returns:
            Dict with cache info
        """
        return {
            "cached_samples": len(self._analysis_cache),
            "cache_size_mb": sum(
                len(str(v)) for v in self._analysis_cache.values()
            ) / (1024 * 1024)
        }
