"""
Create an image combination algroithm that chooses between a minimum
or median image based up bad pixel identification.

:Authors: Christopher Hanley

:License: :doc:`LICENSE`

"""
#   PROGRAM: numcombine.py
#   AUTHOR: Christopher Hanley (based upon original code of Anton Koekemoer)
#   DATE:    February 11, 2004
#   PURPOSE: Create an image combination algroithm that chooses between
#            a minimum or median image based up bad pixel identification.
#   HISTORY:
#      Version 0.1.0: Initial Development -- CJH -- 02/11/04
#      Version 0.1.1: Cast boxsize as type int.  This should always be the case
#        anyways but this protects against errors in the
#        MDRIZTAB -- CJH/WJH -- 07/08/04
#      Version 0.1.2: Improve error message handing in the case where
#        the boxcar convolution step fails.  --CJH -- 10/13/04
#      Version 0.2.0: The creation of the median image will now more closesly
#        replicate the IRAF IMCOMBINE behavior of nkeep = 1 and nhigh = 1.
#        -- CJH -- 03/29/05
#      Version 0.3.0: Rewritten to optimize the code and to bring
#        code up to modern standards.-- Mihai Cara -- 02/19/2018
import warnings
import numpy as np
from scipy import signal
from stsci.image.numcombine import numCombine, num_combine
from .version import *

class minmed:
    """ **DEPRECATED** Create a median array, rejecting the highest pixel and
    computing the lowest valid pixel after mask application

        # In this case we want to calculate two things:
        #   1) the median array, rejecting the highest pixel (thus running
        #      imcombine with nlow=0, nhigh=1, nkeep=1, using the masks)
        #   2) the lowest valid pixel after applying the masks (thus running
        #      imcombine with nlow=0, nhigh=3, nkeep=1, using the masks)
        #
        # We also calculate the sum of the weight files (to produce the total
        # effective exposure time for each pixel).
        #
        # The total effective background in the final image is calculated as follows:
        #   - convert background for each input image to counts/s (divide by exptime)
        #   - multiply this value by the weight image, to obtain the effective background
        #     counts (in DN) for each pixel, for each image
        #   - Add these images together, to obtain the total effective background
        #     for the combined image.
        #
        # Once we've made these two files, then calculate the SNR based on the
        # median-pixel image, and compare with the minimum.

    In this version of the mimmed algorithm we assume that the units of all
    input data is electons.
    """
    def __init__(self,
            imageList,              # list of input data to be combined.
            weightImageList,        # list of input data weight images to be combined.
            readnoiseList,          # list of readnoise values to use for the input images.
            exposureTimeList,       # list of exposure times to use for the input images.
            backgroundValueList,    # list of image background values to use for the input images
            weightMaskList= None,   # list of imput data weight masks to use for pixel rejection.
            combine_grow = 1,       # Radius (pixels) for neighbor rejection
            combine_nsigma1 = 4,    # Significance for accepting minimum instead of median
            combine_nsigma2 = 3,     # Significance for accepting minimum instead of median
            fillval = False         # Turn on use of imedian/imean

            ):

        warnings.warn("The 'minmed' class is deprecated and may be removed"
                      " in a future version. Use 'min_med()' instead.",
                      DeprecationWarning)

        # Define input variables
        self._imageList = imageList
        self._weightImageList = weightImageList
        self._weightMaskList = weightMaskList
        self._exposureTimeList = exposureTimeList
        self._readnoiseList = readnoiseList
        self._backgroundValueList = backgroundValueList
        self._numberOfImages = len( self._imageList)
        self._combine_grow = combine_grow
        self._combine_nsigma1 = combine_nsigma1
        self._combine_nsigma2 = combine_nsigma2

        if fillval:
            combtype_mean = 'imean'
            combtype_median = 'imedian'
        else:
            combtype_mean = 'mean'
            combtype_median = 'median'

        # Create a different median image based upon the number of images in the input list.
        median_file = np.zeros(self._imageList[0].shape,dtype=self._imageList[0].dtype)
        if (self._numberOfImages == 2):
            tmp = numCombine(self._imageList,numarrayMaskList=self._weightMaskList,
                                 combinationType=combtype_mean,nlow=0,nhigh=0,
                                 nkeep=1,upper=None,lower=None)
            median_file = tmp.combArrObj
        else:
            # The value of NHIGH=1 will cause problems when there is only 1 valid
            # unmasked input image for that pixel due to a difference in behavior
            # between 'numcombine' and 'iraf.imcombine'.
            # This value may need to be adjusted on the fly based on the number of
            # inputs and the number of masked values/pixel.
            #
            tmp = numCombine(self._imageList,numarrayMaskList=self._weightMaskList,
                                 combinationType=combtype_median,nlow=0,nhigh=1,
                                 nkeep=1,upper=None,lower=None)
            median_file = tmp.combArrObj

            if self._weightMaskList in [None,[]]:
                self._weightMaskList = [np.zeros(self._imageList[0].shape,dtype=self._imageList[0].dtype)]*len(self._imageList)
            # The following section of code will address the problem caused by having
            # a value of nhigh = 1.  This will behave in a way similar to the way the
            # IRAF task IMCOMBINE behaves.  In order to accomplish this, the following
            # procedure will be followed:
            # 1) The input masks will be summed.
            # 2) The science data will be summed.
            # 3) In the locations of the summed mask where the sum is 1 less than the
            #    the total number of images, the value of that location in the summed
            #    sceince image will be used to replace the existing value in the
            #    existing median_file.
            #
            # This procuedure is being used to prevent too much data from being thrown
            # out of the image.  Take for example the case of 3 input images.  In two
            # of the images the pixel locations have been masked out.  Now, if nhigh
            # is applied there will be no value to use for that position.  However,
            # if this new procedure is used that value in the resulting images will
            # be the value that was rejected by the nhigh rejection step.
            #

            # We need to make certain that "bad" pixels in the sci data are set to 0.  That way,
            # when the sci images are summed, the value of the sum will only come from the "good"
            # pixels.
            tmpList = []
            for image in range(len(self._imageList)):
                tmp =np.where(self._weightMaskList[image] == 1, 0, self._imageList[image])
                tmpList.append(tmp)

            # Sum the mask files
            maskSum = self._sumImages(self._weightMaskList)
            # Sum the science images
            sciSum = self._sumImages(tmpList)
            del(tmpList)
            # Use the summed sci image values in locations where the maskSum indicates
            # that there is only 1 good pixel to use.  The value will be used in the
            # median_file image
            median_file = np.where(maskSum == self._numberOfImages-1,sciSum,median_file)

        if self._weightMaskList in [None,[]]:
            self._weightMaskList = [np.zeros(self._imageList[0].shape,dtype=self._imageList[0].dtype)]*len(self._imageList)
        # Sum the weightMaskList elements
        maskSum = self._sumImages(self._weightMaskList)

        # Create the minimum image from the stack of input images.
        # Find the maximum pixel value for the image stack.
        maxValue = -1e+9
        for image in self._imageList:
            newMax = image.max()
            if (newMax > maxValue):
                maxValue = newMax

        # For each image, set pixels masked as "bad" to the "super-maximum" value.
        for image in range(len(self._imageList)):
            self._imageList[image] = np.where(self._weightMaskList[image] == 1,maxValue+1,self._imageList[image])

        # Call numcombine throwing out the highest N - 1 pixels.
        tmp = numCombine(self._imageList,numarrayMaskList=None,
                                 combinationType=combtype_median,nlow=0,nhigh=self._numberOfImages-1,
                                 nkeep=1,upper=None,lower=None)
        minimum_file = tmp.combArrObj
        # Reset any pixl at maxValue + 1 to 0.
        minimum_file = np.where(maskSum == self._numberOfImages, 0, minimum_file)

        # Scale the weight images by the background values and add them to the bk
        backgroundFileList = []
        for image in range(len(self._weightImageList)):
            tmp = self._weightImageList[image] * (self._backgroundValueList[image]/(self._exposureTimeList[image]))
            backgroundFileList.append(tmp)

        # Create an image of the total effective background (in DN) per pixel:
        # (which is the sum of all the background-scaled weight files)
        #
        bkgd_file = self._sumImages(backgroundFileList)
        del(backgroundFileList)

        #
        # Scale the weight mask images by the square of the readnoise values
        #
        readnoiseFileList = []
        for image in range(len(self._weightMaskList)):
            tmp = (np.logical_not(self._weightMaskList[image]) *
                   (self._readnoiseList[image] * self._readnoiseList[image]))
            readnoiseFileList.append(tmp)

        # Create an image of the total readnoise**2 per pixel:
        # (which is the sum of all the input readnoise values)
        #
        readnoise_file = self._sumImages(readnoiseFileList)
        del(readnoiseFileList)

        # Create an image of the total effective exposure time per pixel:
        # (which is simply the sum of all the drizzle output weight files)
        #
        weight_file = self._sumImages(self._weightImageList)

        # Scale up both the median and minimum arrays by the total effective exposure time
        # per pixel.
        #
        minimum_file_weighted = minimum_file * weight_file
        median_file_weighted = median_file * weight_file
        del(weight_file)

        # Calculate the 1-sigma r.m.s.:
        #   variance = median_electrons + bkgd_electrons + readnoise**2
        #   rms = sqrt(variance)
        #   This image has units of electrons.
        #
        # make this the abs value so that negative numbers dont throw an exception?
        rms_file2 = np.fmax(
            median_file_weighted + bkgd_file + readnoise_file,
            np.zeros_like(median_file_weighted)
        )
        rms_file = np.sqrt(rms_file2)

        del bkgd_file
        del readnoise_file
        # For the median array, calculate the n-sigma lower threshold to the array
        # and incorporate that into the pixel values.
        #
        median_rms_file = median_file_weighted - (rms_file * self._combine_nsigma1)

        if self._combine_grow != 0:
            #
            # Do a more sophisticated rejection: For all cases where the minimum pixel will
            # be accepted instead of the median, set a lower threshold for that pixel and the
            # ones around it (ie become less conservative in rejecting the median). This is
            # because in cases of triple-incidence cosmic rays, quite often the low-lying
            # outliers of the CRs can influence the median for the initial relatively high
            # value of sigma, so a lower threshold must be used to mnake sure that the minimum
            # is selected.
            #
            # This is done as follows:
            # 1) make an image which is zero everywhere except where the minimum will be accepted
            # 2) box-car smooth this image, to make these regions grow.
            # 3) In the file "median_rms_file_electrons", replace these pixels
            #     by median - combine_nsigma2 * rms
            #
            # Then use this image in the final replacement, in the same way as for the
            # case where this option is not selected.

            minimum_flag_file = np.where(np.less(minimum_file_weighted,median_rms_file), 1, 0)

            # The box size value must be an integer.  This is not a problem since __combine_grow should always
            # be an integer type.  The combine_grow column in the MDRIZTAB should also be an integer type.
            boxsize = int(2 * self._combine_grow + 1)
            boxshape = (boxsize, boxsize)
            minimum_grow_file = np.zeros(self._imageList[0].shape,dtype=self._imageList[0].dtype)

            # If the boxcar convolution has failed it is potentially for two reasons:
            #   1) The kernel size for the boxcar is bigger than the actual image.
            #   2) The grow parameter was specified with a value < 0.  This would result
            #      in an illegal boxshape kernel.  The dimensions of the kernel box *MUST*
            #      be integer and greater than zero.
            #
            #   If the boxcar convolution has failed, try to give a meaningfull explanation
            #   as to why based upon the conditionals described above.

            if (boxsize <= 0):
                errormsg1 =  "############################################################\n"
                errormsg1 += "# The boxcar convolution in minmed has failed.  The 'grow' #\n"
                errormsg1 += "# parameter must be greater than or equal to zero. You     #\n"
                errormsg1 += "# specified an input value for the 'grow' parameter of:    #\n"
                errormsg1 += "        combine_grow: " + str(self._combine_grow)+'\n'
                errormsg1 += "############################################################\n"
                raise ValueError(errormsg1)
            if (boxsize > self._imageList[0].shape[0]):
                errormsg2 =  "############################################################\n"
                errormsg2 += "# The boxcar convolution in minmed has failed.  The 'grow' #\n"
                errormsg2 += "# parameter specified has resulted in a boxcar kernel that #\n"
                errormsg2 += "# has dimensions larger than the actual image.  You        #\n"
                errormsg2 += "# specified an input value for the 'grow' parameter of:    #\n"
                errormsg2 += "        combine_grow: " +str(self._combine_grow)+'\n'
                errormsg2 += "############################################################\n"
                print(self._imageList[0].shape)
                raise ValueError(errormsg2)

            # Attempt the boxcar convolution using the boxshape based upon the user input value of "grow"
            ker = np.ones(boxshape) / float(boxsize**2)
            minimum_grow_file = signal.convolve2d(minimum_flag_file, ker,
                                                  boundary='fill', mode='same')

            del(minimum_flag_file)

            temp1 = (median_file_weighted - (rms_file * self._combine_nsigma1))
            temp2 = (median_file_weighted - (rms_file * self._combine_nsigma2))
            median_rms2_file = np.where(np.equal(minimum_grow_file, 0), temp1, temp2)
            del(temp1)
            del(temp2)
            del(rms_file)
            del(minimum_grow_file)

            # Finally decide whether to use the minimim or the median (in counts/s),
            # based on whether the median is more than 3 sigma above the minimum.
            #
            self.combArrObj = np.where(
                np.less(minimum_file_weighted, median_rms2_file),
                minimum_file,
                median_file
            )

        else:
            # Finally decide whether to use the minimim or the median (in counts/s),
            # based on whether the median is more than 3 sigma above the minimum.
            #
            self.combArrObj = np.where(
                np.less(minimum_file_weighted, median_rms_file),
                minimum_file,
                median_file
            )

        # Set fill regions to a pixel value of 0.
        self.combArrObj = np.where(maskSum == self._numberOfImages, 0, self.combArrObj)

