"""
Bias Calibration
================

This simple script can be used to study the number of bias frames required to meet the VIS calibration requirements.

The following requirements related to the bias calibration has been taken from GDPRD.

R-GDP-CAL-052:
The contribution of the residuals of VIS bias subtraction to the *error on the determination of each ellipticity
component* of the local PSF shall not exceed 3x10-5 (one sigma).

R-GDP-CAL-062:
The contribution of the residuals of VIS bias subtraction to the *relative error* \sigma(R2)/R2 on the determination of
the local PSF R2 shall not exceed 1x10-4 (one sigma).

:requires: PyFITS
:requires: NumPy
:requires: matplotlib
:requires: VISsim-Python

:version: 0.95

:author: Sami-Matias Niemi
:contact: smn2@mssl.ucl.ac.uk
"""
import matplotlib
matplotlib.rc('text', usetex=True)
matplotlib.rcParams['font.size'] = 17
matplotlib.rc('xtick', labelsize=14)
matplotlib.rc('axes', linewidth=1.1)
matplotlib.rcParams['legend.fontsize'] = 11
matplotlib.rcParams['legend.handlelength'] = 3
matplotlib.rcParams['xtick.major.size'] = 5
matplotlib.rcParams['ytick.major.size'] = 5
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
import pyfits as pf
import numpy as np
import math, datetime, cPickle, sys
from scipy import interpolate
from analysis import shape
from support import logger as lg
from support import surfaceFitting as sf
from support import bleedingtest as write
from support import files as fileIO


