import os
import warnings
import numpy as np
from astropy.io import fits

from src.objects import File
from src.utils import scale_image_data


def rescale_image_data(img_header, img_data):
    """ Rescales (uncompresses) the data in line with the PS1 documentation.

    The image flux scaling is also non-standard. The images have been
    non-linearly scaled using an arc hyperbolic sine (asinh) transformation that
    converts them to a pseudo-magnitude scale_image_data related to the asinh
    magnitudes (aka "luptitudes") that are used in the Sloan Digital Sky
    Survey. The scaling is determined by the BSOFTEN/BOFFSET keywords in the
    FITS header.

    Here are the relevant lines from the header of the skycell 1725.051
    g-band image:

    .. math::

        BZERO   & =  3.283630371094E+00 / Scaling: TRUE = BZERO + BSCALE *
        DISK \\

        BSCALE  & =  2.008622776974E-04 / Scaling: TRUE = BZERO + BSCALE *
        DISK \\

        BSOFTEN & =  8.739592975627E+01 / Scaling: LINEAR = 2 * BSOFTEN * \sinh(
        TRUE/a) \\

        BOFFSET & =  2.654963016510E+00 / Scaling: UNCOMP = BOFFSET + LINEAR

    The comments describe the transformation to convert the pixel values to
    fluxes, but they are not very clear. BZERO and BSCALE are the standard FITS
    keywords for converting an integer image to a floating point value. The
    constant a is defined below. "UNCOMP" is the linear flux value. If
    :math:`j` is the original integer pixel value, these equations convert to
    the float pixel value :math:`v` and then to a standard linear flux:

    .. math::

        v & = BZERO + BSCALE * j \\

        a & = 2.5/\ln(10) = 1.0857362 \\

        x & = v/a \\

        flux & = BOFFSET + BSOFTEN * \sinh(x) \\

             & = BOFFSET + BSOFTEN * (e^{x} - e^{-x}) \\

             & = BOFFSET + BSOFTEN * (10^{0.4v} - 10^{-0.4v})

    Parameters
    ----------
    img_header :
        Header from the HDU list of the image file.

    img_data : array-like
        Image data from the HDU list of the image file.

    Returns
    ----------
    uncomp: The uncompressed data.

    """
    # Set variables for scaling / rescaling to header values.
    bsoften, boffset = img_header['BSOFTEN'], img_header['BOFFSET']
    a = 2.5 / np.log(10)
    true = scale_image_data(img_header, img_data)
    linear = 2 * bsoften * np.sinh(true / a)
    uncomp = boffset + linear
    return uncomp


def uncompress(file, data_out='none'):
    """ Uncompress an image file using the non-linear scaling formula
    described in `src.data.extract.uncompress.rescale_image`.

    For each directory generated by the get_data function, each file is
    rescaled (uncompressed) and then written out to the same directory with the
    original file name prepended with "uncomp_" to indicate it is uncompressed.

    Parameters
    ----------
    file : File
        A compressed PS1 science or weight image downloaded with the `PS1 Image
        Cutout Service <https://outerspace.stsci.edu/display/PANSTARRS/PS1
        +Image+Cutout+Service>`_.

    Returns
    ----------
    None

    Note
    ----------
    The following warning is to be expected: "VerifyWarning:  Invalid 'BLANK'
    keyword in header."
    """
    try:
        if isinstance(file, File):
            file = file.abspath

        data_out = file
        data_in = file.replace('interim', 'raw').replace('uncomp.', '')

        # Remove truncated (corrupt) compressed image files automatically.
        with warnings.catch_warnings(record=True):
            # 'error' turns matching warnings into exceptions
            trunc = 'File may have been truncated: actual file length*'
            warnings.filterwarnings(action='error', message=trunc, append=True)
            with fits.open(
                    data_in,
                    do_not_scale_image_data=True,
                    ignore_blank=True,
                    checksum=True) as hdu_list:
                img_header, img_data = hdu_list[1].header, hdu_list[1].data

        # Uncompress data and write to new file.
        file.log.info(f"Attempting to uncompress {data_in}")
        uncomp_data = rescale_image_data(img_header, img_data)
        hdu_output = fits.PrimaryHDU(data=uncomp_data, header=img_header)
        hdu_list = fits.HDUList([hdu_output])
        # checksum=True adds both ``DATASUM`` and ``CHECKSUM`` cards to headers.
        hdu_list.writeto(data_out, overwrite=True, checksum=True)
        # Stamp ``CHECKED`` if ``DATASUM`` and ``CHECKSUM`` are verified.
        file.stamp_check()

        file.log.info(f"Successfully uncompressed {data_out}")

    except Warning as e:
        file.log.exception('A warning occurred')
        if 'truncated' in str(e):
            file.log.warn('Deleting corrupt raw fits file, download it again.')
            os.remove(data_in)
            if file.exists:
                os.remove(data_out)

    except Exception as e:
        file.log.exception('An exception occurred')