#        self.out_file1 = median_rms2_file.copy()
#        self.out_file2 = minimum_file_weighted.copy()

    def _sumImages(self,numarrayObjectList):
        """ Sum a list of numarray objects. """
        if numarrayObjectList in [None, []]:
            return None

        tsum = np.zeros(numarrayObjectList[0].shape, dtype=numarrayObjectList[0].dtype)

        for image in numarrayObjectList:
            tsum += image

        return tsum

def min_med(images, weight_images, readnoise_list, exptime_list,
            background_values, weight_masks=None, combine_grow=1,
            combine_nsigma1=4, combine_nsigma2=3, fillval=False):
    """ Create a median array, rejecting the highest pixel and
    computing the lowest valid pixel after mask application.

    .. note::
        In this version of the mimmed algorithm we assume that the units of
        all input data is electons.

    Parameters
    ----------
    images : list of numpy.ndarray
        List of input data to be combined.

    weight_images : list of numpy.ndarray
        List of input data weight images to be combined.

    readnoise_list : list
        List of readnoise values to use for the input images.

    exptime_list : list
        List of exposure times to use for the input images.

    background_values : list
        List of image background values to use for the input images.

    weight_masks : list of numpy.ndarray, None
        List of imput data weight masks to use for pixel rejection.
        (Default: `None`)

    combine_grow : int
        Radius (pixels) for neighbor rejection. (Default: 1)

    combine_nsigma1 : float
        Significance for accepting minimum instead of median. (Default: 4)

    combine_nsigma2 : float
        Significance for accepting minimum instead of median. (Default: 3)

    fillval : bool
        Turn on use of imedian/imean. (Default: `False`)

    Returns
    -------
    combined_array : numpy.ndarray
        Combined array.

    """
    # In this case we want to calculate two things:
    #   1) the median array, rejecting the highest pixel (thus running
    #      imcombine with nlow=0, nhigh=1, nkeep=1, using the masks)
    #   2) the lowest valid pixel after applying the masks (thus running
    #      imcombine with nlow=0, nhigh=3, nkeep=1, using the masks)
    #
    # We also calculate the sum of the weight files (to produce the total
    # effective exposure time for each pixel).
    #
    # The total effective background in the final image is calculated as
    # follows:
    #   - convert background for each input image to counts/s
    #     (divide by exptime)
    #   - multiply this value by the weight image, to obtain the effective
    #     background counts (in DN) for each pixel, for each image
    #   - Add these images together, to obtain the total effective background
    #     for the combined image.
    #
    # Once we've made these two files, then calculate the SNR based on the
    # median-pixel image, and compare with the minimum.

    nimages = len(images)
    combtype_median = 'imedian' if fillval else 'median'
    images = np.asarray(images)
    weight_images = np.asarray(weight_images)

    if weight_masks == [] or weight_masks is None:
        weight_masks = None
        mask_sum = np.zeros(images.shape[1:], dtype=np.int16)
        all_bad_idx = np.array([], dtype=np.int)
        all_bad_idy = np.array([], dtype=np.int)
    else:
        weight_masks = np.asarray(weight_masks, dtype=np.bool)
        mask_sum = np.sum(weight_masks, axis=0, dtype=np.int16)
        all_bad_idx, all_bad_idy = np.where(mask_sum == nimages)

    # Create a different median image based upon the number of images in the
    # input list.
    if nimages == 2:
        median_file = num_combine(
            images,
            masks=weight_masks,
            combination_type='imean' if fillval else 'mean',
            nlow=0, nhigh=0, lower=None, upper=None
        )

    else:
        # The value of NHIGH=1 will cause problems when there is only 1 valid
        # unmasked input image for that pixel due to a difference in behavior
        # between 'num_combine' and 'iraf.imcombine'.
        # This value may need to be adjusted on the fly based on the number of
        # inputs and the number of masked values/pixel.
        #
        median_file = num_combine(
            images,
            masks=weight_masks,
            combination_type=combtype_median,
            nlow=0, nhigh=1, lower=None, upper=None
        )

        # The following section of code will address the problem caused by
        # having a value of nhigh = 1.  This will behave in a way similar to
        # the way the IRAF task IMCOMBINE behaves.  In order to accomplish
        # this, the following procedure will be followed:
        # 1) The input masks will be summed.
        # 2) The science data will be summed.
        # 3) In the locations of the summed mask where the sum is 1 less than
        #    the total number of images, the value of that location in the
        #    summed science image will be used to replace the existing value
        #    in the existing median_file.
        #
        # This procedure is being used to prevent too much data from being
        # thrown out of the image. Take for example the case of 3 input images.
        # In two of the images the pixel locations have been masked out.
        # Now, if nhigh is applied there will be no value to use for that
        # position.  However, if this new procedure is used that value in
        # the resulting images will be the value that was rejected by the
        # nhigh rejection step.

        # We need to make certain that "bad" pixels in the sci data are set to
        # 0. That way, when the sci images are summed, the value of the sum
        # will only come from the "good" pixels.
        if weight_masks is None:
            sci_sum = np.sum(images, axis=0)
            if nimages == 1:
                median_file = sci_sum

        else:
            sci_sum = np.sum(images * np.logical_not(weight_masks), axis=0)
            # Use the summed sci image values in locations where the mask_sum
            # indicates that there is only 1 good pixel to use. The value will
            # be used in the median_file image
            idx = np.where(mask_sum == (nimages - 1))
            median_file[idx] = sci_sum[idx]

    # Create the minimum image from the stack of input images.
    if weight_masks is not None:
        # make a copy of images to avoid side-effect of modifying input
        # argument:
        images = images.copy()
        images[weight_masks] = np.nan
        images[:, all_bad_idx, all_bad_idy] = 0
        minimum_file = np.nanmin(images, axis=0)
    else:
        minimum_file = np.amin(images, axis=0)

    # Scale the weight images by the background values and add them to the bk
    # Create an image of the total effective background (in DN) per pixel:
    # (which is the sum of all the background-scaled weight files)
    s = np.asarray([bv / et for bv, et in
                    zip(background_values, exptime_list)])
    bkgd_file = np.sum(weight_images * s[:, None, None], axis=0)

    # Scale the weight mask images by the square of the readnoise values.
    # Create an image of the total readnoise**2 per pixel
    # (which is the sum of all the input readnoise values).
    if weight_masks is None:
        rdn2 = sum((r**2 for r in readnoise_list))
        readnoise_file = rdn2 * np.ones_like(images[0])

    else:
        readnoise_file = np.sum(
            np.logical_not(weight_masks) *
            (np.asarray(readnoise_list)**2)[:, None, None],
            axis=0
        )

    # Create an image of the total effective exposure time per pixel:
    # (which is simply the sum of all the drizzle output weight files)
    weight_file = np.sum(weight_images, axis=0)

    # Scale up both the median and minimum arrays by the total effective
    # exposure time per pixel.
    minimum_file_weighted = minimum_file * weight_file
    median_file_weighted = median_file * weight_file
    del weight_file

    # Calculate the 1-sigma r.m.s.:
    #   variance = median_electrons + bkgd_electrons + readnoise**2
    #   rms = sqrt(variance)
    #   This image has units of electrons.
    #
    # make this the abs value so that negative numbers dont throw an exception?
    rms_file2 = np.fmax(
        median_file_weighted + bkgd_file + readnoise_file,
        np.zeros_like(median_file_weighted)
    )
    rms_file = np.sqrt(rms_file2)
    del bkgd_file, readnoise_file

    # For the median array, calculate the n-sigma lower threshold to the array
    # and incorporate that into the pixel values.
    median_rms_file = median_file_weighted - rms_file * combine_nsigma1

    if combine_grow != 0:
        # Do a more sophisticated rejection: For all cases where the minimum
        # pixel will be accepted instead of the median, set a lower threshold
        # for that pixel and the ones around it (ie become less conservative
        # in rejecting the median). This is because in cases of
        # triple-incidence cosmic rays, quite often the low-lying outliers
        # of the CRs can influence the median for the initial relatively high
        # value of sigma, so a lower threshold must be used to mnake sure that
        # the minimum is selected.
        #
        # This is done as follows:
        # 1) make an image which is zero everywhere except where the minimum
        #    will be accepted
        # 2) box-car smooth this image, to make these regions grow.
        # 3) In the file "median_rms_file_electrons", replace these pixels
        #     by median - combine_nsigma2 * rms
        #
        # Then use this image in the final replacement, in the same way as for
        # the case where this option is not selected.
        minimum_flag_file = np.less(minimum_file_weighted,
                                    median_rms_file).astype(np.float64)

        # The box size value must be an integer. This is not a problem since
        # __combine_grow should always be an integer type. The combine_grow
        # column in the MDRIZTAB should also be an integer type.
        boxsize = int(2 * combine_grow + 1)
        boxshape = (boxsize, boxsize)
        minimum_grow_file = np.zeros_like(images[0])

        # If the boxcar convolution has failed it is potentially for
        # two reasons:
        #   1) The kernel size for the boxcar is bigger than the actual image.
        #   2) The grow parameter was specified with a value < 0.  This would
        #      result in an illegal boxshape kernel. The dimensions of the
        #      kernel box *MUST* be integer and greater than zero.
        #
        #   If the boxcar convolution has failed, try to give a meaningfull
        #   explanation as to why based upon the conditionals described above.
        if boxsize <= 0:
            errormsg1 = "############################################################\n"
            errormsg1 += "# The boxcar convolution in minmed has failed.  The 'grow' #\n"
            errormsg1 += "# parameter must be greater than or equal to zero. You     #\n"
            errormsg1 += "# specified an input value for the 'grow' parameter of:    #\n"
            errormsg1 += "        combine_grow: " + str(combine_grow)+'\n'
            errormsg1 += "############################################################\n"
            raise ValueError(errormsg1)

        if boxsize > images.shape[1]:
            errormsg2 = "############################################################\n"
            errormsg2 += "# The boxcar convolution in minmed has failed.  The 'grow' #\n"
            errormsg2 += "# parameter specified has resulted in a boxcar kernel that #\n"
            errormsg2 += "# has dimensions larger than the actual image.  You        #\n"
            errormsg2 += "# specified an input value for the 'grow' parameter of:    #\n"
            errormsg2 += "        combine_grow: " + str(combine_grow) + '\n'
            errormsg2 += "############################################################\n"
            print(images.shape[1:])
            raise ValueError(errormsg2)

        # Attempt the boxcar convolution using the boxshape based upon the user
        # input value of "grow"
        ker = np.ones((boxsize, boxsize)) / float(boxsize**2)
        minimum_grow_file = signal.convolve2d(minimum_flag_file, ker,
                                              boundary='fill', mode='same')

        median_rms_file = np.where(
            np.equal(minimum_grow_file, 0),
            median_file_weighted - rms_file * combine_nsigma1,
            median_file_weighted - rms_file * combine_nsigma2
        )
        del rms_file, minimum_grow_file

    # Finally decide whether to use the minimim or the median (in counts/s),
    # based on whether the median is more than 3 sigma above the minimum.
    combined_array = np.where(
        np.less(minimum_file_weighted, median_rms_file),
        minimum_file,
        median_file
    )
    # Set fill regions to a pixel value of 0.
    combined_array[all_bad_idx, all_bad_idy] = 0

    return combined_array


