"""
Implementation the PACO algorithm implementation for VIP
Based on Flasseur+ 2018 https://ui.adsabs.harvard.edu/abs/2018A%26A...618A.138F/abstract

Variable naming is based on the notation of Flasseur+ 2018,
see table 1 in the paper for a description.

Last updated 2022-05-09 by Evert Nasedkin (nasedkinevert@gmail.com).
"""

import sys
#import os
from abc import abstractmethod
# Required so numpy parallelization doesn't conflict with multiprocessing
# os.environ["MKL_NUM_THREADS"] = "1"
# os.environ["NUMEXPR_NUM_THREADS"] = "1"
# os.environ["OMP_NUM_THREADS"] = "1"

#from multiprocessing import Pool
from typing import Tuple, Union, Optional, Callable
import numpy as np
from scipy import ndimage
from scipy.ndimage import filters

from ..config.utils_conf import pool_map, iterable
from ..preproc.rescaling import frame_px_resampling, cube_px_resampling, frame_shift
from ..var.coords import cart_to_pol, pol_to_cart
from ..metrics.detection import detection
from ..fm import normalize_psf
__author__ = "Evert Nasedkin"
__all__ = ['FastPACO',
           'FullPACO']

class PACO:
    """
    This class implements the bulk of the PACO algorithm as described by
    Flasseur et al (2018). In general, the idea is to take in an ADI stack of
    images and statistically determine if there is a signal above the
    background in each 'patch' of the image. This is done by tracing the ark of
    the hypothesized planet through the stack, and comparing this set of patches
    to a set consisting of background only. This is done for each pixel (or
    sub-pixel) location in.
    The output is a signal-to-noise and/or a flux map over the field of view.
    The user can choose to use FullPACO or FastPACO, which are described by
    algorithms 1 and 2 of Flasseur+ 2018. FastPACO has been parallelized,
    and is the recommended usage.

    This output can then be used to compute an unbiased estimate of the flux
    of point sources detected in the image above some user-supplied detection
    threshold.

    Parameters
    ----------
    cube : numpy.ndarray
        3D science frames taken in pupil tracking/ADI mode.
        Dimensions should be (time, x, y), and units should be detector units (ie output
        of SPHERE or GPI reduction pipelines). The data should be centered, and have
        pre-processing already applied (e.g. bad pixel correction).
    angles : numpy.ndarray
        List of parallactic angles for each frame in degrees. Length of this array
        should be the same as the time axis of the science cube. Sign convention
        is the same as in the rest of VIP.
    psf : numpy.ndarray
        Unsaturated PSF image. If a cube is provided, the median of the cube will be used.
    dit_psf : float, optional
        Integration time of the unsaturated PSF in seconds. The PSF is normalised
        to dit_science/dit_psf/nd_transmission.
    dit_science : float, optional
        Integration time of the science frames in seconds.T he PSF is normalised
        to dit_science/dit_psf/nd_transmission.
    nd_transmission : float, optional
        Transmission of an ND filter used to aquire the unsaturated PSF. The PSF is normalised
        to dit_science/dit_psf/nd_transmission.
    fwhm : float, optional
        FWHM of PSF in arcseconds. Default values give 4px radius.
    pixscale : float, optional
        Detector pixel scale in arcseconds per pixel.  Default values give 4px radius.
    rescaling_factor : float, optional
        Scaling for sub/super pixel resolution for PACO. Will rescale both the science
        cube and the PSF.
    verbose : bool, optional
        Sets level of printed outputs.
    """

    def __init__(self,
                 cube: np.ndarray,
                 angles: np.ndarray,
                 psf: np.ndarray,
                 dit_psf: Optional[float] = 1.0,
                 dit_science: Optional[float] = 1.0,
                 nd_transmission: Optional[float] = 1.0,
                 fwhm: Optional[float] = 4.0,
                 pixscale: Optional[float] = 1.0,
                 rescaling_factor: Optional[float] = 1.0,
                 verbose: Optional[bool] = False) -> None:

        # Science image setup
        try:
            self.cube = cube
        except BaseException:
            raise ValueError("You must provide a 3D cube of science data!")

        self.num_frames = self.cube.shape[0]
        self.width = self.cube.shape[2]
        self.height = self.cube.shape[1]

        # Parallactic angles
        try:
            self.angles = angles
        except BaseException:
            raise ValueError("You must provide an array of parallactic angles!")

        # Pixel scaling
        self.pixscale = pixscale
        self.rescaling_factor = rescaling_factor

        # PSF setup
        self.fwhm = int(fwhm/pixscale)
        try:
            # How do we want to deal with stacks of psfs? Median? Just take the first one?
            # Ideally if nPSFs = nImages, use each for each. Need to update!
            if len(psf.shape) > 2:
                psf = np.nanmedian(psf, axis=0)
            self.psf = psf * dit_science/dit_psf/nd_transmission
            self.dit_science = dit_science
            self.dit_psf = dit_psf

            mask = create_boolean_circular_mask(self.cube[0].shape,
                                                radius=self.fwhm)
            self.patch_area_pixels = self.cube[0][mask].ravel().shape[0]
        except BaseException:
            raise ValueError("You must provide an unsaturated PSF image!")

        self.patch_width = 2*int(self.fwhm) + 3
        self.verbose = verbose

        # These are what we're calculating
        self.snr = None
        self.flux = None
        self.std = None

        # Diagnostics
        if self.verbose:
            print("---------------------- ")
            print("Summary of PACO setup: \n")
            print(f"Image Cube shape = {self.cube.shape}")
            print(f"PIXSCALE = {self.pixscale:06}")
            print("PSF |  Area  |  Rad   |  Width | ")
            print(f"    |   {self.patch_area_pixels:01}   |"
                  + f"  {self.fwhm:02}    |  {self.psf.shape[0]:03}   | ")
            print(f"Patch width: {self.patch_width}")
            print("---------------------- \n")
            sys.stdout.flush()

    @abstractmethod
    def PACOCalc(self,
                 phi0s : np.ndarray,
                 use_subpixel_psf_astrometry : Optional[bool] = True,
                 cpu : Optional[int] = 1) -> None:
        """
        This function is algorithm dependant, and sets up the actual
        calculation process.

        Parameters
        ----------
        phi0s : numpy.ndarray
            Array of pixel coordinates to try to search for the planet signal. Typically a grid
            created using numpy.meshgrid.
        use_subpixel_psf_astrometry : bool
            If true, the PSF model for each patch is shifted to the correct
            location as predicted by the starting location and the parallactic
            angles, before being resampled for the patch. If false, the PSF
            model is simply located at the center of each patch. Significantly
            improves performance if set to False, but the SNR is reduced.
        cpu : int, optional
            Number of cpus to use for parallelization.

        Returns
        -------
        a : numpy.ndarray
            a_l from Equation 15 of Flasseur+ 2018
        b : numpy.ndarray
            b_l from Equation 16 of Flasseur+ 2018
        """

    def run(self,
            cpu : Optional[int] = 1,
            imlib : Optional[str] = 'vip-fft',
            interpolation : Optional[str]= 'lanczos4',
            keep_center : Optional[bool] = True,
            use_subpixel_psf_astrometry : Optional[bool] = True) -> Tuple[np.ndarray, np.ndarray]:
        """
        Run method of the PACO class. This function wraps up the PACO
        calculation steps, and returns the snr and flux estimates as
        outputs. The image library arguments are used if a rescaling of the
        data is desired: upsampling the images will result in a better SNR
        detection, but will be slower.

        Parameters
        ----------
        cpu : int, optional
            Number of processers to use
        imlib : str, optional
            See the documentation of the ``vip_hci.preproc.frame_px_resampling``
            function.
        interpolation : str, optional
            See the documentation of the ``vip_hci.preproc.frame_px_resampling``
            function.
        keep_center: bool, opt
            If input dimensions are even and the star centered (i.e. on
            dim//2, dim//2), whether to keep the star centered after scaling, i.e.
            on (new_dim//2, new_dim//2). For a non-centered input cube, better to
            leave it to False.
        use_subpixel_psf_astrometry : bool
            If true, the PSF model for each patch is shifted to the correct
            location as predicted by the starting location and the parallactic
            angles, before being resampled for the patch. If false, the PSF
            model is simply located at the center of each patch. Significantly
            improves performance if set to False, but the SNR is reduced.

        Returns
        -------
        snr : numpy.ndarray
            2D map of the signal-to-noise estimate as computed by PACO.
            This is b/sqrt(a), as in eqn 24 of Flasseur 2018.
        flux : numpy.ndarray
            2D map of the flux estimate as computed by PACO
            This is b/a, as in eqn 21 of Flasseur 2018.
        """

        if self.rescaling_factor != 1:
            self.rescale_cube_and_psf(imlib=imlib,
                                      interpolation=interpolation,
                                      keep_center=keep_center,
                                      verbose=self.verbose)
            if self.verbose:
                print("---------------------- ")
                print(f"Using {cpu} processor(s).")
                print(f"Rescaled Image Cube shape: {self.cube.shape}")
                print("Rescaled PSF:")
                print("PSF |  Area  |  Rad   |  Width | ")
                print(f"    |   {self.patch_area_pixels:01}   |"
                      + f"  {self.fwhm:02}    |  {self.psf.shape[0]:03}   | ")
                print("---------------------- \n")

        # Setup pixel coordinates
        x, y = np.meshgrid(np.arange(0, self.height),
                           np.arange(0, self.width))
        phi0s = np.column_stack((x.flatten(), y.flatten()))
        # Compute a,b
        a, b = self.PACOCalc(np.array(phi0s), cpu=cpu)

        # Reshape into a 2D image, with the same dimensions as the input images
        a = np.reshape(a, (self.height, self.width))
        b = np.reshape(b, (self.height, self.width))
        # Output arrays
        snr = b/np.sqrt(a)
        flux = b/a
        self.snr = snr
        self.flux = flux
        self.std = 1/np.sqrt(a)
        return snr, flux

    """
    Utility Functions
    """
    # Set the image stack to be processed

    def set_cube(self, cube: np.ndarray) -> None:
        """
        Provide a 3D image array to process. This updates
        the science cube, and the associated dimensions.

        Parameters
        ----------
        cube : numpy.ndarray
            3D science frames taken in pupil tracking/ADI mode.
            Dimensions should be (time, x, y), and units should be detector
            units (ie output of SPHERE or GPI reduction pipelines). The data
            should be centered, and have pre-processing already applied (e.g.
            bad pixel correction).
        """
        self.cube = np.array(cube)
        self.num_frames = self.cube.shape[0]
        self.width = self.cube.shape[2]
        self.height = self.cube.shape[1]

    # Set the template PSF
    def set_psf(self, psf: np.ndarray) -> None:
        """
        Read in the PSF template

        Parameters
        ----------
        psf: numpy.ndarray
            An unsaturated psf to use as the template.
        """
        self.psf = psf

    # Set parallactic angles
    def set_angles(self, angles: np.ndarray) -> None:
        """
        Set the rotation angle for each frame

        Parameters
        ----------
        angles: numpy.ndarray
            A list of the parallactic angles for each frame of the science data.
        """

        self.angles = angles

    def get_patch(self,
                  px: Tuple[int, int],
                  width: Optional[int] = None,
                  mask: Optional[np.ndarray] = None) -> np.ndarray:
        """
        Gets patch at given pixel px with size k for the current img sequenc

        Parameters
        ----------
        px : Tuple[int, int]
            Pixel coordinates for center of patch
        width : int
            width of a square patch to be masked

        Returns
        -------
        patch : numpy.ndarray
            A PACO "patch". This is a column through the time dimension of the
            unrotated frames, used to build the background statistics at the
            location of a given pixel.
        """
        if width is None:
            width = self.patch_width
        if mask is None:
            mask = create_boolean_circular_mask(self.cube[0].shape,
                                                radius=self.fwhm,
                                                center=px)
        k = int(width/2)
        if width % 2 != 0:
            k2 = k+1
        else:
            k2 = k
        nx, ny = np.shape(self.cube[0])[:2]
        if px[0]+k2 > nx or px[0]-k < 0 or px[1]+k2 > ny or px[1]-k < 0:
            return np.ones((self.num_frames, self.patch_area_pixels))*np.nan
        patch = self.cube[np.broadcast_to(mask, self.cube.shape)]\
            .reshape(self.num_frames, self.patch_area_pixels)
        return patch

    def set_scale(self, scale: float) -> None:
        """
        Set subpixel scaling factor

        Parameters
        ----------
        scale : float
            Scaling factor. Greater than one will result in an upsampled image,
            less than one will result in a downsampled image.
        """

        self.rescaling_factor = scale

    def rescale_cube_and_psf(self,
                             imlib: Optional[str] = 'vip-fft',
                             interpolation: Optional[str] = 'lanczos4',
                             keep_center: Optional[bool] = True) -> None:
        """
        Rescale each image in the stack by the class level scaling factor
        set during initialization or with set_scale. A scale factor of greater
        than one will upsample the image, a factor of less than one will downsample
        the image. This function wraps the VIP scaling function, and uses the same
        arguments to choose libraries and interpolation methods.

        Parameters
        ----------
        imlib : str, optional
            See the documentation of the ``vip_hci.preproc.frame_px_resampling``
            function.
        interpolation : str, optional
            See the documentation of the ``vip_hci.preproc.frame_px_resampling``
            function.
        keep_center: bool, opt
            If input dimensions are even and the star centered (i.e. on
            dim//2, dim//2), whether to keep the star centered after scaling, i.e.
            on (new_dim//2, new_dim//2). For a non-centered input cube, better to
            leave it to False.
        """

        if self.rescaling_factor == 1:
            if self.verbose:
                print("Scale is 1, no scaling applied.")
            return

        # Resample the science cube
        cube_px_resampling(self.cube,
                           self.rescaling_factor,
                           imlib=imlib,
                           interpolation=interpolation,
                           keep_center=keep_center,
                           verbose=False)

        self.pixscale = self.pixscale/self.rescaling_factor
        self.fwhm = int(self.fwhm*self.rescaling_factor)

        # Resample the PSF
        if self.psf is not None:
            self.psf = frame_px_resampling(self.psf,
                                           self.rescaling_factor,
                                           imlib=imlib,
                                           interpolation=interpolation,
                                           keep_center=keep_center,
                                           verbose=False)
        mask = create_boolean_circular_mask(self.psf.shape, self.fwhm)
        self.patch_area_pixels = self.psf[mask].shape[0]
        self.patch_width = 2*int(self.fwhm) + 3

    """
    Math Functions
    """

    def psf_model_function(self, mean: float, model: Callable,
                           params: dict) -> np.ndarray:
        """
        This function is deprecated in favour of directly supplying
        a PSF. In principle, an analytic model (ie a gaussian or moffat PSF)
        can be used in place of a measured unsaturated PSF.

        Parameters
        ----------
        mean : float
            If using the psfTemplateModel function, the mean
        model : dnc
            numpy statistical model (need to import numpy module for this)
        **kwargs: dict
            additional arguments for model

        Returns
        -------
        self.psf : numpy.ndarray
            Returns the PSF template used by PACO
        """

        if self.psf:
            return self.psf
        if model is None:
            print("Please input either a 2D PSF or a model function.")
            sys.exit(1)
        else:
            if model.__name__ == "psfTemplateModel":
                try:
                    self.psf = model(mean, params)
                    return self.psf
                except ValueError:
                    print("Fix template size")
            self.psf = model(mean, params)
            return self.psf

    def al(self,
           hfl: Union[list, np.ndarray],
           Cfl_inv: Union[list, np.ndarray],
           method: Optional[str] = "") -> np.ndarray:
        """
        a_l
        The sum of a_l is the inverse of the variance of the background at the given pixel.
        Einsum can get slow with large tensors, and may not actually be faster.
        If einsum is used, arguments must be numpy arrays, otherwise lists.

        Parameters
        ----------
        hfl : list
            This is a list of flattened psf templates
        Cfl_inv : list
            This is a list of inverse covariance matrices
        method: string
            Can be empty or "einsum". This determines the method
            used to do the matrix operations. "einsum" is slower for large arrays.

        Returns
        -------
        a : numpy.ndarray
            a_l from equation 15 of Flasseur 2018.
        """
        if method == "einsum":
            d1 = np.einsum('ijk,gj', Cfl_inv, hfl)
            return np.einsum('ml,ml', hfl, np.diagonal(d1).T)

        a = np.sum(np.array([np.dot(hfl[i], np.dot(Cfl_inv[i], hfl[i]).T)
                             for i in range(len(hfl))]), axis=0)
        return a

    def bl(self, hfl: Union[list, np.ndarray],
           Cfl_inv: Union[list, np.ndarray],
           r_fl: Union[list, np.ndarray],
           m_fl: Union[list, np.ndarray],
           method: Optional[str] = "") -> np.ndarray:
        """
        b_l
        The sum of b_l is the flux estimate at the given pixel.
        Einsum can get slow with large tensors, and may not actually be faster.
        If einsum is used, arguments must be numpy arrays, otherwise lists.

        Parameters
        ----------
        hfl : numpy.ndarray
            This is an array of flattened psf templates.
        Cfl_inv : numpy.ndarray
            This is an array of inverse covariance matrices.
        r_fl : numpy.ndarray
            This is an array of flux measurements following the predicted path.
        m_fl : numpy.ndarray
            This is an array of mean background statistics for each location in the path.
        method: string
            Can be empty or "einsum". This determines the method
            used to do the matrix operations. "einsum" is slower for large arrays.

        Returns
        -------
        b : numpy.ndarray
            b_l from equation 16 of Flasseur 2018.

        """
        if method == "einsum":
            d1 = np.einsum('ijk,gj', Cfl_inv, r_fl-m_fl)
            return np.einsum('ml,ml', hfl, np.diagonal(d1).T)

        b = np.sum(np.array([np.dot(np.dot(Cfl_inv[i], hfl[i]).T, (r_fl[i]-m_fl[i]))
                             for i in range(len(hfl))]), axis=0)
        return b

    """
    FluxPACO
    """

    def flux_estimate(self,
                      phi0s: np.ndarray,
                      eps: Optional[float] = 0.1,
                      initial_est: Optional[list] = [0.0]) -> list:
        """
        Unbiased estimate of the flux of a source located at p0
        The estimate of the flux is given by ahat * h, where h is the PSF template.
        This implements algorithm 3 from Flasseur+ 2018.
        TODO: Further testing to ensure that the extracted contrast is actually unbiased.
        Don't trust this estimate without checking!

        Parameters
        ----------
        phi0s : numpy.ndarray
            List of locations of sources to compute unbiased flux estimate in pixel units.
            Origin is at the bottom left. Should be a list of (x,y) tuples, or a 2D numpy
            array.
        eps : float
            Precision requirement for iteration (0,1)
        initial_est : float
            Initial estimate of the flux at p0 in contrast units

        Returns
        -------
        ests : list
            List of a-hat values for each detected source in the SNR map. This is the unbiased
            estimate of the flux at that location. Practically, this is similar to negative PSF
            injection. If the PSF is correctly normalized, this should be in contrast units.
        stds : list
            List of the estimated standard deviation on the flux estimates in contrast units.
        norm : float
            np.nanmax(psf) - Scaling factor for flux estimate and standard deviations.
        """
        print("Computing unbiased flux estimate...")

        if self.verbose:
            print("Initial guesses:")
            print("Positions: ", phi0s)
            print("Contrasts: ", initial_est)

        dim = self.width/2
        # Create arrays needed for storage
        # Store for each image pixel, for each temporal frame an image
        # for patches: for each time, we need to store a column of patches
        normalised_psf,norm,fwhm = normalize_psf(self.psf,
                                            fwhm='fit',
                                            size=None,
                                            threshold=None,
                                            mask_core=None,
                                            model='airy',
                                            imlib='vip-fft',
                                            interpolation='lanczos4',
                                            force_odd=False,
                                            full_output=True,
                                            verbose=self.verbose,
                                            debug=False)

        psf_mask = create_boolean_circular_mask(normalised_psf.shape, radius=self.fwhm)
        hoff = np.zeros((self.num_frames,self.num_frames, self.patch_area_pixels)) # The off axis PSF at each point
        # Create arrays needed for storage
        # Store for each image pixel, for each temporal frame an image
        # for patches: for each time, we need to store a column of patches
        # 2d selection of pixels around a given point
        x, y = np.meshgrid(np.arange(-dim, dim), np.arange(-dim, dim))

        ests = []
        stds = []
        for i, p0 in enumerate(phi0s):
            p0 = (p0[1],p0[0])
            angles_px = np.array(get_rotated_pixel_coords(x, y, p0, self.angles))
            hon = []
            for l, ang in enumerate(angles_px):
                # Get the column of patches at this point
                offax = frame_shift(normalised_psf,
                                    ang[1]-int(ang[1]),
                                    ang[0]-int(ang[0]),
                                    imlib='vip-fft',
                                    interpolation='lanczos4',
                                    border_mode='reflect')[psf_mask]
                hoff[l,l] = offax
                hon.append(offax)

            Cinv, m, patches = self.compute_statistics(np.array(angles_px).astype(int))
            # Get Angles
            # Ensure within image bounds
            # Extract relevant patches and statistics
            Cinlst = []
            mlst = []
            patch = []
            for l, ang in enumerate(angles_px):
                Cinlst.append(Cinv[int(ang[0]), int(ang[1])])
                mlst.append(m[int(ang[0]), int(ang[1])])
                patch.append(patches[int(ang[0]), int(ang[1]),l])
            a = self.al(hon,Cinlst)
            b = self.bl(hon, Cinlst,patch, mlst)
            print(b/a)
            # Fill patches and signal template

            # Unbiased flux estimation
            ahat = initial_est[i]
            aprev = 1e10 # Arbitrary large value so that the loop will run
            while np.abs(ahat - aprev) > np.abs(ahat * eps):
                a = 0.0
                b = 0.0
                # the mean of a temporal column of patches at each pixel
                m = np.zeros((self.num_frames, self.patch_area_pixels))
                # the inverse covariance matrix at each point
                Cinv = np.zeros((self.num_frames, self.patch_area_pixels, self.patch_area_pixels))

                # Patches here are columns in time
                for l,ang in enumerate(angles_px):
                    apatch = self.get_patch(ang.astype(int))
                    m[l], Cinv[l] = self.iterate_flux_calc(ahat, apatch, hoff[l])
                # Patches here are where the planet is expected to be
                a = self.al(hon, Cinv)
                b = self.bl(hon, Cinv, patch, m)
                aprev = ahat
                ahat = b/a
                if self.verbose:
                    print(f"Flux estimate: {ahat/norm}")
            ests.append(np.abs(ahat/norm))
            stds.append(1/np.sqrt(a)/norm)
        print("Extracted contrasts")
        print("-------------------")
        for i in range(len(phi0s)):
            print(
                f"x: {phi0s[i][0]}, y: {phi0s[i][1]}, flux: {ests[i]}±{stds[i]}")
        return ests, stds, norm

    def iterate_flux_calc(self, est: float, patch: np.ndarray,
                          model: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Compute the iterative estimates for the mean and inverse covariance.

        Parameters
        ----------
        est : float
            Current estimate for the magnitude of the flux
        patch : numpy.ndarray
            Column of patches about p0
        model : numpy.ndarray
            Template for PSF

        Returns
        -------
        m : numpy.ndarray
            Mean of background patches.
        Cinv : numpy.ndarray
            List of inverse covariance matrices between the patches.
        """

        if patch is None:
            return None, None

        unbiased = np.array([apatch - est*model[l] for l,apatch in enumerate(patch)])
        m,Cinv = compute_statistics_at_pixel(unbiased)
        return m, Cinv

    def subpixel_threshold_detect(self, snr_map: np.ndarray,
                                  threshold: float,
                                  mode: Optional[str] = 'lpeaks',
                                  bkg_sigma: Optional[float] = 5.0,
                                  matched_filter: Optional[bool] = False,
                                  mask: Optional[bool] = True,
                                  full_output: Optional[bool] = False,
                                  cpu: Optional[int] = 1) -> np.ndarray:
        """ Wraps VIP.metrics.detection.detection, see that function for further documentation.
        Note that the output convention here is different - this function returns xx,yy.

        Finds blobs in a 2d array. The algorithm is designed for automatically
        finding planets in post-processed high contrast final frames. Blob can be
        defined as a region of an image in which some properties are constant or
        vary within a prescribed range of values. See ``Notes`` below to read about
        the algorithm details.
        Parameters
        ----------
        snr_map : numpy ndarray, 2d
            Input frame.
        threshold : float
            S/N threshold for deciding whether the blob is a detection or not. Used
            to threshold the S/N map when ``mode`` is set to 'snrmap' or 'snrmapf'.
        mode : {'lpeaks', 'log', 'dog', 'snrmap', 'snrmapf'}, optional
            Sets with algorithm to use. Each algorithm yields different results. See
            notes for the details of each method.
        bkg_sigma : int or float, optional
            The number standard deviations above the clipped median for setting the
            background level. Used when ``mode`` is either 'lpeaks', 'dog' or 'log'.
        matched_filter : bool, optional
            Whether to correlate with the psf of not. Used when ``mode`` is either
            'lpeaks', 'dog' or 'log'.
        mask : bool, optional
            If True the central region (circular aperture of 2*FWHM radius) of the
            image will be masked out.
        full_output : bool, optional
            Whether to output just the coordinates of blobs that fulfill the SNR
            constraint or a table with all the blobs and the peak pixels and SNR.
        cpu : None or int, optional
            The number of processes for running the ``snrmap`` function.
        verbose : bool, optional
            Whether to print to stdout information about found blobs.
        Returns
        -------
        peaks : np.ndarray
            xx,yy values of the centers of local maxima above the provided threshold
        """
        peaks = detection(snr_map,
                          fwhm=self.fwhm,
                          psf=self.psf/np.nanmax(self.psf),
                          mode=mode,
                          bkg_sigma=bkg_sigma,
                          matched_filter=matched_filter,
                          mask=mask,
                          snr_thresh=threshold,
                          nproc=cpu,
                          plot=False,
                          debug=False,
                          full_output=full_output,
                          verbose=self.verbose)
        return peaks.T

    def pixel_threshold_detection(
            self, snr_map: np.ndarray, threshold: float) -> np.ndarray:
        """
        Returns a list of the pixel coordinates of center of signals above a given threshold

        Parameters
        ----------
        snr_map : numpy.ndarray
            SNR map, b/sqrt(a) as computed by run()
        threshold: float
            Threshold for detection in sigma

        Returns
        -------
        locs : numpy.array
            Array of (x,y) pixel location estimates for the location of point sources
            above the provided threshold.
        """

        data_max = filters.maximum_filter(snr_map, size=self.fwhm)
        maxima = (snr_map == data_max)
        diff = (data_max > threshold)
        maxima[diff == 0] = 0

        labeled, _ = ndimage.label(maxima)
        slices = ndimage.find_objects(labeled)
        x, y = [], []
        for dy, dx in slices:
            x_center = (dx.start + dx.stop - 1)/2
            x.append(x_center)
            y_center = (dy.start + dy.stop - 1)/2
            y.append(y_center)
        return np.array(list(zip(x, y)))

    def compute_statistics(self, phi0s : np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        This function computes the mean and inverse covariance matrix for
        each patch in the image stack in Serial. Used by FastPACO and flux
        estimation.

        Parameters
        ----------
        phi0s : numpy.ndarray
            Array of pixel locations to estimate companion position

        Returns
        -------
        Cinv : numpy.ndarray
            Inverse covariance matrix between the the mean of each of the patches.
            The patches are a column through the time axis of the unrotated
            science images. Together with the mean this provides an empirical
            estimate of the background statistics. An inv covariance matrix is provided
            for each pixel location in ph0s.
        m : numpy.ndarray
            Mean of each of the background patches along the time axis, for each pixel location
            in phi0s.
        patch : numpy.ndarray
            The background column for each test pixel location in phi0s.
        """
        if self.verbose:
            print("Precomputing Statistics...")

        # Store for each image pixel, for each temporal frame an image
        # for patches: for each time, we need to store a column of patches
        patch = np.zeros((self.width, self.height, self.num_frames, self.patch_area_pixels))

        # the mean of a temporal column of patches centered at each pixel
        m = np.zeros((self.height, self.width, self.patch_area_pixels))
        # the inverse covariance matrix at each point
        Cinv = np.zeros((self.height, self.width, self.patch_area_pixels, self.patch_area_pixels))

        # *** SERIAL ***
        # Loop over all pixels
        # i is the same as theta_k in the PACO paper
        for p0 in phi0s:
            apatch = self.get_patch(p0)
            # For some black magic reason this needs to be inverted here.
            m[p0[1]][p0[0]], Cinv[p0[1]][p0[0]] = compute_statistics_at_pixel(apatch)
            patch[p0[1]][p0[0]] = apatch
        return Cinv, m, patch

"""
**************************************************
*                                                *
*                  Fast PACO                     *
*                                                *
**************************************************
"""

class FastPACO(PACO):
    """
    This class implements Algorithm 2 from Flasseur+ 2018.
    """

    def PACOCalc(self,
                 phi0s : np.ndarray,
                 use_subpixel_psf_astrometry : Optional[bool] = True,
                 cpu : Optional[int] = 1) -> None:
        """
        PACOCalc

        This function iterates of a list of test points (phi0) and a list
        of angles between frames to produce 'a' and b', which can be used to
        generate a signal to noise map where SNR = b/sqrt(a) at each pixel.

        Parameters
        ----------
        phi0s : numpy.ndarray
            Array of (x,y) pixel locations to estimate companion position
        use_subpixel_psf_astrometry : bool
            If true, the PSF model for each patch is shifted to the correct
            location as predicted by the starting location and the parallactic
            angles, before being resampled for the patch. If false, the PSF
            model is simply located at the center of each patch. Significantly
            improves performance if set to False, but the SNR is reduced.
        cpu : int
            Number of cores to use for parallel processing

        Returns
        -------
        a : numpy.ndarray
            a_l from Equation 15 of Flasseur+ 2018
        b : numpy.ndarray
            b_l from Equation 16 of Flasseur+ 2018
        """
        npx = len(phi0s)  # Number of pixels in an image
        dim = self.width/2

        a = np.zeros(npx)  # Setup output arrays
        b = np.zeros(npx)
        phi0s = np.array([phi0s[:,1],phi0s[:,0]]).T

        if cpu == 1:
            Cinv, m, patches = self.compute_statistics(phi0s)
        else:
            Cinv, m, patches = self.compute_statistics_parallel(phi0s, cpu=cpu)
        normalised_psf = normalize_psf(self.psf,
                                       fwhm='fit',
                                       size=None,
                                       threshold=None,
                                       mask_core=None,
                                       model='airy',
                                       imlib='vip-fft',
                                       interpolation='lanczos4',
                                       force_odd=False,
                                       full_output=False,
                                       verbose=self.verbose,
                                       debug=False)
        psf_mask = create_boolean_circular_mask(normalised_psf.shape, radius=self.fwhm)

        # Create arrays needed for storage
        # Store for each image pixel, for each temporal frame an image
        # for patches: for each time, we need to store a column of patches

        # Currently forcing integer grid, but meshgrid takes floats as
        # arguments...
        x, y = np.meshgrid(np.arange(-dim, dim), np.arange(-dim, dim))
        if self.verbose:
            print("Running Fast PACO...")

        # Loop over all pixels
        # i is the same as theta_k in the PACO paper
        for i, p0 in enumerate(phi0s):
            # Get Angles
            angles_px = get_rotated_pixel_coords(x, y, p0, self.angles)
            # Ensure within image bounds
            if(int(np.max(angles_px.flatten())) >= self.width or
               int(np.min(angles_px.flatten())) < 0):
                a[i] = np.nan
                b[i] = np.nan
                continue

            # Extract relevant patches and statistics
            Cinlst = []
            mlst = []
            hlst = []
            patch = []
            for l, ang in enumerate(angles_px):
                Cinlst.append(Cinv[int(ang[0]), int(ang[1])])
                mlst.append(m[int(ang[0]), int(ang[1])])
                if use_subpixel_psf_astrometry:
                    offax = frame_shift(normalised_psf,
                                        ang[1]-int(ang[1]),
                                        ang[0]-int(ang[0]),
                                        imlib='vip-fft',
                                        interpolation='lanczos4',
                                        border_mode='reflect')[psf_mask]
                else:
                    offax = normalised_psf[psf_mask]
                hlst.append(offax)
                patch.append(patches[int(ang[0]), int(ang[1]), l])

            # Calculate a and b, matrices
            a[i] = self.al(hlst, Cinlst)
            b[i] = self.bl(hlst, Cinlst, patch, mlst)
        if self.verbose:
            print("Done")
        return a, b

    def compute_statistics_parallel(self,
                                    phi0s: np.ndarray,
                                    cpu: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        This function computes the mean and inverse covariance matrix for
        each patch in the image stack in parallel.

        Parameters
        ----------
        phi0s : int numpy.ndarray
            Array of pixel locations to estimate companion position
        cpu : int
            Number of processors to use

        Returns
        -------
        Cinv : numpy.ndarray
            Inverse covariance matrix between the the mean of each of the patches.
            The patches are a column through the time axis of the unrotated
            science images. Together with the mean this provides an empirical
            estimate of the background statistics. An inv covariance matrix is provided
            for each pixel location in ph0s.
        m : numpy.ndarray
            Mean of each of the background patches along the time axis, for each pixel location
            in phi0s.
        patch : numpy.ndarray
            The background column for each test pixel location in phi0s.

        """

        if self.verbose:
            print("Precomputing Statistics using %d Processes..."%cpu)
        # the mean of a temporal column of patches at each pixel
        m = np.zeros((self.height*self.width*self.patch_area_pixels))
        # the inverse covariance matrix at each point
        Cinv = np.zeros((self.height*
                         self.width*
                         self.patch_area_pixels*
                         self.patch_area_pixels))

        # *** Parallel Processing ***
        p_data = pool_map(cpu, self.get_patch, iterable(phi0s))
        patches = [p for p in p_data]
        data = pool_map(cpu, compute_statistics_at_pixel, iterable(patches))
        # p_pool = Pool(processes=cpu)
        #p_data = p_pool.map(self.get_patch, phi0s, chunksize=int(npx/cpu))
        # p_pool.close()
        # p_pool.join()
        # patches = [p for p in p_data]
        # p = Pool(processes=cpu)
        # data = p.map(compute_statistics_at_pixel, patches, chunksize=int(npx/cpu))
        # p.close()
        # p.join()
        ms, cs = [], []
        for d in data:
            if d[0] is None or d[1] is None:
                ms.append(np.full(self.patch_area_pixels, np.nan))
                cs.append(np.full((self.patch_area_pixels,
                                   self.patch_area_pixels), np.nan))
            else:
                ms.append(d[0])
                cs.append(d[1])
        ms = np.array(ms)
        cs = np.array(cs)
        patches = np.array(patches)

        # Reshape outputs
        patches = patches.reshape((self.width,
                                   self.height,
                                   self.num_frames,
                                   self.patch_area_pixels))
        m = ms.reshape((self.height,
                        self.width,
                        self.patch_area_pixels))
        Cinv = cs.reshape((self.height,
                           self.width,
                           self.patch_area_pixels,
                           self.patch_area_pixels))
        patches = np.swapaxes(patches,0,1)
        m = np.swapaxes(m,0,1)
        Cinv = np.swapaxes(Cinv,0,1)

        return Cinv, m, patches

"""
**************************************************
*                                                *
*                  Full PACO                     *
*                                                *
**************************************************
"""

class FullPACO(PACO):
    """
    Implementation of Algorithm 1 from Flasseur+ 2018
    """

    def PACOCalc(self,
                 phi0s : np.ndarray,
                 use_subpixel_psf_astrometry : Optional[bool] = True,
                 cpu : Optional[int] = 1) -> None:
        """
        PACOCalc

        This function iterates of a list of test points (phi0) and a list
        of angles between frames to produce 'a' and b', which can be used to
        generate a signal to noise map where SNR = b/sqrt(a) at each pixel.

        Parameters
        ----------
        phi0s : numpy.ndarray
            Array of (x,y) pixel locations to estimate companion position
        use_subpixel_psf_astrometry : bool
            If true, the PSF model for each patch is shifted to the correct
            location as predicted by the starting location and the parallactic
            angles, before being resampled for the patch. If false, the PSF
            model is simply located at the center of each patch. Significantly
            improves performance if set to False, but the SNR is reduced.
        cpu : int
            Number of cores to use for parallel processing. TODO: Not yet implemented.

        Returns
        -------
        a : numpy.ndarray
            a_l from Equation 15 of Flasseur+ 2018
        b : numpy.ndarray
            b_l from Equation 16 of Flasseur+ 2018
        """

        npx = len(phi0s)  # Number of pixels in an image
        dim = self.width/2

        a = np.zeros(npx)  # Setup output arrays
        b = np.zeros(npx)
        normalised_psf = normalize_psf(self.psf,
                                       fwhm='fit',
                                       size=None,
                                       threshold=None,
                                       mask_core=None,
                                       model='airy',
                                       imlib='vip-fft',
                                       interpolation='lanczos4',
                                       force_odd=False,
                                       full_output=False,
                                       verbose=self.verbose,
                                       debug=False)
        psf_mask = create_boolean_circular_mask(normalised_psf.shape, radius=self.fwhm)

        if self.verbose:
            print("Running Full PACO...")

        # Set up coordinates so 0 is at the center of the image
        x, y = np.meshgrid(np.arange(-dim, dim), np.arange(-dim, dim))

        if cpu > 1:
            print("Multiprocessing for full PACO is not yet implemented!")

        # Store intermediate results
        patch = np.zeros((self.width, self.height, self.num_frames, self.patch_area_pixels))
        # the mean of a temporal column of patches centered at each pixel
        m = np.zeros((self.height, self.width, self.patch_area_pixels))
        # the inverse covariance matrix at each point
        Cinv = np.zeros((self.height, self.width, self.patch_area_pixels, self.patch_area_pixels))

        # Loop over all pixels
        # i is the same as theta_k in the PACO paper
        for i, p0 in enumerate(phi0s):
            # Get list of pixels for each rotation angle
            angles_px = get_rotated_pixel_coords(x, y, (p0[1],p0[0]), self.angles)

            # Ensure within image bounds
            if(int(np.max(angles_px.flatten())) >= self.width or
               int(np.min(angles_px.flatten())) < 0):
                a[i] = np.nan
                b[i] = np.nan
                continue

            # Iterate over each temporal frame/each angle
            # Same as iterating over phi_l
            current_patch = []
            mlst = []
            h = []
            clst = []
            for l, ang in enumerate(angles_px):
                # Get the column of patches at this point
                if np.max(patch[int(ang[0]),int(ang[1])]) == 0:
                    apatch = self.get_patch((int(ang[1]),int(ang[0])))
                    patch[int(ang[0]),int(ang[1])] = apatch
                    m[int(ang[0]),int(ang[1])], Cinv[int(ang[0]),int(ang[1])] = compute_statistics_at_pixel(apatch)
                else:
                    apatch = patch[int(ang[0]),int(ang[1])]
                if apatch is None:
                    continue
                mlst.append(m[int(ang[0]),int(ang[1])])
                clst.append(Cinv[int(ang[0]),int(ang[1])])

                current_patch.append(apatch)
                if use_subpixel_psf_astrometry:
                    offax = frame_shift(normalised_psf,
                                        ang[1]-int(ang[1]),
                                        ang[0]-int(ang[0]),
                                        imlib='vip-fft',
                                        interpolation='lanczos4',
                                        border_mode='reflect')[psf_mask]
                else:
                    offax = normalised_psf[psf_mask]

                h.append(offax)
            current_patch = np.array(current_patch)
            patches = np.array([current_patch[l,l] for l in range(len(angles_px))])
            h = np.array(h)
            mlst = np.array(mlst)
            clst = np.array(clst)
            # Calculate a and b, matrices
            a[i] = self.al(h, clst)
            b[i] = self.bl(h, clst, patches, mlst)
        if self.verbose:
            print("Done")
        return a, b

"""
Math functions for computing patch covariance
"""

def compute_statistics_at_pixel(
        patch: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """
    Calculate the mean and inverse covariance within a patch
    Reimplemented in PACO class, can probably be deleted

    Parameters
    ----------
    patch : numpy.ndarray
        Array of circular (flattened) patches centered on the same physical
        pixel vertically throughout the image stack
    """

    if patch is None:
        return None, None
    T = patch.shape[0]
    #size = patch.shape[1]

    # Calculate the mean of the column
    m = np.mean(patch, axis=0)
    # Calculate the covariance matrix
    S = sample_covariance(patch, m, T)
    rho = shrinkage_factor(S, T)
    F = diagsample_covariance(S)
    C = covariance(rho, S, F)
    Cinv = np.linalg.inv(C)
    return m, Cinv

def covariance(rho: np.ndarray, S: np.ndarray, F: np.ndarray) -> np.ndarray:
    """
    Ĉ: Shrinkage covariance matrix

    Parameters
    ----------
    rho : float
        Shrinkage factor weight
    S : numpy.ndarray
        Sample covariance matrix
    F : numpy.ndarray
        Diagonal of sample covariance matrix

    Returns
    -------
    m : numpy.ndarray
        Mean of each of the background patches along the time axis.
    Cinv : numpy.ndarray
        Inverse covariance matrix between the the mean of each of the patches.
        The patches are a column through the time axis of the unrotated
        science images. Together with the mean this provides an empirical
        estimate of the background statistics.
    """

    C = (1.0-rho)*S + rho*F
    return C

def sample_covariance(r: np.ndarray, m: np.ndarray,
                      T: np.ndarray) -> np.ndarray:
    """
    Ŝ: Sample covariance matrix

    Parameters
    ----------
    r : numpy.ndarray
        Observed intensity at position θk and time tl
    m : numpy.ndarray
        Mean of all background patches at position θk
    T : int
        Number of temporal frames

    Returns
    -------
    S : numpy.ndarray
        Sample covariance
    """

    #S = (1.0/T)*np.sum([np.outer((p-m).ravel(),(p-m).ravel().T) for p in r], axis=0)
    S = (1.0/T)*np.sum([np.cov(np.stack((p, m)),
                               rowvar=False, bias=False) for p in r], axis=0)
    return S

def diagsample_covariance(S: np.ndarray) -> np.ndarray:
    """
    F: Diagonal elements of the sample covariance matrix

    Parameters
    ----------
    S : arr
        Sample covariance matrix

    Returns
    -------
    F : numpy.ndarray
        Diagonal elements of the sample covariance matrix
    """

    return np.diag(np.diag(S))

def shrinkage_factor(S: np.ndarray, T: np.ndarray) -> float:
    """
    ρ: Shrinkage factor to regularize covariant matrix

    Parameters
    ----------
    S : numpy.ndarray
        Sample covariance matrix
    T : int
        Number of temporal frames

    Returns
    -------
    ρ : float
        Shrinkage factor to regularize covariant matrix
    """

    top = (np.trace(np.dot(S, S)) + np.trace(S)**2 -
           2.0*np.sum(S**2.0))
    bot = ((T+1.0)*(np.trace(np.dot(S, S)) -
                    np.sum(np.diag(S)**2.0)))
    p = top/bot
    return max(min(p, 1.0), 0.0)

def get_rotated_pixel_coords(x: np.ndarray,
                             y: np.ndarray,
                             p0: Tuple[int, int],
                             angles: np.ndarray,
                             astro_convention: Optional[bool] = False) -> np.ndarray:
    """
    For a given pixel, find the new pixel location after a rotation for each angle in angles

    Parameters
    ----------
    x : numpy.ndarrayr
        Grid of x components of pixel coordinates
    y : numpy.ndarrayr
        Grid of y components of pixel coordinates
    p0 : (int,int)
        Initial pixel location
    angles : numpy.ndarrayr
        List of angles for which to compute the new pixel location
    Returns
    -------
    nx : numpy.ndarray
        New array of x pixels coordinates following the rotation
    ny : numpy.ndarray
        New array of y pixels coordinates following the rotation

    """
    # Current pixel
    phi0 = np.array([x[int(p0[0]), int(p0[1])], y[int(p0[0]), int(p0[1])]])

    # Convert to polar coordinates
    rad, theta = cart_to_pol(
        phi0[0], phi0[1], astro_convention=astro_convention)

    # Rotate by parallactic angles
    angles_rad = -1*angles + theta

    # Rotate the polar coordinates by each frame angle
    angles_pol = np.array([rad*np.ones_like(angles_rad), angles_rad])

    # Find the new pixel coordinates after rotation
    nx, ny = pol_to_cart(
        angles_pol[0], angles_pol[1], astro_convention=astro_convention)

    # Shift to center coordinates (central pixel is 0)
    # TODO - use vip cx, cy arguments rather than shifting after?
    nx += +int(x.shape[0]/2)
    ny += +int(x.shape[0]/2)
    return np.array([nx, ny]).T

def create_boolean_circular_mask(shape: np.ndarray,
                                 radius: Optional[int] = 4,
                                 center: Optional[Tuple[int, int]] = None) -> np.ndarray:
    """
    Returns a 2D boolean mask given some radius and location

    Parameters
    ----------
    shape : numpy.ndarray
        Shape of a 2D numpy array
    radius : int
        Radius of the mask in pixels
    center : (int,int)
        Pixel coordinates denoting the center of the mask,
        None defaults to center of shape

    Returns
    -------
    mask : numpy.ndarray
        A boolean mask of the the same shape as the science
        input data (provided by shape argument). The mask is 0
        outside of a circular region located at center, with a specified radius.
    """

    w = shape[0]
    h = shape[1]
    if center is None:
        center = [int(w/2), int(h/2)]
    if radius is None:
        radius = min(center[0], center[1], w-center[0], h-center[1])
    X, Y = np.ogrid[:w, :h]
    dist2 = (X - center[0])**2 + (Y-center[1])**2
    mask = dist2 <= radius**2
    return mask


import pkg_resources
import os
import pdb
from astropy.io import fits
import astropy.constants as const
import astropy.units as u
import numpy as np
import matplotlib.pyplot as plt
import sys
import tayph.util as ut
from tayph.vartests import typetest,dimtest
import tayph.tellurics as mol
import tayph.system_parameters as sp
import tayph.functions as fun
import tayph.operations as ops
from tayph.ccf import xcor
import copy
import scipy.interpolate as interp
import pickle
from pathlib import Path
import warnings
import glob
from scipy import interpolate
import tayph.masking as masking
import subprocess
import textwrap
from astropy.utils.data import download_file
from .phoenix import get_phoenix_wavelengths, get_phoenix_model_spectrum
from scipy.ndimage import uniform_filter1d

__all__ = [
    "read_harpslike",
    "read_espresso",
    "read_uves",
    "read_spirou",
    "read_gianob"
]

def read_harpslike(inpath,filelist,mode,read_s1d=True):
    """
    This reads a folder of HARPS or HARPSN data. Input is a list of filepaths and the mode (HARPS
    or HARPSN).
    """

    if mode=='HARPS':
        catkeyword = 'HIERARCH ESO DPR CATG'
        bervkeyword = 'HIERARCH ESO DRS BERV'
        thfilekeyword = 'HIERARCH ESO DRS CAL TH FILE'
        Zstartkeyword = 'HIERARCH ESO TEL AIRM START'
        Zendkeyword = 'HIERARCH ESO TEL AIRM END'
    elif mode=='HARPSN':
        catkeyword = 'OBS-TYPE'
        bervkeyword = 'HIERARCH TNG DRS BERV'
        thfilekeyword = 'HIERARCH TNG DRS CAL TH FILE'
        Zstartkeyword = 'AIRMASS'
        Zendkeyword = 'AIRMASS'#These are the same because HARPSN doesnt have start and end keywords.
        #Down there, the airmass is averaged, so there is no problem in taking the average of the same number.
    else:
        raise ValueError(f"Error in read_harpslike: mode should be set to HARPS or HARPSN ({mode})")

    #The following variables define lists in which all the necessary data will be stored.
    framename=[]
    header=[]
    s1dhdr=[]
    obstype=[]
    texp=np.array([])
    date=[]
    mjd=np.array([])
    s1dmjd=np.array([])
    npx=np.array([])
    norders=np.array([])
    e2ds=[]
    s1d=[]
    wave1d=[]
    airmass=np.array([])
    berv=np.array([])
    wave=[]
    # wavefile_used = []
    for i in range(len(filelist)):
        if filelist[i].endswith('e2ds_A.fits'):
            print(f'------{filelist[i]}', end="\r")
            hdul = fits.open(inpath/filelist[i])
            data = copy.deepcopy(hdul[0].data)
            hdr = hdul[0].header
            hdul.close()
            del hdul[0].data
            if hdr[catkeyword] == 'SCIENCE':
                framename.append(filelist[i])
                header.append(hdr)
                obstype.append(hdr[catkeyword])
                texp=np.append(texp,hdr['EXPTIME'])
                date.append(hdr['DATE-OBS'])
                mjd=np.append(mjd,hdr['MJD-OBS'])
                npx=np.append(npx,hdr['NAXIS1'])
                norders=np.append(norders,hdr['NAXIS2'])
                e2ds.append(data)
                berv=np.append(berv,hdr[bervkeyword])
                airmass=np.append(airmass,0.5*(hdr[Zstartkeyword]+hdr[Zendkeyword]))#This is an approximation where we take the mean airmass.
                # if nowave == True:
                # wavefile_used.append(hdr[thfilekeyword])
                #Record which wavefile was used by the pipeline to
                #create the wavelength solution.
                wavedata=ut.read_wave_from_e2ds_header(hdr,mode=mode)/10.0#convert to nm.
                wave.append(wavedata)
                # if filelist[i].endswith('wave_A.fits'):
                #     print(filelist[i]+' (wave)')
                #     if nowave == True:
                #         warnings.warn(" in read_e2ds: nowave was set to True but a wave_A file was detected. This wave file is now ignored in favor of the header.",RuntimeWarning)
                #     else:
                #         wavedata=fits.getdata(inpath/filelist[i])
                #         wave.append(wavedata)

                if read_s1d:
                    s1d_path=inpath/Path(str(filelist[i]).replace('e2ds_A.fits','s1d_A.fits'))
                    ut.check_path(s1d_path,exists=True)#Crash if the S1D doesn't exist.
        # if filelist[i].endswith('s1d_A.fits'):
                    hdul = fits.open(s1d_path)
                    data_1d = copy.deepcopy(hdul[0].data)
                    hdr1d = hdul[0].header
                    hdul.close()
                    del hdul
            # if hdr[catkeyword] == 'SCIENCE':
                    s1d.append(data_1d)
                    if mode == 'HARPSN':#In the case of HARPS-N we need to convert the units of the
                        #elevation and provide a UTC keyword.
                        hdr1d['TELALT'] = np.degrees(float(hdr1d['EL']))
                        hdr1d['UTC'] = (float(hdr1d['MJD-OBS'])%1.0)*86400.0
                    s1dhdr.append(hdr1d)
                    s1dmjd=np.append(s1dmjd,hdr1d['MJD-OBS'])
                    berv1d = hdr1d[bervkeyword]
                    if berv1d != hdr[bervkeyword]:
                        wrn_msg = ('WARNING in read_harpslike(): BERV correction of s1d file is not'
                        f'equal to that of the e2ds file. {berv1d} vs {hdr[bervkeyword]}')
                        ut.tprint(wrn_msg)
                    gamma = (1.0-(berv1d*u.km/u.s/const.c).decompose().value)#Doppler factor BERV.
                    wave1d.append((hdr1d['CDELT1']*np.arange(len(data_1d), dtype=float)+hdr1d['CRVAL1'])*gamma)

    #Check that all exposures have the same number of pixels, and clip s1ds if needed.
    # min_npx1d = int(np.min(np.array(npx1d)))
    # if np.sum(np.abs(np.array(npx1d)-npx1d[0])) != 0:
    #     warnings.warn("in read_e2ds when reading HARPS data: Not all s1d files have the same number of pixels. This could have happened if the pipeline has extracted one or two extra pixels in some exposures but not others. The s1d files will be clipped to the smallest length.",RuntimeWarning)
    #     for i in range(len(s1d)):
    #         wave1d[i]=wave1d[i][0:min_npx1d]
    #         s1d[i]=s1d[i][0:min_npx1d]
    #         npx1d[i]=min_npx1d
    output = {'wave':wave,'e2ds':e2ds,'header':header,'wave1d':wave1d,'s1d':s1d,'s1dhdr':s1dhdr,
    'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
    'norders':norders,'berv':berv,'airmass':airmass,'s1dmjd':s1dmjd}
    return(output)

def read_carmenes(inpath,filelist,channel,construct_s1d=True):
    """
    This reads a folder of CARMENES visible (VIS) or infra-red (NIR) channel data. Input is a list
    of filepaths and the mode ('VIS' or 'NIR').
    """

    catkeyword = 'HIERARCH CAHA INS ICS IMAGETYP'
    bervkeyword = 'HIERARCH CARACAL BERV'
    thfilekeyword = 'HIERARCH CARACAL WAVE FILE'
    Zstartkeyword = 'AIRMASS'
    Zendkeyword = 'AIRMASS'#These are the same because CARMENES doesnt have start and end keywords.

    if channel not in ['vis','nir']:
        raise ValueError(f"Error in read_carmenes: channel should be set to VIS or NIR ({channel})")

    #The following variables define lists in which all the necessary data will be stored.
    framename=[]
    header=[]
    s1dhdr=[]
    obstype=[]
    texp=np.array([])
    date=[]
    mjd=np.array([])
    s1dmjd=np.array([])
    npx=np.array([])
    norders=np.array([])
    e2ds=[]
    blaze=[]
    s1d=[]
    wave1d=[]
    airmass=np.array([])
    berv=np.array([])
    wave=[]
    # wavefile_used = []
    for i in range(len(filelist)):
        if filelist[i].endswith('-'+channel.lower()+'_A.fits'):
            print(f'------{filelist[i]}', end="\r")
            hdul = fits.open(inpath/filelist[i])
            data = copy.deepcopy(hdul[1].data)
            cont = copy.deepcopy(hdul[2].data)
            sigma = copy.deepcopy(hdul[3].data)
            wavedata = copy.deepcopy(hdul[4].data)/10.0
            hdr = hdul[0].header
            spechdr = hdul[1].header
            hdul.close()
            del hdul[1].data
            del hdul[2].data
            del hdul[3].data
            del hdul[4].data
            if hdr[catkeyword] == 'SCIENCE':
                framename.append(filelist[i])
                header.append(hdr)
                obstype.append(hdr[catkeyword])
                texp=np.append(texp,hdr['EXPTIME'])
                date.append(hdr['DATE-OBS'])
                mjd=np.append(mjd,hdr['MJD-OBS'])
                npx=np.append(npx,spechdr['NAXIS1'])
                norders=np.append(norders,spechdr['NAXIS2'])
                e2ds.append(data)
                blaze.append(data/sigma**2)
                berv=np.append(berv,hdr[bervkeyword])
                airmass=np.append(airmass,0.5*(hdr[Zstartkeyword]+hdr[Zendkeyword]))#This is an approximation where we take the mean airmass.
                wave.append(ops.vactoair(wavedata))

                if construct_s1d:
                    wave_1d, data_1d = spec_stich_n_norm(data,wavedata,cont,sigma)

                    s1d.append(data_1d)

                    hdr1d = copy.deepcopy(hdr)
                    hdr1d['UTC'] = (float(hdr1d['MJD-OBS'])%1.0)*86400.0
                    s1dhdr.append(hdr1d)
                    s1dmjd=np.append(s1dmjd,hdr1d['MJD-OBS'])
                    berv1d = hdr1d[bervkeyword]
                    gamma = (1.0-(berv1d*u.km/u.s/const.c).decompose().value)#Doppler factor BERV.
                    gamma = 1 #turning berv off
                    wave1d.append(ops.vactoair(wave_1d)*gamma*10)

    BLAZE_Model = blaze_model(np.nanmean(blaze, axis = 0)) #Calucalting blaze model
    e2ds = list(e2ds*BLAZE_Model[np.newaxis,:]) #Deblazing the data

    if construct_s1d:
        output = {'wave':wave,'e2ds':e2ds,'header':header,'wave1d':wave1d,'s1d':s1d,'s1dhdr':s1dhdr,
        'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
        'norders':norders,'berv':berv,'airmass':airmass,'s1dmjd':s1dmjd}
    else:
        output = {'wave':wave,'e2ds':e2ds,'header':header,
        'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
        'norders':norders,'berv':berv,'airmass':airmass}
    return(output)

def spec_stich_n_norm(spec, wave, cont, sig):
    """This stitches and continuum normalises CARMENES E2DS spectra into 1D spectra for use in
    molecfit.
    N. Borsato - 24-02-2021"""

    import numpy as np
    from astropy.io import fits
    from scipy.interpolate import interp1d

    #These arrays will be filled with the stiched data.
    Total_Specs = np.array([])
    Total_Waves = np.array([])
    Total_Cont = np.array([])

    step_size = np.diff(wave)
    step_size = np.min(step_size[step_size>0])/2

    for i in range(len(spec)-1):

        waves = np.linspace(np.min(wave[i]), np.max(wave[i+1]), 4*wave[i:i+1].size) #Specifiying wavelength grid

        #Interpolating the spectral orders to the new grid for both the current and proceeding order
        I_spectra_1 = interp1d(wave[i], spec[i], bounds_error = False) #Intep for spectra
        I_spectra_2 = interp1d(wave[i+1], spec[i+1], bounds_error = False)

        I_sig_1 = interp1d(wave[i], sig[i], bounds_error = False) #Interp for sig vals
        I_sig_2 = interp1d(wave[i+1], sig[i+1], bounds_error = False)

        I_cont_1 = interp1d(wave[i], cont[i], bounds_error = False) #Interp for continuum vals
        I_cont_2 = interp1d(wave[i+1], cont[i+1], bounds_error = False)

        #Using interpolator to create a vecotr of the same length for both orders.
        ## Note: If the order doesn't span the wavelength range it creates a nan
        I_spectra_1 = I_spectra_1(waves)
        I_spectra_2 = I_spectra_2(waves)

        I_sig_1 = I_sig_1(waves)
        I_sig_2 = I_sig_2(waves)

        I_cont_1 = I_cont_1(waves)
        I_cont_2 = I_cont_2(waves)

        #If a nan value is present, replace the value with the corresponding value of the other interpolated vecotor
        ## This creates an overlapping vecotr of the same value, which means you can take an average and get the same value
        I_spectra_1 = np.nan_to_num(I_spectra_1, nan = I_spectra_2)
        I_spectra_2 = np.nan_to_num(I_spectra_2, nan = I_spectra_1)

        I_sig_1 = np.nan_to_num(I_sig_1, nan = I_sig_2)
        I_sig_2 = np.nan_to_num(I_sig_2, nan = I_sig_1)

        I_cont_1 = np.nan_to_num(I_cont_1, nan = I_cont_2)
        I_cont_2 = np.nan_to_num(I_cont_2, nan = I_cont_1)

        Spec_Combo = np.array([I_spectra_1, I_spectra_2]) #Combing spectra as 2d array
        Spec_Combo = Spec_Combo.T #Take transpose to pair the data values

        #Do the same for the sig values and continuum
        Sig_Combo = np.array([I_sig_1, I_sig_2])
        Sig_Combo = Sig_Combo.T

        Cont_Combo = np.array([I_cont_1, I_cont_2])
        Cont_Combo = Cont_Combo.T

        #Money is made here, a weighted average is taken using the sig values than normalise by dividing by the continuu,
        Ave_Spec = np.average(Spec_Combo, weights = 1/(Sig_Combo**2), axis = 1)/np.average(Cont_Combo, axis = 1)
        Ave_Spec = Ave_Spec.T

        #Averages are append with their corresponding wavlength
        Total_Specs = np.append(Total_Specs, Ave_Spec)
        Total_Waves = np.append(Total_Waves, waves)
        #Total_Cont = np.append(Total_Cont, np.average(Cont_Combo, axis = 1))

    #Reaorder values to a new grid to create the final stiched spectra
    waves = np.arange(np.min(Total_Waves), np.max(Total_Waves), step_size)
    I_T_Spectra = interp1d(Total_Waves, Total_Specs)

    return waves, I_T_Spectra(waves) #returns wavelength grid and the normalised flux values

def blaze_model(blaze,sdev=3):
    """Applies running average fit of the blaze order data.
        args:
            blaze: the average escelle blaze data
            sdev: stanrard deviation cutoff

        returns:
            a: the mean blaze model
    """

    def nan_helper(data): #Function which allows interpolation though nans
        return lambda z: z.nonzero()[0]

    blaze = blaze.copy()
    data = blaze.copy()

    nan_mask = np.isnan(blaze)
    no_nan = nan_helper(blaze)

    #This will re-create the dataset but will interpolate though the nans giving it a place holder average value
    blaze[nan_mask]= np.interp(no_nan(nan_mask), no_nan(~nan_mask), blaze[~nan_mask], period = 1)

    a = uniform_filter1d(blaze,size=300,mode="nearest")#Applies moving averages on data

    #Takes difference the replaces datavalues which fall less 3 std of mean trend with mean
    diff = blaze - a
    cleaned_blaze_1 = blaze
    cleaned_blaze_1[diff<-sdev*np.std(diff)] = a[diff<-sdev*np.std(diff)]

    #Repeats process on the new data but masks out datavalues which fall outside 3std in both directions
    a = uniform_filter1d(cleaned_blaze_1,size=300,mode="nearest")
    diff = blaze - a
    cleaned_blaze_2 = cleaned_blaze_1
    cleaned_blaze_2[np.absolute(diff)>sdev*np.std(diff)] = a[np.absolute(diff)>sdev*np.std(diff)]
    a = uniform_filter1d(cleaned_blaze_1,size=300,mode="nearest")

    #Replace all positions that contanined nan values initially with nans again
    a[np.isnan(data)] = data[np.isnan(data)]

    return a

def read_uves(inpath,filelist,mode):
    """This reads a folder of UVES-blue or UVES-red intermediate pipeline products. Input is a list of filepaths and the mode."""
    #The following variables define lists in which all the necessary data will be stored.
    framename=[]
    header=[]
    s1dhdr=[]
    obstype=[]
    texp=np.array([])
    date=[]
    mjd=np.array([])
    s1dmjd=np.array([])
    npx=np.array([])
    npx1d=np.array([])
    norders=np.array([])
    e2ds=[]
    s1d=[]
    wave1d=[]
    airmass=np.array([])
    berv=np.array([])
    wave=[]
    catkeyword = 'HIERARCH ESO DPR CATG'
    bervkeyword = 'HIERARCH ESO DRS BERV'
    thfilekeyword = 'HIERARCH ESO DRS CAL TH FILE'
    Zstartkeyword = 'HIERARCH ESO TEL AIRM START'
    Zendkeyword = 'HIERARCH ESO TEL AIRM END'
    for i in range(len(filelist)):
        print(f'------{filelist[i]}', end="\r")
        if (inpath/Path(filelist[i])).is_dir():
            tmp_products = [i for i in (inpath/Path(filelist[i])).glob('resampled_science_*.fits')]
            tmp_products1d = [i for i in (inpath/Path(filelist[i])).glob('red_science_*.fits')]
            if mode == 'UVES-red' and len(tmp_products) != 2:
                raise ValueError(f"in read_e2ds: When mode=UVES-red there should be 2 resampled_science files (redl and redu), but {len(tmp_products)} were detected in {str(inpath/Path(filelist[i]))}.")
            if mode == 'UVES-blue' and len(tmp_products) != 1:
                raise ValueError(f"in read_e2ds: When mode=UVES-rblue there should be 1 resampled_science files (blue), but {len(tmp_products)} were detected in {str(inpath/Path(filelist[i]))}.")
            if mode == 'UVES-red' and len(tmp_products1d) != 2:
                raise ValueError(f"in read_e2ds: When mode=UVES-red there should be 2 red_science files (redl and redu), but {len(tmp_products1d)} were detected in {str(inpath/Path(filelist[i]))}.")
            if mode == 'UVES-blue' and len(tmp_products1d) != 1:
                raise ValueError(f"in read_e2ds: When mode=UVES-rblue there should be 1 red_science files (blue), but {len(tmp_products1d)} were detected in {str(inpath/Path(filelist[i]))}.")

            data_combined = []#This will store the two chips (redu and redl) in case of UVES_red, or simply the blue chip if otherwise.
            wave_combined = []
            wave1d_combined=[]
            data1d_combined=[]
            norders_tmp = 0
            for tmp_product in tmp_products:
                hdul = fits.open(tmp_product)
                data = copy.deepcopy(hdul[0].data)
                hdr = hdul[0].header
                hdul.close()
                del hdul[0].data
                if not hdr['HIERARCH ESO PRO SCIENCE']:#Only add if it's actually a science product:#I force the user to supply only science exposures in the input  folder. No BS allowed... UVES is hard enough as it is.
                    raise ValueError(f' in read_e2ds: UVES file {tmp_product} is not classified as a SCIENCE file, but should be. Remove it from the folder?')
                wavedata=ut.read_wave_from_e2ds_header(hdr,mode='UVES')/10.0#Convert to nm.
                data_combined.append(data)
                wave_combined.append(wavedata)
                norders_tmp+=np.shape(data)[0]

            for tmp_product in tmp_products1d:
                hdul = fits.open(tmp_product)
                data_1d = copy.deepcopy(hdul[0].data)
                hdr1d = hdul[0].header
                hdul.close()
                del hdul[0].data
                if not hdr1d['HIERARCH ESO PRO SCIENCE']:#Only add if it's actually a science product:#I force the user to supply only science exposures in the input  folder. No BS allowed... UVES is hard enough as it is.
                    raise ValueError(f' in read_e2ds: UVES file {tmp_product} is not classified as a SCIENCE file, but should be. Remove it from the folder?')
                npx_1d = hdr1d['NAXIS1']
                wavedata = np.arange(npx_1d, dtype=float)*hdr1d['CDELT1']+hdr1d['CRVAL1']
                data1d_combined.append(data_1d)
                wave1d_combined.append(wavedata)

            if len(data_combined) < 1 or len(data_combined) > 2:#Double-checking that length here...
                raise ValueError(f'in read_e2ds(): Expected 1 or 2 chips, but {len(data_combined)} files were somehow read.')
            #The chips generally don't give the same size. Therefore I will pad the smaller one with NaNs to make it fit:
            if len(data_combined) != len(data1d_combined):
                raise ValueError(f'in read_e2ds(): The number of chips in the 1d and 2d spectra is not the same {len(data1d_combined)} vs {len(data_combined)}.')

            if len(data_combined) == 2:
                chip1 = data_combined[0]
                chip2 = data_combined[1]
                wave1 = wave_combined[0]
                wave2 = wave_combined[1]
                npx_1 = np.shape(chip1)[1]
                npx_2 = np.shape(chip2)[1]
                no_1 = np.shape(chip1)[0]
                no_2 = np.shape(chip2)[0]
                npx_max = np.max([npx_1,npx_2])
                npx_min = np.min([npx_1,npx_2])
                diff = npx_max-npx_min
                #Pad the smaller one with NaNs to match the wider one:
                if npx_1 < npx_2:
                    chip1=np.hstack([chip1,np.zeros((no_1,diff))*np.nan])
                    wave1=np.hstack([wave1,np.zeros((no_1,diff))*np.nan])
                else:
                    chip2=np.hstack([chip2,np.zeros((no_2,diff))*np.nan])
                    wave2=np.hstack([wave2,np.zeros((no_2,diff))*np.nan])
                #So now they can be stacked:
                e2ds_stacked = np.vstack((chip1,chip2))
                wave_stacked = np.vstack((wave1,wave2))
                if np.shape(e2ds_stacked)[1] != np.shape(wave_stacked)[1]:
                    raise ValueError("Width of stacked e2ds and stacked wave frame are not the same. Is the wavelength solution in the header of this file correct?")
                npx=np.append(npx,np.shape(e2ds_stacked)[1])

                e2ds.append(e2ds_stacked)
                wave.append(wave_stacked)
                chip1_1d = data1d_combined[0]
                chip2_1d = data1d_combined[1]
                wave1_1d = wave1d_combined[0]
                wave2_1d = wave1d_combined[1]
                if np.nanmean(wave1_1d) < np.nanmean(wave2_1d):
                    combined_data_1d = np.concatenate((chip1_1d,chip2_1d))
                    combined_wave_1d = np.concatenate((wave1_1d,wave2_1d))
                else:
                    combined_data_1d = np.concatenate((chip2_1d,chip1_1d))
                    combined_wave_1d = np.concatenate((wave2_1d,wave1_1d))
                wave1d.append(combined_wave_1d)
                s1d.append(combined_data_1d)
                npx1d=np.append(npx1d,len(combined_wave_1d))
            else:
                e2ds.append(data_combined[0])
                wave.append(wave_combined[0])
                npx=np.append(npx,np.shape(data_combined[0])[1])
                wave1d.append(wave1d_combined[0])
                s1d.append(data1d_combined[0])
                npx1d=np.append(npx1d,len(combined_wave_1d))
            #Only using the keyword from the second header in case of redl,redu.
            s1dmjd=np.append(s1dmjd,hdr1d['MJD-OBS'])
            framename.append(hdr['ARCFILE'])
            header.append(hdr)
            obstype.append('SCIENCE')
            texp=np.append(texp,hdr['EXPTIME'])
            date.append(hdr['DATE-OBS'])
            mjd=np.append(mjd,hdr['MJD-OBS'])
            norders=np.append(norders,norders_tmp)
            airmass=np.append(airmass,0.5*(hdr[Zstartkeyword]+hdr[Zendkeyword]))#This is an approximation where we take the mean airmass.
            berv_i=sp.calculateberv(hdr['MJD-OBS'],hdr['HIERARCH ESO TEL GEOLAT'],hdr['HIERARCH ESO TEL GEOLON'],hdr['HIERARCH ESO TEL GEOELEV'],hdr['RA'],hdr['DEC'])
            berv = np.append(berv,berv_i)
            hdr1d['HIERARCH ESO QC BERV']=berv_i#Append the berv here using the ESPRESSO berv keyword, so that it can be used in molecfit later.
            s1dhdr.append(hdr1d)

    #Check that all exposures have the same number of pixels, and clip orders if needed.
    min_npx = int(np.min(np.array(npx)))
    min_npx1d = int(np.min(np.array(npx1d)))
    if np.sum(np.abs(np.array(npx)-npx[0])) != 0:
        warnings.warn("in read_e2ds when reading UVES data: Not all e2ds files have the same number of pixels. This could have happened if the pipeline has extracted one or two extra pixels in some exposures but not others. The e2ds files will be clipped to the smallest width.",RuntimeWarning)
        for i in range(len(e2ds)):
            wave[i]=wave[i][:,0:min_npx]
            e2ds[i]=e2ds[i][:,0:min_npx]
            npx[i]=min_npx
    if np.sum(np.abs(np.array(npx1d)-npx1d[0])) != 0:
        warnings.warn("in read_e2ds when reading UVES data: Not all s1d files have the same number of pixels. This could have happened if the pipeline has extracted one or two extra pixels in some exposures but not others. The s1d files will be clipped to the smallest width.",RuntimeWarning)
        for i in range(len(s1d)):
            wave1d[i]=wave1d[i][0:min_npx1d]
            s1d[i]=s1d[i][0:min_npx1d]
            npx1d[i]=min_npx1d
    output = {'wave':wave,'e2ds':e2ds,'header':header,'wave1d':wave1d,'s1d':s1d,'s1dhdr':s1dhdr,'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,'npx1d':npx1d,'norders':norders,'berv':berv,'airmass':airmass,'s1dmjd':s1dmjd}
    return(output)

def read_espresso(inpath,filelist,read_s1d=True,skysub=True):
    #The following variables define lists in which all the necessary data will be stored.
    framename=[]
    header=[]
    s1dhdr=[]
    obstype=[]
    texp=np.array([])
    date=[]
    mjd=np.array([])
    s1dmjd=np.array([])
    npx=np.array([])
    norders=np.array([])
    e2ds=[]
    s1d=[]
    wave1d=[]
    airmass=np.array([])
    berv=np.array([])
    wave=[]
    catkeyword = 'EXTNAME'
    bervkeyword = 'HIERARCH ESO QC BERV'
    airmass_keyword1 = 'HIERARCH ESO TEL'
    airmass_keyword2 = ' AIRM '
    airmass_keyword3_start = 'START'
    airmass_keyword3_end = 'END'

    if skysub:
        type_suffix = 'S2D_SKYSUB_A.fits'
    else:
        type_suffix = 'S2D_BLAZE_A.fits'

    for i in range(len(filelist)):
        if filelist[i].endswith(type_suffix):
            hdul = fits.open(inpath/filelist[i])
            data = copy.deepcopy(hdul[1].data)
            hdr = hdul[0].header
            hdr2 = hdul[1].header
            wavedata=copy.deepcopy(hdul[5].data)
            hdul.close()
            del hdul

            if hdr2[catkeyword] == 'SCIDATA':
                # print('science keyword found')
                print(f'------{filelist[i]}', end="\r")
                framename.append(filelist[i])
                header.append(hdr)
                obstype.append('SCIENCE')
                texp=np.append(texp,hdr['EXPTIME'])
                date.append(hdr['DATE-OBS'])
                mjd=np.append(mjd,hdr['MJD-OBS'])
                npx=np.append(npx,hdr2['NAXIS1'])
                norders=np.append(norders,hdr2['NAXIS2'])
                e2ds.append(data)
                berv=np.append(berv,hdr[bervkeyword])#in km.s.
                telescope = hdr['TELESCOP'][-1]
                airmass = np.append(airmass,0.5*(hdr[airmass_keyword1+telescope+' AIRM START']+hdr[airmass_keyword1+telescope+' AIRM END']))
                wave.append(wavedata/10.0)#*(1.0-(hdr[bervkeyword]*u.km/u.s/const.c).decompose().value))
                #Ok.! So unlike HARPS, ESPRESSO wavelengths are actually BERV corrected in the S2Ds.
                #WHY!!!?. WELL SO BE IT. IN ORDER TO HAVE E2DSes THAT ARE ON THE SAME GRID, AS REQUIRED, WE UNDO THE BERV CORRECTION HERE.
                #WHEN COMPARING WAVE[0] WITH WAVE[1], YOU SHOULD SEE THAT THE DIFFERENCE IS NILL.
                #THATS WHY LATER WE CAN JUST USE WAVE[0] AS THE REPRESENTATIVE GRID FOR ALL.
                #BUT THAT IS SILLY. JUST SAVE THE WAVELENGTHS!

                if read_s1d:
                    s1d_path=inpath/Path(str(filelist[i]).replace('_'+type_suffix,'_S1D_A.fits'))
                    #Need the blazed files. Not the S2D_A's by themselves.
                    ut.check_path(s1d_path,exists=True)#Crash if the S1D doesn't exist.
                    hdul = fits.open(s1d_path)
                    data_table = copy.deepcopy(hdul[1].data)
                    hdr1d = hdul[0].header
                    hdul.close()
                    del hdul
                    s1d.append(data_table.field(2))

                    berv1d = hdr1d[bervkeyword]
                    if berv1d != hdr[bervkeyword]:
                        wrn_msg = ('WARNING in read_espresso(): BERV correction of S1D file is not'
                        f'equal to that of the S2D file. {berv1d} vs {hdr[bervkeyword]}')
                        ut.tprint(wrn_msg)
                    gamma = (1.0-(berv1d*u.km/u.s/const.c).decompose().value)
                    wave1d.append(data_table.field(1)*gamma)#This is in angstroms.
                    #We need to check to which UT ESPRESSO was connected, so that we can read
                    #the weather information (which is UT-specific) and parse them into the
                    #header using UT-agnostic keywords that are in the ESPRESSO.par file.
                    TELESCOP = hdr1d['TELESCOP'].split('U')[1]#This is the number of the UT, either 1, 2, 3 or 4.
                    if TELESCOP not in ['1','2','3','4']:
                        raise ValueError(f"in read_e2ds when reading ESPRESSO data. The UT telescope is not recognised. (TELESCOP={hdr['TELESCOP']})")
                    else:
                        hdr1d['TELALT']     = hdr1d[f'ESO TEL{TELESCOP} ALT']
                        hdr1d['RHUM']       = hdr1d[f'ESO TEL{TELESCOP} AMBI RHUM']
                        hdr1d['PRESSURE']   = (hdr1d[f'ESO TEL{TELESCOP} AMBI PRES START']+
                                            hdr1d[f'ESO TEL{TELESCOP} AMBI PRES END'])/2.0
                        hdr1d['AMBITEMP']   = hdr1d[f'ESO TEL{TELESCOP} AMBI TEMP']
                        hdr1d['M1TEMP']     = hdr1d[f'ESO TEL{TELESCOP} TH M1 TEMP']
                    s1dhdr.append(hdr1d)
                    s1dmjd=np.append(s1dmjd,hdr1d['MJD-OBS'])
    if read_s1d:
        output = {'wave':wave,'e2ds':e2ds,'header':header,'wave1d':wave1d,'s1d':s1d,'s1dhdr':s1dhdr,
        'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
        'norders':norders,'berv':berv,'airmass':airmass,'s1dmjd':s1dmjd}
    else:
        output = {'wave':wave,'e2ds':e2ds,'header':header,
        'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
        'norders':norders,'berv':berv,'airmass':airmass}
    return(output)

def read_spirou(inpath,filelist,read_s1d=False):
    """
    This reads a folder of SPIROU spectra expecting *t.fits for telluric reduce spectra to exists

    As this data is already telluric reduced no effort is done for creating the s1d data structures
    """

    if read_s1d:
        wrn_msg = ('read s1d files not implemented yet for SPIROU, read_s1d reset to false!')
        ut.tprint(wrn_msg)
        read_s1d = False

    #The following variables define lists in which all the necessary data will be stored.
    framename=[]
    header=[]
    s1dhdr=[]
    obstype=[]
    texp=np.array([])
    date=[]
    mjd=np.array([])
    s1dmjd=np.array([])
    npx=np.array([])
    norders=np.array([])
    e2ds=[]
    s1d=[]
    wave1d=[]
    airmass=np.array([])
    berv=np.array([])
    wave=[]
    for i in range(len(filelist)):
        if filelist[i].endswith('t.fits'):
            print(f'------{filelist[i]}', end="\r")
            hdul = fits.open(inpath/filelist[i])
            fluxdata = copy.deepcopy(hdul[1].data)
            wavedata = ops.vactoair(copy.deepcopy(hdul[2].data))
            blazedata = copy.deepcopy(hdul[3].data)
            teldata = copy.deepcopy(hdul[4].data)
            hdr = hdul[0].header
            fluxhdr = hdul[1].header
            wavehdr = hdul[2].header
            blazehdr = hdul[3].header
            telhdr = hdul[4].header
            hdul.close()
            del hdul[1].data
            del hdul[2].data
            del hdul[3].data
            del hdul[4].data

            __npx = fluxhdr['NAXIS1']
            __norders = fluxhdr['NAXIS2']
            idxfilter = []
            for idx in range(__norders):
                numberofNaN = np.count_nonzero(np.isnan(fluxdata[idx]))
                idxfilter.append(numberofNaN != fluxdata[idx].size)
#                __ratio = float(numberofNaN) / float(fluxdata[idx].size)
#                idxfilter.append(__ratio < 1.0)
#                print("numberofNaN"+str(numberofNaN)+"  fluxdata:"+str(fluxdata[idx].size)+"   ratio:"+str(__ratio))
            fluxdata = fluxdata[idxfilter]
            wavedata = wavedata[idxfilter]
            blazedata = blazedata[idxfilter]
            teldata = teldata[idxfilter]
#            print("dropped count of orders: "+str(np.size(idxfilter) - np.sum(idxfilter)))
            __norders = __norders - np.size(idxfilter) + np.sum(idxfilter)  # subtracts number of False

            framename.append(filelist[i])
            header.append(hdr)
            obstype.append(hdr['OBSTYPE'])
            texp=np.append(texp,hdr['EXPTIME'])
            date.append(hdr['DATE-OBS']+'T'+hdr['UTC-OBS'])
            mjd=np.append(mjd,hdr['MJD-OBS'])
            npx=np.append(npx,__npx)
            norders=np.append(norders,__norders)
            e2ds.append(fluxdata)
            wave.append(wavedata)
            berv=np.append(berv,fluxhdr['BERV'])
            airmass=np.append(airmass,hdr['AIRMASS'])

    if read_s1d:
        output = {'wave':wave,'e2ds':e2ds,'header':header,'wave1d':wave1d,'s1d':s1d,'s1dhdr':s1dhdr,
              'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
              'norders':norders,'berv':berv,'airmass':airmass,'s1dmjd':s1dmjd}
    else:
        output = {'wave':wave,'e2ds':e2ds,'header':header,
              'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
              'norders':norders,'berv':berv,'airmass':airmass}
    return(output)

def read_gianob(inpath,filelist,read_s1d=True):
    """
    This reads a folder of SPIROU spectra expecting *t.fits for telluric reduce spectra to exists

    As this data is already telluric reduced no effort is done for creating the s1d data structures
    """
    #The following variables define lists in which all the necessary data will be stored.
    framename=[]
    header=[]
    s1dhdr=[]
    obstype=[]
    texp=np.array([])
    date=[]
    mjd=np.array([])
    s1dmjd=np.array([])
    npx=np.array([])
    norders=np.array([])
    e2ds=[]
    s1d=[]
    wave1d=[]
    airmass=np.array([])
    berv=np.array([])
    wave=[]

    for i in range(len(filelist)):
        if filelist[i].endswith('AB_ms1d.fits'):
            hdul = fits.open(inpath/filelist[i])
            fitsdata = copy.deepcopy(hdul[1].data)
            hdr = hdul[0].header
            hdul.close()
            del hdul[1].data

            if hdr['OBS-TYPE'] == 'SCIENCE':
                # print('science keyword found')
                print(f'------{filelist[i]}', end="\r")
                framename.append(filelist[i])
                header.append(hdr)
                obstype.append('SCIENCE')
                texp=np.append(texp,hdr['EXPTIME'])
                date.append(hdr['DATE-OBS'])
                mjd=np.append(mjd,hdr['MJD-OBS'])
                berv=np.append(berv,hdr['HIERARCH TNG DRS BERV'])#in km.s.
                airmass = np.append(airmass,hdr['AIRMASS'])
                ordernumbers = []
                wavedata = []
                fluxdata = []
                snrdata = []
                __npx = len(fitsdata[0][1])
                __norders = len(fitsdata)
                npx=np.append(npx,__npx)
                norders=np.append(norders,__norders)
                for idx in range(len(fitsdata)):
                    ordernumbers.append(fitsdata[idx][0])
                    wavedata.append(fitsdata[idx][1])
                    fluxdata.append(fitsdata[idx][2])
                    snrdata.append(fitsdata[idx][3])
                wave.append(ops.vactoair(np.array(wavedata)))
                e2ds.append(np.array(fluxdata))
#                print("---wave then flux---")
#                print(np.array(wavedata))
#                print(np.array(fluxdata))

                if read_s1d:
                    s1d_path=inpath/Path(str(filelist[i]).replace('_ms1d.fits','_s1d.fits'))
                    #Need the blazed files. Not the S2D_A's by themselves.
                    ut.check_path(s1d_path,exists=True)#Crash if the S1D doesn't exist.
                    hdul1d = fits.open(s1d_path)
                    hdr1d = hdul1d[0].header
                    fluxdata1d = copy.deepcopy(hdul1d[0].data)
                    hdul1d.close()
                    del hdul1d[0].data

                    s1d.append(fluxdata1d)

                    bervkeyword = 'HIERARCH TNG DRS BERV'
                    berv1d = hdr1d[bervkeyword]
                    if berv1d != hdr[bervkeyword]:
                        wrn_msg = ('WARNING in read_gianob(): BERV correction of S1D file is not'
                        f'equal to that of the S2D file. {berv1d} vs {hdr[bervkeyword]}')
                        ut.tprint(wrn_msg)
                    gamma = (1.0-(berv1d*u.km/u.s/const.c).decompose().value)
                    crval = hdr1d['CRVAL1']
                    cdelt = hdr1d['CDELT1']
                    mywave = crval
                    wavedata1d = [mywave*gamma*10]
                    for idx in range(len(fluxdata1d)-1):
                        mywave += cdelt
                        wavedata1d.append(mywave*gamma*10)
                    wave1d.append(ops.vactoair(np.array(wavedata1d)))
                    s1dhdr.append(hdr1d)
                    s1dmjd=np.append(s1dmjd,hdr1d['MJD-OBS'])
#                    print("wave:"+str(len(wavedata1d))+" flux:"+str(len(fluxdata1d)))

    if read_s1d:
        output = {'wave':wave,'e2ds':e2ds,'header':header,'wave1d':wave1d,'s1d':s1d,'s1dhdr':s1dhdr,
        'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
        'norders':norders,'berv':berv,'airmass':airmass,'s1dmjd':s1dmjd}
    else:
        output = {'wave':wave,'e2ds':e2ds,'header':header,
        'mjd':mjd,'date':date,'texp':texp,'obstype':obstype,'framename':framename,'npx':npx,
        'norders':norders,'berv':berv,'airmass':airmass}
    return(output)

import numpy,import scipy,from scipy,import astropy,from astropy