#!/usr/bin/env python3
"""
Origami Decoder - Reconstruct tensors from origami-compressed format.

Reverses the origami encoding:
  1. Decompress payload
  2. Reconstruct anchors
  3. Dequantize deltas
  4. Apply motif expansions
  5. Add deltas to anchors to recover original rows
"""

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 OrigamiDecodeResult:
    """Result of origami decoding."""
    tensor: np.ndarray
    tensor_name: str
    shape: Tuple[int, ...]
    n_anchors: int
    n_motifs: int


class OrigamiDecoder:
    """Decodes origami-compressed tensors."""

    def decode(self, input_path: Path) -> OrigamiDecodeResult:
        """
        Decode an origami-compressed file back to a tensor.

        Args:
            input_path: Path to .orig file

        Returns:
            OrigamiDecodeResult with reconstructed tensor
        """
        header, payload = self._read_file(input_path)

        shape = tuple(header["shape"])
        rows, cols = shape

        # Unpack payload
        (
            anchors,
            anchor_indices,
            delta_anchor_map,
            quant_deltas,
            delta_codebook,
            motifs,
            motif_refs,
        ) = self._unpack_payload(payload, shape, header)

        # Dequantize deltas
        if len(delta_codebook) > 0:
            deltas = self._dequantize_deltas(quant_deltas, delta_codebook)
        else:
            # Float16 deltas were stored directly
            deltas = quant_deltas.astype(np.float32)

        # Apply motif expansions (copy motif patterns to their reference locations)
        for start_row, motif_idx, length in motif_refs:
            if motif_idx < len(motifs):
                motif = motifs[motif_idx]
                end_row = min(start_row + length, rows)
                actual_len = end_row - start_row
                deltas[start_row:end_row] = motif[:actual_len].astype(np.float32)

        # Reconstruct tensor: anchor + delta for each row
        tensor = np.zeros((rows, cols), dtype=np.float32)

        for i in range(rows):
            anchor_idx = delta_anchor_map[i]
            anchor_row = anchors[anchor_idx].astype(np.float32)
            tensor[i] = anchor_row + deltas[i]

        return OrigamiDecodeResult(
            tensor=tensor,
            tensor_name=header.get("tensor_name", "unknown"),
            shape=shape,
            n_anchors=len(anchor_indices),
            n_motifs=len(motifs),
        )

    def _read_file(self, path: Path) -> Tuple[Dict[str, Any], bytes]:
        """Read origami file header and payload."""
        with open(path, "rb") as f:
            magic = f.read(4)
            if magic != ORIGAMI_MAGIC:
                raise ValueError(f"Invalid origami magic: {magic!r}, expected {ORIGAMI_MAGIC!r}")

            version = struct.unpack(">H", f.read(2))[0]
            if version != ORIGAMI_VERSION:
                raise ValueError(f"Unsupported origami version: {version}")

            header_len = struct.unpack(">H", f.read(2))[0]
            header = json.loads(f.read(header_len).decode("utf-8"))

            payload = f.read()

        return header, payload

    def _unpack_payload(
        self,
        compressed: bytes,
        shape: Tuple[int, int],
        header: Dict[str, Any],
    ) -> Tuple:
        """Unpack compressed payload into components."""
        raw = zlib.decompress(compressed)
        rows, cols = shape

        offset = 0

        def read_array(dtype, count=None):
            nonlocal offset
            length = struct.unpack_from(">I", raw, offset)[0]
            offset += 4
            data = raw[offset:offset + length]
            offset += length
            return np.frombuffer(data, dtype=dtype)

        # Part 1: Anchors
        anchor_len = struct.unpack_from(">I", raw, offset)[0]
        offset += 4
        # Determine anchor dtype from header
        n_anchors = header.get("n_anchors", 0)
        anchor_dtype = np.float16  # default
        anchor_bytes_expected = n_anchors * cols * 2  # float16
        if anchor_len != anchor_bytes_expected:
            anchor_dtype = np.float32
        anchors_flat = np.frombuffer(raw[offset:offset + anchor_len], dtype=anchor_dtype)
        offset += anchor_len
        if n_anchors > 0:
            anchors = anchors_flat.reshape(n_anchors, cols)
        else:
            anchors = anchors_flat.reshape(-1, cols) if len(anchors_flat) > 0 else np.zeros((0, cols), dtype=anchor_dtype)

        # Part 2: Anchor indices
        idx_len = struct.unpack_from(">I", raw, offset)[0]
        offset += 4
        anchor_indices = np.frombuffer(raw[offset:offset + idx_len], dtype=np.uint32)
        offset += idx_len

        # Part 3: Delta-to-anchor map
        map_len = struct.unpack_from(">I", raw, offset)[0]
        offset += 4
        delta_anchor_map = np.frombuffer(raw[offset:offset + map_len], dtype=np.uint16)
        offset += map_len

        # Part 4: Quantized deltas
        delta_len = struct.unpack_from(">I", raw, offset)[0]
        offset += 4
        use_float16 = header.get("use_float16_deltas", False)
        delta_bits = header.get("delta_bits", 8)
        # Use float16 if explicitly set or if bits > 8
        if use_float16 or delta_bits > 8:
            delta_dtype = np.float16
        else:
            delta_dtype = np.uint8
        quant_deltas_flat = np.frombuffer(raw[offset:offset + delta_len], dtype=delta_dtype)
        offset += delta_len
        # Reshape to (rows, cols)
        quant_deltas = quant_deltas_flat.reshape(rows, cols)

        # Part 5: Delta codebook
        cb_len = struct.unpack_from(">I", raw, offset)[0]
        offset += 4
        delta_codebook = np.frombuffer(raw[offset:offset + cb_len], dtype=np.float32)
        offset += cb_len

        # Part 6: Motifs
        n_motifs = struct.unpack_from(">H", raw, offset)[0]
        offset += 2
        motifs = []
        for _ in range(n_motifs):
            m_len = struct.unpack_from(">I", raw, offset)[0]
            offset += 4
            motif = np.frombuffer(raw[offset:offset + m_len], dtype=np.float16)
            offset += m_len
            # Reshape: each motif is (length, cols)
            motif_rows = len(motif) // cols
            if motif_rows > 0:
                motifs.append(motif.reshape(motif_rows, cols))
            else:
                motifs.append(motif.reshape(-1, cols))

        # Part 7: Motif references
        n_refs = struct.unpack_from(">I", raw, offset)[0]
        offset += 4
        motif_refs = []
        for _ in range(n_refs):
            start_row, motif_idx, length = struct.unpack_from(">IHH", raw, offset)
            offset += 8
            motif_refs.append((start_row, motif_idx, length))

        return (
            anchors,
            anchor_indices,
            delta_anchor_map,
            quant_deltas,
            delta_codebook,
            motifs,
            motif_refs,
        )

    def _dequantize_deltas(
        self, quant_deltas: np.ndarray, codebook: np.ndarray
    ) -> np.ndarray:
        """Dequantize deltas using codebook."""
        if len(codebook) == 0:
            return quant_deltas.astype(np.float32)

        # Simple lookup
        deltas = codebook[quant_deltas]
        return deltas


def decode_origami_tensor(input_path: str) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Convenience function to decode an origami file.

    Args:
        input_path: Path to .orig file

    Returns:
        (tensor, metadata_dict)
    """
    decoder = OrigamiDecoder()
    result = decoder.decode(Path(input_path))

    metadata = {
        "tensor_name": result.tensor_name,
        "shape": list(result.shape),
        "n_anchors": result.n_anchors,
        "n_motifs": result.n_motifs,
    }

    return result.tensor, metadata