#!/usr/bin/env python

"""
.. module:: nllFast
   :synopsis: This module provides methods to access the nllfast grid and
              compute k-factors (when available) to SUSY pair
              production cross sections.

.. moduleauthor:: Suchita Kulkarni <suchita.kulkarni@gmail.com>
.. moduleauthor:: Andre Lessa <lessa.a.p@gmail.com>

"""
#
try:
    import commands as executor
except ImportError:
    import subprocess as executor
import os
from smodels.tools import toolBox
import pyslha
from smodels.tools.physicsUnits import TeV
import numpy
import operator

from smodels import installation
from smodels.tools.smodelsLogging import logger

squarks = [1000001,
           2000001,
           1000002,
           2000002,
           1000003,
           2000003,
           1000004,
           2000004]
antisquarks = map(operator.neg, squarks)
third = [1000005,
         2000005,
         1000006,
         2000006]
gluinos = [1000021]

def getKfactorsFor(pIDs, sqrts, slhafile, pdf='cteq'):
    """
    Read the NLLfast grid and returns a pair of k-factors (NLO and NLL) for the
    pair.

    :returns: k-factors = None, if NLLfast does not contain the process; uses
              the slhafile to obtain the SUSY spectrum.

    """
    if not os.path.isfile(slhafile):
        logger.error("SLHA file %s not found", slhafile)
        return False

    # Get process name (in NLLfast notation)
    process = getProcessName(pIDs)
    if not process:
        # Return k-factors = None, if NLLfast does not have the process
        return (None, None)

    # Obtain relevant masses
    readfile = pyslha.readSLHAFile(slhafile)
    masses=readfile.blocks['MASS']
    check_pids=squarks+gluinos+third
    for check in check_pids:
        if not check in masses.entries:
            logger.error ( "cannot compute k factor for pdgid %d: " \
              " no particle mass given. will set mass to inf." % check )
            masses.entries[check]=1.e10

    gluinomass = abs(masses.entries[1000021])
    squarkmass = sum([abs(masses.entries[pid])
                      for pid in squarks]) / 8.
    pid1, pid2 = sorted(pIDs)
    if pid1 in antisquarks and pid2 in squarks:
        squarkmass = (abs(masses.entries[abs(pid1)]) +
                      abs(masses.entries[pid2])) / 2.
    elif pid1 in squarks and pid2 in squarks:
        squarkmass = (abs(masses.entries[pid1]) + abs(masses.entries[pid2])) / 2.
    elif abs(pid1) == pid2 and pid2 in third:
        squarkmass = abs(masses.entries[abs(pid1)])

    # Set up NLLfast run, the old way
    sqrtS = float(sqrts/TeV)
    energy = str(int(sqrtS)) + 'TeV'
    toolname = "nllfast%d" % int(sqrtS)
    box = toolBox.ToolBox()
    tool = box.get(toolname)
    if tool == None:
        logger.warning("No NLLfast data for sqrts = " + str(sqrts))
        return (None, None)
    nllpath = tool.installDirectory()
    tool.pathOfExecutable()
    tool.checkInstallation()
    if process == "st":
        nll_run = "./nllfast_" + energy + " %s %s %s" % \
                  (process, pdf, squarkmass)
    else:
        nll_run = "./nllfast_" + energy + " %s %s %s %s" % \
                  (process, pdf, squarkmass, gluinomass)

    # Run NLLfast
    nll_output = runNLLfast(nll_run, nllpath)

    # If run was successful, return k-factors:
    if "K_NLO" in nll_output:
        # NLLfast ran ok, try to get the k-factors
        kFacs = getKfactorsFrom(nll_output)
        if not kFacs or min(kFacs) <= 0.:
            logger.warning("Error obtaining k-factors")
            return (None, None)
        else:
            return kFacs
    # If run was not successful, check for decoupling error messages:
    elif not "too low/high" in nll_output.lower():
        logger.warning("Error running NLLfast")
        return (None, None)

    # Check for decoupling cases with a decoupling grid (only for sb and gg)
    doDecoupling = False
    if "too low/high gluino" in nll_output.lower():
        if gluinomass > 500. and process == 'sb':
            doDecoupling = True
            dcpl_mass = gluinomass
    elif "too low/high squark" in nll_output.lower():
        if squarkmass > 500. and process == 'gg':
            doDecoupling = True
            dcpl_mass = squarkmass

    # If process do not have decoupled grids, return None:
    if not doDecoupling:
        logger.warning("Masses out of NLLfast grid for " + process)
        return (None, None)

    # Obtain k-factors from the NLLfast decoupled grid
    kfacs = getDecoupledKfactors(nllpath,process,energy,pdf,min(gluinomass,squarkmass))
    # Decoupling limit is satisfied, do not interpolate
    if not kfacs:
        logger.warning("Error obtaining k-factors from the NLLfast decoupled grid for " + process)
        return (None, None)
    elif dcpl_mass/min(gluinomass,squarkmass) > 10.:
        return kfacs
    # Interpolate between the non-decoupled and decoupled grids
    else:
        kFacsVector = [[10.*min(gluinomass,squarkmass),kfacs]]  #First point for interpolation (decoupled grid)
        kfacs = None
        while not kfacs and dcpl_mass > 500.:
            dcpl_mass -= 100.  # Reduce decoupled mass, until NLLfast produces results
            if process == 'sb': nllinput = (process, pdf, squarkmass, dcpl_mass)
            else:  nllinput = (process, pdf, dcpl_mass, gluinomass)
            nll_run = "./nllfast_" + energy + " %s %s %s %s" % nllinput
            nll_output = runNLLfast(nll_run, nllpath)
            kfacs = getKfactorsFrom(nll_output)
        kFacsVector.append([dcpl_mass, kfacs]) #Second point for interpolation (non-decoupled grid)

    if len(kFacsVector) < 2:
        logger.warning("Not enough points for interpolation in the decoupling "
                       "limit")
        return (None, None)
    else:
        # Interpolate k-factors
        kFacs = interpolateKfactors(kFacsVector,
                        max(squarkmass, gluinomass))
    return kFacs

def getProcessName(pIDs):
    """
    Return the process name (in NLLfast notation) for the pair production of
    pIDs.

    :returns: None, if the particle ID pair is not contained in NLLfast

    """
    pid1, pid2 = sorted(pIDs)
    process = None

    # Obtain the type of process:
    #  - gluino-gluino production = gg
    #  - squark-antisquark = sb
    #  - squark-squark = ss
    #  - squark-gluino = sg
    #  - antistop-stop
    #  - antisbottom-sbottom = st
    if pid1 in antisquarks and pid2 in squarks:
        process = 'sb'
    elif abs(pid1) in squarks and abs(pid2) in squarks:
        process = 'ss'
    elif pid1 == pid2 and pid1 in gluinos:
        process = 'gg'
    elif abs(pid1) in squarks and pid2 == 1000021 or \
            abs(pid2) in squarks and pid1 == 1000021:
        process = 'sg'
    elif abs(pid1) == pid2 and pid2 in third:
        process = 'st'

    return process

def runNLLfast(nll_run, nllpath):
    """
    Execute NLLfast with command nll_run at nllpath.

    :returns: NLLfast output as a string

    """
    current_dir = os.getcwd()
    os.chdir(nllpath)
    nll_output = executor.getoutput(nll_run)
    # nll_output = subprocess.check_output ( nll_run, shell=True, universal_newlines=True )
    os.chdir(current_dir)
    return nll_output

