"""
Audio compression module using multiple ML techniques.

This module extends Sempress's semantic compression to audio by:
1. Segmenting audio into frames/chunks
2. Learning frame-level codebooks via various ML techniques:
   - K-Means clustering (baseline)
   - Gaussian Mixture Models (GMM)
   - PCA-based compression
   - DBSCAN clustering
3. Storing quantized indices + optional residuals
4. Supporting multiple audio formats (WAV, MP3, FLAC)
"""

from __future__ import annotations
from dataclasses import dataclass
from typing import Tuple, Dict, Optional, Literal
import numpy as np
from sklearn.cluster import KMeans, DBSCAN
from sklearn.mixture import GaussianMixture
from sklearn.decomposition import PCA
from .container import pack_container
import warnings
warnings.filterwarnings('ignore')

try:
    import librosa
    LIBROSA_AVAILABLE = True
except ImportError:
    LIBROSA_AVAILABLE = False
    print("Warning: librosa not available. Install with: pip install librosa")

try:
    import soundfile as sf
    SOUNDFILE_AVAILABLE = True
except ImportError:
    SOUNDFILE_AVAILABLE = False
    print("Warning: soundfile not available. Install with: pip install soundfile")


CompressionMethod = Literal["kmeans", "gmm", "pca", "dbscan"]


@dataclass
class AudioEncodeConfig:
    """Configuration for audio compression."""
    frame_size: int = 2048  # Number of samples per frame
    hop_length: int = 512  # Overlap between frames
    n_components: int = 256  # Codebook size / number of components
    method: CompressionMethod = "kmeans"  # Compression method
    residual_threshold: float = 0.1  # Store residuals for high-error frames
    use_spectral: bool = True  # Use spectral features (MFCC) vs raw waveform
    n_mfcc: int = 20  # Number of MFCC coefficients if use_spectral=True
    random_state: int = 42


def _load_audio(audio_path: str) -> Tuple[np.ndarray, int]:
    """
    Load audio file and return waveform + sample rate.
    
    Returns:
        audio: (n_samples,) or (n_samples, n_channels)
        sample_rate: int
    """
    if not LIBROSA_AVAILABLE:
        raise RuntimeError("librosa is required for audio processing. Install with: pip install librosa")
    
    # Load audio with librosa (automatically handles various formats)
    audio, sr = librosa.load(audio_path, sr=None, mono=False)
    
    # If stereo, convert to (n_samples, 2) format
    if audio.ndim == 2:
        audio = audio.T
    
    return audio, sr


def _extract_features(audio: np.ndarray, sr: int, cfg: AudioEncodeConfig) -> Tuple[np.ndarray, Dict]:
    """
    Extract features from audio (either raw waveform frames or spectral features).
    
    Returns:
        features: (n_frames, feature_dim)
        metadata: dict with extraction parameters
    """
    if cfg.use_spectral:
        # Extract MFCC features (more compact representation)
        mfcc = librosa.feature.mfcc(
            y=audio,
            sr=sr,
            n_mfcc=cfg.n_mfcc,
            n_fft=cfg.frame_size,
            hop_length=cfg.hop_length
        )
        # Transpose to (n_frames, n_mfcc)
        features = mfcc.T
        
        metadata = {
            "feature_type": "mfcc",
            "n_mfcc": cfg.n_mfcc,
            "n_fft": cfg.frame_size,
            "hop_length": cfg.hop_length
        }
    else:
        # Use raw waveform frames
        frames = librosa.util.frame(
            audio,
            frame_length=cfg.frame_size,
            hop_length=cfg.hop_length
        )
        # Transpose to (n_frames, frame_size)
        features = frames.T
        
        metadata = {
            "feature_type": "waveform",
            "frame_size": cfg.frame_size,
            "hop_length": cfg.hop_length
        }
    
    return features, metadata


def _fit_codebook_kmeans(features: np.ndarray, k: int, random_state: int) -> Tuple[np.ndarray, Dict]:
    """Fit K-Means codebook."""
    k_actual = min(k, len(features))
    
    kmeans = KMeans(
        n_clusters=k_actual,
        n_init=10,
        random_state=random_state,
        max_iter=300
    )
    kmeans.fit(features)
    
    return kmeans.cluster_centers_, {
        "method": "kmeans",
        "n_clusters": k_actual
    }


def _fit_codebook_gmm(features: np.ndarray, k: int, random_state: int) -> Tuple[np.ndarray, Dict]:
    """Fit Gaussian Mixture Model codebook."""
    k_actual = min(k, len(features))
    
    gmm = GaussianMixture(
        n_components=k_actual,
        covariance_type='diag',  # Diagonal covariance for efficiency
        random_state=random_state,
        max_iter=100
    )
    gmm.fit(features)
    
    # Use GMM means as codebook centers
    return gmm.means_, {
        "method": "gmm",
        "n_components": k_actual,
        "covariance_type": "diag"
    }


