#!/usr/bin/env python3
"""
Origami Encoder - Structure-Aware Tensor Compression

Exploits adjacency correlation (C_excess_z) for compression.

Algorithm:
  1. Anchor Selection: Pick representative rows (stride-based or k-means)
  2. Delta Encoding: Store differences from nearest anchor
  3. Motif Detection: Find repeated delta patterns (LZ77-style back-refs)
  4. Quantization: Compress deltas with adaptive precision

File Format (HXZO - Origami):
  [0:4]   Magic: b"HXZO" (reusing sidecar magic for origami v2)
  [4:6]   Version: uint16 BE (2 = origami format)
  [6:8]   Header length: uint16 BE
  [8:N]   Header JSON (tensor_name, shape, anchors_stride, etc.)
  [N:...]  Payload:
           - Anchors (float16 or float32)
           - Delta indices (uint16 back-refs + new deltas)
           - Delta values (quantized or float16)
"""

from __future__ import annotations

import json
import struct
import zlib
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

ORIGAMI_MAGIC = b"ORIG"
ORIGAMI_VERSION = 1


@dataclass
class OrigamiConfig:
    """Configuration for origami encoding."""
    anchor_stride: int = 32           # Store every Nth row as anchor
    delta_bits: int = 8               # Quantization bits for deltas
    motif_window: int = 128           # Lookback window for motif matching
    motif_min_length: int = 4         # Minimum motif length (rows)
    motif_threshold: float = 0.01     # Max L2 error for motif match
    use_float16_anchors: bool = True  # Anchor precision
    use_float16_deltas: bool = False  # Delta precision (False = quantized)


@dataclass
class OrigamiEncodeResult:
    """Result of origami encoding."""
    output_path: Path
    tensor_name: str
    shape: Tuple[int, ...]
    original_bytes: int
    encoded_bytes: int
    compression_ratio: float
    n_anchors: int
    n_motifs: int
    delta_entropy: float
    config: OrigamiConfig


