"""
skiss_block_codec.py
================================================================================
Block-wise Skiss-DCT Image Codec (Production Variant)

Theory:
    - Global Fourier on natural images is skiss-incomplete (K_critical exceeds
      practical bounds due to non-periodic edges and texture).
    - Localizing the basis to B×B blocks makes each patch skiss-complete at
      very low K (e.g. K=8 captures >95% of block energy).
    - We use the DCT-II cosine basis: better edge behaviour than raw Fourier.
    - Coefficients are quantized to int16 and compressed with zlib inside npz,
      producing a true high-compression parametric image.

Expected size for 2752×1536 RGB:
    K=8,  qstep=4.0   →  ~80–150 KB
    K=16, qstep=2.0   →  ~400–800 KB
================================================================================
"""

import numpy as np
from pathlib import Path
from typing import Union, Tuple, Dict, Optional


# ==============================================================================
#  Stationary Basis: DCT-II (Separable Cosine)
# ==============================================================================

def _build_dct_basis_1d(coords: np.ndarray, K: int, N: int) -> np.ndarray:
    """
    DCT-II basis vectors evaluated at arbitrary continuous coordinates.
    
    B[n, k] = cos( π * k * (coords[n] + 0.5) / N ),   k = 0 .. K-1
    """
    coords = np.asarray(coords, dtype=np.float64)
    k = np.arange(K, dtype=np.float64)
    return np.cos(np.pi * k * (coords[:, None] + 0.5) / N)   # (len(coords), K)


# ==============================================================================
#  Encoder: Explicit Pixels → Quantized Skiss Residue θ
# ==============================================================================

class SkissBlockEncoder:
    """
    Retro-compresses an image into block-wise DCT coefficient residues.
    """
    
    def __init__(self, block_size: int = 64, K: int = 8,
                 qstep: float = 4.0, zero_threshold: float = 0.0):
        """
        Args:
            block_size: Size of local blocks (e.g. 64).
            K: Number of low-frequency modes retained per axis (K×K coeffs).
            qstep: Quantization step. Larger = smaller file, lower fidelity.
            zero_threshold: Coefficients below this are forced to 0 before
                            quantization, increasing zlib compressibility.
        """
        self.block_size = block_size
        self.K = K
        self.qstep = qstep
        self.zero_threshold = zero_threshold
        self.theta: Optional[Dict] = None

    def _fit_block(self, patch: np.ndarray) -> np.ndarray:
        """Fit separable DCT surface to one B×B×C patch. Returns (C, K, K)."""
        B = patch.shape[0]
        y_coords = np.arange(B, dtype=np.float64)
        x_coords = np.arange(B, dtype=np.float64)
        
        U = _build_dct_basis_1d(y_coords, self.K, B)   # (B, K)
        V = _build_dct_basis_1d(x_coords, self.K, B)   # (B, K)
        A = U.T @ U
        Bmat = V.T @ V
        
        C = patch.shape[2]
        coeffs = np.zeros((C, self.K, self.K), dtype=np.float64)
        for c in range(C):
            RHS = U.T @ patch[:, :, c] @ V          # (K, K)
            Z = np.linalg.solve(A, RHS)              # A Z = RHS
            Theta = np.linalg.solve(Bmat, Z.T).T   # B Θᵀ = Zᵀ
            coeffs[c] = Theta
        return coeffs

    def encode(self, image: np.ndarray, verbose: bool = True) -> Dict:
        """
        Compress image into block-wise spectral residue.
        
        Args:
            image: ndarray (H, W) or (H, W, C), uint8 or float.
            verbose: Print compression stats.
        
        Returns:
            theta: Serializable dict.
        """
        if image.ndim == 2:
            image = image[:, :, None]
        H, W, C = image.shape
        B = self.block_size
        
        # Pad to multiple of B using edge replication (preserves boundary continuity)
        pad_h = (B - H % B) % B
        pad_w = (B - W % B) % B
        padded = np.pad(image.astype(np.float64), ((0, pad_h), (0, pad_w), (0, 0)), mode='edge')
        H_pad, W_pad = padded.shape[:2]
        
        n_by = H_pad // B
        n_bx = W_pad // B
        
        # Collapse each block into its skiss residue
        coefficients = np.zeros((n_by, n_bx, C, self.K, self.K), dtype=np.float64)
        for by in range(n_by):
            for bx in range(n_bx):
                y0, x0 = by * B, bx * B
                patch = padded[y0:y0 + B, x0:x0 + B]
                coefficients[by, bx] = self._fit_block(patch)
        
        # Quantization: path erasure into low-entropy integer lattice
        if self.zero_threshold > 0:
            coefficients[np.abs(coefficients) < self.zero_threshold] = 0.0
        coefficients_q = np.round(coefficients / self.qstep).astype(np.int16)
        
        # Metrics
        orig_pixels = H * W * C
        param_vals = coefficients_q.size
        sparsity = float(np.mean(coefficients_q == 0))
        
        self.theta = {
            'coefficients_q': coefficients_q,
            'block_size': self.block_size,
            'K': self.K,
            'shape_orig': np.array([H, W, C]),
            'pad_shape': np.array([H_pad, W_pad]),
            'qstep': float(self.qstep),
            'zero_threshold': float(self.zero_threshold),
        }
        
        if verbose:
            ratio = orig_pixels / param_vals
            print(f"[SkissBlockEncoder] Image {H}×{W}×{C}  →  {n_by}×{n_bx} blocks")
            print(f"[SkissBlockEncoder] K={self.K}  |  Params={param_vals}  "
                  f"|  Compression={ratio:.0f}×  |  Sparsity={sparsity*100:.1f}%")
        
        return self.theta

    def save(self, path: Union[str, Path]):
        """Store as compressed npz. The int16 payload compresses massively with npz/ZIP."""
        if self.theta is None:
            raise RuntimeError("Encode first.")
        np.savez_compressed(str(path), **self.theta)
    
    @staticmethod
    def load_theta(path: Union[str, Path]) -> Dict:
        """Load residue dict from disk."""
        data = np.load(str(path), allow_pickle=False)
        return {k: (v.item() if v.ndim == 0 else v) for k, v in data.items()}


