import numpy as np               # for numerical operations, e.g., np.float64, np.abs
import cv2                       # for DCT and IDCT transforms (cv2.dct, cv2.idct)
import pywt                      # for DWT and inverse DWT (pywt.wavedec2, pywt.waverec2)
import sys

project_path = "..."
sys.path.append(project_path)
from utils import (              # assuming the utils file is in the same folder
    is_region_allowed, 
    find_closest_allowed_color, 
    inverse_dwt, 
    get_idx
)


def embed_dwt(sub_block, diag, dwt_lvl, random_sequence, scale, interval_length, 
              adjust_to_center, max_change):

    if dwt_lvl == 1:
        LL, (LH, HL, HH) = pywt.wavedec2(sub_block, wavelet='haar', level=dwt_lvl)
    else:
        coeffs = pywt.wavedec2(sub_block, wavelet='haar', level=dwt_lvl)
        LL = coeffs[0]  # Approximation coefficients
        details = coeffs[1:]  # Detail coefficients (LH, HL, HH) for each level

    if diag >0:
        for i in range(0, min(diag,LL.shape[0])):  
            for j in range(0, min(diag,LL.shape[1])-i):
                coeff_value = LL[i, j]
                if not is_region_allowed(coeff_value, random_sequence, scale, interval_length):
                    new_coeff = find_closest_allowed_color(coeff_value, random_sequence, 
                                                        interval_length, scale, adjust_to_center, max_change)
                    if new_coeff != LL[i, j]:
                        assert is_region_allowed(new_coeff, random_sequence, scale, interval_length)
                    LL[i, j] = new_coeff
                else: # if its allowed, still adjust to center
                    region_idx = get_idx(coeff_value, scale, interval_length)
                    left_border = -(scale/2) + region_idx * interval_length
                    right_border = -(scale/2) + (region_idx + 1) * interval_length
                    center_of_cube = (left_border + right_border)/2
                    LL[i, j] = center_of_cube
    elif diag == -1:
        for i in range(0, LL.shape[0]):  
            for j in range(0, LL.shape[1]-i):
                coeff_value = LL[i, j]
                if not is_region_allowed(coeff_value, random_sequence, scale, interval_length):
                    new_coeff = find_closest_allowed_color(coeff_value, random_sequence, 
                                                        interval_length, scale, adjust_to_center, max_change)
                    if new_coeff != LL[i, j]:
                        assert is_region_allowed(new_coeff, random_sequence, scale, interval_length)
                    LL[i, j] = new_coeff
                else: # if its allowed, still adjust to center
                    region_idx = get_idx(coeff_value, scale, interval_length)
                    left_border = -(scale/2) + region_idx * interval_length
                    right_border = -(scale/2) + (region_idx + 1) * interval_length
                    center_of_cube = (left_border + right_border)/2
                    LL[i, j] = center_of_cube

    if dwt_lvl == 1:
        dwt_reconstructed = inverse_dwt(LL, LH, HL, HH)
    else: 
        dwt_reconstructed = pywt.waverec2([LL] + details, wavelet='haar')
    
        
    return dwt_reconstructed
