import numpy as np
import sys
import random

project_path = "..."
sys.path.append(project_path)

from embed_utils import *
from utils import *
from entropy_utils import compute_entropy

def calculate_seed(mean_color, round=30):
    return np.round(mean_color / round) * round

def embed_watermark(method, image_array, diag = -1, interval_length=1, dwt_lvl=-1, entropy_threshold=0, image_block_size=96, dct_block_size=8, 
                        max_change=5, adjust_to_center=False, scale=4000, max_iter=2, round=30, num_keys=0):
    """
    Divide the image into 96x96 blocks. Each pixel in a 96x96 block will use the mean color of this large block 
    to determine the allowed and disallowed functions. 
    Iterate over each large block in 8x8 blocks and:
    1. Apply DWT (2 or 3 levels) on each 8x8 block
    2. Watermark the LL band
    3. Apply inverse DWT to the watermarked block
    4. Assemble the transformed blocks back into a 96x96 image block 
    """
    print(f"Watermarking with {method} method")

    image_array = image_array.astype(np.float64) 

    height, width = image_array.shape
    num_regions_per_channel = int(scale / interval_length)

    # Counter to keep track of processed blocks
    block_counter = 0

    # Loop over the image in larger 96x96 blocks
    for y in range(0, height, image_block_size):
        for x in range(0, width, image_block_size):
            # Get size of large block (defined as 96 x 96 but may be smaller when we are at the image borders) 
            if num_keys > 0:
                keys = [2468, 123, 7890, 55555, 999] # 42, 2468 , 1357, 8888, 76543, 101010]
                selected_key = random.choice(keys[:num_keys])
            else:
                selected_key = None
            block_height = min(image_block_size, height - y)
            block_width = min(image_block_size, width - x)
            
            # Extract the larger 96x96 block
            block = image_array[y:y + block_height, x:x + block_width]

            if compute_entropy(block) >= entropy_threshold:

                # Prepare the random sequence for watermarking 
                mean_color = np.mean(block)
                rounded_mean_color = calculate_seed(mean_color, round)
                random_sequence = generate_sequence_from_seed(int(rounded_mean_color), num_regions_per_channel, key = selected_key)

                # Prepare an empty array to hold the DCT-transformed 96x96 block
                IMAGE_block = np.zeros((block_height, block_width), dtype=np.float64)

                sub_block_counter = -1

                num_iter = 0 
                while(num_iter < max_iter): # we want to try again in case the mean color has changed
                # Iterate over each 8x8 block within the 96x96 block
                    for sub_y in range(0, block_height, dct_block_size):
                        for sub_x in range(0, block_width, dct_block_size):
                            sub_block_counter += 1
                            sub_block_height = min(dct_block_size, block_height - sub_y)
                            sub_block_width = min(dct_block_size, block_width - sub_x)

                            # Extract the 8x8 sub-block
                            sub_block = block[sub_y:sub_y + sub_block_height, sub_x:sub_x + sub_block_width]

                            rows, cols = sub_block.shape
                            if rows == cols == dct_block_size: ### embed watermark only if the block size is enough (possibly not at borders)
                              
                                watermarked_block = embed_dwt(sub_block, diag, dwt_lvl, random_sequence, scale, interval_length, 
                                                                adjust_to_center, max_change)
                                
                                # Place small watermarked block in the large image block
                                IMAGE_block[sub_y:sub_y + sub_block_height, sub_x:sub_x + sub_block_width] = watermarked_block

                    # Done with the large block - if mean color changed try again  
                    clamped_IMAGE_block = np.clip(IMAGE_block, 0, 255)
                    new_mean_color = np.mean(clamped_IMAGE_block)
                    new_rounded_mean_color = calculate_seed(new_mean_color, round) 
                    
                    if new_rounded_mean_color ==  rounded_mean_color:
                        break
                    else:
                        print(f"Iteration {num_iter}/{max_iter}: Mean color changed for block {block_counter} \
                            from {rounded_mean_color} to {new_rounded_mean_color}.")
                    num_iter += 1
                    # adjust now to the new mean color
                    mean_color = new_mean_color
                    rounded_mean_color = new_rounded_mean_color
                    random_sequence = generate_sequence_from_seed(int(rounded_mean_color), num_regions_per_channel, key=selected_key)

                # Place large watermarked block in the original image 
                image_array[y:y + block_height, x:x + block_width] = IMAGE_block 

            # Update and print progress
            block_counter += 1

    print("Watermark embedding complete.")
    
    return image_array
