import numpy as np
import matplotlib.pyplot as plt

def function(x, y, a=1, b=1, c=1, d=0, e=0):
    return -a * np.sqrt((b * x)**2 + (c * y)**2 + d) + e

def inverse_function(f, t, a, b, c, d, e):
    term = ((f - e) / (-a))**2 - d
    if term < 0:
        return None, None

    x = (1/np.sqrt(b)) * np.sqrt(term) * np.cos(t)
    y = (1/np.sqrt(c)) * np.sqrt(term) * np.sin(t)

    return x, y

def create_gradient_vector_field(gx, gy, image_shape, step=20, reverse=False):
    rows, cols = image_shape
    y, x = np.mgrid[step/2:rows:step, step/2:cols:step].reshape(2, -1).astype(int)

    max_dim = max(rows, cols)
    scale = max_dim / 1000 

    direction = -1 if reverse else 1

    fig, ax = plt.subplots(figsize=(cols/50, rows/50), dpi=100)
    ax.quiver(x, y, direction * gx[y, x], direction * -gy[y, x],
            scale=scale,
            scale_units='width',
            width=0.002 * max_dim / 500,
            headwidth=8,
            headlength=12,
            headaxislength=0,
            color='black',
            minshaft=2,
            minlength=0,
            pivot='tail')
    ax.set_xlim(0, cols)
    ax.set_ylim(rows, 0)
    ax.set_aspect('equal')
    ax.axis('off')

    fig.tight_layout(pad=0)
    fig.canvas.draw()
    vector_field = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    vector_field = vector_field.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    plt.close(fig)

    return vector_field


def apply_inverse_vector_field_transform(image, func, radius, center=(0.5, 0.5), strength=1, edge_smoothness=0.1, center_smoothness=0.20):

    rows, cols = image.shape[:2]
    max_dim = max(rows, cols)

    center_y = int(center[1] * rows)
    center_x = int(center[0] * cols)
    center_y = abs(rows - center_y)

    pixel_radius = int(max_dim * radius)

    y, x = np.ogrid[:rows, :cols]
    y = (y - center_y) / max_dim
    x = (x - center_x) / max_dim

    dist_from_center = np.sqrt(x**2 + y**2)

    z = func(x, y)
    gy, gx = np.gradient(z)

    def sigmoid(x, center, steepness):
        return 1 / (1+ np.exp(-steepness * (x - center)))

    mask = edge_mask * center_mask

    gx = gx * mask
    gy = gy * mask

    magnitude = np.sqrt(gx**2 + gy**2)
    magnitude[magnitude == 0] = 1
    gx = gx / magnitude
    gy = gy / magnitude

    scale_factor = strength * np.log(max_dim) / 100
    gx = gx * scale_factor * mask
    gy = gy * scale_factor * mask

    x_new = x + gx
    y_new = y + gy

    x_new = y_new * max_dim + center_x
    y_new = y_new * max_dim + center_y

    x_new = np.clip(x_new, 0, cols - 1)
    y_new = np.clip(y_new, 0, rows - 1)





if __name__ == '__main__':
    x = 3
    y = 2

    t = np.arctan2(y, x)

    a, b, c, d, e = 1, 1, 1, 0, 0

    print(x, y)
    function = function(3, 2)
    print(function)
    inverse_function = inverse_function(function, t, a, b, c, d, e)
    print(inverse_function)