#!/usr/bin/env python3
"""
svd_capsule.py - SVD Capsule Format for Helix64 Streaming Matmul

Stores low-rank SVD factors (U, S, Vt) in a compact capsule format.
Enables streaming matmul via C04_SVD_MATMUL without materializing W.

File Format (HXZS v1 - SVD Capsule):
  [0:4]     Magic: b"HXZS"
  [4:6]     Version: uint16 BE (1)
  [6:8]     Header length: uint16 BE
  [8:N]     Header JSON with:
              - tensor_name, original_shape, rank
              - U_shape, S_shape, Vt_shape
              - dtype (fp16 or fp32)
              - compression stats
  [N:...]   Payload (may be zlib-compressed):
              - S values [rank] float32 (always uncompressed for precision)
              - U values [in_features × rank] fp16/fp32
              - Vt values [rank × out_features] fp16/fp32

Key insight:
  - W (m×n) requires m*n*4 bytes
  - SVD capsule: U (m×r)*2 + S (r)*4 + Vt (r×n)*2
  - At rank 64, for 4096×4096: 64MB → 1MB (64x compression)

Work Order: WO-HELIX64-SVD-MATMUL
"""

from __future__ import annotations

import json
import struct
import zlib
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, Literal, Optional, Tuple

import numpy as np


# HXZS format constants (SVD capsule)
HXZS_MAGIC = b"HXZS"
HXZS_VERSION = 1

# Layout convention type
LayoutType = Literal["W", "WT"]


