"""
Experimental image compression with multiple ML techniques.

Extends the base image encoder to support:
- K-Means (baseline)
- GMM (Gaussian Mixture Models)
- PCA + Vector Quantization
- DBSCAN clustering
"""

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 sklearn.neighbors import NearestNeighbors
from .container import pack_container
from PIL import Image
import warnings
warnings.filterwarnings('ignore')

# Import helpers from base encoder
from .image_encoder import (
    _extract_patches, 
    _rgb_to_ycbcr,
    ColorSpace
)

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


@dataclass
class ImageEncodeConfigExperimental:
    """Configuration for experimental image compression."""
    patch_size: int = 8
    k: int = 256
    color_space: ColorSpace = "YCbCr"
    residual_threshold: float = 0.1
    method: CompressionMethod = "kmeans"  # NEW: method selection
    quality_mode: Literal["fast", "balanced", "high"] = "balanced"
    random_state: int = 42


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


def _fit_patch_codebook_gmm(patches_flat: np.ndarray, k: int, random_state: int) -> Tuple[np.ndarray, Dict]:
    """Fit Gaussian Mixture Model codebook."""
    k_actual = min(k, len(patches_flat))
    
    gmm = GaussianMixture(
        n_components=k_actual,
        covariance_type='diag',
        random_state=random_state,
        max_iter=50,
        reg_covar=1e-4  # Regularization for stability
    )
    
    try:
        gmm.fit(patches_flat)
        return gmm.means_, {
            "method": "gmm",
            "n_components": k_actual,
            "covariance_type": "diag"
        }
    except Exception as e:
        # Fallback to K-Means if GMM fails
        print(f"Warning: GMM failed ({e}), falling back to K-Means")
        return _fit_patch_codebook_kmeans(patches_flat, k, random_state)


