"""Spectral coherence using MFCC embeddings."""
import os
import librosa
import numpy as np
from typing import List, Tuple, Dict
from sklearn.metrics.pairwise import cosine_similarity
import logging

logger = logging.getLogger(__name__)

class SpectralCoherence:
    """Computes and compares spectral embeddings using MFCCs."""
    
    def __init__(self, n_mfcc: int = 13, n_fft: int = 2048, hop_length: int = 512):
        self.n_mfcc = n_mfcc
        self.n_fft = n_fft
        self.hop_length = hop_length
    
    def compute_embedding(self, audio_path: str, duration: float = 30.0) -> np.ndarray:
        """
        Compute MFCC-based spectral embedding.
        
        Returns:
            Normalized embedding vector (n_mfcc,)
        """
        try:
            y, sr = librosa.load(audio_path, duration=duration)
            
            # Compute MFCCs
            mfcc = librosa.feature.mfcc(
                y=y, sr=sr, n_mfcc=self.n_mfcc,
                n_fft=self.n_fft, hop_length=self.hop_length
            )
            
            # Get mean across time (spectral profile)
            embedding = np.mean(mfcc, axis=1)
            
            # Normalize
            norm = np.linalg.norm(embedding)
            if norm > 0:
                embedding = embedding / norm
            
            return embedding
            
        except Exception as e:
            logger.error(f"Error computing embedding for {audio_path}: {e}")
            return np.zeros(self.n_mfcc)
    
    def compute_similarity(self, emb1: np.ndarray, emb2: np.ndarray) -> float:
        """Compute cosine similarity between two embeddings (0.0-1.0)."""
        return float(cosine_similarity([emb1], [emb2])[0][0])
    
    def find_similar_samples(
        self, 
        target_path: str, 
        library_embeddings: Dict[str, np.ndarray],
        top_k: int = 10,
        min_similarity: float = 0.7
    ) -> List[Tuple[str, float]]:
        """
        Find most similar samples to target.
        
        Returns:
            List of (path, similarity_score) sorted by similarity
        """
        target_emb = self.compute_embedding(target_path)
        
        similarities = []
        for path, emb in library_embeddings.items():
            if path == target_path:
                continue
            sim = self.compute_similarity(target_emb, emb)
            if sim >= min_similarity:
                similarities.append((path, sim))
        
        # Sort by similarity descending
        similarities.sort(key=lambda x: x[1], reverse=True)
        
        return similarities[:top_k]
    
    def compute_all_embeddings(
        self, 
        library_path: str,
        progress_callback=None
    ) -> Dict[str, np.ndarray]:
        """
        Compute embeddings for all samples in library.
        
        Returns:
            Dict mapping {path: embedding_vector}
        """
        embeddings = {}
        
        audio_exts = ('.wav', '.aif', '.aiff', '.mp3', '.flac')
        audio_files = []
        
        for root, dirs, files in os.walk(library_path):
            for f in files:
                if f.lower().endswith(audio_exts):
                    audio_files.append(os.path.join(root, f))
        
        total = len(audio_files)
        
        for i, path in enumerate(audio_files):
            emb = self.compute_embedding(path)
            embeddings[path] = emb
            
            if progress_callback:
                progress_callback(i + 1, total)
        
        return embeddings
    
    def get_coherence_score(self, sample_paths: List[str]) -> float:
        """Compute average pairwise coherence for a set of samples."""
        if len(sample_paths) < 2:
            return 1.0
        
        embeddings = [self.compute_embedding(p) for p in sample_paths]
        
        total_sim = 0.0
        count = 0
        
        for i in range(len(embeddings)):
            for j in range(i + 1, len(embeddings)):
                sim = self.compute_similarity(embeddings[i], embeddings[j])
                total_sim += sim
                count += 1
        
        return total_sim / count if count > 0 else 0.0


# Convenience function
def get_sample_similarity(path1: str, path2: str) -> float:
    """Quick similarity check between two samples."""
    coherence = SpectralCoherence()
    emb1 = coherence.compute_embedding(path1)
    emb2 = coherence.compute_embedding(path2)
    return coherence.compute_similarity(emb1, emb2)