def getKfactorsFrom(output):
    """
    Read NLLfast output and return the k-factors.

    """
    if not output:
        return False
    else:
        lines = output.split('\n')
        il = 0
        line = lines[il]
        process = False
        while not "K_NLO" in line and il < len(lines) - 2:
            if "process" in line:
                process = line[line.find("process:") + 8:].replace(" ", "")
            il += 1
            line = lines[il]
        if not process:
            return False
        # Line with header
        line = lines[il]
        # Count number of mass entries
        nmass = line.count('GeV')
        # Line with values
        line = lines[il + 2]
        data = [eval(x) for x in line.split()]
        if len(data) != nmass + 11:
            return False
        else:
            kFacs = data[9 + nmass:]

    return kFacs

def interpolateKfactors(kFacsVector, xval):
    """
    Interpolate a list of k-factor  values from
    kFacsVector = [[x0,[k1,k2,..]], [x1,[k1,k2,..],...].

    :returns: list of interpolated k-factor values at x-value xval

    """
    kFacs = []

    xpts = [x[0] for x in kFacsVector]
    ypts = [x[1] for x in kFacsVector]
    coeffs = numpy.matrix.transpose(numpy.polyfit(xpts, ypts, len(xpts) - 1))
    for ik in range(len(ypts[0])):
        kfac = 0.
        for ip, coeff in enumerate(coeffs[ik]):
            kfac += coeff * xval ** (len(xpts) - 1 - ip)
        if kfac <= 0.:
            kfac = 1.
        kFacs.append(kfac)

    return kFacs

def getDecoupledKfactors(nllpath,process,energy,pdf,mass):
    """
    Compute k-factors in the decoupled (squark or gluino) regime for the process.
    If a decoupled grid does not exist for the process, return None
    """

    if process != 'sb' and process != 'gg': return None
    elif process == 'sb': process_dcpl = 'sdcpl'
    elif process == 'gg': process_dcpl = 'gdcpl'
    nll_run = "./nllfast_" + energy + " %s %s %s" % \
                      (process_dcpl, pdf, mass)
    nll_output = runNLLfast(nll_run, nllpath)
    if "K_NLO" in nll_output:
        return getKfactorsFrom(nll_output)
    else: return None

if __name__ == "__main__":
    """
    Calculate k factors for a pid pair.

    """
    slhaF = installation.installDirectory() + "inputFiles/slha/T1.slha"
    kNLO, kNLL = getKfactorsFor((1000021, 1000021), 13.*TeV, slhaF)
    print("nlo, nll = " + str(kNLO) + ", " + str(kNLL))


# Licensed under a 3-clause BSD style license - see LICENSE.rst
import numpy as np
from astropy.utils import lazyproperty
import astropy.units as u
from ..utils.fitting import Fit, Parameters
from ..stats import cash
from ..maps import Map, MapAxis

__all__ = ["MapEvaluator", "MapDataset"]

class MapDataset:
    """Perform sky model likelihood fit on maps.

    Parameters
    ----------
    model : `~gammapy.cube.models.SkyModel`
        Fit model
    counts : `~gammapy.maps.WcsNDMap`
        Counts cube
    exposure : `~gammapy.maps.WcsNDMap`
        Exposure cube
    mask : `~gammapy.maps.WcsNDMap`
        Mask to apply to the likelihood.
    psf : `~gammapy.cube.PSFKernel`
        PSF kernel
    edisp : `~gammapy.irf.EnergyDispersion`
        Energy dispersion
    background_model: `~gammapy.cube.models.BackgroundModel`
        Background model to use for the fit.
    likelihood : {"cash"}
	    Likelihood function to use for the fit.
    """

    def __init__(
        self,
        model,
        counts=None,
        exposure=None,
        mask=None,
        psf=None,
        edisp=None,
        background_model=None,
    ):
        if mask is not None and mask.data.dtype != np.dtype("bool"):
            raise ValueError("mask data must have dtype bool")

        self.model = model
        self.counts = counts
        self.exposure = exposure
        self.mask = mask
        self.psf = psf
        self.edisp = edisp
        self.background_model = background_model
        if background_model:
            self.parameters = Parameters(
                self.model.parameters.parameters +
                self.background_model.parameters.parameters
            )
        else:
            self.parameters = Parameters(self.model.parameters.parameters)

        self.evaluator = MapEvaluator(
            model=self.model, exposure=exposure, psf=self.psf, edisp=self.edisp
        )

    @property
    def data_shape(self):
        """Shape of the counts data"""
        return self.counts.data.shape

    def npred(self):
        """Returns npred map (model + background)"""
        model_npred = self.evaluator.compute_npred()
        back_npred = self.background_model.evaluate()
        total_npred = model_npred.data + back_npred.data
        return back_npred.copy(data=total_npred)
        # TODO: return model_npred + back_npred
        # There is some bug: edisp.e_reco.unit is dimensionless
        # thus map arithmetic does not work.

    def likelihood_per_bin(self):
        """Likelihood per bin given the current model parameters"""
        return cash(n_on=self.counts.data, mu_on=self.npred().data)

    def likelihood(self, parameters, mask=None):
        """Total likelihood given the current model parameters.

        Parameters
        ----------
        mask : `~numpy.ndarray`
            Mask to be combined with the dataset mask.
        """
        if self.mask is None and mask is None:
            stat = self.likelihood_per_bin()
        elif self.mask is None:
            stat = self.likelihood_per_bin()[mask]
        elif mask is None:
            stat = self.likelihood_per_bin()[self.mask.data]
        else:
            stat = self.likelihood_per_bin()[mask & self.mask.data]
        return np.sum(stat, dtype=np.float64)

class MapEvaluator:
    """Sky model evaluation on maps.

    This evaluates a sky model on a 3D map and convolves with the IRFs,
    and returns a map of the predicted counts.
    Note that background counts are not added.

    For now, we only make it work for 3D WCS maps with an energy axis.
    No HPX, no other axes, those can be added later here or via new
    separate model evaluator classes.

    Parameters
    ----------
    model : `~gammapy.cube.models.SkyModel`
        Sky model
    exposure : `~gammapy.maps.Map`
        Exposure map
    background : `~gammapy.maps.Map`
        Background map
    psf : `~gammapy.cube.PSFKernel`
        PSF kernel
    edisp : `~gammapy.irf.EnergyDispersion`
        Energy dispersion
    """

    def __init__(self, model=None, exposure=None, psf=None, edisp=None):
        self.model = model
        self.exposure = exposure
        self.psf = psf
        self.edisp = edisp

        self.parameters = Parameters(self.model.parameters.parameters)

    @lazyproperty
    def geom(self):
        """This will give the energy axes in e_true"""
        return self.exposure.geom

    @lazyproperty
    def geom_image(self):
        return self.geom.to_image()

    @lazyproperty
    def energy_center(self):
        """True energy axis bin centers (`~astropy.units.Quantity`)"""
        energy_axis = self.geom.get_axis_by_name("energy")
        energy = energy_axis.center * energy_axis.unit
        return energy[:, np.newaxis, np.newaxis]

    @lazyproperty
    def energy_edges(self):
        """Energy axis bin edges (`~astropy.units.Quantity`)"""
        energy_axis = self.geom.get_axis_by_name("energy")
        energy = energy_axis.edges * energy_axis.unit
        return energy[:, np.newaxis, np.newaxis]

    @lazyproperty
    def energy_bin_width(self):
        """Energy axis bin widths (`astropy.units.Quantity`)"""
        return np.diff(self.energy_edges, axis=0)

    @lazyproperty
    def lon_lat(self):
        """Spatial coordinate pixel centers.

        Returns ``lon, lat`` tuple of `~astropy.units.Quantity`.
        """
        lon, lat = self.geom_image.get_coord()
        return (u.Quantity(lon, "deg", copy=False), u.Quantity(lat, "deg", copy=False))

    @lazyproperty
    def lon(self):
        return self.lon_lat[0]

    @lazyproperty
    def lat(self):
        return self.lon_lat[1]

    @lazyproperty
    def solid_angle(self):
        """Solid angle per pixel"""
        return self.geom.solid_angle()

    @lazyproperty
    def bin_volume(self):
        """Map pixel bin volume (solid angle times energy bin width)."""
        omega = self.solid_angle
        de = self.energy_bin_width
        return omega * de

    def compute_dnde(self):
        """Compute model differential flux at map pixel centers.

        Returns
        -------
        model_map : `~gammapy.maps.Map`
            Sky cube with data filled with evaluated model values.
            Units: ``cm-2 s-1 TeV-1 deg-2``
        """
        coord = (self.lon, self.lat, self.energy_center)
        dnde = self.model.evaluate(*coord)
        return dnde

    def compute_flux(self):
        """Compute model integral flux over map pixel volumes.

        For now, we simply multiply dnde with bin volume.
        """
        dnde = self.compute_dnde()
        volume = self.bin_volume
        flux = dnde * volume
        return flux

    def apply_exposure(self, flux):
        """Compute npred cube

        For now just divide flux cube by exposure
        """
        npred = (flux * self.exposure.quantity).to_value("")
        return self.exposure.copy(data=npred)

    def apply_psf(self, npred):
        """Convolve npred cube with PSF"""
        return npred.convolve(self.psf)

    def apply_edisp(self, npred):
        """Convolve map data with energy dispersion.

        Parameters
        ----------
        npred : `~gammapy.maps.Map`
            Predicted counts in true energy bins

        Returns
        ---------
        npred_reco : `~gammapy.maps.Map`
            Predicted counts in reco energy bins
        """
        loc = npred.geom.get_axis_index_by_name("energy")
        data = np.rollaxis(npred.data, loc, len(npred.data.shape))
        data = np.dot(data, self.edisp.pdf_matrix)
        data = np.rollaxis(data, -1, loc)
        e_reco_axis = MapAxis.from_edges(
            self.edisp.e_reco.bins, unit=self.edisp.e_reco.unit, name="energy"
        )
        geom_ereco = self.exposure.geom.to_image().to_cube(axes=[e_reco_axis])
        npred = Map.from_geom(geom_ereco, unit="")
        npred.data = data
        return npred

    def compute_npred(self):
        """
        Evaluate model predicted counts.

        Returns
        -------
        npred : `~gammapy.maps.Map`
            Predicted counts on the map (in reco energy bins)
        """
        flux = self.compute_flux()
        npred = self.apply_exposure(flux)
        if self.psf is not None:
            npred = self.apply_psf(npred)
        if self.edisp is not None:
            npred = self.apply_edisp(npred)
        return npred

import numpy,import astropy,from astropy


# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
import warnings
from copy import deepcopy

import numpy as np

from qutip import (tensor, identity, destroy, sigmax, sigmaz, basis, Qobj,
                    QobjEvo)
from ..circuit import QubitCircuit
from ..operations import Gate
from .processor import Processor
from .modelprocessor import ModelProcessor
from ..operations import expand_operator
from ..pulse import Pulse
from ..compiler import (GateCompiler, CavityQEDCompiler)

