"""
Image compression module using patch-based vector quantization.

This module extends Sempress's semantic compression approach to images by:
1. Dividing images into patches (blocks)
2. Learning patch-level codebooks via MiniBatchKMeans
3. Storing quantized indices + optional residuals
4. Supporting multiple color spaces (RGB, YCbCr)
"""

from __future__ import annotations
from dataclasses import dataclass
from typing import Tuple, Dict, Optional, Literal
import numpy as np
from sklearn.cluster import MiniBatchKMeans
from .container import pack_container
from PIL import Image
import io

ColorSpace = Literal["RGB", "YCbCr", "Grayscale"]

@dataclass
class ImageEncodeConfig:
    """Configuration for image compression."""
    patch_size: int = 8  # 8x8 or 16x16 patches
    k: int = 256  # Codebook size per channel
    color_space: ColorSpace = "YCbCr"  # YCbCr often compresses better
    residual_threshold: float = 0.1  # Store residuals for high-error patches
    quality_mode: Literal["fast", "balanced", "high"] = "balanced"
    random_state: int = 42


def _rgb_to_ycbcr(img: np.ndarray) -> np.ndarray:
    """Convert RGB to YCbCr color space."""
    # Standard BT.601 conversion
    transform = np.array([
        [0.299, 0.587, 0.114],
        [-0.168736, -0.331264, 0.5],
        [0.5, -0.418688, -0.081312]
    ])
    offset = np.array([0, 128, 128])

    # Reshape for matrix multiplication
    original_shape = img.shape
    img_flat = img.reshape(-1, 3).astype(np.float32)

    # Apply transformation
    ycbcr = img_flat @ transform.T + offset
    return ycbcr.reshape(original_shape).clip(0, 255).astype(np.uint8)


def _ycbcr_to_rgb(img: np.ndarray) -> np.ndarray:
    """Convert YCbCr back to RGB color space."""
    # Inverse BT.601 conversion
    inv_transform = np.array([
        [1.0, 0.0, 1.402],
        [1.0, -0.344136, -0.714136],
        [1.0, 1.772, 0.0]
    ])
    offset = np.array([0, 128, 128])

    original_shape = img.shape
    img_flat = img.reshape(-1, 3).astype(np.float32)

    # Apply inverse transformation
    rgb = (img_flat - offset) @ inv_transform.T
    return rgb.reshape(original_shape).clip(0, 255).astype(np.uint8)


def _extract_patches(img: np.ndarray, patch_size: int) -> Tuple[np.ndarray, Tuple[int, int]]:
    """
    Extract non-overlapping patches from image.

    Returns:
        patches: (n_patches, patch_size, patch_size, channels)
        grid_shape: (n_patches_h, n_patches_w)
    """
    h, w = img.shape[:2]
    channels = img.shape[2] if img.ndim == 3 else 1

    # Pad image to be divisible by patch_size
    pad_h = (patch_size - h % patch_size) % patch_size
    pad_w = (patch_size - w % patch_size) % patch_size

    if pad_h > 0 or pad_w > 0:
        if img.ndim == 3:
            img = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), mode='edge')
        else:
            img = np.pad(img, ((0, pad_h), (0, pad_w)), mode='edge')

    h_padded, w_padded = img.shape[:2]
    n_patches_h = h_padded // patch_size
    n_patches_w = w_padded // patch_size

    # Reshape into patches
    if img.ndim == 3:
        patches = img.reshape(n_patches_h, patch_size, n_patches_w, patch_size, channels)
        patches = patches.transpose(0, 2, 1, 3, 4)  # (n_h, n_w, p, p, c)
        patches = patches.reshape(-1, patch_size, patch_size, channels)
    else:
        patches = img.reshape(n_patches_h, patch_size, n_patches_w, patch_size)
        patches = patches.transpose(0, 2, 1, 3)  # (n_h, n_w, p, p)
        patches = patches.reshape(-1, patch_size, patch_size, 1)

    return patches, (n_patches_h, n_patches_w)


def _reconstruct_from_patches(patches: np.ndarray, grid_shape: Tuple[int, int],
                              original_shape: Tuple[int, int]) -> np.ndarray:
    """Reconstruct image from patches."""
    n_patches_h, n_patches_w = grid_shape
    patch_size = patches.shape[1]
    channels = patches.shape[3] if patches.ndim == 4 else 1

    # Reshape patches back to grid
    if channels > 1:
        patches = patches.reshape(n_patches_h, n_patches_w, patch_size, patch_size, channels)
        patches = patches.transpose(0, 2, 1, 3, 4)  # (n_h, p, n_w, p, c)
        img = patches.reshape(n_patches_h * patch_size, n_patches_w * patch_size, channels)
    else:
        patches = patches.reshape(n_patches_h, n_patches_w, patch_size, patch_size)
        patches = patches.transpose(0, 2, 1, 3)  # (n_h, p, n_w, p)
        img = patches.reshape(n_patches_h * patch_size, n_patches_w * patch_size)

    # Crop to original dimensions
    h_orig, w_orig = original_shape[:2]
    if channels > 1:
        return img[:h_orig, :w_orig, :]
    else:
        return img[:h_orig, :w_orig]


