r"""
Telluric Spectrum
-----------------

An experimental container for a single grid-point of a pre-computed telluric atmosphere model.  The spectrum is a vector with coordinates wavelength and flux :math:`F(\lambda)`.

TelluricSpectrum
################
"""

import warnings
import logging
from gollum.precomputed_spectrum import PrecomputedSpectrum
import numpy as np
import astropy
import pandas as pd
from astropy import units as u
from specutils import SpectrumCollection
import os
from specutils.spectra.spectrum1d import Spectrum1D
from tqdm import tqdm

log = logging.getLogger(__name__)

#  See Issue: https://github.com/astropy/specutils/issues/779
warnings.filterwarnings(
    "ignore", category=astropy.utils.exceptions.AstropyDeprecationWarning
)
# See Issue: https://github.com/astropy/specutils/issues/800
warnings.filterwarnings("ignore", category=RuntimeWarning)


## TODO: rename this to SkyCalc??
class TelluricSpectrum(PrecomputedSpectrum):
    r"""
    A container for a single Telluric precomputed synthetic spectrum, currently from the skycalc website at eso.org. 

    Args:
        path (str): The path to a whitespace-delimited ASCII file generated by skycalc on eso.org.  You can either provide the file locally or a weblink.  'default' uses a demo spectrum
        wl_lo (float): the bluest wavelength of the models to keep (Angstroms)
        wl_hi (float): the reddest wavelength of the models to keep (Angstroms)
        """

    def __init__(self, *args, path=None, wl_lo=8038, wl_hi=12849, **kwargs):

        if path == "default":
            path = "https://gist.githubusercontent.com/gully/b93a04ede17f617d36b2a8ffc32a60e6/raw/01c603ee201509b79b1a464093e6200ae2032198/demo_skycalc_telluric_transmission_spectrum.txt"

        if path is not None:

            # Units: nm, transmittance
            df_native = (
                pd.read_csv(
                    path,
                    delim_whitespace=True,
                    names=["wavelength_nm", "transmittance"],
                )
                .sort_values("wavelength_nm")
                .reset_index(drop=True)
            )

            # convert to Angstrom
            df_native["wavelength"] = df_native["wavelength_nm"] * 10.0
            mask = (df_native.wavelength > wl_lo) & (df_native.wavelength < wl_hi)
            df_trimmed = df_native[mask].reset_index(drop=True)

            super().__init__(
                spectral_axis=df_trimmed.wavelength.values * u.Angstrom,
                flux=df_trimmed.transmittance.values * u.dimensionless_unscaled,
                **kwargs,
            )

        else:
            super().__init__(*args, **kwargs)


class TelFitSpectrum(PrecomputedSpectrum):
    r"""
    A container for a single TelFit precomputed synthetic spectrum. 

    Args:
        path (str): The path to a whitespace-delimited ASCII file generated by skycalc on eso.org.  You can either provide the file locally or a weblink.  'default' uses a demo spectrum
        wl_lo (float): the bluest wavelength of the models to keep (Angstroms)
        wl_hi (float): the reddest wavelength of the models to keep (Angstroms)
        """

    def __init__(self, *args, path=None, **kwargs):

        if path is not None:

            names = ["wavelength_nm", "transmission", "continuum", "err"]
            df_native = pd.read_csv(
                path,
                delim_whitespace=True,
                names=names,
                usecols=names[0:2],
                encoding="unicode_escape",  # or UTF-8?
                low_memory=False,
            )

            # Units: nm, transmittance
            # convert to Angstrom
            df_native["wavelength"] = (
                pd.to_numeric(df_native["wavelength_nm"], errors="coerce") * 10.0
            )
            flux_out = pd.to_numeric(df_native["transmission"], errors="coerce")

            good_mask = (df_native.wavelength.values == df_native.wavelength.values) & (
                flux_out.values == flux_out.values
            )

            super().__init__(
                spectral_axis=df_native.wavelength.values[good_mask] * u.Angstrom,
                flux=flux_out.values[good_mask] * u.dimensionless_unscaled,
                **kwargs,
            )

        else:
            super().__init__(*args, **kwargs)

    def air_to_vacuum(self):
        """Converts a spectrum from air to vacuum  

        Based on Morton, D. C. 1991, ApJS, 77, 119
        """
        wave_A = self.wavelength.to(u.Angstrom).value
        n_air = 1.0 + 2.735182e-4 + 131.4182 / wave_A ** 2 + 2.76249e8 / wave_A ** 4
        vacuum_wavelengths = self.wavelength * n_air
        return self._copy(spectral_axis=vacuum_wavelengths, wcs=None)