def _fit_patch_codebook_pca(patches_flat: np.ndarray, k: int, random_state: int) -> Tuple[np.ndarray, Dict]:
    """
    Fit PCA + quantization codebook.
    Reduce dimensionality with PCA, then cluster in reduced space.
    """
    # Determine PCA components
    n_pca_components = min(k // 2, patches_flat.shape[1], len(patches_flat))
    
    pca = PCA(n_components=n_pca_components, random_state=random_state)
    patches_pca = pca.fit_transform(patches_flat)
    
    # K-Means in PCA space
    k_actual = min(k, len(patches_pca))
    kmeans = KMeans(
        n_clusters=k_actual,
        n_init=3,
        random_state=random_state,
        max_iter=50
    )
    kmeans.fit(patches_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": float(np.sum(pca.explained_variance_ratio_))
    }


def _fit_patch_codebook_dbscan(patches_flat: np.ndarray, k: int, random_state: int) -> Tuple[np.ndarray, Dict]:
    """
    Fit DBSCAN clustering codebook.
    """
    # Sample for eps estimation (faster)
    sample_size = min(1000, len(patches_flat))
    sample_idx = np.random.RandomState(random_state).choice(len(patches_flat), sample_size, replace=False)
    patches_sample = patches_flat[sample_idx]
    
    # Estimate eps
    k_neighbors = min(10, sample_size - 1)
    nbrs = NearestNeighbors(n_neighbors=k_neighbors).fit(patches_sample)
    distances, _ = nbrs.kneighbors(patches_sample)
    eps = np.mean(distances[:, -1]) * 1.5
    
    # Apply DBSCAN
    dbscan = DBSCAN(eps=eps, min_samples=5)
    labels = dbscan.fit_predict(patches_flat)
    
    # Get cluster centers
    unique_labels = np.unique(labels[labels >= 0])
    
    if len(unique_labels) == 0:
        print("Warning: DBSCAN found no clusters, falling back to K-Means")
        return _fit_patch_codebook_kmeans(patches_flat, k, random_state)
    
    centroids = []
    for label in unique_labels:
        cluster_points = patches_flat[labels == label]
        centroids.append(np.mean(cluster_points, axis=0))
    
    centroids = np.array(centroids)
    
    # If too many clusters, reduce with K-Means
    if len(centroids) > k:
        kmeans = KMeans(n_clusters=k, random_state=random_state, n_init=3)
        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_patches(patches: np.ndarray, codebook: np.ndarray, 
                   residual_threshold: float) -> Dict:
    """Encode patches using codebook."""
    n_patches = patches.shape[0]
    patch_dim = np.prod(patches.shape[1:])
    patches_flat = patches.reshape(n_patches, patch_dim).astype(np.float32)
    
    # Find nearest centroid
    distances = np.sum((patches_flat[:, None, :] - codebook[None, :, :]) ** 2, axis=2)
    indices = np.argmin(distances, axis=1).astype(np.uint16)
    
    # Reconstruct
    recon_flat = codebook[indices]
    recon = recon_flat.reshape(patches.shape)
    
    # Calculate per-patch error
    patch_errors = np.sqrt(np.mean((patches_flat - recon_flat) ** 2, axis=1))
    patch_errors_normalized = patch_errors / 255.0
    
    # High-error patches need residuals
    high_error_mask = patch_errors_normalized > residual_threshold
    high_error_indices = np.where(high_error_mask)[0]
    
    residuals = {}
    if len(high_error_indices) > 0:
        residual_data = patches_flat[high_error_indices] - recon_flat[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": patch_errors_normalized,
        "residuals": residuals,
        "mean_error": float(np.mean(patch_errors_normalized)),
        "max_error": float(np.max(patch_errors_normalized)),
        "residual_count": len(high_error_indices)
    }


def encode_image_experimental(image_path: str, cfg: ImageEncodeConfigExperimental) -> bytes:
    """
    Encode an image using experimental ML techniques.
    
    Args:
        image_path: Path to image file
        cfg: Compression configuration with method selection
        
    Returns:
        Compressed image blob
    """
    # Load image
    img = Image.open(image_path)
    original_format = img.format
    original_mode = img.mode
    
    # Convert to RGB
    if img.mode not in ['RGB', 'L']:
        img = img.convert('RGB')
    
    img_array = np.array(img)
    original_shape = img_array.shape
    
    # Convert color space
    if cfg.color_space == "YCbCr" and img_array.ndim == 3:
        img_array = _rgb_to_ycbcr(img_array)
    elif cfg.color_space == "Grayscale" and img_array.ndim == 3:
        img_array = np.mean(img_array, axis=2, keepdims=False).astype(np.uint8)
    
    # Extract patches
    patches, grid_shape = _extract_patches(img_array, cfg.patch_size)
    
    # Process each channel
    channels = patches.shape[3] if patches.ndim == 4 else 1
    channel_data = {}
    
    for ch in range(channels):
        if channels > 1:
            ch_patches = patches[:, :, :, ch:ch+1]
        else:
            ch_patches = patches
        
        # Flatten for clustering
        n_patches = ch_patches.shape[0]
        patch_dim = np.prod(ch_patches.shape[1:])
        patches_flat = ch_patches.reshape(n_patches, patch_dim).astype(np.float32)
        
        # Fit codebook based on method
        if cfg.method == "kmeans":
            codebook, model_info = _fit_patch_codebook_kmeans(patches_flat, cfg.k, cfg.random_state)
        elif cfg.method == "gmm":
            codebook, model_info = _fit_patch_codebook_gmm(patches_flat, cfg.k, cfg.random_state)
        elif cfg.method == "pca":
            codebook, model_info = _fit_patch_codebook_pca(patches_flat, cfg.k, cfg.random_state)
        elif cfg.method == "dbscan":
            codebook, model_info = _fit_patch_codebook_dbscan(patches_flat, cfg.k, cfg.random_state)
        else:
            raise ValueError(f"Unknown method: {cfg.method}")
        
        # Encode patches
        encoded = _encode_patches(ch_patches, codebook, cfg.residual_threshold)
        
        channel_data[f"ch{ch}"] = {
            "codebook": codebook.astype(np.float32).tobytes(),
            "indices": encoded["indices"],
            "residuals": encoded["residuals"],
            "mean_error": encoded["mean_error"],
            "max_error": encoded["max_error"],
            "model_info": model_info
        }
    
    # Package
    payload = {
        "domain": "image",
        "version": "0.2.0-experimental",
        "original_shape": original_shape,
        "original_format": original_format,
        "original_mode": original_mode,
        "patch_size": cfg.patch_size,
        "grid_shape": grid_shape,
        "color_space": cfg.color_space,
        "k": cfg.k,
        "channels": channels,
        "channel_data": channel_data,
        "model": {
            "type": f"patch-{cfg.method}",
            "quality_mode": cfg.quality_mode,
            "random_state": cfg.random_state
        },
        "meta": {
            "residual_threshold": cfg.residual_threshold
        }
    }
    
    return pack_container(payload)
