import numpy as np
import zlib

def compress(data: np.ndarray) -> bytes:
    """
    Compress a 1D numpy array using double cumsum + zlib.
    
    Parameters
    ----------
    data : np.ndarray
        Input array of any numeric type (converted to float64 internally).
    
    Returns
    -------
    bytes
        Compressed byte string.
    """
    # Ensure float64 for precision
    x = data.astype(np.float64)
    
    # Scale down to expose byte patterns
    scaled = x * 1e-9
    
    # Double cumulative sum
    first_cumsum = np.cumsum(scaled)
    second_cumsum = np.cumsum(first_cumsum)
    
    # Serialize and compress
    raw_bytes = second_cumsum.tobytes()
    compressed = zlib.compress(raw_bytes, level=9)
    return compressed


def decompress(compressed: bytes, original_shape: tuple = None) -> np.ndarray:
    """
    Decompress and restore the original data.
    
    Parameters
    ----------
    compressed : bytes
        Compressed byte string produced by `compress()`.
    original_shape : tuple, optional
        If the original array was multi‑dimensional, provide its shape.
        For 1D arrays, can be None (the flat length is inferred).
    
    Returns
    -------
    np.ndarray
        Restored array (float64). Use `.astype(original_dtype)` if needed.
    """
    # Decompress and restore float64 sequence
    raw_bytes = zlib.decompress(compressed)
    c = np.frombuffer(raw_bytes, dtype=np.float64)
    
    # First inverse: difference of c -> b
    # Use prepend=0 to keep length
    b = np.diff(c, prepend=c[0])   # b[0] = c[0], b[i] = c[i] - c[i-1]
    
    # Second inverse: difference of b -> scaled a
    a = np.diff(b, prepend=b[0])   # a[0] = b[0], a[i] = b[i] - b[i-1]
    
    # Undo scaling
    x = a * 1e9
    
    # Reshape if original shape was provided
    if original_shape is not None:
        x = x.reshape(original_shape)
    return x


# ========== Demonstration ==========
if __name__ == "__main__":
    # Generate some pseudo‑random test data
    np.random.seed(42)
    original = np.random.normal(0, 1, size=10000).astype(np.float32)
    
    # Compress
    compressed_bytes = compress(original)
    
    # Decompress
    restored = decompress(compressed_bytes, original_shape=original.shape)
    
    # Verify
    max_error = np.max(np.abs(original - restored))
    print(f"Original size:      {original.nbytes} bytes")
    print(f"Compressed size:    {len(compressed_bytes)} bytes")
    print(f"Compression ratio:  {original.nbytes / len(compressed_bytes):.2f}x")
    print(f"Maximum error:      {max_error:.2e}")
    print(f"Perfect match:      {np.allclose(original, restored, atol=1e-9)}")