"""spectral_engine.py - Análisis espectral para selección por similitud tímbrica y síntesis granular."""
import numpy as np
import logging
import json
import os
import wave
import struct
from typing import Dict, List, Optional, Tuple, Any
from dataclasses import dataclass, asdict
from pathlib import Path
from collections import defaultdict

logger = logging.getLogger("SpectralEngine")

LIBROSA_AVAILABLE = False
try:
    import librosa
    LIBROSA_AVAILABLE = True
    logger.info("[SPECTRAL] librosa disponible para síntesis granular")
except ImportError:
    logger.warning("[SPECTRAL] librosa no disponible, síntesis granular limitada")

@dataclass
class SpectralProfile:
    """Perfil espectral de un sample de audio."""
    path: str
    centroid_mean: float
    centroid_std: float
    rolloff_85: float
    flux_mean: float
    mfcc: List[float]
    rms: float
    spectral_flatness: float
    duration: float
    genre_hints: List[str]

class SpectralEngine:
    def __init__(self):
        self._cache: Dict[str, SpectralProfile] = {}
        self._librosa = None
        self._np = np
        self._init_librosa()
        self._load_cached_index()
    
    def _init_librosa(self):
        try:
            import librosa
            self._librosa = librosa
            logger.info("[SPECTRAL] librosa disponible")
        except ImportError:
            logger.warning("[SPECTRAL] librosa no disponible, usando análisis básico")
    
    def _load_cached_index(self):
        INDEX_PATH = os.path.join(os.path.dirname(__file__), "spectral_index.json")
        if os.path.exists(INDEX_PATH):
            try:
                with open(INDEX_PATH) as fh:
                    data = json.load(fh)
                for path, d in data.items():
                    self._cache[path] = SpectralProfile(
                        path=path,
                        centroid_mean=d.get("centroid", 0.0),
                        centroid_std=d.get("centroid_std", 100.0),
                        rolloff_85=d.get("rolloff", 0.0),
                        flux_mean=d.get("flux", 0.1),
                        mfcc=d.get("mfcc", [0.0]*13),
                        rms=d.get("rms", 0.3),
                        spectral_flatness=d.get("flatness", 0.5),
                        duration=d.get("duration", 2.0),
                        genre_hints=d.get("genre_hints", ["unknown"])
                    )
                logger.info(f"[SPECTRAL] Índice cargado: {len(self._cache)} samples")
            except Exception as e:
                logger.warning(f"[SPECTRAL] Error cargando índice: {e}")
    
    def analyze(self, path: str) -> Optional[SpectralProfile]:
        if path in self._cache:
            return self._cache[path]
        
        if self._librosa and os.path.exists(path):
            profile = self._analyze_librosa(path)
        else:
            profile = self._analyze_basic(path)
        if profile:
            self._cache[path] = profile
        return profile
    
    def similarity(self, a: SpectralProfile, b: SpectralProfile) -> float:
        """Retorna similitud 0.0-1.0 entre dos perfiles espectrales."""
        if not a or not b:
            return 0.0
        centroid_sim = 1.0 - min(abs(a.centroid_mean - b.centroid_mean) / max(a.centroid_mean + 1, 1), 1.0)
        rolloff_sim  = 1.0 - min(abs(a.rolloff_85 - b.rolloff_85) / max(a.rolloff_85 + 1, 1), 1.0)
        flux_sim     = 1.0 - min(abs(a.flux_mean - b.flux_mean) / max(a.flux_mean + 1, 1), 1.0)
        mfcc_sim = 0.0
        if a.mfcc and b.mfcc and len(a.mfcc) == len(b.mfcc):
            diff = sum((x-y)**2 for x,y in zip(a.mfcc, b.mfcc))
            mfcc_sim = 1.0 / (1.0 + diff**0.5)
        return 0.35*centroid_sim + 0.25*rolloff_sim + 0.15*flux_sim + 0.25*mfcc_sim

    def find_most_similar(self, reference_path: str, candidates: List[str], top_n: int = 5) -> List[Tuple[str, float]]:
        """Dado un sample de referencia, retorna los N candidatos más similares."""
        ref = self.analyze(reference_path)
        if not ref:
            return []
        scored = []
        for c in candidates:
            prof = self.analyze(c)
            if prof:
                score = self.similarity(ref, prof)
                scored.append((c, score))
        return sorted(scored, key=lambda x: x[1], reverse=True)[:top_n]

    def _analyze_librosa(self, path: str) -> Optional[SpectralProfile]:
        try:
            lib = self._librosa
            y, sr = lib.load(path, sr=None, mono=True, duration=30.0)
            centroid = lib.feature.spectral_centroid(y=y, sr=sr)[0]
            rolloff  = lib.feature.spectral_rolloff(y=y, sr=sr, roll_percent=0.85)[0]
            if hasattr(lib.feature, 'spectral_flux'):
                flux = lib.feature.spectral_flux(y=y, sr=sr)[0]
            else:
                S = np.abs(lib.stft(y))
                flux = np.mean(np.abs(np.diff(S, axis=1)), axis=0)
            mfccs    = lib.feature.mfcc(y=y, sr=sr, n_mfcc=13)
            rms      = lib.feature.rms(y=y)[0]
            flatness = lib.feature.spectral_flatness(y=y)[0]
            duration = lib.get_duration(y=y, sr=sr)
            return SpectralProfile(
                path=path,
                centroid_mean=float(np.mean(centroid)),
                centroid_std=float(np.std(centroid)),
                rolloff_85=float(np.mean(rolloff)),
                flux_mean=float(np.mean(flux)),
                mfcc=[float(np.mean(mfccs[i])) for i in range(13)],
                rms=float(np.mean(rms)),
                spectral_flatness=float(np.mean(flatness)),
                duration=float(duration),
                genre_hints=self._infer_genre_hints(float(np.mean(centroid)), float(np.mean(rms)))
            )
        except Exception as e:
            logger.warning(f"[SPECTRAL] Error analizando {path}: {e}")
            return None

    def _analyze_basic(self, path: str) -> Optional[SpectralProfile]:
        name = os.path.basename(path).lower()
        centroid = 5000.0 if any(k in name for k in ['hat','shaker','top']) else (200.0 if 'bass' in name or 'sub' in name else 2000.0)
        return SpectralProfile(
            path=path, centroid_mean=centroid, centroid_std=100.0,
            rolloff_85=centroid*2, flux_mean=0.1, mfcc=[0.0]*13,
            rms=0.3, spectral_flatness=0.5 if 'noise' in name else 0.1,
            duration=2.0, genre_hints=self._infer_genre_hints(centroid, 0.3)
        )

    def _infer_genre_hints(self, centroid: float, rms: float) -> List[str]:
        hints = []
        if centroid < 500 and rms > 0.4: hints.append('reggaeton_bass')
        if 500 < centroid < 3000: hints.append('reggaeton_perc')
        if centroid > 6000: hints.append('hi_freq_perc')
        return hints or ['unknown']

    def build_similarity_matrix(self, paths: List[str]) -> np.ndarray:
        """T041: Construye matriz de similitud NxN entre samples."""
        n = len(paths)
        matrix = np.zeros((n, n), dtype=np.float32)
        profiles = [self.analyze(p) for p in paths]
        for i in range(n):
            for j in range(n):
                if i == j:
                    matrix[i, j] = 1.0
                elif profiles[i] and profiles[j]:
                    matrix[i, j] = self.similarity(profiles[i], profiles[j])
        return matrix

    def cluster_by_role(self, paths: List[str], n_clusters: int = 5) -> Dict[int, List[str]]:
        """T042: Agrupa samples en N familias tímbricas usando K-means manual."""
        profiles = [self.analyze(p) for p in paths]
        valid_indices = [i for i, p in enumerate(profiles) if p is not None]
        if len(valid_indices) < n_clusters:
            return {0: paths}
        
        centroids_list = [profiles[i].centroid_mean for i in valid_indices]
        rolloffs_list = [profiles[i].rolloff_85 for i in valid_indices]
        features = np.array([[c, r] for c, r in zip(centroids_list, rolloffs_list)], dtype=np.float32)
        
        min_vals = features.min(axis=0)
        max_vals = features.max(axis=0)
        range_vals = max_vals - min_vals + 1e-6
        features_norm = (features - min_vals) / range_vals
        
        np.random.seed(42)
        cluster_centers = features_norm[np.random.choice(len(features_norm), n_clusters, replace=False)]
        
        for _ in range(50):
            distances = np.sqrt(np.sum((features_norm[:, np.newaxis] - cluster_centers) ** 2, axis=2))
            assignments = np.argmin(distances, axis=1)
            new_centers = np.array([
                features_norm[assignments == k].mean(axis=0) if np.sum(assignments == k) > 0 else cluster_centers[k]
                for k in range(n_clusters)
            ])
            if np.allclose(cluster_centers, new_centers, rtol=1e-4):
                break
            cluster_centers = new_centers
        
        clusters: Dict[int, List[str]] = defaultdict(list)
        for idx, cluster_id in enumerate(assignments):
            original_idx = valid_indices[idx]
            clusters[int(cluster_id)].append(paths[original_idx])
        
        return dict(clusters)

    def extract_grain(self, path: str, position_ratio: float = 0.5, grain_ms: float = 50.0) -> Optional[np.ndarray]:
        """
        T136: Extrae un grano de audio de un archivo en una posición relativa.
        
        Args:
            path: Ruta al archivo de audio
            position_ratio: Posición relativa (0.0-1.0) dentro del archivo
            grain_ms: Duración del grano en milisegundos
        
        Returns:
            np.ndarray con el grano extraído, o None si falla
        """
        if not LIBROSA_AVAILABLE:
            logger.warning("[GRANULAR] librosa no disponible para extract_grain")
            return None
        
        try:
            lib = self._librosa
            y, sr = lib.load(path, sr=None, mono=True, duration=30.0)
            
            total_samples = len(y)
            grain_samples = int(sr * grain_ms / 1000.0)
            
            center_sample = int(total_samples * position_ratio)
            start_sample = max(0, center_sample - grain_samples // 2)
            end_sample = min(total_samples, start_sample + grain_samples)
            
            grain = y[start_sample:end_sample]
            
            fade_len = min(len(grain) // 10, 100)
            if fade_len > 0:
                fade_in = np.linspace(0.0, 1.0, fade_len)
                fade_out = np.linspace(1.0, 0.0, fade_len)
                grain[:fade_len] *= fade_in
                grain[-fade_len:] *= fade_out
            
            return grain.astype(np.float32)
            
        except Exception as e:
            logger.warning(f"[GRANULAR] Error extrayendo grano de {path}: {e}")
            return None

    def stretch_grain(self, grain: np.ndarray, target_duration_ms: float, sr: int = 44100) -> Optional[np.ndarray]:
        """
        T137: Estira o comprime un grano a una duración objetivo.
        
        Args:
            grain: Array de audio del grano
            target_duration_ms: Duración objetivo en milisegundos
            sr: Sample rate
        
        Returns:
            np.ndarray con el grano estirado, o None si falla
        """
        if not LIBROSA_AVAILABLE or grain is None or len(grain) == 0:
            return None
        
        try:
            lib = self._librosa
            target_samples = int(sr * target_duration_ms / 1000.0)
            
            if target_samples == len(grain):
                return grain
            
            stretch_ratio = target_samples / len(grain)
            
            stretched = lib.effects.time_stretch(grain, rate=1.0 / stretch_ratio)
            
            if len(stretched) < target_samples:
                padding = np.zeros(target_samples - len(stretched), dtype=np.float32)
                stretched = np.concatenate([stretched, padding])
            elif len(stretched) > target_samples:
                stretched = stretched[:target_samples]
            
            return stretched.astype(np.float32)
            
        except Exception as e:
            logger.warning(f"[GRANULAR] Error estirando grano: {e}")
            return None

    def create_granular_texture(self, path: str, duration_s: float = 4.0, density: float = 0.5, 
                                 output_path: Optional[str] = None) -> Optional[str]:
        """
        T138: Crea una textura granular desde un sample fuente.
        
        Args:
            path: Ruta al archivo de audio fuente
            duration_s: Duración objetivo en segundos
            density: Densidad de granos (0.0-1.0)
            output_path: Ruta de salida opcional
        
        Returns:
            Ruta del archivo generado, o None si falla
        """
        if not LIBROSA_AVAILABLE:
            logger.warning("[GRANULAR] librosa no disponible para create_granular_texture")
            return None
        
        try:
            lib = self._librosa
            sr = 44100
            
            y, file_sr = lib.load(path, sr=sr, mono=True, duration=30.0)
            
            target_samples = int(sr * duration_s)
            output = np.zeros(target_samples, dtype=np.float32)
            
            grain_sizes_ms = [20, 30, 50, 80, 120]
            min_grain_ms = min(grain_sizes_ms)
            max_grain_ms = max(grain_sizes_ms)
            
            base_interval_ms = 50.0
            interval_ms = base_interval_ms / max(density, 0.1)
            
            num_grains = int(duration_s * 1000.0 / interval_ms)
            
            logger.info(f"[GRANULAR] Creando textura: {num_grains} granos, densidad={density}")
            
            for i in range(num_grains):
                position_ratio = np.random.random()
                grain_ms = np.random.choice(grain_sizes_ms)
                
                grain = self.extract_grain(path, position_ratio, grain_ms)
                if grain is None or len(grain) == 0:
                    continue
                
                position_samples = int(target_samples * (i / num_grains))
                position_samples = min(position_samples, target_samples - len(grain))
                
                if position_samples < 0:
                    continue
                
                end_pos = min(position_samples + len(grain), target_samples)
                actual_len = end_pos - position_samples
                output[position_samples:end_pos] += grain[:actual_len] * (0.3 + 0.2 * np.random.random())
            
            rms = np.sqrt(np.mean(output ** 2))
            if rms > 0:
                output = output / (rms * 3)
            
            if output_path is None:
                base_dir = Path(__file__).parents[3] / "libreria" / "reggaeton" / "textures"
                base_dir.mkdir(parents=True, exist_ok=True)
                
                base_name = Path(path).stem
                grain_id = np.random.randint(1000, 9999)
                output_path = str(base_dir / f"{base_name}_granular_{grain_id}.wav")
            
            self._save_wav(output, output_path, sr)
            
            logger.info(f"[GRANULAR] Textura creada: {output_path}")
            return output_path
            
        except Exception as e:
            logger.warning(f"[GRANULAR] Error creando textura granular: {e}")
            return None

    def _save_wav(self, audio: np.ndarray, path: str, sr: int = 44100) -> bool:
        """Guarda un array de audio como archivo WAV."""
        try:
            audio_int = np.clip(audio * 32767, -32768, 32767).astype(np.int16)
            
            with wave.open(path, 'wb') as wav_file:
                wav_file.setnchannels(1)
                wav_file.setsampwidth(2)
                wav_file.setframerate(sr)
                wav_file.writeframes(audio_int.tobytes())
            
            return True
        except Exception as e:
            logger.warning(f"[GRANULAR] Error guardando WAV {path}: {e}")
            return False


class GranularSynthesizer:
    """
    T139: Sintetizador granular para crear pads atmosféricos.
    """
    
    def __init__(self):
        self.engine = get_spectral_engine()
        self._np = np
        self._librosa = None
        if LIBROSA_AVAILABLE:
            import librosa
            self._librosa = librosa
    
    def generate_granular_pad(self, source_path: str, duration_s: float = 8.0, 
                               base_density: float = 0.4, 
                               variation_factor: float = 0.3,
                               output_path: Optional[str] = None) -> Optional[str]:
        """
        T139: Genera un pad granular atmosférico desde una fuente.
        
        Args:
            source_path: Ruta al archivo de audio fuente
            duration_s: Duración objetivo en segundos
            base_density: Densidad base de granos (0.0-1.0)
            variation_factor: Factor de variación tímbrica (0.0-1.0)
            output_path: Ruta de salida opcional
        
        Returns:
            Ruta del archivo generado, o None si falla
        """
        if not LIBROSA_AVAILABLE:
            logger.warning("[GRANULAR] librosa no disponible para generate_granular_pad")
            return None
        
        try:
            sr = 44100
            target_samples = int(sr * duration_s)
            output = np.zeros(target_samples, dtype=np.float32)
            
            if self._librosa:
                y, _ = self._librosa.load(source_path, sr=sr, mono=True, duration=30.0)
            else:
                return None
            
            grain_count = int(duration_s * base_density * 20)
            
            layer_configs = [
                {'density_mult': 1.0, 'grain_ms_range': (30, 80), 'amp_range': (0.25, 0.35)},
                {'density_mult': 0.5, 'grain_ms_range': (80, 150), 'amp_range': (0.15, 0.25)},
                {'density_mult': 0.25, 'grain_ms_range': (150, 300), 'amp_range': (0.08, 0.15)},
            ]
            
            for layer_config in layer_configs:
                layer_density = base_density * layer_config['density_mult']
                layer_grains = int(grain_count * layer_config['density_mult'])
                
                grain_ms_min, grain_ms_max = layer_config['grain_ms_range']
                amp_min, amp_max = layer_config['amp_range']
                
                for i in range(layer_grains):
                    position_ratio = np.random.random()
                    
                    grain_ms = np.random.uniform(grain_ms_min, grain_ms_max)
                    grain_samples = int(sr * grain_ms / 1000.0)
                    
                    center_sample = int(len(y) * position_ratio)
                    start_sample = max(0, center_sample - grain_samples // 2)
                    end_sample = min(len(y), start_sample + grain_samples)
                    
                    grain = y[start_sample:end_sample].copy()
                    
                    fade_len = min(len(grain) // 8, 50)
                    if fade_len > 0 and len(grain) > fade_len * 2:
                        grain[:fade_len] *= np.linspace(0, 1, fade_len)
                        grain[-fade_len:] *= np.linspace(1, 0, fade_len)
                    
                    out_position = int(target_samples * (i / layer_grains))
                    out_position += int(np.random.uniform(-0.1, 0.1) * target_samples / layer_grains)
                    out_position = max(0, min(out_position, target_samples - len(grain)))
                    
                    end_pos = min(out_position + len(grain), target_samples)
                    actual_len = end_pos - out_position
                    
                    if actual_len > 0:
                        amplitude = np.random.uniform(amp_min, amp_max)
                        output[out_position:end_pos] += grain[:actual_len] * amplitude
            
            rms = np.sqrt(np.mean(output ** 2))
            if rms > 0:
                output = output / (rms * 2.5)
            
            if output_path is None:
                base_dir = Path(__file__).parents[3] / "libreria" / "reggaeton" / "textures"
                base_dir.mkdir(parents=True, exist_ok=True)
                
                base_name = Path(source_path).stem
                pad_id = np.random.randint(1000, 9999)
                output_path = str(base_dir / f"{base_name}_pad_{pad_id}.wav")
            
            if self.engine._save_wav(output, output_path, sr):
                logger.info(f"[GRANULAR] Pad generado: {output_path}")
                return output_path
            
            return None
            
        except Exception as e:
            logger.warning(f"[GRANULAR] Error generando pad granular: {e}")
            return None


_engine_instance: Optional[SpectralEngine] = None

def get_spectral_engine() -> SpectralEngine:
    global _engine_instance
    if _engine_instance is None:
        _engine_instance = SpectralEngine()
    return _engine_instance

def get_granular_synthesizer() -> GranularSynthesizer:
    """Factory para obtener instancia del sintetizador granular."""
    return GranularSynthesizer()