# ==============================================================================
#  Decoder: Skiss Residue → On-Demand Pixel Generation
# ==============================================================================

class SkissBlockDecoder:
    """
    Generates pixels from the block-wise skiss residue at arbitrary coordinates.
    """
    
    def __init__(self, theta: Dict):
        self.coefficients_q = theta['coefficients_q']
        self.block_size = int(theta['block_size'])
        self.K = int(theta['K'])
        self.shape_orig = tuple(theta['shape_orig'])
        self.pad_shape = tuple(theta['pad_shape'])
        self.C = int(self.shape_orig[2]) if len(self.shape_orig) == 3 else 1
        self.qstep = float(theta['qstep'])
        self.B = self.block_size
        
        # De-quantize back to continuous spectral residue
        self.coeffs = self.coefficients_q.astype(np.float64) * self.qstep
        self.n_by, self.n_bx = self.coefficients_q.shape[:2]

    def _generate_grid(self, x_grid: np.ndarray, y_grid: np.ndarray) -> np.ndarray:
        """
        Core engine: evaluate the block-wise parametric surface on any regular grid.
        x_grid, y_grid: 2D arrays of global continuous coordinates.
        """
        H_out, W_out = x_grid.shape
        out = np.zeros((H_out, W_out, self.C), dtype=np.float64)
        
        x_flat = x_grid.ravel()
        y_flat = y_grid.ravel()
        
        # Map global coords to block indices and local block coordinates
        bx_flat = (x_flat // self.B).astype(int)
        by_flat = (y_flat // self.B).astype(int)
        nx_flat = x_flat - bx_flat * self.B
        ny_flat = y_flat - by_flat * self.B
        
        # Clip to padded bounds
        bx_flat = np.clip(bx_flat, 0, self.n_bx - 1)
        by_flat = np.clip(by_flat, 0, self.n_by - 1)
        nx_flat = np.clip(nx_flat, 0, self.B - 1)
        ny_flat = np.clip(ny_flat, 0, self.B - 1)
        
        for by in range(self.n_by):
            for bx in range(self.n_bx):
                mask = (bx_flat == bx) & (by_flat == by)
                if not np.any(mask):
                    continue
                
                nx = nx_flat[mask]
                ny = ny_flat[mask]
                flat_idx = np.flatnonzero(mask)
                
                V = _build_dct_basis_1d(nx, self.K, self.B)   # (n, K)
                U = _build_dct_basis_1d(ny, self.K, self.B)   # (n, K)
                Theta = self.coeffs[by, bx]                    # (C, K, K)
                
                # Batch evaluate all channels: I_n = Σ U[n,k] Theta[k,l] V[n,l]
                vals = np.zeros((mask.sum(), self.C), dtype=np.float64)
                for c in range(self.C):
                    vals[:, c] = np.sum((U @ Theta[c]) * V, axis=1)
                
                out.reshape(-1, self.C)[flat_idx] = vals
        
        return out

    def generate(self, x: np.ndarray, y: np.ndarray, channel: Optional[int] = None) -> np.ndarray:
        """
        Evaluate at arbitrary real-valued coordinates (scalar or ndarray).
        """
        x = np.asarray(x, dtype=np.float64)
        y = np.asarray(y, dtype=np.float64)
        if x.shape != y.shape:
            raise ValueError("x and y must have identical shape.")
        
        orig_shape = x.shape
        x_grid, y_grid = x, y  # already 2D or any shape
        
        out = self._generate_grid(x_grid, y_grid)
        
        if channel is not None:
            return out[..., channel].reshape(orig_shape)
        return out.reshape(orig_shape + (self.C,))

    def reconstruct(self, shape: Optional[Tuple[int, int]] = None) -> np.ndarray:
        """
        Reconstruct a full pixel grid at any resolution.
        If shape is None, uses the original image dimensions.
        """
        H0, W0 = int(self.shape_orig[0]), int(self.shape_orig[1])

        if shape is None:
            # Native resolution: sample the EXACT integer pixel positions the
            # basis was fit on. No interpolation, no padding region touched.
            xs = np.arange(W0, dtype=np.float64)
            ys = np.arange(H0, dtype=np.float64)
        else:
            # Arbitrary resolution (e.g. super-res): map the output grid across
            # the ORIGINAL (unpadded) domain [0, W0-1] x [0, H0-1] so every
            # sample stays aligned to the block it was fit on and never lands
            # in the edge-replicated padding.
            H_out, W_out = shape
            xs = np.linspace(0, W0 - 1, W_out)
            ys = np.linspace(0, H0 - 1, H_out)

        x_grid, y_grid = np.meshgrid(xs, ys, indexing='xy')
        return self._generate_grid(x_grid, y_grid)

    def reconstruct_patch(self, x_min: float, x_max: float,
                          y_min: float, y_max: float,
                          res_x: int, res_y: int) -> np.ndarray:
        """
        Generate a continuous viewport patch at arbitrary resolution.
        Useful for region-of-interest decoding without loading the full grid.
        """
        xs = np.linspace(x_min, x_max, res_x)
        ys = np.linspace(y_min, y_max, res_y)
        x_grid, y_grid = np.meshgrid(xs, ys, indexing='xy')
        return self._generate_grid(x_grid, y_grid)


# ==============================================================================
#  Diagnostics
# ==============================================================================

def to_display(arr: np.ndarray, value_range: float = 255.0) -> np.ndarray:
    """
    Convert a raw float reconstruction into a display/save-ready uint8 image.

    The decoder returns float64 in the original pixel range (≈0–255), and can
    slightly overshoot (e.g. -9 .. 256) because the cosine surface is smooth.
    matplotlib/PIL treat float RGB as 0–1, so passing the raw array straight to
    imshow() clamps almost everything to white. Always run reconstructions
    through this before showing or saving them.
    """
    if value_range <= 1.0:
        arr = arr * 255.0
    return np.clip(np.round(arr), 0, 255).astype(np.uint8)


def skiss_metrics(original: np.ndarray, reconstructed: np.ndarray) -> Dict[str, float]:
    """MSE / PSNR / skiss-phase classification."""
    orig = original.astype(np.float64)
    rec = reconstructed.astype(np.float64)
    mse = float(np.mean((orig - rec) ** 2))
    max_val = 255.0 if orig.max() > 1.0 else 1.0
    psnr = float('inf') if mse == 0 else 20.0 * np.log10(max_val / np.sqrt(mse))
    phase = "✅ Solid" if psnr > 35 else ("⚠️ Liquid" if psnr > 20 else "❌ Gas")
    return {'mse': mse, 'psnr': psnr, 'phase': phase}


# ==============================================================================
#  Usage Example (mirrors your previous script)
# ==============================================================================

if __name__ == "__main__":
    from PIL import Image
    import pylab as plt
        
    # 1. Load a real photograph (e.g. 2752×1536 from your example)
    img = np.array(Image.open("photo.png").convert("RGB"))
    print(f"Original image: {img.shape}  ({img.nbytes / 1e6:.2f} MB in memory)")
    
    # 2. ENCODE: collapse into block-wise skiss residue
    #    K=8, qstep=4.0  →  expect ~100-200 KB file for a typical photo
    #encoder = SkissBlockEncoder(block_size=64, K=8, qstep=4.0, zero_threshold=1.0)
    encoder = SkissBlockEncoder(block_size=32, K=8, qstep=4.0, zero_threshold=1.0)
    theta = encoder.encode(img)
    encoder.save("photo.skiss.npz")
    
    # 3. Verify file size
    fsize = Path("photo.skiss.npz").stat().st_size
    orig_size = img.nbytes
    print(f"\nFile size: {fsize / 1024:.1f} KB  ({orig_size / fsize:.0f}× smaller than raw)")
    
    # 4. DECODE: generate from residue
    theta_loaded = SkissBlockEncoder.load_theta("photo.skiss.npz")
    decoder = SkissBlockDecoder(theta_loaded)
    
    # Full reconstruction at original resolution
    rec = decoder.reconstruct()
    m = skiss_metrics(img, rec)
    print(f"\nFull reconstruct: {m}")

    # rec is float64 in 0–255 range. Clip+cast before showing/saving,
    # otherwise imshow/PIL treat it as 0–1 and blow it out to black & white.
    rec_img = to_display(rec)
    Image.fromarray(rec_img).save("photo.reconstructed.png")
    plt.imshow(rec_img)
    plt.show()
    """
    # 5. On-demand: 2× super-resolution at half-pixel offsets (native to skiss theory)
    fine = decoder.reconstruct(shape=(img.shape[0] * 2, img.shape[1] * 2))
    print(f"Super-resolution shape: {fine.shape}")
    
    # 6. Arbitrary viewport (pay work only for the requested region)
    patch = decoder.reconstruct_patch(100.0, 400.0, 100.0, 400.0, res_x=800, res_y=800)
    print(f"Viewport patch shape: {patch.shape}")
    """
    """
    # 7. Progressive transmission demo (send low-K residue first)
    print("\n--- Progressive Skiss (fixed qstep=4.0) ---")
    for K_test in [2, 4, 6, 8, 12]:
        enc_tmp = SkissBlockEncoder(block_size=64, K=K_test, qstep=4.0)
        th_tmp = enc_tmp.encode(img, verbose=False)
        # Emulate file size by counting non-zero int16 coefficients
        nz = np.count_nonzero(th_tmp['coefficients_q'])
        dec_tmp = SkissBlockDecoder(th_tmp)
        rec_tmp = dec_tmp.reconstruct()
        m_tmp = skiss_metrics(img, rec_tmp)
        print(f"  K={K_test:2d}  |  Non-zero coeffs={nz:>7,}  |  "
              f"PSNR={m_tmp['psnr']:>6.2f} dB  |  {m_tmp['phase']}")
    """