"""
ExpansiveCoherenceValidator - Validates coherence for 12+ samples simultaneously.

This engine provides comprehensive coherence validation for large sample kits,
supporting up to and beyond 12 samples with efficient similarity calculations
and intelligent replacement suggestions.

Example usage:
    validator = ExpansiveCoherenceValidator(metadata_store)
    matrix = validator.calculate_similarity_matrix(samples)
    is_coherent = validator.validate_kit_coherence(samples, threshold=0.90)
    replacements = validator.suggest_replacements(samples, bad_indices, all_available)
"""

from typing import List, Dict, Tuple, Optional, Set
import math


class ExpansiveCoherenceValidator:
    """
    Validates coherence for large sample kits (12+ samples).
    
    This class provides methods to:
    - Calculate similarity matrices for all sample pairs
    - Validate overall kit coherence against thresholds
    - Identify incoherent sample pairs
    - Suggest better sample replacements
    - Validate transitions between song sections
    
    Attributes:
        metadata_store: Optional metadata store for sample lookups
        max_samples: Maximum number of samples to validate (default 24)
        feature_weights: Weights for different feature types in similarity calc
    """
    
    def __init__(self, metadata_store=None):
        """
        Initialize the ExpansiveCoherenceValidator.
        
        Args:
            metadata_store: Optional metadata store for sample feature lookups.
                          If provided, enables richer similarity calculations.
        """
        self.metadata_store = metadata_store
        self.max_samples = 24
        self.feature_weights = {
            'spectral_centroid': 0.25,
            'spectral_rolloff': 0.20,
            'zero_crossing_rate': 0.15,
            'rms_energy': 0.25,
            'bpm': 0.10,
            'key': 0.05
        }
    
    def _normalize_features(self, features: Dict) -> Dict:
        """
        Normalize feature values to 0.0-1.0 range for fair comparison.
        
        Args:
            features: Dictionary of raw feature values
            
        Returns:
            Dictionary of normalized feature values
        """
        normalized = {}
        
        # Spectral centroid: typically 0-8000 Hz
        if 'spectral_centroid' in features:
            normalized['spectral_centroid'] = min(1.0, features['spectral_centroid'] / 8000.0)
        
        # Spectral rolloff: typically 0-8000 Hz
        if 'spectral_rolloff' in features:
            normalized['spectral_rolloff'] = min(1.0, features['spectral_rolloff'] / 8000.0)
        
        # Zero crossing rate: typically 0-0.5
        if 'zero_crossing_rate' in features:
            normalized['zero_crossing_rate'] = min(1.0, features['zero_crossing_rate'] / 0.5)
        
        # RMS energy: typically 0-1.0 already
        if 'rms_energy' in features:
            normalized['rms_energy'] = features['rms_energy']
        
        # BPM: normalize around typical range 60-180
        if 'bpm' in features:
            bpm = features['bpm']
            if bpm > 0:
                normalized['bpm'] = max(0.0, min(1.0, (bpm - 60) / 120.0))
            else:
                normalized['bpm'] = 0.5  # Default middle value
        
        return normalized
    
    def _key_similarity(self, key1: str, key2: str) -> float:
        """
        Calculate similarity between two musical keys.
        
        Uses circle of fifths distance for harmonic similarity.
        
        Args:
            key1: First key (e.g., "Am", "C", "F#m")
            key2: Second key (e.g., "Am", "C", "F#m")
            
        Returns:
            Similarity score 0.0-1.0 where 1.0 = identical keys
        """
        if not key1 or not key2:
            return 0.5  # Unknown keys get neutral score
        
        key1 = key1.strip().lower()
        key2 = key2.strip().lower()
        
        if key1 == key2:
            return 1.0
        
        # Circle of fifths order for major keys
        major_circle = ['c', 'g', 'd', 'a', 'e', 'b', 'f#', 'db', 
                       'ab', 'eb', 'bb', 'f']
        
        # Circle of fifths order for minor keys
        minor_circle = ['am', 'em', 'bm', 'f#m', 'c#m', 'g#m', 'd#m', 'bbm',
                       'fm', 'cm', 'gm', 'dm']
        
        # Extract root and mode
        def parse_key(k):
            if k.endswith('m'):
                return k[:-1], 'minor'
            return k, 'major'
        
        root1, mode1 = parse_key(key1)
        root2, mode2 = parse_key(key2)
        
        # Same mode - use appropriate circle
        if mode1 == mode2:
            circle = major_circle if mode1 == 'major' else minor_circle
            if root1 in circle and root2 in circle:
                idx1 = circle.index(root1)
                idx2 = circle.index(root2)
                distance = min(abs(idx1 - idx2), 12 - abs(idx1 - idx2))
                return 1.0 - (distance / 6.0)  # Max distance is 6 (tritone)
        
        # Different modes - check relative major/minor
        relative_majors = {'am': 'c', 'em': 'g', 'bm': 'd', 'f#m': 'a',
                          'c#m': 'e', 'g#m': 'b', 'd#m': 'f#', 'bbm': 'db',
                          'fm': 'ab', 'cm': 'eb', 'gm': 'bb', 'dm': 'f'}
        relative_minors = {v: k for k, v in relative_majors.items()}
        
        if mode1 == 'minor' and relative_majors.get(key1) == key2:
            return 0.85  # Relative major
        if mode2 == 'minor' and relative_majors.get(key2) == key1:
            return 0.85  # Relative major
        
        return 0.5  # Default moderate similarity for different keys
    
    def _euclidean_distance(self, features1: Dict, features2: Dict) -> float:
        """
        Calculate weighted Euclidean distance between two feature sets.
        
        Args:
            features1: First sample's normalized features
            features2: Second sample's normalized features
            
        Returns:
            Distance value (0.0 = identical, higher = more different)
        """
        distance = 0.0
        total_weight = 0.0
        
        for feature, weight in self.feature_weights.items():
            if feature == 'key':
                # Key similarity handled separately
                continue
            
            val1 = features1.get(feature, 0.5)
            val2 = features2.get(feature, 0.5)
            
            diff = val1 - val2
            distance += weight * (diff * diff)
            total_weight += weight
        
        if total_weight > 0:
            distance = math.sqrt(distance / total_weight)
        
        return distance
    
    def calculate_similarity_matrix(self, samples: List[Dict]) -> List[List[float]]:
        """
        Calculate a complete similarity matrix for all sample pairs.
        
        Creates an NxN matrix where each cell [i][j] represents the
        similarity between sample i and sample j. Diagonal is always 1.0.
        
        Args:
            samples: List of sample dictionaries with 'features', 'bpm', 'key'
                    Expected format: {"path": str, "features": Dict, "bpm": float, "key": str}
        
        Returns:
            NxN similarity matrix where values range 0.0-1.0
            
        Raises:
            ValueError: If samples list is empty or exceeds max_samples
        """
        if not samples:
            raise ValueError("Cannot calculate similarity matrix for empty samples list")
        
        if len(samples) > self.max_samples:
            raise ValueError(f"Too many samples: {len(samples)} (max: {self.max_samples})")
        
        n = len(samples)
        matrix = [[0.0 for _ in range(n)] for _ in range(n)]
        
        # Pre-normalize all features
        normalized_features = []
        for sample in samples:
            features = sample.get('features', {})
            features['bpm'] = sample.get('bpm', 0)
            features['key'] = sample.get('key', '')
            normalized = self._normalize_features(features)
            normalized_features.append(normalized)
        
        # Calculate pairwise similarities
        for i in range(n):
            matrix[i][i] = 1.0  # Self-similarity is perfect
            
            for j in range(i + 1, n):
                # Calculate Euclidean distance
                distance = self._euclidean_distance(
                    normalized_features[i], 
                    normalized_features[j]
                )
                
                # Add key similarity component
                key_sim = self._key_similarity(
                    samples[i].get('key', ''),
                    samples[j].get('key', '')
                )
                
                # Combine distance and key similarity
                # Convert distance to similarity (1.0 - distance)
                feature_sim = 1.0 - min(1.0, distance)
                
                # Weighted combination
                combined_sim = (0.85 * feature_sim) + (0.15 * key_sim)
                
                # Symmetric matrix
                matrix[i][j] = combined_sim
                matrix[j][i] = combined_sim
        
        return matrix
    
    def validate_kit_coherence(self, samples: List[Dict], threshold: float = 0.90) -> bool:
        """
        Validate if a complete kit meets coherence requirements.
        
        Calculates the average coherence across all sample pairs and
        compares against the threshold.
        
        Args:
            samples: List of sample dictionaries
            threshold: Minimum average coherence required (default 0.90)
        
        Returns:
            True if average coherence >= threshold, False otherwise
        """
        if len(samples) < 2:
            return True  # Single sample is always coherent
        
        score = self.get_coherence_score(samples)
        return score >= threshold
    
    def get_coherence_score(self, samples: List[Dict]) -> float:
        """
        Calculate the overall coherence score for a kit.
        
        Computes the average of all pairwise similarities in the matrix.
        
        Args:
            samples: List of sample dictionaries
        
        Returns:
            Average coherence score from 0.0 to 1.0
        """
        if len(samples) < 2:
            return 1.0  # Perfect coherence for single sample
        
        matrix = self.calculate_similarity_matrix(samples)
        
        # Calculate average of upper triangle (excluding diagonal)
        n = len(samples)
        total = 0.0
        count = 0
        
        for i in range(n):
            for j in range(i + 1, n):
                total += matrix[i][j]
                count += 1
        
        return total / count if count > 0 else 1.0
    
    def find_incoherent_pairs(self, matrix: List[List[float]], 
                             threshold: float = 0.85) -> List[Tuple[int, int]]:
        """
        Identify pairs of samples with low coherence.
        
        Scans the similarity matrix and returns indices of all pairs
        with similarity below the threshold.
        
        Args:
            matrix: NxN similarity matrix from calculate_similarity_matrix
            threshold: Minimum acceptable similarity (default 0.85)
        
        Returns:
            List of (i, j) tuples where matrix[i][j] < threshold
        """
        incoherent_pairs = []
        n = len(matrix)
        
        for i in range(n):
            for j in range(i + 1, n):
                if matrix[i][j] < threshold:
                    incoherent_pairs.append((i, j))
        
        return incoherent_pairs
    
    def _calculate_sample_coherence(self, matrix: List[List[float]], 
                                    index: int) -> float:
        """
        Calculate average coherence for a single sample against all others.
        
        Args:
            matrix: Similarity matrix
            index: Sample index to evaluate
            
        Returns:
            Average coherence score for the sample
        """
        n = len(matrix)
        total = 0.0
        count = 0
        
        for j in range(n):
            if j != index:
                total += matrix[index][j]
                count += 1
        
        return total / count if count > 0 else 0.0
    
    def suggest_replacements(self, samples: List[Dict], 
                            bad_indices: List[int],
                            all_available: List[Dict]) -> List[Dict]:
        """
        Suggest better sample replacements for incoherent samples.
        
        For each sample index in bad_indices, searches all_available
        for the most coherent alternative that maintains kit consistency.
        
        Args:
            samples: Current kit samples (including incoherent ones)
            bad_indices: Indices of samples needing replacement
            all_available: Pool of all available samples to choose from
        
        Returns:
            List of suggested replacement samples, one per bad_index
        """
        if not bad_indices or not all_available:
            return []
        
        suggestions = []
        
        for bad_idx in bad_indices:
            # Get current kit excluding the bad sample
            good_samples = [s for i, s in enumerate(samples) if i != bad_idx]
            
            if not good_samples:
                # If all samples are bad, pick any sample as anchor
                best_replacement = all_available[0]
                best_score = 0.0
                
                for candidate in all_available:
                    if candidate.get('path') != samples[bad_idx].get('path'):
                        # Calculate solo coherence
                        test_kit = [samples[bad_idx], candidate]
                        score = self.get_coherence_score(test_kit)
                        if score > best_score:
                            best_score = score
                            best_replacement = candidate
                
                suggestions.append(best_replacement)
                continue
            
            # Find best replacement that maximizes coherence with good samples
            best_replacement = None
            best_coherence = 0.0
            
            for candidate in all_available:
                # Skip if candidate is already in the kit
                candidate_path = candidate.get('path')
                if any(s.get('path') == candidate_path for s in samples):
                    continue
                
                # Calculate coherence with good samples
                test_kit = good_samples + [candidate]
                coherence = self.get_coherence_score(test_kit)
                
                if coherence > best_coherence:
                    best_coherence = coherence
                    best_replacement = candidate
            
            if best_replacement:
                suggestions.append(best_replacement)
            else:
                # Fallback: return first available
                suggestions.append(all_available[0])
        
        return suggestions
    
    def validate_section_transition(self, section1_samples: List[Dict],
                                    section2_samples: List[Dict]) -> float:
        """
        Validate coherence between two song sections.
        
        Calculates cross-section similarity to ensure smooth transitions.
        Useful for verifying that verse and chorus samples work together.
        
        Args:
            section1_samples: Samples used in first section (e.g., verse)
            section2_samples: Samples used in second section (e.g., chorus)
        
        Returns:
            Transition coherence score from 0.0 to 1.0
        """
        if not section1_samples or not section2_samples:
            return 1.0  # Empty section has no transition issues
        
        # Calculate all pairwise similarities between sections
        total_similarity = 0.0
        count = 0
        
        # Pre-normalize features for efficiency
        norm_section1 = []
        for sample in section1_samples:
            features = sample.get('features', {})
            features['bpm'] = sample.get('bpm', 0)
            features['key'] = sample.get('key', '')
            norm_section1.append(self._normalize_features(features))
        
        norm_section2 = []
        for sample in section2_samples:
            features = sample.get('features', {})
            features['bpm'] = sample.get('bpm', 0)
            features['key'] = sample.get('key', '')
            norm_section2.append(self._normalize_features(features))
        
        # Cross-section similarities
        for i, s1 in enumerate(section1_samples):
            for j, s2 in enumerate(section2_samples):
                distance = self._euclidean_distance(norm_section1[i], norm_section2[j])
                key_sim = self._key_similarity(s1.get('key', ''), s2.get('key', ''))
                feature_sim = 1.0 - min(1.0, distance)
                combined_sim = (0.85 * feature_sim) + (0.15 * key_sim)
                
                total_similarity += combined_sim
                count += 1
        
        return total_similarity / count if count > 0 else 1.0
    
    def get_coherence_breakdown(self, samples: List[Dict]) -> Dict:
        """
        Get detailed coherence analysis for a kit.
        
        Returns comprehensive breakdown including:
        - Overall coherence score
        - Per-sample coherence scores
        - Incoherent pairs identified
        - Weakest samples
        
        Args:
            samples: List of sample dictionaries
        
        Returns:
            Dictionary with detailed coherence analysis
        """
        if len(samples) < 2:
            return {
                'overall_coherence': 1.0,
                'per_sample_scores': [1.0] if samples else [],
                'incoherent_pairs': [],
                'weakest_sample_index': None,
                'weakest_sample_path': None
            }
        
        matrix = self.calculate_similarity_matrix(samples)
        overall = self.get_coherence_score(samples)
        incoherent = self.find_incoherent_pairs(matrix, threshold=0.85)
        
        # Calculate per-sample scores
        per_sample = []
        for i in range(len(samples)):
            score = self._calculate_sample_coherence(matrix, i)
            per_sample.append({
                'index': i,
                'path': samples[i].get('path', 'unknown'),
                'coherence': score
            })
        
        # Find weakest sample
        weakest = min(per_sample, key=lambda x: x['coherence']) if per_sample else None
        
        return {
            'overall_coherence': overall,
            'per_sample_scores': per_sample,
            'incoherent_pairs': incoherent,
            'weakest_sample_index': weakest['index'] if weakest else None,
            'weakest_sample_path': weakest['path'] if weakest else None,
            'is_valid': overall >= 0.90
        }
