# Path: pyqstem/util.py
# def spatial_frequencies(shape,sampling,return_polar=False,return_nyquist=False,wavelength=None):
# 
#     if not isinstance(shape, collections.Iterable):
#         shape = (shape,)*2
# 
#     if not isinstance(sampling, collections.Iterable):
#         sampling = (sampling,)*2
#     
#     dkx=1/(shape[0]*sampling[0])
#     dky=1/(shape[1]*sampling[1])
# 
#     if shape[0]%2==0:
#         kx = np.fft.fftshift(dkx*np.arange(-shape[0]/2,shape[0]/2,1))
#     else:
#         kx = np.fft.fftshift(dkx*np.arange(-shape[0]/2-.5,shape[0]/2-.5,1))
# 
#     if shape[1]%2==0:
#         ky = np.fft.fftshift(dky*np.arange(-shape[1]/2,shape[1]/2,1))
#     else:
#         ky = np.fft.fftshift(dky*np.arange(-shape[1]/2-.5,shape[1]/2-.5,1))
# 
#     ky,kx = np.meshgrid(ky,kx)
# 
#     k2 = kx**2+ky**2
# 
#     ret = (kx,ky,k2)
#     if return_nyquist:
#         knx = 1/(2*sampling[0])
#         kny = 1/(2*sampling[1])
#         Kx=kx/knx
#         Ky=ky/knx
#         K2 = Kx**2+Ky**2
#         ret += (Kx,Ky,K2)
#     if return_polar:
#         theta = np.sqrt(k2*wavelength**2)
#         phi = np.arctan2(ky,kx)
#         ret += (theta,phi)
# 
#     return ret


import numpy as np
import scipy.ndimage
import matplotlib.pyplot as plt
import warnings
from .util import spatial_frequencies
from skimage.filters import gaussian
    
def detect(img,sampling,dose=None,MTF_param=None,MTF_func=None,blur=None,resample=None):

    if resample is not None:
        if not isinstance(resample, (list, tuple)):
            resample=(resample,)*2
        zoom=(sampling[0]/resample[0],sampling[1]/resample[1])
        sampling=resample
        warnings.filterwarnings('ignore')
        img = scipy.ndimage.interpolation.zoom(img, zoom)
        warnings.filterwarnings('always')
        
    
    if blur is not None:
        img=gaussian(img,blur)
    
    if dose is not None:
        img = img/np.sum(img)*dose*np.product(sampling)*np.product(img.shape)
        
        img[img<0]=0
        #vals = len(np.unique(img))
        #vals = 2**np.ceil(np.log2(vals))
        #img = np.random.poisson(img * vals) / float(vals)

        img = np.random.poisson(img).astype(np.int64)
    
    if MTF_param is not None:
        if MTF_func is None:
            MTF_func=MTF_a
    

# Path: pyqstem/util.py
# def spatial_frequencies(shape,sampling,return_polar=False,return_nyquist=False,wavelength=None):
# 
#     if not isinstance(shape, collections.Iterable):
#         shape = (shape,)*2
# 
#     if not isinstance(sampling, collections.Iterable):
#         sampling = (sampling,)*2
#     
#     dkx=1/(shape[0]*sampling[0])
#     dky=1/(shape[1]*sampling[1])
# 
#     if shape[0]%2==0:
#         kx = np.fft.fftshift(dkx*np.arange(-shape[0]/2,shape[0]/2,1))
#     else:
#         kx = np.fft.fftshift(dkx*np.arange(-shape[0]/2-.5,shape[0]/2-.5,1))
# 
#     if shape[1]%2==0:
#         ky = np.fft.fftshift(dky*np.arange(-shape[1]/2,shape[1]/2,1))
#     else:
#         ky = np.fft.fftshift(dky*np.arange(-shape[1]/2-.5,shape[1]/2-.5,1))
# 
#     ky,kx = np.meshgrid(ky,kx)
# 
#     k2 = kx**2+ky**2
# 
#     ret = (kx,ky,k2)
#     if return_nyquist:
#         knx = 1/(2*sampling[0])
#         kny = 1/(2*sampling[1])
#         Kx=kx/knx
#         Ky=ky/knx
#         K2 = Kx**2+Ky**2
#         ret += (Kx,Ky,K2)
#     if return_polar:
#         theta = np.sqrt(k2*wavelength**2)
#         phi = np.arctan2(ky,kx)
#         ret += (theta,phi)
# 
#     return ret


# Path: pyqstem/detection.py
import numpy as np
import scipy.ndimage
import matplotlib.pyplot as plt
import warnings
from .util import spatial_frequencies
from skimage.filters import gaussian

    
def detect(img,sampling,dose=None,MTF_param=None,MTF_func=None,blur=None,resample=None):

    if resample is not None:
        if not isinstance(resample, (list, tuple)):
            resample=(resample,)*2
        zoom=(sampling[0]/resample[0],sampling[1]/resample[1])
        sampling=resample
        warnings.filterwarnings('ignore')
        img = scipy.ndimage.interpolation.zoom(img, zoom)
        warnings.filterwarnings('always')
        
    
    if blur is not None:
        img=gaussian(img,blur)
    
    if dose is not None:
        img = img/np.sum(img)*dose*np.product(sampling)*np.product(img.shape)
        
        img[img<0]=0
        #vals = len(np.unique(img))
        #vals = 2**np.ceil(np.log2(vals))
        #img = np.random.poisson(img * vals) / float(vals)

        img = np.random.poisson(img).astype(np.int64)
    
    if MTF_param is not None:
        if MTF_func is None:
            MTF_func=MTF_a
    

        kx,ky,k2=spatial_frequencies(img.shape,sampling)