"""
Loader for SDSS individual spectrum files: spec_ files.

.. _spec: https://data.sdss.org/datamodel/files/BOSS_SPECTRO_REDUX/RUN2D/spectra/PLATE4/spec.html
"""
import os
import re
import _io

from astropy.io import fits
from astropy.table import Table
from astropy.wcs import WCS
from astropy.units import Unit, def_unit
from astropy.nddata import StdDevUncertainty, InverseVariance

import numpy as np

from ...spectra import Spectrum1D
from ..registers import data_loader, custom_writer
from ..parsing_utils import read_fileobj_or_hdulist

__all__ = ['spec_identify', 'spSpec_identify',
           'spec_loader', 'spSpec_loader']

_spSpec_pattern = re.compile(r'spSpec-\d{5}-\d{4}-\d{3}\.fit')
_spec_pattern = re.compile(r'spec-\d{4,5}-\d{5}-\d{4}\.fits')


def _sdss_wcs_to_log_wcs(old_wcs):
    """
    The WCS in the SDSS files does not appear to follow the WCS standard - it
    claims to be linear, but is logarithmic in base-10.
    The wavelength is given by:
    λ = 10^(w0 + w1 * i)
    with i being the pixel index starting from 0.

    The FITS standard uses a natural log with a sightly different formulation,
    see WCS Paper 3 (which discusses spectral WCS).

    This function does the conversion from the SDSS WCS to FITS WCS.
    """
    w0 = old_wcs.wcs.crval[0]
    w1 = old_wcs.wcs.cd[0,0]
    crval = 10 ** w0
    cdelt = crval * w1 * np.log(10)
    cunit = old_wcs.wcs.cunit[0] or Unit('Angstrom')
    ctype = "WAVE-LOG"

    w = WCS(naxis=1)
    w.wcs.crval[0] = crval
    w.wcs.cdelt[0] = cdelt
    w.wcs.ctype[0] = ctype
    w.wcs.cunit[0] = cunit
    w.wcs.set()

    return w


def spec_identify(origin, *args, **kwargs):
    """
    Check whether given input is FITS and has SDSS-III/IV spec type
    BINTABLE in first extension. This is used for Astropy I/O Registry.
    """
    # Test if fits has extension of type BinTable and check for spec-specific keys
    with read_fileobj_or_hdulist(*args, **kwargs) as hdulist:
        return (hdulist[0].header.get('TELESCOP') == 'SDSS 2.5-M' and
                hdulist[0].header.get('FIBERID', 0) > 0 and
                len(hdulist) > 1 and
                (isinstance(hdulist[1], fits.BinTableHDU) and
                 hdulist[1].header.get('TTYPE3').lower() == 'ivar'))


def spSpec_identify(origin, *args, **kwargs):
    """
    Check whether given input is FITS with SDSS-I/II spSpec tyamepe data.
    This is used for Astropy I/O Registry.
    """
    # Test telescope keyword and check if primary HDU contains data
    # consistent with spSpec format
    with read_fileobj_or_hdulist(*args, **kwargs) as hdulist:
        return (hdulist[0].header.get('TELESCOP') == 'SDSS 2.5-M' and
                hdulist[0].header.get('FIBERID', 0) > 0 and
                isinstance(hdulist[0].data, np.ndarray) and
                hdulist[0].data.shape[0] == 5)


def spPlate_identify(origin, *args, **kwargs):
    """
    Check whether given input is FITS with SDSS spPlate fibre spectral data.
    This is used for Astropy I/O Registry.
    """
    # Test telescope keyword and check if primary HDU contains data
    # consistent with spSpec format
    with read_fileobj_or_hdulist(*args, **kwargs) as hdulist:
        return (hdulist[0].header.get('TELESCOP') == 'SDSS 2.5-M' and
                hdulist[0].header.get('FIBERID', 0) <= 0 and
                isinstance(hdulist[0].data, np.ndarray) and
                hdulist[0].data.shape[0] > 5)