class OrigamiEncoder:
    """Encodes tensors using origami (adjacency-aware) compression."""

    def __init__(self, config: OrigamiConfig = None):
        self.config = config or OrigamiConfig()

    def encode(
        self,
        tensor: np.ndarray,
        tensor_name: str,
        output_path: Path,
    ) -> OrigamiEncodeResult:
        """
        Encode a 2D tensor using origami compression.

        Args:
            tensor: 2D numpy array (rows x cols)
            tensor_name: Name for metadata
            output_path: Where to write the encoded file

        Returns:
            OrigamiEncodeResult with compression metrics
        """
        if tensor.ndim != 2:
            raise ValueError(f"Origami only supports 2D tensors, got {tensor.ndim}D")

        rows, cols = tensor.shape
        original_bytes = tensor.nbytes

        # 1. Select anchor rows
        anchors, anchor_indices = self._select_anchors(tensor)

        # 2. Compute deltas from nearest anchor
        deltas, delta_anchor_map = self._compute_deltas(tensor, anchors, anchor_indices)

        # 3. Find motifs (repeated delta patterns)
        motifs, motif_refs = self._find_motifs(deltas)

        # 4. Quantize remaining deltas
        quant_deltas, delta_codebook = self._quantize_deltas(deltas, motifs)

        # 5. Pack everything
        payload = self._pack_payload(
            anchors=anchors,
            anchor_indices=anchor_indices,
            delta_anchor_map=delta_anchor_map,
            quant_deltas=quant_deltas,
            delta_codebook=delta_codebook,
            motifs=motifs,
            motif_refs=motif_refs,
            shape=(rows, cols),
        )

        # 6. Write file
        # Determine if we used float16 deltas
        use_float16 = self.config.use_float16_deltas or self.config.delta_bits > 8

        header = {
            "tensor_name": tensor_name,
            "shape": [rows, cols],
            "dtype": str(tensor.dtype),
            "anchor_stride": self.config.anchor_stride,
            "n_anchors": len(anchor_indices),
            "n_motifs": len(motifs),
            "delta_bits": self.config.delta_bits,
            "use_float16_deltas": use_float16,
        }

        self._write_file(output_path, header, payload)

        encoded_bytes = output_path.stat().st_size
        compression_ratio = original_bytes / encoded_bytes if encoded_bytes > 0 else 0

        # Compute delta entropy
        delta_entropy = self._compute_entropy(quant_deltas)

        return OrigamiEncodeResult(
            output_path=output_path,
            tensor_name=tensor_name,
            shape=(rows, cols),
            original_bytes=original_bytes,
            encoded_bytes=encoded_bytes,
            compression_ratio=compression_ratio,
            n_anchors=len(anchor_indices),
            n_motifs=len(motifs),
            delta_entropy=delta_entropy,
            config=self.config,
        )

    def _select_anchors(
        self, tensor: np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Select anchor rows using stride sampling."""
        rows = tensor.shape[0]
        stride = self.config.anchor_stride

        # Stride-based selection
        anchor_indices = np.arange(0, rows, stride)

        # Always include last row if not already
        if anchor_indices[-1] != rows - 1:
            anchor_indices = np.append(anchor_indices, rows - 1)

        anchors = tensor[anchor_indices].astype(np.float32)

        if self.config.use_float16_anchors:
            anchors = anchors.astype(np.float16)

        return anchors, anchor_indices

    def _compute_deltas(
        self,
        tensor: np.ndarray,
        anchors: np.ndarray,
        anchor_indices: np.ndarray,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Compute delta from nearest anchor for each row."""
        rows, cols = tensor.shape
        deltas = np.zeros_like(tensor, dtype=np.float32)
        delta_anchor_map = np.zeros(rows, dtype=np.uint16)

        # For each row, find nearest anchor
        anchor_set = set(anchor_indices)
        current_anchor_idx = 0

        for i in range(rows):
            if i in anchor_set:
                # This row is an anchor - delta is zero
                # Find the index in anchor array
                current_anchor_idx = np.searchsorted(anchor_indices, i)
                delta_anchor_map[i] = current_anchor_idx
                deltas[i] = 0
            else:
                # Find nearest anchor (before or after)
                prev_anchor = anchor_indices[current_anchor_idx]

                if current_anchor_idx + 1 < len(anchor_indices):
                    next_anchor = anchor_indices[current_anchor_idx + 1]
                    # Choose closer anchor
                    if (i - prev_anchor) <= (next_anchor - i):
                        nearest_idx = current_anchor_idx
                    else:
                        nearest_idx = current_anchor_idx + 1
                else:
                    nearest_idx = current_anchor_idx

                delta_anchor_map[i] = nearest_idx
                # Convert anchor back to float32 for delta computation
                anchor_row = anchors[nearest_idx].astype(np.float32)
                deltas[i] = tensor[i].astype(np.float32) - anchor_row

        return deltas, delta_anchor_map

    def _find_motifs(
        self, deltas: np.ndarray
    ) -> Tuple[List[np.ndarray], List[Tuple[int, int, int]]]:
        """
        Find repeated delta patterns (LZ77-style).

        Returns:
            motifs: List of unique motif patterns
            motif_refs: List of (start_row, motif_idx, length) tuples
        """
        rows = deltas.shape[0]
        window = self.config.motif_window
        min_len = self.config.motif_min_length
        threshold = self.config.motif_threshold

        motifs = []
        motif_refs = []
        motif_cache = {}  # hash -> motif_idx

        i = 0
        while i < rows:
            # Try to find a matching motif in the lookback window
            best_match = None
            best_length = 0

            # Look back in window
            start = max(0, i - window)
            for j in range(start, i):
                # Try to match starting at j
                match_len = 0
                while (i + match_len < rows and
                       j + match_len < i and
                       match_len < window):
                    # Check if rows match within threshold
                    diff = np.abs(deltas[i + match_len] - deltas[j + match_len])
                    if np.max(diff) < threshold:
                        match_len += 1
                    else:
                        break

                if match_len >= min_len and match_len > best_length:
                    best_match = j
                    best_length = match_len

            if best_match is not None:
                # Found a motif reference
                # Store the pattern from best_match
                pattern = deltas[best_match:best_match + best_length].copy()

                # Check if this pattern is already stored
                pattern_hash = hash(pattern.tobytes())
                if pattern_hash in motif_cache:
                    motif_idx = motif_cache[pattern_hash]
                else:
                    motif_idx = len(motifs)
                    motifs.append(pattern)
                    motif_cache[pattern_hash] = motif_idx

                motif_refs.append((i, motif_idx, best_length))
                i += best_length
            else:
                i += 1

        return motifs, motif_refs

    def _quantize_deltas(
        self, deltas: np.ndarray, motifs: List[np.ndarray]
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Quantize deltas using uniform quantization or float16."""
        bits = self.config.delta_bits

        # Use float16 for higher fidelity (bits > 8 or explicit flag)
        if self.config.use_float16_deltas or bits > 8:
            # Just use float16 - high fidelity mode
            return deltas.astype(np.float16), np.array([], dtype=np.float32)

        # Uniform quantization with uint8 (bits <= 8)
        levels = 2 ** bits

        # Find min/max excluding zeros (anchor rows)
        nonzero_mask = np.any(deltas != 0, axis=1)
        if np.sum(nonzero_mask) == 0:
            # All deltas are zero
            return np.zeros_like(deltas, dtype=np.uint8), np.array([0.0], dtype=np.float32)

        nonzero_deltas = deltas[nonzero_mask]
        d_min = nonzero_deltas.min()
        d_max = nonzero_deltas.max()

        # Build codebook
        codebook = np.linspace(d_min, d_max, levels, dtype=np.float32)

        # Quantize
        # Scale to [0, levels-1]
        if d_max > d_min:
            scaled = (deltas - d_min) / (d_max - d_min) * (levels - 1)
        else:
            scaled = np.zeros_like(deltas)

        quant = np.clip(np.round(scaled), 0, levels - 1).astype(np.uint8)

        return quant, codebook

    def _compute_entropy(self, quant_deltas: np.ndarray) -> float:
        """Compute Shannon entropy of quantized deltas."""
        flat = quant_deltas.flatten()
        unique, counts = np.unique(flat, return_counts=True)
        probs = counts / counts.sum()
        entropy = -np.sum(probs * np.log2(probs + 1e-10))
        return float(entropy)

    def _pack_payload(
        self,
        anchors: np.ndarray,
        anchor_indices: np.ndarray,
        delta_anchor_map: np.ndarray,
        quant_deltas: np.ndarray,
        delta_codebook: np.ndarray,
        motifs: List[np.ndarray],
        motif_refs: List[Tuple[int, int, int]],
        shape: Tuple[int, int],
    ) -> bytes:
        """Pack all components into compressed payload."""
        parts = []

        # Part 1: Anchors
        anchor_bytes = anchors.tobytes()
        parts.append(struct.pack(">I", len(anchor_bytes)))
        parts.append(anchor_bytes)

        # Part 2: Anchor indices
        anchor_idx_bytes = anchor_indices.astype(np.uint32).tobytes()
        parts.append(struct.pack(">I", len(anchor_idx_bytes)))
        parts.append(anchor_idx_bytes)

        # Part 3: Delta-to-anchor map
        map_bytes = delta_anchor_map.tobytes()
        parts.append(struct.pack(">I", len(map_bytes)))
        parts.append(map_bytes)

        # Part 4: Quantized deltas
        delta_bytes = quant_deltas.tobytes()
        parts.append(struct.pack(">I", len(delta_bytes)))
        parts.append(delta_bytes)

        # Part 5: Delta codebook
        cb_bytes = delta_codebook.tobytes()
        parts.append(struct.pack(">I", len(cb_bytes)))
        parts.append(cb_bytes)

        # Part 6: Motifs (count + each motif)
        parts.append(struct.pack(">H", len(motifs)))
        for motif in motifs:
            m_bytes = motif.astype(np.float16).tobytes()
            parts.append(struct.pack(">I", len(m_bytes)))
            parts.append(m_bytes)

        # Part 7: Motif references (count + each ref)
        parts.append(struct.pack(">I", len(motif_refs)))
        for start_row, motif_idx, length in motif_refs:
            parts.append(struct.pack(">IHH", start_row, motif_idx, length))

        # Compress
        raw = b"".join(parts)
        compressed = zlib.compress(raw, level=9)

        return compressed

    def _write_file(self, path: Path, header: Dict[str, Any], payload: bytes):
        """Write origami file."""
        header_json = json.dumps(header).encode("utf-8")

        with open(path, "wb") as f:
            f.write(ORIGAMI_MAGIC)
            f.write(struct.pack(">H", ORIGAMI_VERSION))
            f.write(struct.pack(">H", len(header_json)))
            f.write(header_json)
            f.write(payload)


def encode_tensor_origami(
    tensor: np.ndarray,
    tensor_name: str,
    output_path: str,
    config: OrigamiConfig = None,
) -> Dict[str, Any]:
    """
    Convenience function to encode a tensor with origami compression.

    Args:
        tensor: 2D numpy array
        tensor_name: Name for metadata
        output_path: Output file path
        config: Optional configuration

    Returns:
        Receipt dict with compression metrics
    """
    encoder = OrigamiEncoder(config)
    result = encoder.encode(tensor, tensor_name, Path(output_path))

    return {
        "status": "OK",
        "tensor_name": result.tensor_name,
        "shape": list(result.shape),
        "original_bytes": result.original_bytes,
        "encoded_bytes": result.encoded_bytes,
        "compression_ratio": result.compression_ratio,
        "n_anchors": result.n_anchors,
        "n_motifs": result.n_motifs,
        "delta_entropy": result.delta_entropy,
        "output_path": str(result.output_path),
        "config": {
            "anchor_stride": result.config.anchor_stride,
            "delta_bits": result.config.delta_bits,
            "motif_window": result.config.motif_window,
        }
    }
