"""
Generate PSF at an arbitrary position in a drizzled image using the WFC3/IR
effective PSFs.
"""
import os
from collections import OrderedDict

import numpy as np

import astropy.io.fits as pyfits
import astropy.wcs as pywcs

#from research.dash.epsf import DrizzlePSF

DEMO_LIST = ['ib6o23rtq_flt.fits', 'ib6o23rwq_flt.fits', 'ib6o23rzq_flt.fits', 'ib6o23s2q_flt.fits']
DEMO_IMAGE = 'wfc3-ersii-g01-b6o-23-119.0-f140w_drz_sci.fits'

try:
    from .. import utils, model
except:
    from grizli import utils, model
    
class DrizzlePSF(object):
    def __init__(self, flt_files=DEMO_LIST, info=None, driz_image=DEMO_IMAGE, driz_hdu=None, beams=None):
        """
        Object for making drizzled PSFs
        
        Parameters
        ----------
        flt_files : list
            List of FLT files that were used to create the drizzled image.
        
        driz_image : str
            Filename of the drizzled image.
            
        """
        if info is None:
            if beams is None:
                info = self._get_flt_wcs(flt_files)
            else:
                info = self._get_wcs_from_beams(beams)
                
        self.flt_keys, self.wcs, self.footprint = info
        self.flt_files = list(np.unique([key[0] for key in self.flt_keys]))

        self.ePSF = utils.EffectivePSF()
        
        if driz_hdu is None:
            self.driz_image = driz_image
            self.driz_header = pyfits.getheader(driz_image)
            self.driz_wcs = pywcs.WCS(self.driz_header)
            self.driz_pscale = utils.get_wcs_pscale(self.driz_wcs)
        else:
            self.driz_image = driz_image
            self.driz_header = driz_hdu.header
            self.driz_wcs = pywcs.WCS(self.driz_header)
            self.driz_pscale = utils.get_wcs_pscale(self.driz_wcs)
        
        self.driz_wcs.pscale = self.driz_pscale
        
    @staticmethod
    def _get_flt_wcs(flt_files):
        """
        TBD
        """
        from shapely.geometry import Polygon, Point
        
        from grizli import utils
        wcs = OrderedDict()
        footprint = OrderedDict()
        
        flt_keys = []
        for file in flt_files:
            flt_j = pyfits.open(file)
            for ext in range(1,5):
                if ('SCI',ext) in flt_j:
                    key = file, ext
                    wcs[key] = pywcs.WCS(flt_j['SCI', ext], relax=True,
                                          fobj=flt_j)
                    
                    wcs[key].pscale = utils.get_wcs_pscale(wcs[key])
                    
                    footprint[key] = Polygon(wcs[key].calc_footprint())
                    flt_keys.append(key)
                    
        return flt_keys, wcs, footprint
    
    @staticmethod
    def _get_wcs_from_beams(beams):
        """
        TBD
        """
        from shapely.geometry import Polygon, Point
        
        from grizli import utils
        wcs = OrderedDict()
        footprint = OrderedDict()
        
        flt_keys = []
        for beam in beams:
            #flt_j = pyfits.open(file)
            file = beam.direct.parent_file
            ext = beam.direct.sci_extn
            
            key = file, ext
            wcs[key] = beam.direct.wcs.copy()
            
            wcs[key].pscale = utils.get_wcs_pscale(wcs[key])
            
            footprint[key] = Polygon(wcs[key].calc_footprint())
            flt_keys.append(key)
                    
        return flt_keys, wcs, footprint
        
    def get_driz_cutout(self, ra=53.06967306, dec=-27.72333015, size=15, get_cutout=False, N=None):
        """
        TBD
        """