"""
See LICENSE.txt
"""
import os
import math
from copy import deepcopy
import logging

from stwcs.wcsutil import HSTWCS
from astropy.io import fits
from astropy.table import Table

from hstaxe.axeerror import aXeError

# make sure there is a logger
_log = logging.getLogger(__name__)

class ProjectionList:
    def __init__(self, filename='', dim_info=''):

        """
        The central class for the new object lists that will
        be generated by projecting out positions from a
        drizzled image catalog.

        Parameters
        ----------
        filename : str
            The name of the drizzle input image.
            This is used to extract the D*DATA
            values saved to the header by astrodrizzle
        dim_info : str
            information on extra pixel rows/columns to be added
            to the 'natural' image size

        Returns
        -------
        Nothing

        Notes:
        ------
        This init-method extracts all necessary instance data
        from the input image expression in the header of the
        drizzled data. Also the dimensions of the input image
        are determined, and with together with the input parameter
        the x/y-min/max of the area allowed in the IOL is defined.
        """
        if not filename:
            raise aXeError("IOLPREP: No filename specified!")
        self.filename = filename
        self.iol_info = {}
        self.dim_info = []

        # get the information related to the IOL
        self._set_iol_info()

        # the the dimension of the input image
        self._set_indims()

        # compute the dimension information
        self._compute_diminfo(dim_info)

    def _set_iol_info(self):
        """Get all info on the extension."""

        # get the bracket information
        exte_str = self.filename.split('.fits')[1]

        # get the inside of the brackets
        exte_data = exte_str[1:len(exte_str)-1]

        # collect fits name, extension name and extension
        # version in the dictionary
        self.iol_info['root'] = self.filename.split('.fits')[0]
        self.iol_info['fits'] = self.filename.split('.fits')[0] + '.fits'
        # make sure the fits file exists
        if not os.path.isfile(self.iol_info['fits']):
            err_msg = ("IOLPREP: File {0:s} does not exist!"
                       .format(self.iol_info['fits']))
            raise aXeError(err_msg)

        self.iol_info['ext_nam'] = exte_data.split(',')[0]
        self.iol_info['ext_ver'] = int(exte_data.split(',')[1])
        self.iol_name = '%s_%i.cat' % (self.iol_info['root'],
                                       self.iol_info['ext_ver'])

    def _compute_diminfo(self, dim_info):
        """Compute the absolute dimension information"""

        # possible add rows/columns on either side.
        # the '0.5' generally have to be there
        # given an image with 1000 pix, SExtractor
        # can distribute object position in [0.5, 1000.5]
        self.dim_info.append(0.5-dim_info[0])
        self.dim_info.append(0.5+dim_info[1]+self.indims[0])
        self.dim_info.append(0.5-dim_info[2])
        self.dim_info.append(0.5+dim_info[3]+self.indims[1])

    def _set_indims(self):
        """Get the input dimensions of the image"""

        in_img = fits.open(self.iol_info['fits'], 'readonly')
        in_head = in_img[self.iol_info['ext_nam'], self.iol_info['ext_ver']].header

        # extract the keywords for the image size from
        # the header
        self.indims = [in_head['NAXIS1'], in_head['NAXIS2']]
        in_img.close()
        del in_head

    def make_grismcat(self, drizzle_image="", catalog=None, hard_angle=False, hard_angle_value=90.):

        """Make the grism catalog.

        The method creates a new input object list. The positional
        information on objects in a drizzled image are projected
        back into the coordinate system of one input image.
        A selection is done on the basis of the projected coordinates,
        and the selected objects are stored to a new IOL file with an
        updated object angle that is in the projected coordinate system.

        Parameters
        ----------
        drizzle_image : str
            The name of the drizzled mosaic image that the catalog
            was constructed from.

        catalog : astropy.table.Table
            A the master catalog which will be used to
            create the dither image catalog with updated angles

        hard_angle : bool
            This will always set the extraction angle to 90degrees.

        hard_angle_value : float
            If a specified angle for extraction is preferred then
            this is the value to use and will replace THETA_IMAGE
            in the output catalog. It's currently specified in
            degrees, and is converted appropriately for whatever
            units are used in the catalog itself.

        Returns
        -------
        Nothing

        Notes
        -----
        This method previous took the names of two text files that contained the
        displayed position and angles calculated from the catalog entry. That
        functionality has been moved to memory since awtran is no longer in use.
        The catalog is read directly.
        """

        if not drizzle_image:
            raise aXeError("IOLPREP: No drizzle image specified.")

        if catalog is None:
            raise aXeError("IOLPREP: No input catalog provided.")
        if not isinstance(catalog, Table):
            raise aXeError("IOLPREP: Expected input catalog to be an astropy table")

        _log.info("\n >>>> Working on Input Object List: {0:s} >>>>\n"
              .format(self.iol_name))

        # now translate ra dec to new image pixel points
        # this must go through the wcs of the mosaic image
        # and then through the wcs for the individual image
        _log.info("Converting coordinates using wcs from grism image {0}\n".format(self.filename))
        trad = catalog['THETA_IMAGE']
        xcat = catalog['X_IMAGE']
        ycat = catalog['Y_IMAGE']

        # translate to degrees if necessary
        if catalog['THETA_IMAGE'].unit.name == 'deg':
            translate = True
        elif catalog['THETA_IMAGE'].unit.name == 'rad':
            translate = False
            hard_angle_value = math.radians(hard_angle_value)
        else:
            raise aXeError("Unknown unit on THETA_IMAGE in input catalog")

        # 10.0 is a made up scaling length to use to
        # get the angle precision
        shifted_x = []
        shifted_y = []
        for t, x, y in zip(trad, xcat, ycat):
            if translate:
                angle = math.radians(t)

            shifted_x.append(x + 10.0 * math.cos(angle))
            shifted_y.append(y + 10.0 * math.sin(angle))

        # translate the catalog (x, y) to (ra, dec)
        mosaic_image_wcs = HSTWCS(drizzle_image, ext=1)
        mosaic_image_ra, mosaic_image_dec = mosaic_image_wcs.wcs_pix2world(shifted_x,
                                                                           shifted_y,
                                                                           1)
        # compute  the location in the dithered image using the shifted coords
        dither_image_wcs = HSTWCS(self.filename)
        dither_image_x, dither_image_y = dither_image_wcs.all_world2pix(mosaic_image_ra,
                                                                        mosaic_image_dec,
                                                                        1)
        trans_ra, trans_dec = mosaic_image_wcs.wcs_pix2world(xcat, ycat, 1)
        trans_x, trans_y = dither_image_wcs.all_world2pix(trans_ra, trans_dec, 1)

        output_catalog = deepcopy(catalog)
        for row in range(len(catalog)-1, -1, -1):
            x = trans_x[row]
            y = trans_y[row]
            # check whether the object position is
            # in the range to be stored
            if (( self.dim_info[0] <= x <= self.dim_info[1])  and
                ( self.dim_info[2] <= y <= self.dim_info[3]) ):

                # compute the new object angle
                dx = xcat[row] - x
                dy = ycat[row] - y
                angle = math.atan2(dy, dx)  # compute local angle change

                # return to degrees for catalog if necessary
                if translate:
                    angle = math.degrees(angle)

                # _log.info("dx {} dy {}  angle {} x,y: ({},{})\n".format(dx, dy, angle, x, y))
                # fill in the new position and angle
                output_catalog['X_IMAGE'][row] = trans_x[row]
                output_catalog['Y_IMAGE'][row] = trans_y[row]
                if hard_angle:
                    output_catalog['THETA_IMAGE'][row] = hard_angle_value
                else:
                    output_catalog['THETA_IMAGE'][row] = angle
            else:
                _log.info(f"{x}\t{y}\t{self.dim_info}\t{output_catalog['NUMBER'][row]}")
                output_catalog.remove_row(row)

        # save the new IOL, this is done especially for the C
        # code which is expecting Source Extractor style catalog
        # files. Decided to keep the output here consistent with
        # the C code so that it can continue to be used separately.
        # numbers start at 1 not zero. The output formatting allows
        # the astropy.sextractor formatter to read the catalog file.
        # There isn't currently an astropy writer for that format.
        if os.access(self.iol_name, os.F_OK):
            os.remove(self.iol_name)
        of = open(self.iol_name, 'w')
        for num, name in zip(range(len(output_catalog.colnames)+1), output_catalog.colnames):
            of.write("# {0:d} {1:s}\t\t{2:s}\t\t[{3:s}]\n".format(num+1,
                                                                  name,
                                                                  output_catalog[name].description,
                                                                  str(output_catalog[name].unit))
                                                                 )
        output_catalog.write(of, format='ascii.no_header', overwrite=False)
        of.close()

        _log.info(f"\n >>>> Catalog: {self.iol_name} written with {len(catalog)} entries.>>>> \n")