__all__ = ['DispersiveCavityQED']

class DispersiveCavityQED(ModelProcessor):
    """
    The processor based on the physical implementation of
    a dispersive cavity QED system.
    The available Hamiltonian of the system is predefined.
    For a given pulse amplitude matrix, the processor can
    calculate the state evolution under the given control pulse,
    either analytically or numerically.
    (Only additional attributes are documented here, for others please
    refer to the parent class :class:`.ModelProcessor`)

    Parameters
    ----------
    num_qubits: int
        The number of qubits in the system.

    correct_global_phase: float, optional
        Save the global phase, the analytical solution
        will track the global phase.
        It has no effect on the numerical solution.

    num_levels: int, optional
        The number of energy levels in the resonator.

    t1: list or float, optional
        Characterize the decoherence of amplitude damping for
        each qubit. A list of size `num_qubits` or a float for all qubits.

    t2: list of float, optional
        Characterize the decoherence of dephasing for
        each qubit. A list of size `num_qubits` or a float for all qubits.

    **params:
        Keyword argument for hardware parameters, in the unit of GHz.
        Qubit parameters can either be a float or a list of the length
        ``num_qubits``.

        - ``deltamax``: the pulse strength of sigma-x control, default ``1.0``
        - ``epsmax``: the pulse strength of sigma-z control, default ``9.5``
        - ``eps``: the bare transition frequency for each of the qubits,
          default ``9.5``
        - ``delta``: the coupling between qubit states, default ``0.0``
        - ``g``: the coupling strength between the resonator and the qubit,
          default ``1.0``
        - ``w0``: the bare frequency of the resonator. Should only be a float,
          default ``0.01``

        The dressed qubit frequency is `wq` is computed by
        :math:`w_q=\sqrt{\epsilon^2+\delta^2}`

    Attributes
    ----------
    wq: list of float
        The frequency of the qubits calculated from
        eps and delta for each qubit.

    Delta: list of float
        The detuning with respect to w0 calculated
        from wq and w0 for each qubit.
    """

    def __init__(self, num_qubits, correct_global_phase=True,
                 num_levels=10, t1=None, t2=None, **params):
        super(DispersiveCavityQED, self).__init__(
            num_qubits, correct_global_phase=correct_global_phase,
            t1=t1, t2=t2)
        self.correct_global_phase = correct_global_phase
        self.spline_kind = "step_func"
        self.num_levels = num_levels
        self.params = {  # default parameters
            "deltamax": 1.0,
            "epsmax": 9.5,
            "w0": 10,
            "eps": 9.5,
            "delta": 0.0,
            "g": 0.01,
        }
        if params is not None:
            self.params.update(params)
        for key, value in self.params.items():
            if key != "w0":
                # if float, make it an array
                self.params[key] = self.to_array(value, self.num_qubits)
        self.set_up_params()
        self.set_up_ops(num_qubits)
        self.dims = [num_levels] + [2] * num_qubits
        self.pulse_dict = self.get_pulse_dict()
        self.native_gates = ["SQRTISWAP", "ISWAP", "RX", "RZ"]

    def set_up_ops(self, num_qubits):
        """
        Generate the Hamiltonians for the spinchain model and save them in the
        attribute `ctrls`.

        Parameters
        ----------
        num_qubits: int
            The number of qubits in the system.
        """
        # single qubit terms
        for m in range(num_qubits):
            self.add_control(2*np.pi*sigmax(), [m+1], label="sx" + str(m))
        for m in range(num_qubits):
            self.add_control(2*np.pi*sigmaz(), [m+1], label="sz" + str(m))
        # coupling terms
        a = tensor(
            [destroy(self.num_levels)] +
            [identity(2) for n in range(num_qubits)])
        for n in range(num_qubits):
            sm = tensor([identity(self.num_levels)] +
                        [destroy(2) if m == n else identity(2)
                         for m in range(num_qubits)])
            self.add_control(
                2*np.pi * a.dag() * sm + 2*np.pi * a * sm.dag(),
                list(range(num_qubits+1)), label="g" + str(n)
            )

    def set_up_params(self):
        """
        Compute the qubit frequency and detune.
        """
        # backward compatibility
        self.params["sz"] = self.params["epsmax"]
        self.params["sx"] = self.params["deltamax"]
        eps = self.params["eps"]
        delta = self.params["delta"]
        w0 = self.params["w0"]
        g = self.params["g"]
        # computed
        self.wq = np.sqrt(eps**2 + delta**2)
        self.Delta = self.wq - w0

        # rwa/dispersive regime tests
        if any(g / (w0 - self.wq) > 0.05):
            warnings.warn("Not in the dispersive regime")

        if any((w0 - self.wq)/(w0 + self.wq) > 0.05):
            warnings.warn(
                "The rotating-wave approximation might not be valid.")
        print(self.params)

    @property
    def sx_ops(self):
        """
        list: A list of sigmax Hamiltonians for each qubit.
        """
        return self.ctrls[0: self.num_qubits]

    @property
    def sz_ops(self):
        """
        list: A list of sigmaz Hamiltonians for each qubit.
        """
        return self.ctrls[self.num_qubits: 2*self.num_qubits]

    @property
    def cavityqubit_ops(self):
        """
        list: A list of interacting Hamiltonians between cavity and each qubit.
        """
        return self.ctrls[2*self.num_qubits: 3*self.num_qubits]

    @property
    def sx_u(self):
        """array-like: Pulse matrix for sigmax Hamiltonians."""
        return self.coeffs[: self.num_qubits]

    @property
    def sz_u(self):
        """array-like: Pulse matrix for sigmaz Hamiltonians."""
        return self.coeffs[self.num_qubits: 2*self.num_qubits]

    @property
    def g_u(self):
        """
        array-like: Pulse matrix for interacting Hamiltonians
        between cavity and each qubit.
        """
        return self.coeffs[2*self.num_qubits: 3*self.num_qubits]

    def get_operators_labels(self):
        """
        Get the labels for each Hamiltonian.
        It is used in the method method :meth:`.Processor.plot_pulses`.
        It is a 2-d nested list, in the plot,
        a different color will be used for each sublist.
        """
        return ([[r"$\sigma_x^%d$" % n for n in range(self.num_qubits)],
                 [r"$\sigma_z^%d$" % n for n in range(self.num_qubits)],
                 [r"$g_{%d}$" % (n) for n in range(self.num_qubits)]])

    def eliminate_auxillary_modes(self, U):
        """
        Eliminate the auxillary modes like the cavity modes in cqed.
        """
        psi_proj = tensor(
            [basis(self.num_levels, 0)] +
            [identity(2) for n in range(self.num_qubits)])
        return psi_proj.dag() * U * psi_proj

    def load_circuit(
            self, qc, schedule_mode="ASAP", compiler=None):
        if compiler is None:
            compiler = CavityQEDCompiler(
                self.num_qubits, self.params, global_phase=0.)
        tlist, coeff = super().load_circuit(
            qc, schedule_mode=schedule_mode, compiler=compiler)
        self.global_phase = compiler.global_phase
        return tlist, coeff


import warnings

import astropy.units as u
import numpy as np
from astropy.time import Time

from sora.config.decorators import deprecated_function, deprecated_alias
from sora.prediction import occ_params, PredictionTable
from .fitting import fit_ellipse as ellipse_fitting

__all__ = ['Occultation']
warnings.simplefilter('always', UserWarning)