def factors_for_matmul(
    U: np.ndarray,
    S: np.ndarray,
    Vt: np.ndarray,
    stored_layout: LayoutType,
    target_layout: LayoutType = "WT",
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Return factors for the desired matmul convention.

    PyTorch linear layers compute Y = X @ W.T where W is stored as [out, in].
    This function converts between W and W.T conventions without re-computing SVD.

    Math identity:
        If A = U @ diag(S) @ Vt, then A.T = Vt.T @ diag(S) @ U.T

    Args:
        U: Left singular vectors from SVD
        S: Singular values
        Vt: Right singular vectors (V transposed)
        stored_layout: What was decomposed ("W" or "WT")
        target_layout: What matmul needs (usually "WT" for Y = X @ W.T)

    Returns:
        (U_out, S_out, Vt_out) for the target layout

    Example:
        # If capsule stores SVD of W but we need Y = X @ W.T:
        U_t, S_t, Vt_t = factors_for_matmul(U, S, Vt, stored_layout="W", target_layout="WT")
        # Now: (U_t * S_t) @ Vt_t ≈ W.T
    """
    if stored_layout == target_layout:
        return U, S, Vt

    # Swap: A.T = Vt.T @ diag(S) @ U.T
    # So U_new = Vt.T, Vt_new = U.T, S stays same
    return Vt.T.copy(), S, U.T.copy()


def compute_svd_truncated(
    W: np.ndarray,
    rank: int,
    full_svd_fallback: bool = True,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, Any]]:
    """
    Compute rank-r truncated SVD of weight matrix W.

    Args:
        W: Weight matrix [out_features, in_features] or [in_features, out_features]
        rank: Desired rank (will be clipped to min(W.shape) - 1)
        full_svd_fallback: Use full SVD if sparse SVD fails

    Returns:
        (U, S, Vt, stats) where W ≈ U @ diag(S) @ Vt
    """
    from scipy.sparse.linalg import svds
    from scipy.linalg import svd as full_svd_fn

    W = np.asarray(W, dtype=np.float64)  # SVD needs float64 for numerical stability
    m, n = W.shape

    # Clip rank to valid range
    max_rank = min(m, n) - 1
    actual_rank = min(rank, max_rank)
    if actual_rank < 1:
        actual_rank = 1

    try:
        # Sparse SVD (faster for low rank)
        U, S, Vt = svds(W, k=actual_rank)
        # svds returns ascending order, we want descending
        idx = np.argsort(S)[::-1]
        U = U[:, idx]
        S = S[idx]
        Vt = Vt[idx, :]
    except Exception as e:
        if full_svd_fallback:
            # Fall back to full SVD
            U, S, Vt = full_svd_fn(W, full_matrices=False)
            U = U[:, :actual_rank]
            S = S[:actual_rank]
            Vt = Vt[:actual_rank, :]
        else:
            raise RuntimeError(f"SVD failed: {e}")

    # Compute reconstruction metrics
    W_approx = (U * S) @ Vt
    residual = W - W_approx

    frobenius_error = np.linalg.norm(residual) / np.linalg.norm(W)
    energy_captured = np.sum(S**2) / np.sum(np.linalg.svd(W, compute_uv=False)**2)

    # Cosine similarity of flattened matrices
    W_flat = W.flatten()
    W_approx_flat = W_approx.flatten().astype(W.dtype)
    cosine = np.dot(W_flat, W_approx_flat) / (np.linalg.norm(W_flat) * np.linalg.norm(W_approx_flat))

    stats = {
        "requested_rank": rank,
        "actual_rank": actual_rank,
        "original_shape": list(W.shape),
        "frobenius_error": float(frobenius_error),
        "energy_captured": float(energy_captured),
        "cosine_similarity": float(cosine),
        "singular_values_top5": S[:5].tolist(),
        "singular_values_sum": float(np.sum(S)),
    }

    return U.astype(np.float32), S.astype(np.float32), Vt.astype(np.float32), stats


def encode_svd_capsule(
    U: np.ndarray,
    S: np.ndarray,
    Vt: np.ndarray,
    tensor_name: str,
    original_shape: Tuple[int, int],
    output_path: str,
    use_fp16: bool = True,
    compress: bool = True,
    decomposed: LayoutType = "WT",
    extra_meta: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Encode SVD factors into HXZS capsule format.

    Args:
        U: Left singular vectors [m × rank] float32
        S: Singular values [rank] float32
        Vt: Right singular vectors [rank × n] float32
        tensor_name: Name of the original tensor
        original_shape: Original tensor shape (m, n)
        output_path: Where to write .hxzs file
        use_fp16: Store U, Vt as float16 (saves 2x space)
        compress: Apply zlib compression to payload
        decomposed: Layout convention ("W" or "WT"). "WT" is default for
                    PyTorch linear layers where Y = X @ W.T.
        extra_meta: Additional metadata

    Returns:
        Receipt with encoding stats
    """
    U = np.asarray(U, dtype=np.float32)
    S = np.asarray(S, dtype=np.float32).ravel()
    Vt = np.asarray(Vt, dtype=np.float32)

    rank = len(S)
    m, n = original_shape

    # Convert U, Vt to fp16 if requested
    store_dtype = np.float16 if use_fp16 else np.float32
    U_store = U.astype(store_dtype)
    Vt_store = Vt.astype(store_dtype)

    # S always float32 for precision
    S_store = S.astype(np.float32)

    # Build header
    header = {
        "schema": "hxzs_svd_capsule_v1",
        "tensor_name": tensor_name,
        "original_shape": list(original_shape),
        "rank": rank,
        "decomposed": decomposed,  # "W" or "WT" - which matrix was SVD'd
        "in_features": original_shape[1] if decomposed == "W" else original_shape[0],
        "out_features": original_shape[0] if decomposed == "W" else original_shape[1],
        "U_shape": list(U_store.shape),
        "S_shape": [rank],
        "Vt_shape": list(Vt_store.shape),
        "factor_dtype": "float16" if use_fp16 else "float32",
        "S_dtype": "float32",
        "compressed": compress,
        "created_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    }
    if extra_meta:
        header.update(extra_meta)

    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")

    # Build payload: S (always first, always fp32) + U + Vt
    S_bytes = S_store.tobytes()
    U_bytes = U_store.tobytes()
    Vt_bytes = Vt_store.tobytes()

    payload = S_bytes + U_bytes + Vt_bytes
    payload_uncompressed_size = len(payload)

    if compress:
        payload = zlib.compress(payload, level=9)

    # Build file
    output = bytearray()
    output.extend(HXZS_MAGIC)
    output.extend(struct.pack(">H", HXZS_VERSION))
    output.extend(struct.pack(">H", len(header_bytes)))
    output.extend(header_bytes)
    output.extend(payload)

    # Write
    out_path = Path(output_path)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_bytes(bytes(output))

    # Storage calculations
    original_bytes = m * n * 4  # Original fp32
    capsule_bytes = len(output)

    # Without compression, theoretical bytes
    bytes_per_factor = 2 if use_fp16 else 4
    theoretical_bytes = (
        rank * 4 +                    # S (always fp32)
        m * rank * bytes_per_factor + # U
        rank * n * bytes_per_factor   # Vt
    )

    receipt = {
        "schema": "encode_svd_capsule_v1",
        "tensor_name": tensor_name,
        "original_shape": list(original_shape),
        "rank": rank,
        "decomposed": decomposed,
        "in_features": original_shape[1] if decomposed == "W" else original_shape[0],
        "out_features": original_shape[0] if decomposed == "W" else original_shape[1],
        "factor_dtype": "float16" if use_fp16 else "float32",
        "compressed": compress,
        "original_bytes": original_bytes,
        "theoretical_bytes": theoretical_bytes,
        "payload_uncompressed_bytes": payload_uncompressed_size,
        "payload_compressed_bytes": len(payload) if compress else payload_uncompressed_size,
        "capsule_bytes": capsule_bytes,
        "compression_vs_original": original_bytes / capsule_bytes,
        "compression_vs_theoretical": theoretical_bytes / capsule_bytes if compress else 1.0,
        "output_path": str(out_path),
    }

    return receipt


def decode_svd_capsule(
    capsule_path: str,
    return_fp32: bool = True,
    target_layout: Optional[LayoutType] = None,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, Any]]:
    """
    Decode SVD factors from HXZS capsule.

    Args:
        capsule_path: Path to .hxzs file
        return_fp32: Convert factors to float32 (default True for compute)
        target_layout: If specified, convert factors to this layout ("W" or "WT").
                       If None, returns factors as stored.

    Returns:
        (U, S, Vt, metadata) where:
          - U: [m × rank]
          - S: [rank]
          - Vt: [rank × n]
          - metadata includes 'decomposed' field indicating stored layout
    """
    path = Path(capsule_path)
    data = path.read_bytes()

    # Parse header
    if data[:4] != HXZS_MAGIC:
        raise ValueError(f"Not an HXZS file (magic: {data[:4]!r})")

    version = struct.unpack(">H", data[4:6])[0]
    if version != HXZS_VERSION:
        raise ValueError(f"Unsupported HXZS version: {version}")

    header_len = struct.unpack(">H", data[6:8])[0]
    header = json.loads(data[8:8 + header_len].decode("utf-8"))

    # Decompress payload if needed
    payload = data[8 + header_len:]
    if header.get("compressed", False):
        payload = zlib.decompress(payload)

    # Parse factors
    rank = header["rank"]
    U_shape = tuple(header["U_shape"])
    Vt_shape = tuple(header["Vt_shape"])

    factor_dtype = np.float16 if header.get("factor_dtype") == "float16" else np.float32

    # S is always float32
    S_bytes = rank * 4
    S = np.frombuffer(payload[:S_bytes], dtype=np.float32).copy()

    # U
    U_elements = U_shape[0] * U_shape[1]
    U_bytes = U_elements * (2 if factor_dtype == np.float16 else 4)
    U = np.frombuffer(payload[S_bytes:S_bytes + U_bytes], dtype=factor_dtype).copy()
    U = U.reshape(U_shape)

    # Vt
    Vt = np.frombuffer(payload[S_bytes + U_bytes:], dtype=factor_dtype).copy()
    Vt = Vt.reshape(Vt_shape)

    # Convert to fp32 if requested
    if return_fp32 and factor_dtype == np.float16:
        U = U.astype(np.float32)
        Vt = Vt.astype(np.float32)

    # Get stored layout (default "WT" for backward compatibility)
    stored_layout: LayoutType = header.get("decomposed", "WT")  # type: ignore

    # Convert layout if requested
    output_layout = stored_layout
    if target_layout is not None and target_layout != stored_layout:
        U, S, Vt = factors_for_matmul(U, S, Vt, stored_layout, target_layout)
        output_layout = target_layout

    metadata = {
        "tensor_name": header.get("tensor_name", "unknown"),
        "original_shape": tuple(header.get("original_shape", [])),
        "rank": rank,
        "decomposed": stored_layout,  # What was originally stored
        "output_layout": output_layout,  # What's being returned
        "in_features": header.get("in_features"),
        "out_features": header.get("out_features"),
        "factor_dtype": header.get("factor_dtype", "float32"),
        "file_size_bytes": len(data),
        "header": header,
    }

    return U, S, Vt, metadata


def inspect_svd_capsule(capsule_path: str) -> Dict[str, Any]:
    """
    Inspect HXZS capsule header without loading factors.

    Args:
        capsule_path: Path to .hxzs file

    Returns:
        Header metadata and storage stats
    """
    path = Path(capsule_path)
    file_size = path.stat().st_size

    with path.open("rb") as f:
        magic = f.read(4)
        if magic != HXZS_MAGIC:
            return {"valid": False, "error": f"Invalid magic: {magic!r}"}

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

    # Compute storage efficiency
    original_shape = header.get("original_shape", [0, 0])
    m, n = original_shape if len(original_shape) == 2 else (0, 0)
    rank = header.get("rank", 0)

    original_bytes = m * n * 4
    compression_ratio = original_bytes / file_size if file_size > 0 else 0

    return {
        "valid": True,
        "version": version,
        "file_size_bytes": file_size,
        "tensor_name": header.get("tensor_name", "unknown"),
        "original_shape": original_shape,
        "rank": rank,
        "decomposed": header.get("decomposed", "WT"),  # Default for backward compat
        "in_features": header.get("in_features"),
        "out_features": header.get("out_features"),
        "U_shape": header.get("U_shape", []),
        "Vt_shape": header.get("Vt_shape", []),
        "factor_dtype": header.get("factor_dtype", "float32"),
        "compressed": header.get("compressed", False),
        "original_bytes": original_bytes,
        "compression_ratio": compression_ratio,
        "header": header,
    }


def svd_matmul_from_capsule(
    capsule_path: str,
    X: np.ndarray,
    matmul_convention: LayoutType = "WT",
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Compute X @ W (or X @ W.T) using SVD capsule WITHOUT materializing W.

    This is the streaming matmul entry point.

    Args:
        capsule_path: Path to .hxzs file
        X: Input batch [batch × in_features]
        matmul_convention: What operation to compute:
            - "WT": Y = X @ W.T (PyTorch linear default, where W is [out, in])
            - "W":  Y = X @ W (direct matmul with stored matrix)

    Returns:
        (output [batch × out_features], stats)
    """
    # Decode with layout conversion to match matmul convention
    U, S, Vt, meta = decode_svd_capsule(capsule_path, target_layout=matmul_convention)

    X = np.asarray(X, dtype=np.float32)

    # Streaming SVD matmul: X @ A = ((X @ U) * S) @ Vt
    # where A is the target matrix (W or W.T depending on convention)
    XU = np.matmul(X, U)
    XU_S = XU * S
    result = np.matmul(XU_S, Vt)

    # Stats
    batch_size = X.shape[0]
    in_features = X.shape[1] if X.ndim > 1 else X.shape[0]
    rank = len(S)
    out_features = Vt.shape[1]

    # FLOPs
    flops_step1 = 2 * batch_size * in_features * rank
    flops_step2 = batch_size * rank
    flops_step3 = 2 * batch_size * rank * out_features
    total_flops = flops_step1 + flops_step2 + flops_step3

    # Memory comparison
    materialized_W_bytes = in_features * out_features * 4
    factors_bytes = U.nbytes + S.nbytes + Vt.nbytes

    stats = {
        "tensor_name": meta.get("tensor_name"),
        "rank": rank,
        "batch_size": batch_size,
        "in_features": in_features,
        "out_features": out_features,
        "matmul_convention": matmul_convention,
        "stored_layout": meta.get("decomposed"),
        "layout_converted": meta.get("decomposed") != matmul_convention,
        "total_flops": total_flops,
        "W_materialized_bytes": materialized_W_bytes,
        "factors_bytes": factors_bytes,
        "memory_savings": materialized_W_bytes / factors_bytes,
        "W_materialized": False,
    }

    return result, stats


def reconstruct_W_from_capsule(capsule_path: str) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Reconstruct full W from SVD capsule (for verification/comparison).

    WARNING: This defeats the purpose of streaming matmul!
    Only use for verification.

    Args:
        capsule_path: Path to .hxzs file

    Returns:
        (W_approx [m × n], metadata)
    """
    U, S, Vt, meta = decode_svd_capsule(capsule_path)
    W_approx = (U * S) @ Vt
    return W_approx, meta


# CLI test
if __name__ == "__main__":
    import tempfile

    print("=== SVD Capsule Format Test ===")
    print()

    # Create test weight matrix
    rng = np.random.default_rng(42)
    m, n = 4096, 4096
    rank_true = 64  # Low-rank ground truth

    # Generate low-rank matrix W = A @ B where A is m×rank, B is rank×n
    A = rng.standard_normal((m, rank_true)).astype(np.float32)
    B = rng.standard_normal((rank_true, n)).astype(np.float32)
    W = A @ B

    print(f"Original W: {W.shape}, {W.nbytes:,} bytes ({W.nbytes/1024/1024:.1f} MB)")
    print(f"True rank: {rank_true}")
    print()

    # Compute truncated SVD
    print("Computing truncated SVD...")
    U, S, Vt, svd_stats = compute_svd_truncated(W, rank=64)
    print(f"  Rank: {svd_stats['actual_rank']}")
    print(f"  Frobenius error: {svd_stats['frobenius_error']:.6f}")
    print(f"  Cosine similarity: {svd_stats['cosine_similarity']:.6f}")
    print(f"  Energy captured: {svd_stats['energy_captured']:.4f}")
    print()

    # Encode to capsule
    with tempfile.NamedTemporaryFile(suffix=".hxzs", delete=False) as f:
        tmp_path = f.name

    print("Encoding SVD capsule...")
    receipt = encode_svd_capsule(
        U=U, S=S, Vt=Vt,
        tensor_name="test_weight",
        original_shape=(m, n),
        output_path=tmp_path,
        use_fp16=True,
        compress=True,
    )
    print(f"  Capsule size: {receipt['capsule_bytes']:,} bytes ({receipt['capsule_bytes']/1024:.1f} KB)")
    print(f"  Compression vs original: {receipt['compression_vs_original']:.1f}x")
    print()

    # Decode and verify
    print("Decoding SVD capsule...")
    U_dec, S_dec, Vt_dec, meta = decode_svd_capsule(tmp_path)
    print(f"  U shape: {U_dec.shape}")
    print(f"  S shape: {S_dec.shape}")
    print(f"  Vt shape: {Vt_dec.shape}")
    print()

    # Streaming matmul test
    print("Streaming matmul test...")
    batch_size = 32
    X = rng.standard_normal((batch_size, m)).astype(np.float32)

    # Ground truth (materializing W)
    Y_gt = X @ W

    # Streaming (no W materialization)
    Y_stream, stats = svd_matmul_from_capsule(tmp_path, X)

    # Compare
    error = np.linalg.norm(Y_gt - Y_stream) / np.linalg.norm(Y_gt)
    cosine = np.dot(Y_gt.flatten(), Y_stream.flatten()) / (np.linalg.norm(Y_gt) * np.linalg.norm(Y_stream))

    print(f"  Batch size: {batch_size}")
    print(f"  Relative L2 error: {error:.6f}")
    print(f"  Cosine similarity: {cosine:.6f}")
    print(f"  W materialized: {stats['W_materialized']}")
    print(f"  Memory savings: {stats['memory_savings']:.1f}x")
    print()

    # Cleanup
    Path(tmp_path).unlink()

    # Summary
    print("=" * 50)
    if error < 0.01 and cosine > 0.999:
        print("PASS: SVD capsule streaming matmul works correctly")
    else:
        print("FAIL: SVD capsule streaming matmul has issues")
