import numpy as np
from scipy.signal import convolve, fftconvolve


def generate_random_binary_array_with_squares(shape, num_squares, square_size_mean, square_size_var=0.5, width=2, intersect=False, FG=False):
    '''
    Generate a random binary array with the borders of squares. The squares are placed randomly in the array.
    When FG is true, the function returns the array with the borders and the array with the filled squares.
    
    agrs:
        shape: tuple or int, size of the array
        num_squares: int, number of squares to place in the array
        avg_square_size: int, average size of the squares
        sq  uare_size_var: float, variance of the square sizes
        width: int, width of the squares' border
        intersect: bool, if True, the squares can intersect
        FG: bool, if True, the function returns the array with the filled squares
    '''

    if isinstance(shape, int):
        shape = (shape, shape)
    if shape[0] < square_size_mean or shape[1] < square_size_mean:
        raise ValueError("The square size is larger than the array size")

    array = np.zeros(shape, dtype=int)

    if FG:
        fg_array = np.zeros(shape, dtype=int)
    
    for _ in range(num_squares):

        feasible = False
        while not feasible:
            sq_size = int(np.round(np.random.normal(square_size_mean, square_size_mean*square_size_var))) 

            if sq_size > 0 and (shape[0]-sq_size+1) > 0 and (shape[1]-sq_size+1) > 0:
                # Randomly choose the top-left corner of the square
                row = np.random.randint(0, shape[0] - sq_size + 1)
                col = np.random.randint(0, shape[1] - sq_size + 1)

                # Check if the square can be placed
                if (row + sq_size)<= shape[0] and (col + sq_size)<= shape[1]:
                    feasible = True
            
        # Place the square of 1s
        array[row:row + sq_size, col:col + sq_size] += 1
        if FG:
            fg_array[row:row + sq_size, col:col + sq_size] = 1
        if intersect:
            array[(row+width):(row + sq_size-width), (col+width):(col + sq_size-width)] -= 1
        else:     
            array[(row+width):(row + sq_size-width), (col+width):(col + sq_size-width)] = 0
    if FG:
        return array.clip(0, 1), fg_array
    return array.clip(0, 1)


def generate_random_array_with_correlated_noise(shape, correlation_scale):
#TODO: find out why this only works with gausian noise !

    '''
    Generate a random array with spatially correlated noise. The noise is generated by convolving a random array with a Gaussian kernel.
    Finally, the noise is normalized to [0, 1].

    args:
        shape: tuple or int, size of the array
        correlation_scale: float, scale of the Gaussian kernel
    '''
    if isinstance(shape, int):
        shape = (shape, shape)


    x = np.arange(-correlation_scale, correlation_scale)
    y = np.arange(-correlation_scale, correlation_scale)
    X, Y = np.meshgrid(x, y)
    dist = np.sqrt(X*X + Y*Y)
    #filter_kernel = np.exp(-dist**2/(2*correlation_scale))
    filter_kernel = (1/(correlation_scale**2))*np.exp(-dist**2/(2*correlation_scale))


    noise = np.random.rand(*shape)
    noise = fftconvolve(noise, filter_kernel, mode='same')
    #noise = (noise-np.min(noise))/(np.max(noise)-np.min(noise))

    return noise


def generate_two_dim_sin_cos_wave(x_max, y_max, freq_x=1, freq_y=1, phase_x=0, phase_y=0):
    '''
    Generate a 2D array with a sin and cos wave in x and y direction

    params:
    x_max, y_max : int, x and y dimensions of the array
    freq_x, freq_y : float, frequency of the sin wave in x direction (the cos wave in y direction)
    phase_x, phase_y : float, phase of the sin wave in x direction (the cos wave in y direction)
    '''

    array = np.zeros((x_max, y_max))
    x = np.arange(x_max)
    y = np.arange(y_max)

    xx, yy = np.meshgrid(x, y)

    array = np.sin(2 * np.pi * freq_x * xx + phase_x) + np.cos(2 * np.pi * freq_y * yy + phase_y)

    return array, f'${{\sin}}$(${{2\pi}}$*{freq_x}*x + {phase_x}) + ${{\cos}}$(${{2\pi}}$*{freq_y}*y + {phase_y})'


def generate_binary_quadrant_array(size):
    """
    Generate a (size, size) array where the top-right and bottom-left quadrants are 1s, else 0s.
    
    :param size: The size of the square array (size x size).
    :return: A 2D NumPy array with the specified block pattern.
    """
    array = np.zeros((size, size), dtype=int)  # Initialize with zeros

    # Define quadrant boundaries
    half = size // 2
    
    # Set top-right and bottom-left quadrants to 1
    array[:half, half:] = 1   # Top-right
    array[half:, :half] = 1   # Bottom-left
    return array


def generate_checkerboard_array(size):
    return np.indices((size, size)).sum(axis=0) % 2 