class Occultation:
    """Instantiates the Occultation object and performs the reduction of the
    occultation.

    Attributes
    ----------
    star : `sora.Star`, `str`, required
        the coordinate of the star in the same reference frame as the ephemeris.
        It must be a Star object or a string with the coordinates of the object
        to search on Vizier.

    body : `sora.Body`, `str`
        Object that will occult the star. It must be a Body object or its name
        to search in the Small Body Database.

    ephem : `sora.Ephem`, `list`
        Object ephemeris. It must be an Ephemeris object or a list.

    time : `str`, `astropy.time.Time`, required
        Reference time of the occultation. Time does not need to be exact, but
        needs to be within approximately 50 minutes of the occultation closest
        approach to calculate occultation parameters.

    reference_center : `str`, `sora.Observer`, `sora.Spacecraft`
        A SORA observer object or a string 'geocenter'.
        The occultation parameters will be calculated in respect
        to this reference as center of projection.

    Important
    ---------
    When instantiating with "body" and "ephem", the user may define the
    Occultation in 3 ways:

    1. With `body` and `ephem`.

    2. With only "body". In this case, the "body" parameter must be a Body
    object and have an ephemeris associated (see Body documentation).

    3. With only `ephem`. In this case, the `ephem` parameter must be one of the
    Ephem Classes and have a name (see Ephem documentation) to search for the
    body in the Small Body Database.

    """
    def __init__(self, star, body=None, ephem=None, time=None, reference_center='geocenter'):

        from sora.body import Body
        from sora.star import Star
        from .chordlist import ChordList

        if body is None and ephem is None:
            raise ValueError('"body" and/or "ephem" must be given.')
        if time is None:
            raise ValueError('"time" parameter must be given.')
        if isinstance(star, str):
            star = Star(coord=star)
        elif not isinstance(star, Star):
            raise ValueError('"star" must be a Star object or a string with coordinates of the star')
        self._star = star
        if body is not None:
            if not isinstance(body, (str, Body)):
                raise ValueError('"body" must be a string with the name of the object or a Body object')
            if isinstance(body, str):
                body = Body(name=body)
            self._body = body
        if ephem is not None:
            if body is not None:
                self.body.ephem = ephem
            else:
                if hasattr(ephem, 'name'):
                    self._body = Body(name=ephem.name, ephem=ephem)
                else:
                    raise ValueError('When only "ephem" is given, "ephem" must have a name for search.')
        try:
            ephem = self.body.ephem
        except AttributeError:
            raise ValueError('An Ephem object must be defined in Body object.')
        self._reference_center = reference_center

        tca, ca, pa, vel, dist = occ_params(self.star, ephem, time, reference_center=reference_center)
        self.ca = ca   # Closest Approach distance
        self.pa = pa   # Position Angle at CA
        self.vel = vel  # Shadow velocity at CA
        self.dist = dist  # object distance at CA
        self.tca = tca   # Instant of CA
        self.star_diam = self.star.apparent_diameter(self.dist, verbose=False)

        meta = {
            'name': self.body.name, 'radius': self.body.radius.to(u.km).value,
            'error_ra': self.body.ephem.error_ra.to(u.mas).value, 'error_dec': self.body.ephem.error_dec.to(u.mas).value}
        self.predict = PredictionTable(
            time=[tca], coord_star=[self.star.get_position(tca, observer=reference_center)],
            coord_obj=[self.body.ephem.get_position(tca, observer=reference_center)],
            ca=[ca.value], pa=[pa.value], vel=[vel.value], dist=[dist.value],
            mag=[self.star.mag['G']], source=[self.star.code], meta=meta)

        self.__observations = []
        self._chords = ChordList(star=self.star, body=self._body, time=self.tca)
        self._chords._shared_with['occultation'] = {"vel": np.absolute(self.vel), "dist": float(self.dist.AU),
                                                    "star_diam": float(self.star_diam.km)}

    @property
    def star(self):
        return self._star

    @property
    def body(self):
        return self._body

    @property
    def chords(self):
        return self._chords

    # remove this block for v1.0
    @deprecated_function(message="Please use chords.add_chord to add new observations.")
    def add_observation(self, obs, lightcurve):
        """Adds observations to the Occultation object.

        Parameters
        ----------
        obs : `sora.Observer`
            The Observer object to be added.

        lightcurve : `sora.LightCurve`
            The LightCurve object to be added.
        """
        self.chords.add_chord(name=lightcurve.name, observer=obs, lightcurve=lightcurve)

    @deprecated_function(message="Please use chords.remove_chord to remove observations.")
    def remove_observation(self, key, key_lc=None):
        """Removes an observation from the Occultation object.

        Parameters
        ----------
        key : `str`
            The name given to `Observer` or `LightCurve` to remove from the list.

        key_lc : `str`
            In the case where repeated names are present for different
            observations, `key_lc` must be given for the name of the `LightCurve`
            and key will be used for the name of the `Observer`.
        """
        try:
            self.chords.remove_chord(name=key)
        except KeyError:
            rm_list = []
            for name, chord in self.chords.items():
                if chord.observer.name == key and (key_lc is None or chord.lightcurve.name == key_lc):
                    rm_list.append(name)
            if len(rm_list) != 1:
                if len(rm_list) == 0:
                    err_mes = "No observation was identified with given keys."
                else:
                    err_mes = "More than one chord was identified. Please provide unique keys."
                raise ValueError(err_mes)
            else:
                self.chords.remove_chord(name=rm_list[0])

    @deprecated_function(message="Please use chords.")
    def observations(self):
        """ Print all the observations added to the Occultation object
        Pair (`Observer`, `LightCurve`)
        """
        print(self.chords.__repr__())

    # end of block removal

    def fit_ellipse(self, **kwargs):
        chisquare = ellipse_fitting(self, **kwargs)
        return chisquare

    fit_ellipse.__doc__ = ellipse_fitting.__doc__

    # remove this block for v1.0
    @property
    @deprecated_function(message="Please use chords.summary()")
    def positions(self):
        """ Calculates the position and velocity for all chords.
        Saves it into an `_PositionDict` object.
        """
        from functools import partial
        from .meta import _PositionDict

        if not hasattr(self, '_position'):
            self._position = _PositionDict()
        position = self._position
        if len(self.chords) == 0:
            raise ValueError('There is no observation defined for this occultation')

        pair = []
        for name, chord in self.chords.items():
            obs = chord.observer
            obs_name = obs.name
            if obs_name == '':
                obs_name = '{}_obs'.format(chord.name)
            lc = chord.lightcurve
            lc_name = lc.name
            if lc_name == '':
                lc_name = '{}_lc'.format(chord.name)
            pair.append((obs_name, lc_name))

            coord = [obs.lon, obs.lat, obs.height]
            if obs.name not in position.keys():
                position['_occ_'+obs_name] = _PositionDict(lon=obs.lon, lat=obs.lat, height=obs.height)
                position[obs_name]['_occ_lon'] = obs.lon
                position[obs_name]['_occ_lat'] = obs.lat
                position[obs_name]['_occ_height'] = obs.height
            pos_obs = position[obs_name]
            coord2 = [pos_obs['lon'], pos_obs['lat'], pos_obs['height']]
            if obs.lon != pos_obs['lon']:
                position[obs_name]['_occ_lon'] = obs.lon
            if obs.lat != pos_obs['lat']:
                position[obs_name]['_occ_lat'] = obs.lat
            if obs.height != pos_obs['height']:
                position[obs_name]['_occ_height'] = obs.height
            samecoord = (coord == coord2)

            if lc_name not in pos_obs.keys():
                pos_obs['_occ_'+lc_name] = _PositionDict()
            pos_lc = pos_obs[lc_name]

            pos_lc['_occ_status'] = chord.status()

            if hasattr(lc, 'immersion'):
                if 'immersion' not in pos_lc.keys():
                    pos_lc['_occ_immersion'] = _PositionDict(on=chord.is_able['immersion'],
                                                             enable=partial(chord.enable, time='immersion'),
                                                             disable=partial(chord.disable, time='immersion'))
                obs_im = pos_lc['immersion']
                obs_im['_occ_on'] = chord.is_able['immersion']
                do_err = False
                if samecoord and 'time' in obs_im.keys() and obs_im['time'] == lc.immersion:
                    pass
                else:
                    do_err = True
                    f1, g1, vf1, vg1 = chord.get_fg(time='immersion', vel=True)
                    obs_im['_occ_time'] = lc.immersion
                    obs_im['_occ_value'] = (round(f1, 3), round(g1, 3))
                    obs_im['_occ_vel'] = (round(vf1, 3), round(vg1, 3))
                if not do_err and 'time_err' in obs_im.keys() and obs_im['time_err'] == lc.immersion_err:
                    pass
                else:
                    fe1, ge1 = chord.get_fg(time=lc.immersion-lc.immersion_err*u.s)
                    fe2, ge2 = chord.get_fg(time=lc.immersion+lc.immersion_err*u.s)
                    obs_im['_occ_time_err'] = lc.immersion_err
                    obs_im['_occ_error'] = ((round(fe1, 3), round(ge1, 3)), (round(fe2, 3), round(ge2, 3)))

            if hasattr(lc, 'emersion'):
                pos_lc['_occ_emersion'] = _PositionDict(on=chord.is_able['emersion'],
                                                        enable=partial(chord.enable, time='emersion'),
                                                        disable=partial(chord.disable, time='emersion'))
                obs_em = pos_lc['emersion']
                do_err = False
                if samecoord and 'time' in obs_em.keys() and obs_em['time'] == lc.emersion:
                    pass
                else:
                    do_err = True
                    f1, g1, vf1, vg1 = chord.get_fg(time='emersion', vel=True)
                    obs_em['_occ_time'] = lc.emersion
                    obs_em['_occ_value'] = (round(f1, 3), round(g1, 3))
                    obs_em['_occ_vel'] = (round(vf1, 3), round(vg1, 3))
                if not do_err and 'time_err' in obs_em.keys() and obs_em['time_err'] == lc.emersion_err:
                    pass
                else:
                    fe1, ge1 = chord.get_fg(time=lc.emersion-lc.emersion_err*u.s)
                    fe2, ge2 = chord.get_fg(time=lc.emersion+lc.emersion_err*u.s)
                    obs_em['_occ_time_err'] = lc.emersion_err
                    obs_em['_occ_error'] = ((round(fe1, 3), round(ge1, 3)), (round(fe2, 3), round(ge2, 3)))

            if pos_lc['status'] == 'negative':
                if 'start_obs' not in pos_lc.keys():
                    pos_lc['_occ_start_obs'] = _PositionDict()
                obs_start = pos_lc['start_obs']
                if samecoord and 'time' in obs_start.keys() and obs_start['time'] == lc.initial_time:
                    pass
                else:
                    f, g, vf, vg = chord.get_fg(time='start', vel=True)
                    obs_start['_occ_time'] = lc.initial_time
                    obs_start['_occ_value'] = (round(f, 3), round(g, 3))
                    obs_start['_occ_vel'] = (round(vf, 3), round(vg, 3))
                if 'end_obs' not in pos_lc.keys():
                    pos_lc['_occ_end_obs'] = _PositionDict()
                obs_end = pos_lc['end_obs']
                if samecoord and 'time' in obs_end.keys() and obs_end['time'] == lc.end_time:
                    pass
                else:
                    f, g, vf, vg = chord.get_fg(time='end', vel=True)
                    obs_end['_occ_time'] = lc.end_time
                    obs_end['_occ_value'] = (round(f, 3), round(g, 3))
                    obs_end['_occ_vel'] = (round(vf, 3), round(vg, 3))

        for key in list(position):
            n = 0
            for key_lc in list(position[key]):
                if type(position[key][key_lc]) != _PositionDict:
                    continue
                if (key, key_lc) not in pair:
                    del position[key][key_lc]
                else:
                    n += 1
            if n == 0:
                del position[key]

        return self._position

    @positions.setter
    def positions(self, value):
        """
        Note
        ----
        If the users tries to set a value to position, it must be ``'on'`` or
        ``'off'``, and it will be assigned to all chords.
        """
        if value not in ['on', 'off']:
            raise ValueError("Value must be 'on' or 'off' only.")
        pos = self.positions
        for key in pos.keys():
            pos[key] = value

    # end of block removal

    def check_velocities(self):
        """Prints the current velocity used by the LightCurves and its radial velocity.
        """
        if hasattr(self, 'fitted_params'):
            center = np.array([self.fitted_params['center_f'][0], self.fitted_params['center_g'][0]])
        else:
            center = np.array([0, 0])
        for name, chord in self.chords.items():
            im = getattr(chord.lightcurve, 'immersion', None)
            em = getattr(chord.lightcurve, 'emersion', None)
            if im is None and em is None:
                continue
            print('{} - Velocity used: {:.3f}'.format(name, chord.lightcurve.vel))
            if im is not None:
                vals = chord.get_fg(time=im, vel=True)
                delta = np.array(vals[0:2]) - center
                print('    Immersion Radial Velocity: {:.3f}'.
                      format(np.abs(np.dot(np.array(vals[2:]), delta)/np.linalg.norm(delta))))
            if em is not None:
                vals = chord.get_fg(time=em, vel=True)
                delta = np.array(vals[0:2]) - center
                print('    Emersion Radial Velocity: {:.3f}'.
                      format(np.abs(np.dot(np.array(vals[2:]), delta)/np.linalg.norm(delta))))

    @deprecated_alias(log='verbose')  # remove this line in v1.0
    def new_astrometric_position(self, time=None, offset=None, error=None, verbose=True, observer=None):
        """Calculates the new astrometric position for the object given fitted parameters.

        Parameters
        ----------
        time : `str`, `astropy.time.Time`
            Reference time to calculate the position. If not given, it uses the
            instant of the occultation Closest Approach. It can be a string
            in the ISO format (yyyy-mm-dd hh:mm:ss.s) or an astropy Time object.

        offset : `list`
            Offset to apply to the position. If not given, uses the parameters
            from the fitted ellipse.

            Note
            ----
            Must be a list of 3 values being [X, Y, 'unit']. 'unit' must be
            Angular or Distance unit.

            If Distance units for X and Y:
            Ex: [100, -200, 'km'], [0.001, 0.002, 'AU']

            If Angular units fox X [d*a*cos(dec)] and Y [d*dec]:
            Ex: [30.6, 20, 'mas'], or [-15, 2, 'arcsec']

        error : `list`
            Error bar of the given offset. If not given, it uses the 1-sigma
            value of the fitted ellipse.

            Note
            ----
            Error must be a list of 3 values being [dX, dY, 'unit'], similar to
            offset. It does not need to be in the same unit as offset.

        verbose : `bool`
            If true, it Prints text, else it Returns text.

        observer : `str`, `sora.Observer`, `sora.Spacecraft`
            IAU code of the observer (must be present in given list of kernels),
            a SORA observer object or a string: ['geocenter', 'barycenter']
        """
        from astropy.coordinates import SkyCoord, SkyOffsetFrame

        if time is not None:
            time = Time(time)
        else:
            time = self.tca

        tca_diff = np.absolute(time-self.tca)
        if tca_diff > 1*u.day:
            warnings.warn('The difference between the given time and the closest approach instant is {:.1f} days. '
                          'This position could not have a physical meaning.'.format(tca_diff.jd))

        if offset is not None:
            off_ra = offset[0]*u.Unit(offset[2])
            off_dec = offset[1]*u.Unit(offset[2])
            if off_ra.unit.physical_type == 'length':
                dist = True
            elif off_ra.unit.physical_type == 'angle':
                dist = False
            else:
                raise ValueError('Offset unit must be a distance or angular value.')
        elif hasattr(self, 'fitted_params'):
            off_ra = self.fitted_params['center_f'][0]*u.km
            off_dec = self.fitted_params['center_g'][0]*u.km
            dist = True
        else:
            warnings.warn('No offset given or found. Using 0.0 instead.')
            off_ra = 0.0*u.mas
            off_dec = 0.0*u.mas
            dist = False

        if error is not None:
            e_off_ra = error[0]*u.Unit(error[2])
            e_off_dec = error[1]*u.Unit(error[2])
            if e_off_ra.unit.physical_type == 'length':
                e_dist = True
            elif e_off_ra.unit.physical_type == 'angle':
                e_dist = False
            else:
                raise ValueError('Error unit must be a distance or angular value.')
        elif hasattr(self, 'fitted_params'):
            e_off_ra = self.fitted_params['center_f'][1]*u.km
            e_off_dec = self.fitted_params['center_g'][1]*u.km
            e_dist = True
        else:
            warnings.warn('No error given or found. Using 0.0 instead.')
            e_off_ra = 0.0*u.mas
            e_off_dec = 0.0*u.mas
            e_dist = False

        if observer is None:
            observer = self._reference_center

        coord = self.body.ephem.get_position(time, observer=observer)
        distance = coord.distance.to(u.km)
        coord_frame = SkyOffsetFrame(origin=coord)
        if dist:
            off_ra = np.arctan2(off_ra, distance)
            off_dec = np.arctan2(off_dec, distance)
        if e_dist:
            e_off_ra = np.arctan2(e_off_ra, distance)
            e_off_dec = np.arctan2(e_off_dec, distance)
        new_pos = SkyCoord(lon=off_ra, lat=off_dec, frame=coord_frame)
        new_pos = new_pos.icrs

        error_star = self.star.error_at(self.tca)
        error_ra = np.sqrt(error_star[0]**2 + e_off_ra**2)
        error_dec = np.sqrt(error_star[1]**2 + e_off_dec**2)

        out = 'Ephemeris offset (km): X = {:.1f} +/- {:.1f}; Y = {:.1f} +/- {:.1f}\n'.format(
            distance*np.sin(off_ra.to(u.mas)).value, distance*np.sin(e_off_ra.to(u.mas)).value,
            distance*np.sin(off_dec.to(u.mas)).value, distance*np.sin(e_off_dec.to(u.mas)).value)
        out += 'Ephemeris offset (mas): da_cos_dec = {:.3f} +/- {:.3f}; d_dec = {:.3f} +/- {:.3f}\n'.format(
            off_ra.to(u.mas).value, e_off_ra.to(u.mas).value, off_dec.to(u.mas).value, e_off_dec.to(u.mas).value)
        out += '\nAstrometric object position at time {} for reference {}\n'.format(time.iso, observer.__repr__())
        out += 'RA = {} +/- {:.3f} mas; DEC = {} +/- {:.3f} mas'.format(
            new_pos.ra.to_string(u.hourangle, precision=7, sep=' '), error_ra.to(u.mas).value,
            new_pos.dec.to_string(u.deg, precision=6, sep=' '), error_dec.to(u.mas).value)

        if verbose:
            print(out)
        else:
            return out

    # remove this block for v1.0
    @deprecated_function(message="Please use chords.plot_chords to have a better control of the plots")
    def plot_chords(self, all_chords=True, positive_color='blue', negative_color='green', error_color='red',
                    ax=None, lw=2):
        """Plots the chords of the occultation.

        Parameters
        ----------
        all_chords : `bool`, default=True
            If True, it plots all the chords. If False, it sees what was
            deactivated in self.positions and ignores them.

        positive_color : `str`, default='blue'
            Color for the positive chords.

        negative_color : `str`, default='green'
            Color for the negative chords.

        error_color : `str`, default='red'
            Color for the error bars of the chords.

        ax : `maptlotlib.pyplot.Axes`, default=None
            Axis where to plot chords (default: Uses matplotlib pool).

        lw : `int`, `float`, default=2
            Linewidth of the chords.

        """
        self.chords.plot_chords(segment='positive', only_able=not all_chords, color=positive_color, lw=lw, ax=ax)
        self.chords.plot_chords(segment='error', only_able=not all_chords, color=error_color, lw=lw, ax=ax)
        self.chords.plot_chords(segment='negative', color=negative_color, lw=lw, ax=ax, linestyle='--')

    # end of block removal

    def get_map_sites(self):
        """Returns Dictionary with sites in the format required by plot_occ_map function.

        Returns
        -------
        sites : `dict`
            Dictionary with the sites in the format required by `plot_occ_map`
            function.
        """
        sites = {}
        color = {'positive': 'blue', 'negative': 'red'}
        for name, chord in self.chords.items():
            obs = chord.observer
            sites[name] = [obs.lon.deg, obs.lat.deg, 10, 10, color[chord.status()]]
        return sites

    def plot_occ_map(self, **kwargs):
        if 'radius' not in kwargs and hasattr(self, 'fitted_params'):
            r_equa = self.fitted_params['equatorial_radius'][0]
            obla = self.fitted_params['oblateness'][0]
            pos_ang = self.fitted_params['position_angle'][0]
            theta = np.linspace(-np.pi, np.pi, 1800)
            map_pa = self.predict['P/A'][0]
            circle_x = r_equa*np.cos(theta)
            circle_y = r_equa*(1.0-obla)*np.sin(theta)
            ellipse_y = -circle_x*np.sin((pos_ang-map_pa)*u.deg) + circle_y*np.cos((pos_ang-map_pa)*u.deg)
            kwargs['radius'] = ellipse_y.max()
            print('Projected shadow radius = {:.1f} km'.format(kwargs['radius']))
        kwargs['sites'] = kwargs.get('sites', self.get_map_sites())
        if 'offset' not in kwargs and hasattr(self, 'fitted_params'):
            off_ra = self.fitted_params['center_f'][0]*u.km
            off_dec = self.fitted_params['center_g'][0]*u.km
            off_ra = np.arctan2(off_ra, self.dist)
            off_dec = np.arctan2(off_dec, self.dist)
            kwargs['offset'] = [off_ra, off_dec]
        self.predict.plot_occ_map(**kwargs)

    plot_occ_map.__doc__ = PredictionTable.plot_occ_map.__doc__

    def to_log(self, namefile=None):
        """Saves the occultation log to a file.

        Parameters
        ----------
        namefile : `str`
            Filename to save the log.
        """
        if namefile is None:
            namefile = 'occ_{}_{}.log'.format(self.body.shortname.replace(' ', '_'), self.tca.isot[:16])
        f = open(namefile, 'w')
        f.write(self.__str__())
        f.close()

    def check_time_shift(self, time_interval=30, time_resolution=0.001, verbose=False, plot=False, use_error=True, delta_plot=100,
                         ignore_chords=None):
        """Check the needed time offset, so all chords have their center aligned.

        Parameters
        ----------
        time_interval : `int`, `float`
            Time interval to check, default is 30 seconds.

        time_resolution : `int`, `float`
            Time resolution of the search, default is 0.001 seconds.

        verbose : `bool`
            If True, it prints text, default is False.

        plot : `bool`, default=False
            If True, it plots figures as a visual aid.

        use_error : `bool`, default=True
            if True, the linear fit considers the time uncertainty.

        delta_plot : `int`, `float`, default=100
            Value to be added to increase the plot limit, in km.

        ignore_chords : `str`, list, default=None
            Names of the chords to be ignored in the linear fit.

        Returns
        -------
        time_shift : `dict`
            Dictionary with needed time offset to align the chords, each key is
            the name of the chord.
        """
        import matplotlib.pyplot as plt
        fm = np.array([])
        gm = np.array([])
        dfm = np.array([])
        dgm = np.array([])
        vfm = np.array([])
        vgm = np.array([])
        chord_name = np.array([])
        ignore_chords = np.array(ignore_chords, ndmin=1)
        use_chords = np.array([], dtype=bool)
        out_dic = {}

        chords = range(len(self.chords))
        for i in chords:
            chord = self.chords[i]
            if chord.status() == 'positive':
                ffm, ggm, vffm, vggm = chord.get_fg(time=chord.lightcurve.time_mean, vel=True)
                df1, dg1, df2, dg2 = chord.path(segment='error')
                dfm = np.append(dfm, np.sqrt(((df1.max() - df1.min())/2)**2 + ((df2.max() - df2.min())/2)**2))
                dgm = np.append(dgm, np.sqrt(((dg1.max() - dg1.min())/2)**2 + ((dg2.max() - dg2.min())/2)**2))
                fm = np.append(fm, ffm)
                gm = np.append(gm, ggm)
                vfm = np.append(vfm, vffm)
                vgm = np.append(vgm, vggm)
                chord_name = np.append(chord_name, chord.name)
                if chord.name in ignore_chords:
                    use_chords = np.append(use_chords, False)
                else:
                    use_chords = np.append(use_chords, True)
        if len(fm[use_chords]) < 2:
            raise ValueError('The number of fitted chords should be higher than two')
        out = self.__linear_fit_error(x=fm[use_chords], y=gm[use_chords], sx=dfm[use_chords], sy=dgm[use_chords],
                                      verbose=verbose, use_error=use_error)
        fm_fit = np.arange(-delta_plot+np.min([fm.min(), gm.min()]), delta_plot+np.max([fm.max(), gm.max()]),
                           time_resolution*np.absolute(self.vel.value))
        gm_fit = self.__func_line_decalage(out.beta, fm_fit)

        previous_dt = np.array([])
        fm_decalage = np.array([])
        gm_decalage = np.array([])
        time_decalage = np.array([])

        for i, j in enumerate(chord_name):
            previous_dt = np.append(previous_dt, self.chords[j].lightcurve.dt)
            dist_min = np.array([])
            dtt = np.arange(-time_interval, +time_interval, time_resolution)
            for dt in dtt:
                fm_new = fm[i] + dt*vfm[i]
                gm_new = gm[i] + dt*vgm[i]
                dist = np.sqrt((fm_new - fm_fit)**2 + (gm_new - gm_fit)**2)
                dist_min = np.append(dist_min, dist.min())
            print('dt = {:+6.3f} seconds; chord = {}'.format(previous_dt[i] + dtt[dist_min.argmin()], chord_name[i]))
            fm_decalage = np.append(fm_decalage, fm[i] + dtt[dist_min.argmin()]*vfm[i])
            gm_decalage = np.append(gm_decalage, gm[i] + dtt[dist_min.argmin()]*vgm[i])
            time_decalage = np.append(time_decalage, dtt[dist_min.argmin()])

        for i in range(len(chord_name)):
            out_dic[chord_name[i]] = time_decalage[i]
        if plot:
            plt.figure(figsize=(5, 5))
            plt.title('Before', fontsize=15)
            self.chords.plot_chords(color='blue')
            self.chords.plot_chords(color='red', segment='error')
            plt.plot(fm[use_chords], gm[use_chords], linestyle='None', marker='o', color='k')
            plt.plot(fm[np.invert(use_chords)], gm[np.invert(use_chords)], linestyle='None', marker='x', color='r')
            plt.plot(fm_fit, gm_fit, 'k-')
            plt.xlim(-delta_plot + np.min([fm.min(), gm.min()]), delta_plot + np.max([fm.max(), gm.max()]))
            plt.ylim(-delta_plot + np.min([fm.min(), gm.min()]), delta_plot + np.max([fm.max(), gm.max()]))
            plt.show()
            for i, j in enumerate(chord_name):
                self.chords[j].lightcurve.dt = previous_dt[i] + time_decalage[i]
            plt.figure(figsize=(5, 5))
            plt.title('After', fontsize=15)
            self.chords.plot_chords(color='blue')
            self.chords.plot_chords(color='red', segment='error')
            plt.plot(fm_decalage[use_chords], gm_decalage[use_chords], linestyle='None', marker='o', color='k')
            plt.plot(fm_decalage[np.invert(use_chords)], gm_decalage[np.invert(use_chords)], linestyle='None', marker='x',
                     color='r')
            plt.plot(fm_fit, gm_fit, 'k-')
            plt.xlim(-delta_plot + np.min([fm.min(), gm.min()]), delta_plot + np.max([fm.max(), gm.max()]))
            plt.ylim(-delta_plot + np.min([fm.min(), gm.min()]), delta_plot + np.max([fm.max(), gm.max()]))
            plt.show()
            for i, j in enumerate(chord_name):
                self.chords[j].lightcurve.dt = previous_dt[i]
        return out_dic

    def to_file(self):
        """Saves the occultation data to a file.

        Three files are saved containing the positions and velocities for the
        observations. They are for the positive, negative and error bars positions.

        The format of the files are: positions in f and g, velocities in f and
        g, the Julian Date of the observation, light curve name of the
        corresponding position.
        """
        pos = []
        neg = []
        err = []
        for name, chord in self.chords.items():
            status = chord.status()
            l_name = name.replace(' ', '_')
            if status == 'positive':
                im = chord.lightcurve.immersion
                ime = chord.lightcurve.immersion_err
                f, g, vf, vg = chord.get_fg(time=im, vel=True)
                f1, g1 = chord.get_fg(time=im-ime*u.s)
                f2, g2 = chord.get_fg(time=im+ime*u.s)
                pos.append([f, g, vf, vg, im.jd, l_name+'_immersion'])
                err.append([f1, g1, vf, vg, (im-ime*u.s).jd, l_name+'_immersion_err-'])
                err.append([f2, g2, vf, vg, (im+ime*u.s).jd, l_name+'_immersion_err+'])

                em = chord.lightcurve.emersion
                eme = chord.lightcurve.emersion_err
                f, g, vf, vg = chord.get_fg(time=em, vel=True)
                f1, g1 = chord.get_fg(time=em-eme*u.s)
                f2, g2 = chord.get_fg(time=em+eme*u.s)
                pos.append([f, g, vf, vg, em.jd, l_name+'_emersion'])
                err.append([f1, g1, vf, vg, (em-eme*u.s).jd, l_name+'_emersion_err-'])
                err.append([f2, g2, vf, vg, (em+eme*u.s).jd, l_name+'_emersion_err+'])

            if status == 'negative':
                ini = chord.lightcurve.initial_time
                f, g, vf, vg = chord.get_fg(time=ini, vel=True)
                neg.append([f, g, vf, vg, ini.jd, l_name+'_start'])

                end = chord.lightcurve.end_time
                f, g, vf, vg = chord.get_fg(time=end, vel=True)
                neg.append([f, g, vf, vg, end.jd, l_name+'_end'])
        if len(pos) > 0:
            f = open('occ_{}_pos.txt'.format(self.body.shortname.replace(' ', '_')), 'w')
            for line in pos:
                f.write('{:10.3f} {:10.3f} {:-6.2f} {:-6.2f} {:16.8f} {}\n'.format(*line))
            f.close()
            f = open('occ_{}_err.txt'.format(self.body.shortname.replace(' ', '_')), 'w')
            for line in err:
                f.write('{:10.3f} {:10.3f} {:-6.2f} {:-6.2f} {:16.8f} {}\n'.format(*line))
            f.close()
        if len(neg) > 0:
            f = open('occ_{}_neg.txt'.format(self.body.shortname.replace(' ', '_')), 'w')
            for line in neg:
                f.write('{:10.3f} {:10.3f} {:-6.2f} {:-6.2f} {:16.8f} {}\n'.format(*line))
            f.close()

    def __str__(self):
        """ String representation of the Occultation class
        """
        out = ('Stellar occultation of star {} {} by {}.\n\n'
               'Geocentric Closest Approach: {:.3f}\n'
               'Instant of CA: {}\n'
               'Position Angle: {:.2f}\n'
               'Geocentric shadow velocity: {:.2f}\n'
               'Sun-Geocenter-Target angle:  {:.2f} deg\n'
               'Moon-Geocenter-Target angle: {:.2f} deg\n\n\n'.format(
            self.star._catalogue, self.star.code, self.body.name, self.ca, self.tca.iso,
            self.pa, self.vel, self.predict['S-G-T'].data[0], self.predict['M-G-T'].data[0])
        )

        count = {'positive': 0, 'negative': 0, 'visual': 0}
        string = {'positive': '', 'negative': '', 'visual': ''}
        if len(self.chords) > 0:
            for chord in self.chords.values():
                status = chord.status()
                string[status] += chord.__str__() + '\n'
                count[status] += 1

        if len(self.chords) == 0:
            out += 'No observations reported'
        else:
            out += '\n'.join(['{} {} observations'.format(count[k], k) for k in string.keys() if count[k] > 0])

        out += '\n\n'

        out += '#'*79 + '\n{:^79s}\n'.format('STAR') + '#'*79 + '\n'
        out += self.star.__str__() + '\n\n'

        coord = self.star.get_position(self.tca, observer=self._reference_center)
        try:
            error_star = self.star.error_at(self.tca)
        except:
            error_star = [0, 0]*u.mas
        out += 'Geocentric star coordinate at occultation Epoch ({}):\n'.format(self.tca.iso)
        out += 'RA={} +/- {:.4f}, DEC={} +/- {:.4f}\n\n'.format(
            coord.ra.to_string(u.hourangle, sep='hms', precision=5), error_star[0],
            coord.dec.to_string(u.deg, sep='dms', precision=4), error_star[1])

        out += self.body.__str__() + '\n'

        for status in string.keys():
            if count[status] == 0:
                continue
            out += ('#'*79 + '\n{:^79s}\n'.format(status.upper() + ' OBSERVATIONS') + '#'*79)
            out += '\n\n'
            out += string[status]

        if hasattr(self, 'fitted_params'):
            out += '#'*79 + '\n{:^79s}\n'.format('RESULTS') + '#'*79 + '\n\n'
            out += 'Fitted Ellipse:\n'
            out += '\n'.join(['{}: {:.3f} +/- {:.3f}'.format(k, *self.fitted_params[k])
                              for k in self.fitted_params.keys()]) + '\n'
            polar_radius = self.fitted_params['equatorial_radius'][0]*(1.0-self.fitted_params['oblateness'][0])
            equivalent_radius = np.sqrt(self.fitted_params['equatorial_radius'][0]*polar_radius)
            out += 'polar_radius: {:.3f} km \n'.format(polar_radius)
            out += 'equivalent_radius: {:.3f} km \n'.format(equivalent_radius)
            if not np.isnan(self.body.H):
                H_sun = -26.74
                geometric_albedo = (10**(0.4*(H_sun - self.body.H.value))) * ((u.au.to('km')/equivalent_radius)**2)
                out += 'geometric albedo (V): {:.3f} ({:.1%}) \n'.format(geometric_albedo, geometric_albedo)
            else:
                out += 'geometric albedo (V): not calculated, absolute magnitude (H) is unknown \n'
            out += '\nMinimum chi-square: {:.3f}\n'.format(self.chi2_params['chi2_min'])
            out += 'Number of fitted points: {}\n'.format(self.chi2_params['npts'])
            out += 'Number of fitted parameters: {}\n'.format(self.chi2_params['nparam'])
            out += 'Minimum chi-square per degree of freedom: {:.3f}\n'.format(
                self.chi2_params['chi2_min']/(self.chi2_params['npts'] - self.chi2_params['nparam']))
            out += 'Radial dispersion: {:.3f} +/- {:.3f} km\n'.format(
                self.chi2_params['radial_dispersion'].mean(), self.chi2_params['radial_dispersion'].std(ddof=1))
            out += 'Radial error:      {:.3f} +/- {:.3f} km\n'.format(
                self.chi2_params['radial_error'].mean(), self.chi2_params['radial_error'].std(ddof=1))

            out += '\n' + self.new_astrometric_position(verbose=False)

        return out

    def __func_line_decalage(self, p, x):
        """ Private function returns a linear function, intended for fitting inside the self.check_decalage().
        """
        a, b = p
        return a*x + b

    def __linear_fit_error(self, x, y, sx, sy, verbose=False, use_error=True):
        """ Private function returns a linear fit, intended for fitting inside the self.check_decalage().
        """
        import scipy.odr as odr

        model = odr.Model(self.__func_line_decalage)
        if use_error:
            data = odr.RealData(x=x, y=y, sx=sx, sy=sy)
        else:
            data = odr.RealData(x=x, y=y)
        fit = odr.ODR(data, model, beta0=[0., 1.])
        out = fit.run()
        if verbose:
            print('Linear fit procedure')
            out.pprint()
            print('\n')
        return out

import numpy,import scipy,import astropy,from astropy