#!/usr/bin/env python

#
# This builds test data testing the PSF orienation and calls the 'spot'
# program from the command line... This should be included in the
# /branches/fortran/spot/ subdirectory of the svn repository.
#

from sys import argv
import numpy as np
import pyfits, subprocess, os


def rms_from_SNR(snr, image):
    """
    Calculates the rms noise level [sqrt(B) from Appendix A in Bridle et al 08]
    using the input image at a given SNR...
    """
    rootB = np.sqrt(np.sum(image**2)) / snr
    return rootB


def where(*args):
    """
    Performs a function rather like a stripped down IDL where() function.
    Compares multiple input Boolean arrays (np.array objects) and outputs
    the indices where all conditions are simultaneously fulfilled.

    Written B. Rowe, 2011.
    """
    comb = np.where(args[0], 1, 0)
    scomb = np.shape(comb)
    for i in xrange(len(args) - 1):
        if np.shape(args[i + 1]) != scomb:
            raise ValueError('WHERE: Not all input condition vectors'+\
                             ' equal length')
        comb = comb * np.where(args[i + 1], 1, 0)

    q = np.arange(len(comb))[comb == 1]
    return q


if len(argv) < 3 or len(argv) > 4:
    print "usage: ./test_PSF1.py <FWHM [dbl]> <beta [dbl]> <SNR [dbl]>"
    exit(1)
fwhm = float(argv[1])
beta = float(argv[2])
snr = float(argv[3])

print "test_PSF1: Generating run_PSF1.py-testing image and catalogues"
print "test_PSF1: FWHM = "+str(fwhm)
print "test_PSF1: beta = "+str(beta)
print "test_PSF1: SNR = "+str(snr)

# Basic image/postage stamp size parameters
npix = 576
nstamp = 48
nrows = 1 #npix / nstamp
ncols = npix / nstamp

# Set the ellipticity parameters
theta = np.pi * np.arange(ncols, dtype=float) / float(ncols)
cos2t = np.cos(2. * theta)
sin2t = np.sin(2. * theta)
ellip = 0.4

pixscale_DES = 0.27 # DES pixel scale

# Define the arrays for storing the output params
x_true = np.zeros(ncols, dtype=float)
y_true = np.zeros(ncols, dtype=float)
e1 = np.zeros(ncols, dtype=float)
e2 = np.zeros(ncols, dtype=float)

# Setup the random seed
np.random.seed(None)

# Build the image
outimfile = 'test_PSF1_fwhm'+str(fwhm)+'beta'+str(beta)+'im.fits'
image = np.zeros((2 * nstamp, ncols * nstamp))
#    # Setup the random x, y offsets in the truth catalogue
#    dx = np.random.standard_normal(size=(nrows, ncols)) * 0.5
#    dy = np.random.standard_normal(size=(nrows, ncols)) * 0.5
#    xstamp = float(nstamp) / 2. + dx
#    ystamp = float(nstamp) / 2. + dy
for jcol in xrange(ncols):
    e1[jcol] = ellip * cos2t[jcol]
    e2[jcol] = ellip * sin2t[jcol]
    xstamp = 0.5 * np.float(nstamp)
    ystamp = 0.5 * np.float(nstamp)
    subprocess.call(['./spot', 'M', str(nstamp), str(nstamp),
                     str(xstamp), str(ystamp),
                     str(e1[jcol]), str(e2[jcol]),
                     str(fwhm), '0.', str(beta)])
    hdu_spot = pyfits.open('spot.fits')
    spot = hdu_spot[0].data
    hdu_spot.close()
# Define the region within which to put this spot image
    jmin = nstamp * jcol
    jmax = jmin + nstamp
    print 0, nstamp, jmin, jmax
    x_true[jcol] = jmin + xstamp#[irow, jcol]
    y_true[jcol] = 0 + ystamp#[irow, jcol]
    image[0:nstamp, jmin:jmax] = spot

for jcol in xrange(ncols):
    e1[jcol] = ellip * cos2t[jcol]
    e2[jcol] = ellip * sin2t[jcol]
    xstamp = 0.5 * np.float(nstamp)
    ystamp = 0.5 * np.float(nstamp)
    subprocess.call(['./spot', 'M', str(nstamp), str(nstamp),
                     str(xstamp), str(ystamp),
                     str(e1[jcol]), str(e2[jcol]),
                     str(fwhm), '0.', str(beta)])
    hdu_spot = pyfits.open('spot.fits')
    spot = hdu_spot[0].data
    hdu_spot.close()
# Define the region within which to put this spot image
    jmin = nstamp * jcol
    jmax = jmin + nstamp
    print nstamp, 2* nstamp, jmin, jmax
    x_true[jcol] = jmin + xstamp#[irow, jcol]
    y_true[jcol] = 0 + ystamp#[irow, jcol]
    image[nstamp: 2 * nstamp, jmin:jmax] = spot

# Output to FITS image file

print 'test_PSF_orientation: Writing test image to '+outimfile
hdu_image = pyfits.PrimaryHDU(image)
hdu_image.writeto(outimfile, clobber=True)

# Then add the noise...
hdu_image = pyfits.open(outimfile)
image = hdu_image[0].data
hdu_image.close()
print 'test_PSF_orientation.py: Adding noise at SNR = '+str(snr)
for jcol in xrange(ncols):
    # Define the region of each spot within this image
    jmin = nstamp * jcol
    jmax = jmin + nstamp
    spot = image[0:nstamp, jmin:jmax]
    rms = rms_from_SNR(snr, spot)
    noise = np.random.standard_normal((nstamp, nstamp)) * rms
    image[0:nstamp, jmin:jmax]+= noise

for jcol in xrange(ncols):
    # Define the region of each spot within this image
    jmin = nstamp * jcol
    jmax = jmin + nstamp
    spot = image[nstamp: 2 * nstamp, jmin:jmax]
    rms = rms_from_SNR(snr, spot)
    noise = np.random.standard_normal((nstamp, nstamp)) * rms
    image[nstamp: 2 * nstamp, jmin:jmax]+= noise

# Output to FITS image file
outimfile_snr = 'test_PSF_orientation_fwhm'+str(fwhm)+'beta'+\
                str(beta)+'snr'+str(snr)+'im.fits'
print 'test_PSF_orientation: Writing test image to '+outimfile_snr
hdu_image = pyfits.PrimaryHDU(image)
hdu_image.writeto(outimfile_snr, clobber=True)
os.remove('mtf_nopix.fits')
os.remove('mtf_pix.fits')
os.remove('spot.fits')