from Crypto.Random import get_random_bytes
import numpy as np

def write_file(filename, array_name, length, count=10000):
    with open(filename, 'w') as file:
        file.write(f"{array_name} = [\n")
        for _ in range(count):
            random_bytes = get_random_bytes(length)
            # Get the repr of bytes, which Python can later evaluate to bytes
            byte_string = repr(random_bytes)
            file.write(f'    {byte_string},\n')
        file.write("]\n")

# Create the files with specific byte lengths
# write_file('messages_2.py', 'MESSAGES', 32)
# write_file('nonces_2.py', 'NONCES', 12)
# write_file('keys_2.py', 'KEYS', 32,count=10)

def simulate_users(uid_length,num_users,num_images,mean_images,std_images):
    #####draw user ids
    uids = [get_random_bytes(uid_length) for _ in range(num_users)]
    indices = np.arange(num_users)
    #####replicate according to distribution
    counts = np.random.normal(loc=mean_images, scale=std_images, size=num_users)
    # Ensure counts are positive integers and sum to sample_size
    counts = np.maximum(0, np.round(counts)).astype(int)  
    counts = counts / counts.sum() * num_images  # Normalize to sum to sample_size
    counts = np.round(counts).astype(int)

    # Adjust to exactly match sample_size (if rounding causes slight mismatch)
    while counts.sum() < num_images:
        counts[np.random.randint(num_users)] += 1
    while counts.sum() > num_images:
        counts[np.random.randint(num_users)] -= 1

    # Generate the sample
    sample = np.repeat(indices, counts)

    # Shuffle the sample
    np.random.shuffle(sample)
    
    #####check stats of distribution
    # Compute actual mean and std of occurrences
    actual_counts = np.bincount(sample, minlength=num_users)
    actual_mean = np.mean(actual_counts)
    actual_std = np.std(actual_counts)

    print(f"Expected Mean: {mean_images}, Actual Mean: {actual_mean}")
    print(f"Expected Std: {std_images}, Actual Std: {actual_std}")

    ######write to file
    with open("messanges_long_2.py", 'w') as file:
        file.write(f"MESSAGES = [\n")
        for uid in sample:            
            # Get the repr of bytes, which Python can later evaluate to bytes
            byte_string = repr(uids[uid])
            file.write(f'    {byte_string},\n')
        file.write("]\n")

    with open("messnges_short.py","w") as file:
        file.write(f"MESSAGES = [\n")
        for uid in uids:            
            # Get the repr of bytes, which Python can later evaluate to bytes
            byte_string = repr(uid)
            file.write(f'    {byte_string},\n')
        file.write("]\n")

simulate_users(uid_length=32,num_users=1000,num_images=10000,mean_images=10,std_images=3)