"""
Table interpolation class, to allow interpolation of 2D images in any number of other 
dimensions. Reads in an interim FITS table as defined below:

Tables organised as a standard FITS image file with each table as an image HDU. 
---------------
First (primary) HDU must contain the following header entries:

CRPIXx, CRVALx, CRDELTAx: Number, position and pixel spacing of the reference pixel in 
the image. Where x is the axis number (1 or 2).

GRIDVALS: Names of values describing the grid points
e.g. GRIDVALS = ALT,AZ
---------------
Each HDU must contain a header entry containing the grid points described by GRIDVALS
e.g. ALT =70
AZI = 0

It is also strongly recommended to include documentation for each interpolation 
dimension containing at least the units. This entry ust begin with DOC, followed by the 
dimension name
e.g. DOCALT = "Altitude of event (deg)"

TODO:
    - Improve error handling
    - Add option for caching nd interpolated value
    - Better deal with edges of phase space
    - Allow non-linear interpolation
"""

from scipy import interpolate
import numpy as np
from astropy.io import fits


class TableInterpolator:
    """
    This is a simple class for loading lookup tables from a fits file and
    interpolating between them
    """

    def __init__(self, filename, verbose=1):
        """
        Initialisation of class to load templates from a file and create the interpolation
        objects

        Parameters
        ----------
        filename: string
            Location of Template file
        verbose: int
            Verbosity level,
            0 = no logging
            1 = File + interpolation point information
            2 = Detailed description of interpolation points
        """
        self.verbose = verbose
        if self.verbose:
            print("Loading lookup tables from", filename)

        grid, bins, template = self.parse_fits_table(filename)
        x_bins, y_bins = bins

        self.interpolator = interpolate.LinearNDInterpolator(grid, template, fill_value=0)
        self.nearest_interpolator = interpolate.NearestNDInterpolator(grid, template)

        self.grid_interp = interpolate.RegularGridInterpolator((x_bins, y_bins),
                                                               np.zeros([x_bins.shape[0],
                                                                         y_bins.shape[
                                                                             0]]),
                                                               method="linear",
                                                               bounds_error=False,
                                                               fill_value=0)

    def parse_fits_table(self, filename):
        """
        Function opens tables contained within fits files and parses them into a format 
        recognisable by the interpolator.
        
        Parameters
        ----------
        filename: str
            Name of table file

        Returns
        -------
            tuple (grid points, bin centres, images)
        """
        file = fits.open(filename)
        template = list()
        grid = list()

        primHDU = file[0].header  # We require first HDU to be primary

        # Below definitions are standard
        ix, iy = primHDU["CRPIX2"], primHDU["CRPIX1"]
        val_x, val_y = primHDU["CRVAL2"], primHDU["CRVAL1"]
        print(val_x, val_y)
        delta_x, delta_y = primHDU["CRDELTA2"], primHDU["CRDELTA1"]
        nbins_x, nbins_y = primHDU["NAXIS2"], primHDU["NAXIS1"]
        ix *= delta_x
        iy *= delta_y

        x_bins = np.arange(val_x-ix ,val_x+(delta_x*nbins_x)-ix, step=delta_x)
        y_bins = np.arange(val_y-iy ,val_y+(delta_y*nbins_y)-iy, step=delta_y)
        grid_vals = primHDU["GRIDVALS"]
        points = grid_vals.split(",")

        if self.verbose:
            print("Interpolation point source be called in order", points)
        if self.verbose>1:
            for p in points:
                print(p,":", primHDU["DOC"+p])

        for hdu in file:
            template.append(hdu.data)
            
            hdu_pt = list()
            for p in points:
                hdu_pt.append(hdu.header[p])
            grid.append(np.array(hdu_pt))

        print(np.array(grid))
        bins = (x_bins, y_bins)
        return grid, bins, template

    def interpolate(self, params, pixel_pos_x, pixel_pos_y):
        """

        Parameters
        ----------
        params: ndarray
            numpy array of interpolation parameters
            currently [energy, impact distance, xmax]
        pixel_pos_x: ndarray
            pixel position in degrees
        pixel_pos_y: ndarray
            pixel position in degrees

        Returns
        -------
        ndarray of expected intensity for all pixel positions given

        """

        image = self.interpolated_image(params)
        self.grid_interp.values = image

        points = np.array([pixel_pos_x, pixel_pos_y])
        return self.grid_interp(points.T)

    def interpolated_image(self, params):
        """
        Function for creating a ful interpolated image template from the interpolation library

        Parameters
        ----------
        params: ndarray
            numpy array of interpolation parameters
            currently [energy, impact distance, xmax]

        Returns
        -------
        ndarray of a single image template

        """
        image = self.interpolator(params)[0]
        if np.isnan(image).all():
            print("Found a NaN", params)
            image = self.nearest_interpolator(params)[0]

        return image
