import numpy as np
import matplotlib.pyplot as plt
from scipy import special

### FITTING MODEL DEFINITIONS ####
#    """ These functions will output a 2D numpy array containing a model with
#    center position crow, ccol. The center position refers to the corresponding
#    image, while this is the UV-data."""

def optically_thin_sphere_uv((u,v), flux, crow, ccol, d):
    """ Expressions from Pearson (1995) Appendix B.4 and A.3."""
    # d is diameter
    # Calculate position shift. Note, sphere is 
    # independent of rotation, so keep both unrotated.
    # Columns are horizontal so with u, and rows are vertical so with v
    shift = np.exp(-2*np.pi*1j*(ccol*u + crow*v))
    rho = np.sqrt(u**2+v**2)
    # Avoid singularity, manually set to limit later
    rho[0][0] = 1
    z = np.pi*d*rho
    ftsphere = 3.0*(np.sin(z)-z*np.cos(z))/z**3
    ftsphere[0][0] = 1.0 # From integral of image plane expression B14. E.g.
    # Wolfram alpha: 
    # "integrate (6/(pi*d**2))*2*pi*r*sqrt(1-(2r/d)**2) dr from r=0 to r=d/2"
    # gives answer 1.0.
    ftsphere = ftsphere*flux
    return ftsphere * shift

def optically_thin_shell_uv((u,v), flux, crow, ccol, dout):
    # ASSUMES RATIO 0.7 between inner and outer diameter
    # dout is outer diameter in pixels
    # Flux in Jy
    # center pos crow, ccol in pixels
    # First create two spheres with flux densities equal
    # to their respective volumes
    volout = (4.0/3)*np.pi*(0.5*dout)**3
    din = 0.7*dout
    volin = (4.0/3)*np.pi*(0.5*din)**3
    sphout = optically_thin_sphere_uv((u,v), volout, crow, ccol, dout)
    sphin = optically_thin_sphere_uv((u,v), volin, crow, ccol, din)
    # The shell is now the difference between the spheres,
    # Normalised to the specified flux
    shell = sphout-sphin
    shell = shell*flux/(volout-volin)
    return shell

def twoDgaussian_uv((_u,_v), flux, _crow, _ccol, sigma_x, sigma_y, PA, mode = 'normal'):
    # After visual inspection
    PA = 0.5*np.pi-PA
    # From http://en.wikipedia.org/wiki/Fourier_transform#Tables_of_important_Fourier_transforms
    # no 401 with parameter identification from 
    # http://en.wikipedia.org/wiki/Gaussian_function#Two-dimensional_Gaussian_function
    # Calculate position shift. Note,
    # independent of rotation, so keep both unrotated
    shift = np.exp(-2*np.pi*1j*(_ccol*_u + _crow*_v))
    # Rotate u,v-plane according to PA
    u = np.cos(PA) * _u - np.sin(PA) * _v
    v = np.sin(PA) * _u + np.cos(PA) * _v
    # Calculate exponential u,v-factor, i.e. damping in uvplane
    size = np.exp(-2*np.pi**2*(u**2*sigma_x**2 + v**2*sigma_y**2))
    # In most cases we want a source with the specified flux,
    # and no special norm is needed, i.e. norm = 1.
    # But, if for example we want a beam normalised
    # for the peak to be 1, we need a different norm.
    # NOTE: This is adapted to get an image
    # with peak 1 if inverting this result with numpy.fft.ifft2.
    if mode == 'beam':
        norm = 2*sigma_x*sigma_y * np.pi
    else:
        norm = 1
    Z = flux * shift * norm * size
    return Z

def build_uvmesh(nrow, ncol):
    """Construct a meshgrid in the UV-plane of size nrow, ncol. Uses carteesian
    indexing to return array of (us, vs) since u,v conventionally means
    horizontal and vertical respectively."""
    # Calculate corresponding fourier frequencies for UV stamps
    frqu = np.fft.fftfreq(ncol, d = 1.0)
    frqv = np.fft.fftfreq(nrow, d = 1.0)
    us,vs = np.meshgrid(frqu, frqv, indexing = 'xy')
    return (us,vs)

def optf(pars, data, uvmesh, uvbeam, rms):
    flux, crow, ccol, diam  = pars
    uvshell = optically_thin_shell_uv(uvmesh, flux, crow, ccol, diam)
    convuvshell = uvshell * uvbeam
    model = np.real(np.fft.ifft2(convuvshell))
    sigma = np.sqrt((0.1*data)**2 + rms**2)
    return np.ravel(np.divide((model - data)**2,sigma**2))