#!/usr/bin/env python3
"""Decompress the BWT+MTF+LZW equation database files."""
import json, struct, sys
from pathlib import Path

def lzw_decode(data: bytes, max_dict: int = 4096) -> list[int]:
    """Decode LZW-compressed data into a list of integers."""
    result = []
    if not data:
        return result
    
    dict_size = 256
    dictionary = {i: bytes([i]) for i in range(dict_size)}
    
    # Read 12-bit values regardless of max_dict
    bits = []
    for byte in data:
        bits.extend([(byte >> (7 - i)) & 1 for i in range(8)])
    
    pos = 0
    def read_bits(n):
        nonlocal pos
        val = 0
        for _ in range(n):
            val = (val << 1) | bits[pos]
            pos += 1
        return val
    
    code_bits = 9
    prev = read_bits(code_bits)
    result.append(prev)
    prev_bytes = dictionary[prev]
    
    while pos < len(bits):
        code = read_bits(code_bits)
        if code in dictionary:
            entry = dictionary[code]
        elif code == dict_size:
            entry = prev_bytes + prev_bytes[:1]
        else:
            break
        
        result.extend(entry)
        dictionary[dict_size] = prev_bytes + entry[:1]
        dict_size += 1
        if dict_size >= (1 << code_bits) and code_bits < 12:
            code_bits += 1
        prev_bytes = entry
    
    return result


def mtf_decode(data: list[int]) -> list[int]:
    """Move-to-front decoding."""
    alphabet = list(range(256))
    result = []
    for code in data:
        if code < len(alphabet):
            val = alphabet[code]
            result.append(val)
            alphabet.pop(code)
            alphabet.insert(0, val)
    return result


def bwt_inverse(data: list[int], primary_index: int) -> bytes:
    """Inverse Burrows-Wheeler transform."""
    if not data:
        return b''
    n = len(data)
    # Count occurrences
    counts = [0] * 256
    for b in data:
        counts[b] += 1
    
    # Compute cumulative counts (start positions)
    start = [0] * 256
    total = 0
    for i in range(256):
        start[i] = total
        total += counts[i]
    
    # Build next array
    next_idx = [0] * n
    temp_count = [0] * 256
    for i in range(n):
        b = data[i]
        next_idx[i] = start[b] + temp_count[b]
        temp_count[b] += 1
    
    # Reconstruct
    result = bytearray()
    idx = primary_index
    for _ in range(n):
        idx = next_idx[idx]
        result.append(data[idx])
    
    return bytes(result)


def decompress(path: Path) -> dict:
    """Decompress an equation database file."""
    raw = path.read_bytes()
    
    # Read header
    offset = 0
    header_len = struct.unpack('>I', raw[offset:offset+4])[0]
    offset += 4
    header = json.loads(raw[offset:offset+header_len].decode())
    offset += header_len
    
    chain = header.get("chain", ["bwt", "mtf", "lzw"])
    meta = header.get("metadata", {})
    compressed = raw[offset:]
    
    # Apply decompression steps in reverse order
    data = compressed
    
    # Step 1: LZW decode
    lzw_meta = meta.get("lzw", {})
    integers = lzw_decode(data, lzw_meta.get("max_dict", 4096))
    
    # Step 2: MTF decode
    integers = mtf_decode(integers)
    
    # Step 3: BWT inverse
    bwt_meta = meta.get("bwt", {})
    primary = bwt_meta.get("primary_index", 0)
    raw_bytes = bwt_inverse(integers, primary)
    
    # Parse as JSON
    return json.loads(raw_bytes.decode("utf-8"))


def main():
    path = Path(sys.argv[1]) if len(sys.argv) > 1 else \
        "/home/allaun/Research Stack/3-Mathematical-Models/equations_compressed/chain_equations_20260504_134248.compressed"
    
    print(f"Decompressing {path}...")
    data = decompress(path)
    
    if isinstance(data, dict):
        print(f"Type: dict with {len(data)} keys")
        for k, v in data.items():
            if isinstance(v, list):
                print(f"  {k}: list[{len(v)}]")
                if len(v) > 0:
                    print(f"    first: {str(v[0])[:200]}")
            elif isinstance(v, dict):
                print(f"  {k}: dict with {len(v)} keys")
            else:
                print(f"  {k}: {type(v).__name__} = {str(v)[:200]}")
    elif isinstance(data, list):
        print(f"Type: list[{len(data)}]")
        if len(data) > 0:
            print(f"  first: {str(data[0])[:200]}")


if __name__ == "__main__":
    main()