def testBiasCalibrationDelta(log, numdata=2066, floor=995, xsize=2048, ysize=2066, order=3, biases=15, surfaces=100,
                        file='psf1x.fits', psfs=500, psfscale=1.e3, debug=False, plots=False):
    """
    Derive the PSF ellipticities for a given number of random surfaces with random PSF positions
    and a given number of biases median combined and compare to the nominal PSF ellipticity.

    This function can be used to derive the error (delta) in determining ellipticity and size given
    a reference PSF.

    Choices that need to be made and effect the results:

        #. bias surface that is assumed (amplitude, complexity, etc.)
        #. whether the order of the polynomial surface to be fitted is known or not
        #. size of the Gaussian weighting function when calculating the ellipticity components

    There are also other choices such as the number of PSFs and scaling and the random numbers generated for
    the surface that also affect the results, however, to a lesser degree.

    Generates a set of plots that can be used to inspect the simulation.
    """
    log.info('Processing file %s' % file)
    #read in data without noise or bias level and scale it to 20k electrons
    data = pf.getdata(file)
    data /= np.max(data)
    data *= psfscale

    #derive the reference value from the scaled data
    sh = shape.shapeMeasurement(data.copy(), log)
    results = sh.measureRefinedEllipticity()
    #sh.writeFITS(results['GaussianWeighted'], file.replace('.fits', 'Gweighted.fits'))
    #rescale to not loose numerical accuracy
    reference1 = results['e1']
    reference2 = results['e2']
    refR2 = results['R2']
    reference = results['ellipticity']

    print 'Reference Ellipticities and R2 :'
    print reference1, reference2, reference, refR2

    #generate a random quadrant surface representing BIAS without noise
    #modify zclean if a different order surface is needed
    x = np.random.random(numdata)
    y = np.random.random(numdata)
    xx, yy = np.meshgrid(np.linspace(x.min(), x.max(), xsize),
                         np.linspace(y.min(), y.max(), ysize))
    zclean = yy - xx + 0.78*xx**2 + 15.0*yy**2 - 1.75*xx*yy + 10.0*xx**3 + 0.3*yy**3 + floor

    #random positions for the PSFs, these positions are the lower corners
    xpositions = np.random.random_integers(0, zclean.shape[1] - data.shape[1], psfs)
    ypositions = np.random.random_integers(0, zclean.shape[0] - data.shape[0], psfs)

    if plots:
        # generate 2D plot
        im = plt.imshow(zclean, extent=(0, ysize, xsize, 0))
        plt.scatter(xpositions/data.shape[1]/2, ypositions/data.shape[0]/2)
        c1 = plt.colorbar(im)
        c1.set_label('BIAS [ADUs]')
        plt.xlim(0, ysize)
        plt.ylim(0, xsize)
        plt.xlabel('Y [pixels]')
        plt.ylabel('X [pixels]')
        plt.savefig('NoNoise2D.png')
        plt.close()
        #and 3D
        fig = plt.figure()
        ax = Axes3D(fig)
        ax.plot_surface(xx*xsize, yy*ysize, zclean, rstride=100, cstride=100, alpha=0.6, cmap=cm.jet)
        ax.set_xlabel('X [pixels]')
        ax.set_ylabel('Y [pixels]')
        ax.set_zlabel('BIAS [ADUs]')
        plt.savefig('NoNoise.png')
        plt.close()

    out = {}
    #number of biases to median combine
    for a in xrange(biases):
        print 'Number of Biases: %i / %i' % (a+1, biases)

        #data storage
        de1 = []
        de2 = []
        de = []
        R2 = []
        R2abs = []

        #number of random readnoised surfaces to loop over
        for b in xrange(surfaces):

            print 'Surface: %i / %i' % (b+1, surfaces)

            #add readout noise based on a+1 median combined biases
            z = addReadoutNoise(zclean.copy(), number=a+1)

            if plots:
                # generate 2D plot
                im = plt.imshow(z, extent=(0, ysize, xsize, 0))
                c1 = plt.colorbar(im)
                c1.set_label('BIAS [ADUs]')
                plt.xlim(0, ysize)
                plt.ylim(0, xsize)
                plt.xlabel('Y [pixels]')
                plt.ylabel('X [pixels]')
                plt.savefig('Readnoised%i%i.png' % (a+1, b+1))
                plt.close()
                #and 3D
                fig = plt.figure()
                ax = Axes3D(fig)
                ax.plot_surface(xx*xsize, yy*ysize, z, rstride=100, cstride=100, alpha=0.6, cmap=cm.jet)
                ax.set_xlabel('X [pixels]')
                ax.set_ylabel('Y [pixels]')
                ax.set_zlabel('BIAS [ADUs]')
                plt.savefig('Readnoised3D%i%i.png' % (a+1, b+1))
                plt.close()

            # Fit 2d polynomial to the noised data
            m = sf.polyfit2d(xx.ravel(), yy.ravel(), z.ravel(), order=order)
            # Evaluate it on a rectangular grid
            fitted = sf.polyval2d(xx, yy, m)

            if plots:
                # generate 2D plot
                im = plt.imshow(fitted, extent=(0, ysize, xsize, 0))
                c1 = plt.colorbar(im)
                c1.set_label('BIAS [ADUs]')
                plt.xlim(0, ysize)
                plt.ylim(0, xsize)
                plt.xlabel('Y [pixels]')
                plt.ylabel('X [pixels]')
                plt.savefig('Fitted2D%i.png' % (a+1))
                plt.close()
                #and 3D
                fig = plt.figure()
                ax = Axes3D(fig)
                ax.plot_surface(xx*xsize, yy*ysize, fitted, rstride=100, cstride=100, alpha=0.6, cmap=cm.jet)
                ax.set_xlabel('X [pixels]')
                ax.set_ylabel('Y [pixels]')
                ax.set_zlabel('BIAS [ADUs]')
                plt.savefig('Fitted3D%i.png' % (a+1))
                plt.close()

            #subtract the no noise surface from the fit
            fitted -= zclean

            if plots:
                # generate 2D plot
                im = plt.imshow(fitted, extent=(0, ysize, xsize, 0))
                c1 = plt.colorbar(im)
                c1.set_label(r'$\Delta$BIAS [ADUs]')
                plt.xlim(0, ysize)
                plt.ylim(0, xsize)
                plt.xlabel('Y [pixels]')
                plt.ylabel('X [pixels]')
                plt.savefig('Residual2D%i%i.png' % (a+1, b+1))
                plt.close()
                #and 3D
                fig = plt.figure()
                ax = Axes3D(fig)
                ax.plot_surface(xx*xsize, yy*ysize, fitted, rstride=100, cstride=100, alpha=0.6, cmap=cm.jet)
                ax.set_xlabel('X [pixels]')
                ax.set_ylabel('Y [pixels]')
                ax.set_zlabel(r'$\Delta$BIAS [ADUs]')
                ax.set_zlim(-0.01, 0.01)
                plt.savefig('Residual3D%i%i.png' % (a+1, b+1))
                plt.close()

            #loop over the PSFs
            for xpos, ypos in zip(xpositions, ypositions):
                #measure e and R2 from the postage stamp image
                small = fitted[ypos:ypos+data.shape[0], xpos:xpos+data.shape[1]].copy()
                small += data.copy()
                sh = shape.shapeMeasurement(small.copy(), log)
                results = sh.measureRefinedEllipticity()

                #save delta values
                de1.append(results['e1'] - reference1)
                de2.append(results['e2'] - reference2)
                de.append(results['ellipticity'] - reference)
                R2.append((results['R2'] - refR2) / refR2)
                R2abs.append((results['R2'] - refR2))

                if debug:
                    print xpos, ypos
                    write.writeFITSfile(small/data, 'testResidualDelta.fits')
                    print 'DEBUG mode -- exiting now'
                    import sys; sys.exit()
        if plots:
            plotDeltaEs(de1, de2, de, 'MultipleBiases%i.pdf' % (a+1), title='%i Biases median combined' % (a+1))

        out[a+1] = [de1, de2, de, R2, R2abs]

    return out


def testBiasCalibrationSigma(log, numdata=2066, floor=3500, xsize=2048, ysize=2066, order=3, biases=15, surfaces=100,
                             file='psf1x.fits', psfs=500, psfscale=1e5, gain=3.1,
                             debug=False, plots=True):
    """
    Derive the PSF ellipticities for a given number of random surfaces with random PSF positions
    and a given number of biases median combined.

    This function is to derive the the actual values so that the knowledge (variance) can be studied.

    Choices that need to be made and effect the results:

        #. bias surface that is assumed (amplitude, complexity, etc.)
        #. whether the order of the polynomial surface to be fitted is known or not
        #. size of the Gaussian weighting function when calculating the ellipticity components

    There are also other choices such as the number of PSFs and scaling and the random numbers generated for
    the surface that also affect the results, however, to a lesser degree.

    Generates a set of plots that can be used to inspect the simulation.
    """
    log.info('Processing file %s' % file)

    #read in data without noise or bias level and renormalize it
    data = pf.getdata(file)
    data /= np.max(data)
    data  = data * psfscale / gain

    #generate a random quadrant surface representing BIAS without noise
    #modify zclean if a different order surface is needed
    x = np.random.random(numdata)
    y = np.random.random(numdata)