from __future__ import print_function
try:
    import stwcs
    from stwcs import wcsutil
except:
    stwcs = None

from stsci.tools import fileutil

OPUS_WCSKEYS = ['OCRVAL1','OCRVAL2','OCRPIX1','OCRPIX2',
                'OCD1_1','OCD1_2','OCD2_1','OCD2_2',
                'OCTYPE1','OCTYPE2']


def archive_prefix_OPUS_WCS(fobj,extname='SCI'):
    """ Identifies WCS keywords which were generated by OPUS and archived
        using a prefix of 'O' for all 'SCI' extensions in the file

        Parameters
        ----------
        fobj : str or `astropy.io.fits.HDUList`
            Filename or fits object of a file

    """
    if stwcs is None:
        print('=====================')
        print('The STWCS package is needed to convert an old-style OPUS WCS to an alternate WCS')
        print('=====================')
        raise ImportError


    closefits = False
    if isinstance(fobj,str):
        # A filename was provided as input
        fobj = fits.open(fobj,mode='update')
        closefits=True

    # Define the header
    ext = ('sci',1)
    hdr = fobj[ext].header

    numextn = fileutil.countExtn(fobj)
    extlist = []
    for e in range(1,numextn+1):
        extlist.append(('sci',e))

    # Insure that the 'O' alternate WCS is present
    if 'O' not in wcsutil.wcskeys(hdr):
        # if not, archive the Primary WCS as the default OPUS WCS
        wcsutil.archiveWCS(fobj,extlist, wcskey='O', wcsname='OPUS')

    # find out how many SCI extensions are in the image
    numextn = fileutil.countExtn(fobj,extname=extname)
    if numextn == 0:
        extname = 'PRIMARY'

    # create HSTWCS object from PRIMARY WCS
    wcsobj = wcsutil.HSTWCS(fobj,ext=ext,wcskey='O')
    # get list of WCS keywords
    wcskeys = list(wcsobj.wcs2header().keys())

    # For each SCI extension...
    for e in range(1,numextn+1):
        # Now, look for any WCS keywords with a prefix of 'O'
        for key in wcskeys:
            okey = 'O'+key[:7]
            hdr = fobj[(extname,e)].header
            if okey in hdr:
                # Update alternate WCS keyword with prefix-O OPUS keyword value
                hdr[key] = hdr[okey]

    if closefits:
        fobj.close()

def create_prefix_OPUS_WCS(fobj,extname='SCI'):
    """ Creates alternate WCS with a prefix of 'O' for OPUS generated WCS values
        to work with old MultiDrizzle.

        Parameters
        ----------
        fobj : str or `astropy.io.fits.HDUList`
            Filename or fits object of a file

        Raises
        ------
        IOError:
            if input FITS object was not opened in 'update' mode

    """
    # List of O-prefix keywords to create
    owcskeys = OPUS_WCSKEYS

    closefits = False
    if isinstance(fobj,str):
        # A filename was provided as input
        fobj = fits.open(fobj, mode='update')
        closefits=True
    else:
        # check to make sure this FITS obj has been opened in update mode
        if fobj.fileinfo(0)['filemode'] != 'update':
            print('File not opened with "mode=update". Quitting...')
            raise IOError

    # check for existance of O-prefix WCS
    if owcskeys[0] not in fobj['sci',1].header:

        # find out how many SCI extensions are in the image
        numextn = fileutil.countExtn(fobj,extname=extname)
        if numextn == 0:
            extname = ''
        for extn in range(1,numextn+1):
            hdr = fobj[(extname,extn)].header
            for okey in owcskeys:
                hdr[okey] = hdr[okey[1:]+'O']

    # Close FITS image if we had to open it...
    if closefits:
        fobj.close()