import tempfile

import numpy as np
import torch
import gradio as gr


def save_numpy_as_binary_bitwise(array, K, filename):
    """Save a NumPy array as a binary file with bitwise storage."""
    bits_per_value = int(np.ceil(np.log2(K)))  # Number of bits required per value
    bitstring = ''.join(format(val, f'0{bits_per_value}b') for val in array)  # Convert each number to binary

    # Convert bitstring to bytes
    byte_array = int(bitstring, 2).to_bytes((len(bitstring) + 7) // 8, byteorder='big')

    # Write to binary file
    with open(filename, 'wb') as f:
        f.write(byte_array)


def load_numpy_from_binary_bitwise(filename, K, T, model_type, effective_num_values):
    if filename is None:
        return None
    """Load a NumPy array from a binary file stored in bitwise format."""
    bits_per_value = int(np.ceil(np.log2(K)))  # Number of bits required per value

    if f'-K{K}-' not in filename:
        raise gr.Error("Please set the codebook size to match the bitstream file you provided")

    if f'-T{T}-' not in filename:
        raise gr.Error("Please set the number of diffusion timesteps to match the bitstream file you provided")

    if f'-M{model_type}-' not in filename:
        raise gr.Error("Please set the image size to match the bitstream file you provided")

    # Read the binary file as bytes
    with open(filename, 'rb') as f:
        byte_data = f.read()
    # Convert bytes to a binary string
    bitstring = bin(int.from_bytes(byte_data, byteorder='big'))[2:]  # Remove '0b' prefix

    # Pad with leading zeros if needed
    bitstring = bitstring.zfill(effective_num_values * bits_per_value)

    # Extract values from bitstring
    values = [int(bitstring[i:i + bits_per_value], 2) for i in range(0, len(bitstring), bits_per_value)]

    return torch.from_numpy(np.array(values, dtype=np.int32)).squeeze()


def generate_binary_file(np_arr, num_noises, timesteps, model_type):
    temp_file = tempfile.NamedTemporaryFile(delete=False,
                                            suffix=f".bitstream-T{timesteps}-K{num_noises}-M{model_type}-")
    save_numpy_as_binary_bitwise(np_arr, num_noises, temp_file.name)
    return temp_file.name