class IOLMaker:
    """Central class to take the input and to create Input Object Lists

    for the list of images extracted from the header of the
    drizzled image.
    """
    def __init__(self, drizzle_image, input_cat,  dim_term):
        """
        Parameters
        ----------
        drizzle_image: str
            the name of the drizzled image
        input_cat: str
            the name of the input catalogue made
            from the mosaic drizzled image, the master catalog
        dim_term:
            description of the additional rows/column for the
            Input Object Lists

        Returns
        -------
        Nothing

        Notes
        -----
        Basic checks on the input is done. The existence of the images is
        checked, also the data type of the various real
        or integer numbers.
        """
        self.iol_list = []
        self.dim_info = []

        # check whether the drizzled image exists,
        # store the name if it exists
        if not os.path.isfile(drizzle_image):
            err_msg = "File: {0:s} does not exist!".format(drizzle_image)
            raise aXeError(err_msg)
        else:
            self.drizzle_image = drizzle_image

        # check whether the input catalogue exists,
        # store the name if it exists
        if not os.path.isfile(input_cat):
            err_msg = 'File: ' + input_cat + ' does not exist!'
            raise aXeError(err_msg)
        else:
            self.input_cat = Table.read(input_cat,
                                        format='ascii.sextractor')

        # resolve and get the dimension information
        self._set_dimension_info(dim_term)

        # create the list of Input Object Instances
        self._fill_iollist()

    def _set_dimension_info(self, dimension_term):
        """Get the dimension information"""
        # initialize the array

        # check the dimension input
        dim_entries = dimension_term.split(',')

        # is the number of items correct?
        if len(dim_entries) != 4:
            err_msg = ("There must be 4 entries in the term: {0:s},"
                       "not {1:d}!".format(dimension_term, len(dim_entries)))
            raise aXeError(err_msg)

        # check whether each item is an integer
        for item in dim_entries:
            try:
                int_item = int(item.strip())
            except ValueError:
                raise aXeError("Item: {} must be integer!".format(item))

            # store the item
            self.dim_info.append(int_item)

    def _fill_iollist(self):
        """Derive the names for the input parameters from input images.

        The method derives the names for the input
        images from the header of the drizzled image
        and then creates an Input Object List instance
        for each input image.

        Parameters
        ----------
        None

        Returns
        -------
        Saves the list with the IOL instances.

        """

        # determine the number of input images
        n_iols = self._get_niol()

        drizzle_head = fits.getheader(self.drizzle_image, ext=0)

        # for each input image
        for index in range(1, n_iols+1):

            # extract the name of an input file
            keyname = 'D%03iDATA' % (index)

            # create a new IOl instance at the end of the list
            self.iol_list.append(ProjectionList(drizzle_head[keyname],
                                                self.dim_info))

    def _get_niol(self):
        """Return the number of input images from the header
        Parameters
        ----------
        None

        Returns
        -------
        ID: int
            The number of input images

        Notes
        -----
        The method looks in the header of the drizzled image
        for the name of the input images used in the drizzle
        process. This number is returned
        """
        # open the fits file and get the header
        drizzle_head = fits.getheader(self.drizzle_image, ext=0)

        # create the keyname for the first input image
        ID = 1
        keyname = "D{0:03d}DATA".format(ID)

        # create the keyname for subsequent input images
        # and continue until the keynames do not exist
        while keyname in drizzle_head:
            ID = ID+1
            keyname = "D{0:03d}DATA".format(ID)

        # correct the number
        ID = ID-1

        # return the number
        return ID

    def run(self):
        """Create the Input Object Lists.

        This method is responsible for creating the
        Input Object Lists. Other internal methods as well as
        methods of other classes are successively
        called to create the Input Object Lists associated with the
        input images listed in the header of the
        drizzled image.


        Parameters
        ----------
        None

        Returns
        -------
        Nothing
        """

        # create the new Input Object Lists
        for iol in self.iol_list:
            # make a new IOL for each input image
            iol.make_grismcat(self.drizzle_image, self.input_cat)