def _fit_codebook_pca(features: np.ndarray, k: int, random_state: int) -> Tuple[np.ndarray, Dict]:
    """
    Fit PCA + quantization codebook.
    First reduce dimensionality with PCA, then quantize in reduced space.
    """
    # Determine PCA components (use fewer dimensions)
    n_pca_components = min(k // 2, features.shape[1], len(features))
    
    pca = PCA(n_components=n_pca_components, random_state=random_state)
    features_pca = pca.fit_transform(features)
    
    # Now apply K-Means in PCA space
    k_actual = min(k, len(features_pca))
    kmeans = KMeans(
        n_clusters=k_actual,
        n_init=5,
        random_state=random_state,
        max_iter=100
    )
    kmeans.fit(features_pca)
    
    # Transform centroids back to original space
    centroids_original = pca.inverse_transform(kmeans.cluster_centers_)
    
    return centroids_original, {
        "method": "pca",
        "n_pca_components": n_pca_components,
        "n_clusters": k_actual,
        "explained_variance_ratio": float(np.sum(pca.explained_variance_ratio_))
    }


def _fit_codebook_dbscan(features: np.ndarray, k: int, random_state: int) -> Tuple[np.ndarray, Dict]:
    """
    Fit DBSCAN clustering codebook.
    Note: DBSCAN doesn't use 'k' directly, we estimate eps from data.
    """
    # Estimate eps based on k-distance (heuristic)
    from sklearn.neighbors import NearestNeighbors
    
    # Use a sample for faster computation
    sample_size = min(1000, len(features))
    sample_idx = np.random.RandomState(random_state).choice(len(features), sample_size, replace=False)
    features_sample = features[sample_idx]
    
    # Find average distance to k-th nearest neighbor
    k_neighbors = min(10, sample_size - 1)
    nbrs = NearestNeighbors(n_neighbors=k_neighbors).fit(features_sample)
    distances, _ = nbrs.kneighbors(features_sample)
    eps = np.mean(distances[:, -1]) * 1.5  # Heuristic multiplier
    
    # Apply DBSCAN
    dbscan = DBSCAN(eps=eps, min_samples=5)
    labels = dbscan.fit_predict(features)
    
    # Get cluster centers (mean of each cluster)
    unique_labels = np.unique(labels)
    # Remove noise label (-1)
    unique_labels = unique_labels[unique_labels >= 0]
    
    if len(unique_labels) == 0:
        # Fallback to K-Means if DBSCAN finds no clusters
        print("Warning: DBSCAN found no clusters, falling back to K-Means")
        return _fit_codebook_kmeans(features, k, random_state)
    
    centroids = []
    for label in unique_labels:
        cluster_points = features[labels == label]
        centroids.append(np.mean(cluster_points, axis=0))
    
    centroids = np.array(centroids)
    
    # If we have too many clusters, reduce with K-Means
    if len(centroids) > k:
        kmeans = KMeans(n_clusters=k, random_state=random_state, n_init=5)
        kmeans.fit(centroids)
        centroids = kmeans.cluster_centers_
    
    return centroids, {
        "method": "dbscan",
        "eps": float(eps),
        "n_clusters_found": len(unique_labels),
        "n_clusters_final": len(centroids),
        "n_noise_points": int(np.sum(labels == -1))
    }


def _encode_features(features: np.ndarray, codebook: np.ndarray, 
                     residual_threshold: float) -> Dict:
    """
    Encode features using codebook.
    
    Returns:
        Dictionary with indices, reconstruction error, and residuals
    """
    # Find nearest codebook entry for each frame
    distances = np.sum((features[:, None, :] - codebook[None, :, :]) ** 2, axis=2)
    indices = np.argmin(distances, axis=1).astype(np.uint16)
    
    # Reconstruct
    recon = codebook[indices]
    
    # Calculate per-frame error (normalized RMSE)
    frame_errors = np.sqrt(np.mean((features - recon) ** 2, axis=1))
    # Normalize by signal magnitude
    signal_magnitude = np.sqrt(np.mean(features ** 2, axis=1)) + 1e-8
    frame_errors_normalized = frame_errors / signal_magnitude
    
    # Identify high-error frames that need residuals
    high_error_mask = frame_errors_normalized > residual_threshold
    high_error_indices = np.where(high_error_mask)[0]
    
    # Store residuals for high-error frames
    residuals = {}
    if len(high_error_indices) > 0:
        residual_data = features[high_error_indices] - recon[high_error_indices]
        residuals = {
            "indices": high_error_indices.astype(np.uint32).tobytes(),
            "data": residual_data.astype(np.float32).tobytes()
        }
    
    return {
        "indices": indices.tobytes(),
        "recon": recon,
        "errors": frame_errors_normalized,
        "residuals": residuals,
        "mean_error": float(np.mean(frame_errors_normalized)),
        "max_error": float(np.max(frame_errors_normalized)),
        "residual_count": len(high_error_indices)
    }


def encode_audio(audio_path: str, cfg: AudioEncodeConfig) -> bytes:
    """
    Encode an audio file using semantic compression.
    
    Args:
        audio_path: Path to audio file
        cfg: Compression configuration
        
    Returns:
        Compressed audio blob
    """
    # Load audio
    audio, sr = _load_audio(audio_path)
    
    # Handle stereo by processing channels independently
    is_stereo = audio.ndim == 2 and audio.shape[1] == 2
    
    if is_stereo:
        # Process each channel separately
        channel_data = {}
        for ch_idx in range(2):
            ch_audio = audio[:, ch_idx]
            ch_data = _encode_audio_channel(ch_audio, sr, cfg)
            channel_data[f"ch{ch_idx}"] = ch_data
        
        n_channels = 2
        n_samples = audio.shape[0]
    else:
        # Mono audio
        if audio.ndim == 2:
            audio = audio[:, 0]
        
        channel_data = {
            "ch0": _encode_audio_channel(audio, sr, cfg)
        }
        n_channels = 1
        n_samples = len(audio)
    
    # Package everything
    payload = {
        "domain": "audio",
        "version": "0.1.0",
        "sample_rate": sr,
        "n_samples": n_samples,
        "n_channels": n_channels,
        "channel_data": channel_data,
        "config": {
            "frame_size": cfg.frame_size,
            "hop_length": cfg.hop_length,
            "n_components": cfg.n_components,
            "method": cfg.method,
            "use_spectral": cfg.use_spectral,
            "n_mfcc": cfg.n_mfcc if cfg.use_spectral else None
        },
        "meta": {
            "residual_threshold": cfg.residual_threshold,
            "random_state": cfg.random_state
        }
    }
    
    return pack_container(payload)


def _encode_audio_channel(audio: np.ndarray, sr: int, cfg: AudioEncodeConfig) -> Dict:
    """Encode a single audio channel."""
    # Extract features
    features, feature_metadata = _extract_features(audio, sr, cfg)
    
    # Fit codebook based on method
    if cfg.method == "kmeans":
        codebook, model_info = _fit_codebook_kmeans(features, cfg.n_components, cfg.random_state)
    elif cfg.method == "gmm":
        codebook, model_info = _fit_codebook_gmm(features, cfg.n_components, cfg.random_state)
    elif cfg.method == "pca":
        codebook, model_info = _fit_codebook_pca(features, cfg.n_components, cfg.random_state)
    elif cfg.method == "dbscan":
        codebook, model_info = _fit_codebook_dbscan(features, cfg.n_components, cfg.random_state)
    else:
        raise ValueError(f"Unknown compression method: {cfg.method}")
    
    # Encode features
    encoded = _encode_features(features, codebook, cfg.residual_threshold)
    
    return {
        "codebook": codebook.astype(np.float32).tobytes(),
        "codebook_shape": codebook.shape,
        "indices": encoded["indices"],
        "residuals": encoded["residuals"],
        "mean_error": encoded["mean_error"],
        "max_error": encoded["max_error"],
        "residual_count": encoded["residual_count"],
        "feature_metadata": feature_metadata,
        "model_info": model_info
    }


def encode_audio_from_array(audio: np.ndarray, sr: int, cfg: AudioEncodeConfig) -> bytes:
    """
    Encode an audio array using semantic compression.
    
    Args:
        audio: Audio waveform array (n_samples,) or (n_samples, n_channels)
        sr: Sample rate
        cfg: Compression configuration
        
    Returns:
        Compressed audio blob
    """
    # Similar to encode_audio but takes array directly
    is_stereo = audio.ndim == 2 and audio.shape[1] == 2
    
    if is_stereo:
        channel_data = {}
        for ch_idx in range(2):
            ch_audio = audio[:, ch_idx]
            ch_data = _encode_audio_channel(ch_audio, sr, cfg)
            channel_data[f"ch{ch_idx}"] = ch_data
        
        n_channels = 2
        n_samples = audio.shape[0]
    else:
        if audio.ndim == 2:
            audio = audio[:, 0]
        
        channel_data = {
            "ch0": _encode_audio_channel(audio, sr, cfg)
        }
        n_channels = 1
        n_samples = len(audio)
    
    payload = {
        "domain": "audio",
        "version": "0.1.0",
        "sample_rate": sr,
        "n_samples": n_samples,
        "n_channels": n_channels,
        "channel_data": channel_data,
        "config": {
            "frame_size": cfg.frame_size,
            "hop_length": cfg.hop_length,
            "n_components": cfg.n_components,
            "method": cfg.method,
            "use_spectral": cfg.use_spectral,
            "n_mfcc": cfg.n_mfcc if cfg.use_spectral else None
        },
        "meta": {
            "residual_threshold": cfg.residual_threshold,
            "random_state": cfg.random_state
        }
    }
    
    return pack_container(payload)
