# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from datetime import datetime
from tensorflow import keras

from bitsback import *

# +
(x_train, _), (x_test, _) = keras.datasets.mnist.load_data()

orig_x_train = x_train
orig_x_test = x_test

# Scale to be between [-1, 1]
x_train = (x_train.astype('float32') - 127.5) / 127.5
x_test = (x_test.astype('float32') - 127.5) / 127.5

print(x_train.shape)

mnist_digits = x_train
mnist_digits = np.expand_dims(mnist_digits, -1)


# +
def display_digits(X, digit_size=28, n=10):
    figure = np.zeros((digit_size * n, digit_size * n))
    
    for i in range(n):
        for j in range(n):
            index = np.random.randint(0, X.shape[0])
            digit = X[index].reshape(digit_size, digit_size)
            
            x = i * digit_size
            y = j * digit_size
            figure[x:x + digit_size, y:y + digit_size] = digit
    
    plt.figure(figsize=(max(n, 3), max(n, 3)))
    plt.imshow(figure, cmap='Greys_r')
    plt.show()
    

print(mnist_digits[0, 14, 10:20].reshape(10))
display_digits(mnist_digits)
# -

enc = keras.models.load_model('vae-encoder')
dec = keras.models.load_model('vae-decoder')

stack = np.random.randint(low=0, high=(1 << 32) - 1, size=50, dtype=np.uint32).tolist()
stack.append(int(np.random.randint(low=(1 << 32), high=(1 << 63))))
init_stack = stack.copy()
num_images = len(orig_x_test)
quant_bits = 16
stats = {}
for i in range(num_images):
    stack = bitsback_append(orig_x_test[i:(i+1)], stack, enc, dec, quant_bits=quant_bits, stats=stats)
    if (i > 0 and (i+1) % 10 == 0) or (i == num_images - 1):
        orig = 784 * (i + 1)
        compress = 4 * len(stack) + 4
        bits_per_dim = compress * 8 / orig
        print('[%s]: N=%05d, bits per dim=%.4f' % (datetime.now(), i+1, bits_per_dim))

images = []
for i in range(num_images):
    data, stack = bitsback_pop(stack, enc, dec, quant_bits=quant_bits)
    images = [data] + images
print(len(stack))

assert(stack == init_stack)
for i in range(num_images):
    x1 = np.array(images[i]).reshape(28, 28).tolist()
    x2 = orig_x_test[i].tolist()
    assert(x1 == x2)

# +
import pandas as pd

df = pd.DataFrame(stats)
df.hist(bins=10, figsize=(12, 8))
df.describe()

StarcoderdataJupyterScriptsDedupFiltered