def _fit_patch_codebook(patches: np.ndarray, k: int, random_state: int) -> np.ndarray:
    """
    Learn codebook for image patches using MiniBatchKMeans.

    Args:
        patches: (n_patches, patch_size, patch_size, channels)
        k: codebook size

    Returns:
        codebook: (k, patch_size * patch_size * channels)
    """
    # Flatten patches for clustering
    n_patches = patches.shape[0]
    patch_dim = np.prod(patches.shape[1:])
    patches_flat = patches.reshape(n_patches, patch_dim).astype(np.float32)

    # Adjust k if we have fewer unique patches
    k_actual = min(k, n_patches)

    # Fit MiniBatchKMeans (n_init=1 with k-means++ is sufficient)
    kmeans = MiniBatchKMeans(
        n_clusters=k_actual, batch_size=2048, n_init=1,
        random_state=random_state, max_iter=100
    )
    kmeans.fit(patches_flat)

    return kmeans.cluster_centers_


def _encode_patches(patches: np.ndarray, codebook: np.ndarray,
                   residual_threshold: float) -> Dict:
    """
    Encode patches using codebook.

    Uses the dot product identity: ||a-b||^2 = ||a||^2 + ||b||^2 - 2*a.b
    to avoid materializing the full (n_patches, k, patch_dim) distance tensor.

    Returns:
        Dictionary with indices, reconstruction error, and residuals
    """
    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 via dot product trick (avoids O(n*k*d) intermediate)
    patches_sq = np.sum(patches_flat ** 2, axis=1, keepdims=True)   # (n, 1)
    codebook_sq = np.sum(codebook ** 2, axis=1, keepdims=True).T    # (1, k)
    cross = patches_flat @ codebook.T                                # (n, k)
    distances = patches_sq + codebook_sq - 2 * cross                 # (n, k)
    idx_dtype = np.uint8 if codebook.shape[0] <= 256 else np.uint16
    indices = np.argmin(distances, axis=1).astype(idx_dtype)

    # Reconstruct
    recon_flat = codebook[indices]
    recon = recon_flat.reshape(patches.shape)

    # Calculate per-patch error (normalized RMSE)
    patch_errors = np.sqrt(np.mean((patches_flat - recon_flat) ** 2, axis=1))
    patch_errors_normalized = patch_errors / 255.0  # Normalize to [0, 1]

    # Identify high-error patches that need residuals
    high_error_mask = patch_errors_normalized > residual_threshold
    high_error_indices = np.where(high_error_mask)[0]

    # Store residuals for high-error patches (as float16 for compactness)
    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.float16).tobytes(),
            "residual_dtype": "float16",
        }

    return {
        "indices": indices.tobytes(),
        "index_dtype": "uint8" if idx_dtype == np.uint8 else "uint16",
        "recon": recon,
        "errors": patch_errors_normalized,
        "residuals": residuals,
        "mean_error": float(np.mean(patch_errors_normalized)),
        "max_error": float(np.max(patch_errors_normalized))
    }


def encode_image(image_path: str, cfg: ImageEncodeConfig) -> bytes:
    """
    Encode an image file using patch-based vector quantization.

    Args:
        image_path: Path to image file
        cfg: Compression configuration

    Returns:
        Compressed image blob
    """
    # Load image
    img = Image.open(image_path)
    original_format = img.format
    original_mode = img.mode

    # Convert to RGB if needed
    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 requested
    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 separately for better compression
    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

        # Fit codebook
        codebook = _fit_patch_codebook(ch_patches, cfg.k, cfg.random_state)

        # 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"],
            "index_dtype": encoded["index_dtype"],
            "residuals": encoded["residuals"],
            "mean_error": encoded["mean_error"],
            "max_error": encoded["max_error"]
        }

    # Package everything
    payload = {
        "domain": "image",
        "version": "0.2.0",
        "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": "patch-kmeans",
            "quality_mode": cfg.quality_mode,
            "random_state": cfg.random_state
        },
        "meta": {
            "residual_threshold": cfg.residual_threshold
        }
    }

    return pack_container(payload)


def encode_image_from_array(img_array: np.ndarray, cfg: ImageEncodeConfig) -> bytes:
    """
    Encode an image array using patch-based vector quantization.

    Args:
        img_array: Image as numpy array
        cfg: Compression configuration

    Returns:
        Compressed image blob
    """
    original_shape = img_array.shape

    # Convert color space if requested
    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 separately
    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

        # Fit codebook and encode
        codebook = _fit_patch_codebook(ch_patches, cfg.k, cfg.random_state)
        encoded = _encode_patches(ch_patches, codebook, cfg.residual_threshold)

        channel_data[f"ch{ch}"] = {
            "codebook": codebook.astype(np.float32).tobytes(),
            "indices": encoded["indices"],
            "index_dtype": encoded["index_dtype"],
            "residuals": encoded["residuals"],
            "mean_error": encoded["mean_error"],
            "max_error": encoded["max_error"]
        }

    # Package
    payload = {
        "domain": "image",
        "version": "0.2.0",
        "original_shape": original_shape,
        "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": "patch-kmeans",
            "quality_mode": cfg.quality_mode,
            "random_state": cfg.random_state
        },
        "meta": {
            "residual_threshold": cfg.residual_threshold
        }
    }

    return pack_container(payload)