@data_loader(
    label="SDSS-III/IV spec", identifier=spec_identify, extensions=['fits'],
    priority=10,
)
def spec_loader(file_obj, **kwargs):
    """
    Loader for SDSS-III/IV optical spectrum "spec" files.

    Parameters
    ----------
    file_obj: str, file-like, or HDUList
          FITS file name, object (provided from name by Astropy I/O Registry),
          or HDUList (as resulting from astropy.io.fits.open()).

    Returns
    -------
    data: Spectrum1D
        The spectrum that is represented by the 'loglam' (wavelength) and 'flux'
        data columns in the BINTABLE extension of the FITS `file_obj`.
    """
    with read_fileobj_or_hdulist(file_obj, **kwargs) as hdulist:
        header = hdulist[0].header
        name = header.get('NAME')
        meta = {'header': header}

        bunit = header.get('BUNIT', '1e-17 erg / (Angstrom cm2 s)')
        if 'Ang' in bunit and 'strom' not in bunit:
            bunit = bunit.replace('Ang', 'Angstrom')
        flux_unit = Unit(bunit)

        # spectrum is in HDU 1
        flux = hdulist[1].data['flux'] * flux_unit

        uncertainty = InverseVariance(hdulist[1].data['ivar'] / flux_unit**2)

        dispersion = 10**hdulist[1].data['loglam']
        dispersion_unit = Unit('Angstrom')

        mask = hdulist[1].data['and_mask'] != 0

    return Spectrum1D(flux=flux, spectral_axis=dispersion * dispersion_unit,
                      uncertainty=uncertainty, meta=meta, mask=mask)


@data_loader(
    label="SDSS-I/II spSpec", identifier=spSpec_identify,
    extensions=['fit', 'fits'], priority=10,
)
def spSpec_loader(file_obj, **kwargs):
    """
    Loader for SDSS-I/II spSpec files.

    Parameters
    ----------
    file_obj: str, file-like, or HDUList
           FITS file name, object (provided from name by Astropy I/O Registry),
           or HDUList (as resulting from astropy.io.fits.open()).

    Returns
    -------
    data: Spectrum1D
        The spectrum that is represented by the wavelength solution from the
        header WCS and data array of the primary HDU.
    """
    with read_fileobj_or_hdulist(file_obj, **kwargs) as hdulist:
        header = hdulist[0].header
        # name = header.get('NAME')
        meta = {'header': header}
        wcs = WCS(header).dropaxis(1)

        bunit = header.get('BUNIT', '1e-17 erg / (Angstrom cm2 s)')
        # fix mutilated flux unit
        bunit = bunit.replace('/cm/s/Ang', '/ (Angstrom cm2 s)')
        if 'Ang' in bunit and 'strom' not in bunit:
            bunit = bunit.replace('Ang', 'Angstrom')
        flux_unit = Unit(bunit)
        flux = hdulist[0].data[0, :] * flux_unit

        uncertainty = StdDevUncertainty(hdulist[0].data[2, :] * flux_unit)

        # Fix the WCS if it is claimed to be linear
        if header.get('DC-Flag', 1) == 1:
            fixed_wcs = _sdss_wcs_to_log_wcs(wcs)
        else:
            fixed_wcs = wcs

        mask = hdulist[0].data[3, :] != 0

    return Spectrum1D(flux=flux, wcs=fixed_wcs,
                      uncertainty=uncertainty, meta=meta, mask=mask)


@data_loader(label="SDSS spPlate", identifier=spPlate_identify, extensions=['fits'])
def spPlate_loader(file_obj, limit=None, **kwargs):
    """
    Loader for SDSS spPlate files, reading flux spectra from all fibres into single array.

    Parameters
    ----------
    file_obj: str, file-like, or HDUList
           FITS file name, object (provided from name by Astropy I/O Registry),
           or HDUList (as resulting from astropy.io.fits.open()).

    limit : :class:`int`, optional
        If set, only return the first `limit` spectra in `flux` array.

    Returns
    -------
    Spectrum1D
        The spectra represented by the wavelength solution from the header WCS
        and the data array of the primary HDU (typically 640 along dimension 1).
    """
    with read_fileobj_or_hdulist(file_obj, **kwargs) as hdulist:
        header = hdulist[0].header
        meta = {'header': header}
        wcs = WCS(header).dropaxis(1)
        if limit is None:
            limit = header['NAXIS2']

        bunit = header.get('BUNIT', '1e-17 erg / (Angstrom cm2 s)')
        if 'Ang' in bunit and 'strom' not in bunit:
            bunit = bunit.replace('Ang', 'Angstrom')
        flux_unit = Unit(bunit)
        flux = hdulist[0].data[0:limit, :] * flux_unit
        uncertainty = InverseVariance(hdulist[1].data[0:limit, :] / flux_unit**2)

        # Fix the WCS if it is claimed to be linear
        if header.get('DC-Flag', 1) == 1:
            fixed_wcs = _sdss_wcs_to_log_wcs(wcs)
        else:
            fixed_wcs = wcs

        mask = hdulist[2].data[0:limit, :] != 0
        meta['plugmap'] = Table.read(hdulist[5])[0:limit]

    return Spectrum1D(flux=flux, wcs=fixed_wcs,
                      uncertainty=uncertainty, meta=meta, mask=mask)
