import sys
import h5py
import hdf5plugin
import numpy as np

print("=" * 60)
print("=== STARTING C-ZOO VERIFICATION (R/C -> Python) ===")
print("=" * 60)
sys.stdout.flush()

FILENAME = "c_encoding_zoo.h5"
SHAPE = (128, 128) # Matches the C/R chunk dimensions

def get_expected_data(shape, dtype):
    """Recreates the deterministic data generated by the C library."""
    total_elements = np.prod(shape)
    
    if np.issubdtype(dtype, np.floating):
        base = np.linspace(0.0, 1.0, total_elements, dtype=dtype)
        wave = (np.sin(np.linspace(0, 50, total_elements)) * 0.1).astype(dtype)
        return (base + wave).reshape(shape)
    else:
        return np.arange(1, total_elements + 1, dtype=dtype).reshape(shape)

def main():
    try:
        f = h5py.File(FILENAME, 'r')
    except Exception as e:
        print(f"FATAL: Could not open {FILENAME}: {e}")
        sys.exit(1)

    passed_tests = 0
    total_datasets = len(f.keys())

    print(f"\n--- Verifying {total_datasets} datasets in {FILENAME} ---")
    
    for name in f.keys():
        dset = f[name]
        print(f"[{name: <25}] -> Reading... ", end="")
        sys.stdout.flush()
        
        is_lossy = any(kw in name for kw in ["rate", "precision", "accuracy", "expert", "truncprec", "sz_", "zfp"])
        
        try:
            # Trigger decompression via hdf5plugin
            read_data = dset[...]
            expected_data = get_expected_data(SHAPE, read_data.dtype)

            if is_lossy:
                # Use the 0.5 tolerance defined in your C verification logic
                max_diff = np.max(np.abs(read_data - expected_data))
                if max_diff <= 0.5:
                    print("OK (Lossy Bounds Passed)")
                    passed_tests += 1
                else:
                    print(f"FAILED (Exceeded tolerance: Max diff {max_diff})")
            else:
                if np.issubdtype(read_data.dtype, np.floating):
                    # Use the 1e-5 tolerance for bit-exact floats
                    max_diff = np.max(np.abs(read_data - expected_data))
                    if max_diff <= 1e-5:
                         print("OK (Bit-Exact Float)")
                         passed_tests += 1
                    else:
                         print(f"FAILED (Mismatch: Max diff {max_diff})")
                else:
                    if np.array_equal(read_data, expected_data):
                        print("OK (Bit-Exact Int)")
                        passed_tests += 1
                    else:
                        print("FAILED (Data Mismatch)")
                        
        except Exception as e:
            print(f"ERROR ({e})")

    print(f"\n=== VERIFICATION COMPLETE: {passed_tests}/{total_datasets} passed ===")
    
    if passed_tests != total_datasets:
        sys.exit(1)

if __name__ == "__main__":
    main()
