#!/usr/bin/env python3
# Compression testing for image data for expected cases of ours.

# TODO: Implement our own CLI tool that uses our LZW implementation, as it differs.

import imageio.v3 as iio
import ncompress
import numpy as np
import sys

uncompressed = iio.imread(sys.argv[1])
compressed = ncompress.compress(uncompressed.tobytes())
print("Uncompressed:", uncompressed.nbytes)
print("Compressed:", len(compressed), len(compressed) / uncompressed.nbytes)

with open(sys.argv[1] + "-indexed.raw", "rb") as file:
    raw_bytes = file.read()
indexmapping = {}
indexset = set(raw_bytes)
for idx, value in enumerate(indexset):
    indexmapping[value] = idx
print(indexmapping)
indexmap = np.zeros((uncompressed.shape[0] * uncompressed.shape[1], 1), dtype=np.uint8)
for y in range(uncompressed.shape[0]):
    for x in range(uncompressed.shape[1]):
        indexmap[y * uncompressed.shape[0] + x] = indexmapping[raw_bytes[y * uncompressed.shape[0] + x]]
indexedcompressed = ncompress.compress(indexmap.tobytes())
print("Indexed:", len(indexedcompressed), len(indexedcompressed) / len(raw_bytes))

# https://arxiv.org/html/2006.02639v2
def pixel_art_predictive_filter(image_channel):
    """
    Applies a spatial prediction filter to a 2D image matrix 
    to prepare it for lossless compression.
    """
    rows, cols = image_channel.shape
    residual_map = np.zeros((rows, cols), dtype=np.int16)
    
    for r in range(rows):
        for c in range(cols):
            # Handle edge cases (boundaries where neighbors don't exist)
            if r == 0 and c == 0:
                pred = 0
            elif r == 0:
                pred = image_channel[r, c-1] # Only Left neighbor exists
            elif c == 0:
                pred = image_channel[r-1, c] # Only Top neighbor exists
            else:
                # Core DSP Neighborhood definition
                A = image_channel[r, c-1]  # Left
                B = image_channel[r-1, c]  # Top
                C = image_channel[r-1, c-1] # Top-Left
                
                # Median Edge Detector logic
                if C >= max(A, B):
                    pred = min(A, B)
                elif C <= min(A, B):
                    pred = max(A, B)
                else:
                    pred = int(A) + int(B) - int(C)
            
            # Map the actual value to the residual difference
            residual_map[r, c] = int(image_channel[r, c]) - int(pred)
            
    return residual_map

predictivefilter = pixel_art_predictive_filter(indexmap)
predictivefiltercompressed = ncompress.compress(predictivefilter.tobytes())
print("Predictive filter:", len(predictivefiltercompressed), len(predictivefiltercompressed) / len(raw_bytes))

def encode_zorder_2d(x: int, y: int) -> int:
    z = 0
    for i in range(16):
        z |= ((x & (1 << i)) << i) | ((y & (1 << i)) << (i + 1))
    return z

zordered = np.zeros((uncompressed.shape[0] * uncompressed.shape[1], 1), dtype=np.uint8)
for y in range(uncompressed.shape[0]):
    for x in range(uncompressed.shape[1]):
        idx = encode_zorder_2d(x, y)
        zordered[idx] = uncompressed[y, x]

zcompressed = ncompress.compress(zordered.tobytes())
print("Zordered:", len(zcompressed), len(zcompressed) / uncompressed.nbytes)

delta = np.zeros((uncompressed.shape[0] * uncompressed.shape[1], 1), dtype=np.int8)
for y in range(uncompressed.shape[0]):
    for x in range(uncompressed.shape[1]):
        delta[y * uncompressed.shape[0] + x] = int(uncompressed[y, x]) - int(uncompressed[y, (x + 1) % uncompressed.shape[0]])
deltacompressed = ncompress.compress(delta.tobytes())
print("Delta:", len(deltacompressed), len(deltacompressed) / uncompressed.nbytes)

deltaz = np.zeros((uncompressed.shape[0] * uncompressed.shape[1], 1), dtype=np.int8)
for y in range(uncompressed.shape[0]):
    for x in range(uncompressed.shape[1]):
        idx = encode_zorder_2d(x, y)
        deltaz[idx] = int(uncompressed[y, x]) - int(uncompressed[y, (x + 1) % uncompressed.shape[0]])
deltazcompressed = ncompress.compress(deltaz.tobytes())
print("Deltaz:", len(deltazcompressed), len(deltazcompressed) / uncompressed.nbytes)