import numpy as np
from scipy.ndimage import binary_dilation

def compute_complex_phi(phi):
    F = np.fft.fft2(phi)
    phi_c = np.fft.ifft2(F)
    return phi_c

def admissible_edit_mask(phi, imag_grad_threshold=None, percentile=95, min_abs=1e-3, dilate_iters=1):
    """
    Robust mask:
    - If imag_grad_threshold provided, use it.
    - Else compute threshold as max(percentile value, min_abs).
    - Optionally dilate mask to give edit zones some area.
    Returns (mask, mag).
    """
    phi_c = compute_complex_phi(phi)
    imag = np.imag(phi_c)
    gx, gy = np.gradient(imag)
    mag = np.sqrt(gx**2 + gy**2)

    if imag_grad_threshold is None:
        pval = np.percentile(mag, percentile)
        thresh = max(pval, min_abs)
    else:
        thresh = imag_grad_threshold

    mask = mag > thresh

    if dilate_iters > 0:
        mask = binary_dilation(mask, iterations=dilate_iters)

    return mask, mag