import numpy as np
from scipy.optimize import minimize


def get_droplet(radius, norm, shiftx, shifty, gamma, resolution=512):
    dim = int(2*radius+10)
    # print(dim)
    center = dim/2
    x = np.arange(0, dim)
    y = np.arange(0, dim)
    xv, yv = np.meshgrid(x, y)

    densmod= 2.* radius* np.sin(np.arccos(((xv-center)**2 + (yv-center)**2)**0.5/radius))
    densmod[densmod<0]=0
    densmod[np.isnan(densmod)]=0
    # densmod = gaussian_filter(densmod, sigma=sigma)
    densmod = (densmod/np.sum(densmod))**gamma
    densmod = densmod/np.sum(densmod)*norm

    dens = densmod * np.exp(1j*2*np.pi*(shiftx*(xv-center)/resolution + shifty*(yv-center)/resolution))



    densout = np.zeros([resolution, resolution], dtype=complex)
    densout[resolution//2-dens.shape[0]//2: resolution//2+dens.shape[0]//2, resolution//2-dens.shape[1]//2: resolution//2+dens.shape[1]//2] = dens[:dens.shape[0]//2*2,:dens.shape[1]//2*2]


    # print("init:", np.amax(np.abs(densout)))
    return densout


def get_droplet_pattern(radius, norm, shiftx, shifty, gamma, resolution=512):

    dens = get_droplet(radius, norm, shiftx, shifty, gamma, resolution)
    patt = np.abs(np.fft.fftshift(np.fft.fft2(dens)))**2
    return patt