#!/usr/bin/env python
"""
Remove horizontal stripes from ACS WFC post-SM4 data.

For more information, see
`Removal of Bias Striping Noise from Post-SM4 ACS WFC Images <http://www.stsci.edu/hst/acs/software/destripe/>`_.

.. note::

    * Does not work on RAW image.

    * Uses the flatfield specified by the image header keyword PFLTFILE.
      If keyword value is 'N/A', as is the case with biases and darks,
      then unity flatfield is used.

    * Uses post-flash image specified by the image header keyword FLSHFILE.
      If keyword value is 'N/A', then dummy post-flash with zeroes is used.

    * Uses the dark image specified by the image header keyword DARKFILE.
      If keyword value is 'N/A', then dummy dark with zeroes is used.

Examples
--------

In Python:

>>> from acstools import acs_destripe
>>> acs_destripe.clean('uncorrected_flt.fits', 'csck',
...                    mask1='mymask_sci1.fits', mask2='mymask_sci2.fits',
...                    clobber=False, maxiter=15, sigrej=2.0)

From command line::

    % acs_destripe [-h] [--stat STAT] [--lower [LOWER]] [--upper [UPPER]]
                   [--binwidth BINWIDTH] [--mask1 MASK1] [--mask2 MASK2]
                   [--dqbits [DQBITS]] [--rpt_clean RPT_CLEAN]
                   [--atol [ATOL]] [-c] [-q] [--version]
                   input suffix [maxiter] [sigrej]

"""
# STDLIB
import logging

# THIRD-PARTY
import numpy as np

try:
    # This supports PIXVALUE
    from stsci.tools import stpyfits as fits
except ImportError:
    # Falls back to Astropy
    from astropy.io import fits

# LOCAL
from .utils_calib import extract_dark, extract_flash, extract_flatfield, SM4_MJD

__taskname__ = 'acs_destripe'
__version__ = '0.8.2'
__vdate__ = '22-Sep-2016'
__author__ = 'Norman Grogin, STScI, March 2012.'
__all__ = ['clean']

#
# HISTORY:
# .........
# 12MAR2015 (v0.6.3) Cara added capability to use DQ mask;
#           added support for multiple input files and wildcards in the file
#           names. See Ticket #1178.
# 23MAR2015 (v0.7.1) Cara added weighted (by NPix) background computations
#           (especially important for vigneted filters). See Ticket #1180.
# 31MAR2015 (v0.8.0) Cara added repeated de-stripe iterations (to improve
#           corrections in the "RAW" space) and support for various
#           statistics modes. See Ticket #1183.

logging.basicConfig()
LOG = logging.getLogger(__taskname__)
LOG.setLevel(logging.INFO)

class StripeArray:
    """Class to handle data array to be destriped."""

    def __init__(self, image):
        self.hdulist = fits.open(image)
        self.ampstring = self.hdulist[0].header['CCDAMP']
        self.flatcorr = self.hdulist[0].header['FLATCORR']
        self.flshcorr = self.hdulist[0].header['FLSHCORR']
        self.darkcorr = self.hdulist[0].header['DARKCORR']
        self.configure_arrays()

    def configure_arrays(self):
        """Get the SCI and ERR data."""
        self.science = self.hdulist['sci', 1].data
        self.err = self.hdulist['err', 1].data
        self.dq = self.hdulist['dq', 1].data
        if (self.ampstring == 'ABCD'):
            self.science = np.concatenate(
                (self.science, self.hdulist['sci', 2].data[::-1, :]), axis=1)
            self.err = np.concatenate(
                (self.err, self.hdulist['err', 2].data[::-1, :]), axis=1)
            self.dq = np.concatenate(
                (self.dq, self.hdulist['dq', 2].data[::-1, :]), axis=1)
        self.ingest_dark()
        self.ingest_flash()
        self.ingest_flatfield()

    def ingest_flatfield(self):
        """Process flatfield."""

        self.invflat = extract_flatfield(
            self.hdulist[0].header, self.hdulist[1])

        # If BIAS or DARK, set flatfield to unity
        if self.invflat is None:
            self.invflat = np.ones_like(self.science)
            return

        # Apply the flatfield if necessary
        if self.flatcorr != 'COMPLETE':
            self.science = self.science * self.invflat
            self.err = self.err * self.invflat

    def ingest_flash(self):
        """Process post-flash."""

        self.flash = extract_flash(self.hdulist[0].header, self.hdulist[1])

        # Set post-flash to zeros
        if self.flash is None:
            self.flash = np.zeros_like(self.science)
            return

        # Apply the flash subtraction if necessary.
        # Not applied to ERR, to be consistent with ingest_dark()
        if self.flshcorr != 'COMPLETE':
            self.science = self.science - self.flash

    def ingest_dark(self):
        """Process dark."""

        self.dark = extract_dark(self.hdulist[0].header, self.hdulist[1])

        # If BIAS or DARK, set dark to zeros
        if self.dark is None:
            self.dark = np.zeros_like(self.science)
            return

        # Apply the dark subtraction if necessary.
        # Effect of DARK on ERR is insignificant for de-striping.
        if self.darkcorr != 'COMPLETE':
            self.science = self.science - self.dark

    def write_corrected(self, output, clobber=False):
        """Write out the destriped data."""

        # un-apply the flatfield if necessary
        if self.flatcorr != 'COMPLETE':
            self.science = self.science / self.invflat
            self.err = self.err / self.invflat

        # un-apply the post-flash if necessary
        if self.flshcorr != 'COMPLETE':
            self.science = self.science + self.flash

        # un-apply the dark if necessary
        if self.darkcorr != 'COMPLETE':
            self.science = self.science + self.dark

        # reverse the amp merge
        if (self.ampstring == 'ABCD'):
            tmp_1, tmp_2 = np.split(self.science, 2, axis=1)
            self.hdulist['sci', 1].data = tmp_1.copy()
            self.hdulist['sci', 2].data = tmp_2[::-1, :].copy()

            tmp_1, tmp_2 = np.split(self.err, 2, axis=1)
            self.hdulist['err', 1].data = tmp_1.copy()
            self.hdulist['err', 2].data = tmp_2[::-1, :].copy()
        else:
            self.hdulist['sci', 1].data = self.science.copy()
            self.hdulist['err', 1].data = self.err.copy()

        # Write the output
        self.hdulist.writeto(output, overwrite=clobber)

    def close(self):
        """Close open file(s)."""
        self.hdulist.close()

def _read_mask(mask1, mask2):
    if isinstance(mask1, str):
        mask1 = fits.getdata(mask1)
    if isinstance(mask2, str):
        mask2 = fits.getdata(mask2)

    if mask1 is None and mask2 is None:
        return None
    elif mask2 is None:
        mask = mask1
    elif mask1 is None:
        mask = mask2
    else:
        dat1 = mask1
        dat2 = mask2
        mask = np.concatenate((dat1, dat2[::-1, :]), axis=1)

    # Mask must only have binary values
    if mask is not None:
        mask[mask != 0] = 1

    return mask

def clean(input, suffix, stat="pmode1", maxiter=15, sigrej=2.0,
          lower=None, upper=None, binwidth=0.3,
          mask1=None, mask2=None, dqbits=None,
          rpt_clean=0, atol=0.01, clobber=False, verbose=True):
    r"""Remove horizontal stripes from ACS WFC post-SM4 data.

    Parameters
    ----------
    input : str or list of str
        Input filenames in one of these formats:

            * a Python list of filenames
            * a partial filename with wildcards ('\*flt.fits')
            * filename of an ASN table ('j12345670_asn.fits')
            * an at-file (``@input``)

    suffix : str
        The string to use to add to each input file name to
        indicate an output product. This string will be appended
        to the suffix in each input filename to create the
        new output filename. For example, setting `suffix='csck'`
        will create '\*_csck.fits' images.

    stat : { 'pmode1', 'pmode2', 'mean', 'mode', 'median', 'midpt' } (Default = 'pmode1')
        Specifies the statistics to be used for computation of the
        background in image rows:

        * 'pmode1' - SEXTRACTOR-like mode estimate based on a
          modified `Pearson's rule <https://en.wikipedia.org/wiki/Nonparametric_skew#Pearson.27s_rule>`_:
          ``2.5*median-1.5*mean``;
        * 'pmode2' - mode estimate based on
          `Pearson's rule <https://en.wikipedia.org/wiki/Nonparametric_skew#Pearson.27s_rule>`_:
          ``3*median-2*mean``;
        * 'mean' - the mean of the distribution of the "good" pixels (after
          clipping, masking, etc.);
        * 'mode' - the mode of the distribution of the "good" pixels;
        * 'median' - the median of the distribution of the "good" pixels;
        * 'midpt' - estimate of the median of the distribution of the "good"
          pixels based on an algorithm similar to IRAF's ``imagestats`` task
          (``CDF(midpt)=1/2``).

        .. note::
            The midpoint and mode are computed in two passes through the
            image. In the first pass the standard deviation of the pixels
            is calculated and used with the *binwidth* parameter to compute
            the resolution of the data histogram. The midpoint is estimated
            by integrating the histogram and computing by interpolation
            the data value at which exactly half the pixels are below that
            data value and half are above it. The mode is computed by
            locating the maximum of the data histogram and fitting the peak
            by parabolic interpolation.

    maxiter : int
        This parameter controls the maximum number of iterations
        to perform when computing the statistics used to compute the
        row-by-row corrections.

    sigrej : float
        This parameters sets the sigma level for the rejection applied
        during each iteration of statistics computations for the
        row-by-row corrections.

    lower : float, None (Default = None)
        Lower limit of usable pixel values for computing the background.
        This value should be specified in the units of the input image(s).

    upper : float, None (Default = None)
        Upper limit of usable pixel values for computing the background.
        This value should be specified in the units of the input image(s).

    binwidth : float (Default = 0.1)
        Histogram's bin width, in sigma units, used to sample the
        distribution of pixel brightness values in order to compute the
        background statistics. This parameter is aplicable *only* to *stat*
        parameter values of `'mode'` or `'midpt'`.

    clobber : bool
        Specify whether or not to 'clobber' (delete then replace)
        previously generated products with the same names.

    mask1 : str, numpy.ndarray, None, or list of these types
        Mask images for ``SCI,1``, one for each input file.
        Pixels with zero values will be masked out, in addition to clipping.

    mask2 : str, numpy.ndarray, None, or list of these types
        Mask images for ``SCI,2``, one for each input file.
        Pixels with zero values will be masked out, in addition to clipping.
        This is not used for subarrays.

    dqbits : int, str, None (Default = None)
        Integer sum of all the DQ bit values from the input image's DQ array
        that should be considered "good" when building masks for de-striping
        computations. For example, if pixels in the DQ array can be
        combinations of 1, 2, 4, and 8 flags and one wants to consider
        DQ "defects" having flags 2 and 4 as being acceptable for de-striping
        computations, then `dqbits` should be set to 2+4=6. Then a DQ pixel
        having values 2,4, or 6 will be considered a good pixel, while a DQ
        pixel with a value, e.g., 1+2=3, 4+8=12, etc. will be flagged
        as a "bad" pixel.

        Alternatively, one can enter a comma- or '+'-separated list of
        integer bit flags that should be added to obtain the final
        "good" bits. For example, both ``4,8`` and ``4+8`` are equivalent to
        setting `dqbits` to 12.

        | Set `dqbits` to 0 to make *all* non-zero pixels in the DQ
          mask to be considered "bad" pixels, and the corresponding image
          pixels not to be used for de-striping computations.

        | Default value (`None`) will turn off the use of image's DQ array
          for de-striping computations.

        | In order to reverse the meaning of the `dqbits`
          parameter from indicating values of the "good" DQ flags
          to indicating the "bad" DQ flags, prepend '~' to the string
          value. For example, in order not to use pixels with
          DQ flags 4 and 8 for sky computations and to consider
          as "good" all other pixels (regardless of their DQ flag),
          set `dqbits` to ``~4+8``, or ``~4,8``. To obtain the
          same effect with an `int` input value (except for 0),
          enter -(4+8+1)=-9. Following this convention,
          a `dqbits` string value of ``'~0'`` would be equivalent to
          setting ``dqbits=None``.

        .. note::
            DQ masks (if used), *will be* combined with user masks specified
            in the `mask1` and `mask2` parameters (if any).

    rpt_clean : int
        An integer indicating how many *additional* times stripe cleaning
        should be performed on the input image. Default = 0.

    atol : float, None
        The threshold for maximum absolute value of bias stripe correction
        below which repeated cleanings can stop. When `atol` is `None`
        cleaning will be repeated `rpt_clean` number of times.
        Default = 0.01 [e].

    verbose : bool
        Print informational messages. Default = True.

    """
    from stsci.tools import parseinput  # Optional package dependency

    flist = parseinput.parseinput(input)[0]

    if isinstance(mask1, str):
        mlist1 = parseinput.parseinput(mask1)[0]
    elif isinstance(mask1, np.ndarray):
        mlist1 = [mask1.copy()]
    elif mask1 is None:
        mlist1 = []
    elif isinstance(mask1, list):
        mlist1 = []
        for m in mask1:
            if isinstance(m, np.ndarray):
                mlist1.append(m.copy())
            elif isinstance(m, str):
                mlist1 += parseinput.parseinput(m)[0]
            else:
                raise TypeError("'mask1' must be a list of str or "
                                "numpy.ndarray values.")
    else:
        raise TypeError("'mask1' must be either a str, or a "
                        "numpy.ndarray, or a list of the two type of "
                        "values.")

    if isinstance(mask2, str):
        mlist2 = parseinput.parseinput(mask2)[0]
    elif isinstance(mask2, np.ndarray):
        mlist2 = [mask2.copy()]
    elif mask2 is None:
        mlist2 = []
    elif isinstance(mask2, list):
        mlist2 = []
        for m in mask2:
            if isinstance(m, np.ndarray):
                mlist2.append(m.copy())
            elif isinstance(m, str):
                mlist2 += parseinput.parseinput(m)[0]
            else:
                raise TypeError("'mask2' must be a list of str or "
                                "numpy.ndarray values.")
    else:
        raise TypeError("'mask2' must be either a str or a "
                        "numpy.ndarray, or a list of the two type of "
                        "values.")

    n_input = len(flist)
    n_mask1 = len(mlist1)
    n_mask2 = len(mlist2)

    if n_input == 0:
        raise ValueError("No input file(s) provided or "
                         "the file(s) do not exist")

    if n_mask1 == 0:
        mlist1 = [None] * n_input
    elif n_mask1 != n_input:
        raise ValueError('Insufficient masks for [SCI,1]')

    if n_mask2 == 0:
        mlist2 = [None] * n_input
    elif n_mask2 != n_input:
        raise ValueError('Insufficient masks for [SCI,2]')

    for image, maskfile1, maskfile2 in zip(flist, mlist1, mlist2):
        # Skip processing pre-SM4 images
        if (fits.getval(image, 'EXPSTART') <= SM4_MJD):
            LOG.warning(f'{image} is pre-SM4. Skipping...')
            continue

        # Data must be in ELECTRONS
        if (fits.getval(image, 'BUNIT', ext=1) != 'ELECTRONS'):
            LOG.warning(f'{image} is not in ELECTRONS. Skipping...')
            continue

        # Skip processing CTECORR-ed images
        if (fits.getval(image, 'PCTECORR') == 'COMPLETE'):
            LOG.warning(f'{image} already has PCTECORR applied. Skipping...')
            continue

        # generate output filename for each input based on specification
        # of the output suffix
        output = image.replace('.fits', '_' + suffix + '.fits')
        LOG.info('Processing ' + image)

        # verify masks defined (or not) simultaneously:
        if (fits.getval(image, 'CCDAMP') == 'ABCD' and
                ((mask1 is not None and mask2 is None) or
                 (mask1 is None and mask2 is not None))):
            raise ValueError("Both 'mask1' and 'mask2' must be specified "
                             "or not specified together.")

        maskdata = _read_mask(maskfile1, maskfile2)
        perform_correction(image, output, stat=stat, maxiter=maxiter,
                           sigrej=sigrej, lower=lower, upper=upper,
                           binwidth=binwidth, mask=maskdata, dqbits=dqbits,
                           rpt_clean=rpt_clean, atol=atol,
                           clobber=clobber, verbose=verbose)
        LOG.info(output + ' created')

def perform_correction(image, output, stat="pmode1", maxiter=15, sigrej=2.0,
                       lower=None, upper=None, binwidth=0.3,
                       mask=None, dqbits=None,
                       rpt_clean=0, atol=0.01, clobber=False, verbose=True):
    """
    Clean each input image.

    Parameters
    ----------
    image : str
        Input image name.

    output : str
        Output image name.

    mask : `numpy.ndarray`
        Mask array.

    maxiter, sigrej, clobber
        See :func:`clean`.

    dqbits : int, str, or None
        Data quality bits to be considered as "good" (or "bad").
        See :func:`clean` for more details.

    rpt_clean : int
        An integer indicating how many *additional* times stripe cleaning
        should be performed on the input image. Default = 0.

    atol : float, None
        The threshold for maximum absolute value of bias stripe correction
        below which repeated cleanings can stop. When `atol` is `None`
        cleaning will be repeated `rpt_clean` number of times.
        Default = 0.01 [e].

    verbose : bool
        Print informational messages. Default = True.

    """
    # construct the frame to be cleaned, including the
    # associated data stuctures needed for cleaning
    frame = StripeArray(image)

    # combine user mask with image's DQ array:
    mask = _mergeUserMaskAndDQ(frame.dq, mask, dqbits)

    # Do the stripe cleaning
    Success, NUpdRows, NMaxIter, Bkgrnd, STDDEVCorr, MaxCorr, Nrpt = clean_streak(
        frame, stat=stat, maxiter=maxiter, sigrej=sigrej,
        lower=lower, upper=upper, binwidth=binwidth, mask=mask,
        rpt_clean=rpt_clean, atol=atol, verbose=verbose
    )

    if Success:
        if verbose:
            LOG.info('perform_correction - =====  Overall statistics for '
                     'de-stripe corrections:  =====')

        if (STDDEVCorr > 1.5*0.9):
            LOG.warning('perform_correction - STDDEV of applied de-stripe '
                        f'corrections ({STDDEVCorr:.3g}) exceeds\n'
                        'known bias striping STDDEV of 0.9e '
                        '(see ISR ACS 2011-05) more than 1.5 times.')

        elif verbose:
            LOG.info('perform_correction - STDDEV of applied de-stripe '
                     f'corrections {STDDEVCorr:.3g}.')

        if verbose:
            LOG.info('perform_correction - Estimated background: '
                     f'{Bkgrnd:.5g}.')
            LOG.info('perform_correction - Maximum applied correction: '
                     f'{MaxCorr:.3g}.')
            LOG.info('perform_correction - Effective number of clipping '
                     f'iterations: {NMaxIter}.')
            LOG.info('perform_correction - Effective number of additional '
                     f'(repeated) cleanings: {Nrpt}.')
            LOG.info('perform_correction - Total number of corrected rows: '
                     f'{NUpdRows}.')

    frame.write_corrected(output, clobber=clobber)
    frame.close()

def _mergeUserMaskAndDQ(dq, mask, dqbits):
    from astropy.nddata.bitmask import (interpret_bit_flags,
                                        bitfield_to_boolean_mask)

    dqbits = interpret_bit_flags(dqbits)
    if dqbits is None:
        if mask is None:
            return None
        else:
            return mask.copy().astype(dtype=np.uint8)

    if dq is None:
        raise ValueError("DQ array is None while 'dqbits' is not None.")

    dqmask = bitfield_to_boolean_mask(dq, dqbits, good_mask_value=1,
                                      dtype=np.uint8)

    if mask is None:
        return dqmask

    # merge user mask with DQ mask:
    dqmask[mask == 0] = 0

    return dqmask

def clean_streak(image, stat="pmode1", maxiter=15, sigrej=2.0,
                 lower=None, upper=None, binwidth=0.3, mask=None,
                 rpt_clean=0, atol=0.01, verbose=True):
    """
    Apply destriping algorithm to input array.

    Parameters
    ----------
    image : `StripeArray` object
        Arrays are modifed in-place.

    stat : str
        Statistics for background computations
        (see :py:func:`clean` for more details)

    mask : `numpy.ndarray`
        Mask array. Pixels with zero values are masked out.

    maxiter, sigrej : see `clean`

    rpt_clean : int
        An integer indicating how many *additional* times stripe cleaning
        should be performed on the input image. Default = 0.

    atol : float, None
        The threshold for maximum absolute value of bias stripe correction
        below which repeated cleanings can stop. When `atol` is `None`
        cleaning will be repeated `rpt_clean` number of times.
        Default = 0.01 [e].

    verbose : bool
        Print informational messages. Default = True.

    Returns
    -------
    Success : bool
        Indicates successful execution.

    NUpdRows : int
        Number of updated rows in the image.

    NMaxIter : int
        Maximum number of clipping iterations performed on image rows.

    Bkgrnd, STDDEVCorr, MaxCorr : float
        Background, standard deviation of corrections and maximum correction
        applied to the non-flat-field-corrected (i.e., RAW) image rows.

    Nrpt : int
        Number of *additional* (performed *after* initial run) cleanings.

    """
    # Optional package dependency
    try:
        from stsci.imagestats import ImageStats
    except ImportError:
        ImageStats = None

    if mask is not None and image.science.shape != mask.shape:
        raise ValueError('Mask shape does not match science data shape')

    Nrpt = 0
    warn_maxiter = False
    NUpdRows = 0
    NMaxIter = 0
    STDDEVCorr = 0.0
    MaxCorr = 0.0
    wmean = 0.0

    stat = stat.lower().strip()
    if stat not in ['pmode1', 'pmode2', 'mean', 'mode', 'median', 'midpt']:
        raise ValueError("Unsupported value for 'stat'.")

    # array to hold the stripe amplitudes
    corr = np.empty(image.science.shape[0], dtype=np.float64)

    # array to hold cumulative stripe amplitudes and latest row npix:
    cumcorr = np.zeros(image.science.shape[0], dtype=np.float64)
    cnpix = np.zeros(image.science.shape[0], dtype=int)

    # other arrays
    corr_scale = np.empty(image.science.shape[0], dtype=np.float64)
    npix = np.empty(image.science.shape[0], dtype=int)
    sigcorr2 = np.zeros(image.science.shape[0], dtype=np.float64)
    updrows = np.zeros(image.science.shape[0], dtype=int)

    # for speed-up and to reduce rounding errors in ERR computations,
    # keep a copy of the squared error array:
    imerr2 = image.err**2

    # arrays for detecting oscillatory behaviour:
    nonconvi0 = np.arange(image.science.shape[0])
    corr0 = np.zeros(image.science.shape[0], dtype=np.float64)

    if stat == 'pmode1':
        # SExtractor-esque central value statistic; slightly sturdier against
        # skewness of pixel histogram due to faint source flux
        def getcorr(): return (2.5 * SMedian - 1.5 * SMean)

    elif stat == 'pmode2':
        # "Original Pearson"-ian estimate for mode:
        def getcorr(): return (3.0 * SMedian - 2.0 * SMean)

    elif stat == 'mean':
        def getcorr(): return (SMean)

    elif stat == 'median':
        def getcorr(): return (SMedian)

    elif stat == 'mode':
        if ImageStats is None:
            raise ImportError('stsci.imagestats is missing')

        def getcorr():
            imstat = ImageStats(image.science[i][BMask], 'mode',
                                lower=lower, upper=upper, nclip=0)
            if imstat.npix != NPix:
                raise ValueError(f'imstate.npix ({imstat.npix}) != '
                                 f'NPix ({NPix})')
            return (imstat.mode)

    elif stat == 'midpt':
        if ImageStats is None:
            raise ImportError('stsci.imagestats is missing')

        def getcorr():
            imstat = ImageStats(image.science[i][BMask], 'midpt',
                                lower=lower, upper=upper, nclip=0)
            if imstat.npix != NPix:
                raise ValueError(f'imstate.npix ({imstat.npix}) != '
                                 f'NPix ({NPix})')
            return (imstat.midpt)

    nmax_rpt = 1 if rpt_clean is None else max(1, rpt_clean+1)

    for rpt in range(nmax_rpt):
        Nrpt += 1

        if verbose:
            if Nrpt <= 1:
                if nmax_rpt > 1:
                    LOG.info("clean_streak - Performing initial image bias "
                             "de-stripe:")
                else:
                    LOG.info("clean_streak - Performing image bias de-stripe:")
            else:
                LOG.info("clean_streak - Performing repeated image bias "
                         f"de-stripe #{Nrpt - 1}:")

        # reset accumulators and arrays:
        corr[:] = 0.0
        corr_scale[:] = 0.0
        npix[:] = 0

        tcorr = 0.0
        tnpix = 0
        tnpix2 = 0
        NMaxIter = 0

        # loop over rows to fit the stripe amplitude
        mask_arr = None
        for i in range(image.science.shape[0]):
            if mask is not None:
                mask_arr = mask[i]

            # row-by-row iterative sigma-clipped mean;
            # sigma, iters are adjustable
            SMean, SSig, SMedian, NPix, NIter, BMask = djs_iterstat(
                image.science[i], MaxIter=maxiter, SigRej=sigrej,
                Min=lower, Max=upper, Mask=mask_arr, lineno=i+1
            )

            if NPix > 0:
                corr[i] = getcorr()
                npix[i] = NPix
                corr_scale[i] = 1.0 / np.average(image.invflat[i][BMask])
                sigcorr2[i] = corr_scale[i]**2 * \
                    np.sum((image.err[i][BMask])**2)/NPix**2
                cnpix[i] = NPix
                tnpix += NPix
                tnpix2 += NPix*NPix
                tcorr += corr[i] * NPix

            if NIter > NMaxIter:
                NMaxIter = NIter

        if tnpix <= 0:
            LOG.warning('clean_streak - No good data points; cannot de-stripe.')
            return False, 0, 0, 0.0, 0.0, 0

        if NMaxIter >= maxiter:
            warn_maxiter = True

        # require that bias stripe corrections have zero mean:
        # 1. compute weighted background of the flat-fielded image:
        wmean = tcorr / tnpix
        Bkgrnd = wmean
        # 2. estimate corrections:
        corr[npix > 0] -= wmean

        # convert corrections to the "raw" space:
        corr *= corr_scale

        # weighted mean and max value for current corrections
        # to the *RAW* image:
        trim_npix = npix[npix > 0]
        trim_corr = corr[npix > 0]
        cwmean = np.sum(trim_npix * trim_corr) / tnpix
        current_max_corr = np.amax(np.abs(trim_corr - cwmean))
        wvar = np.sum(trim_npix * (trim_corr - cwmean) ** 2) / tnpix
        uwvar = wvar / (1.0 - float(tnpix2) / float(tnpix) ** 2)
        STDDEVCorr = np.sqrt(uwvar)

        # keep track of total corrections:
        cumcorr += corr

        # apply corrections row-by-row
        for i in range(image.science.shape[0]):
            if npix[i] < 1:
                continue
            updrows[i] = 1

            ffdark = (image.dark[i] + image.flash[i]) * image.invflat[i]
            t1 = np.maximum(image.science[i] + ffdark, 0.0)

            # stripe is constant along the row, before flatfielding;
            # afterwards it has the shape of the inverse flatfield
            truecorr = corr[i] * image.invflat[i]
            #truecorr_sig2 = sigcorr2[i] * image.invflat[i]**2  # DEBUG

            # correct the SCI extension
            image.science[i] -= truecorr

            t2 = np.maximum(image.science[i] + ffdark, 0.0)

            T = (t1 - t2) * image.invflat[i]

            # correct the ERR extension
            # NOTE: np.abs() in the err array recomputation is used for safety
            #       only and, in principle, assuming no errors have been made
            #       in the derivation of the formula, np.abs() should not be
            #       necessary.
            imerr2[i] -= T
            image.err[i] = np.sqrt(np.abs(imerr2[i]))
            # NOTE: for debugging purposes, one may want to uncomment
            #       next line:
            #assert( np.all(imerr2 >= 0.0))

        if atol is not None:
            if current_max_corr < atol:
                break

            # detect oscilatory non-convergence:
            nonconvi = np.nonzero(np.abs(corr) > atol)[0]
            nonconvi_int = np.intersect1d(nonconvi, nonconvi0)
            if (nonconvi.shape[0] == nonconvi0.shape[0] and
                    nonconvi.shape[0] == nonconvi_int.shape[0] and
                    np.all(corr0[nonconvi]*corr[nonconvi] < 0.0) and Nrpt > 1):
                LOG.warning("clean_streak - Repeat bias stripe cleaning\n"
                            "process appears to be oscillatory for "
                            f"{nonconvi.shape[0]:d} image "
                            "rows.\nTry to adjust 'sigrej', 'maxiter', and/or "
                            "'dqbits' parameters.\n"
                            "In addition,  consider using masks or adjust "
                            "existing masks.")
                break

            nonconvi0 = nonconvi.copy()
            corr0 = corr.copy()

        if verbose:
            if Nrpt <= 1:
                LOG.info("clean_streak - Image bias de-stripe: Done.")
            else:
                LOG.info(f"clean_streak - Repeated (#{Nrpt - 1}) image bias "
                         "de-stripe: Done.")

    if verbose and Nrpt > 1:
        LOG.info('clean_streak - =====  Repeated de-stripe "residual" '
                 'estimates:  =====')
        LOG.info('clean_streak - STDDEV of the last applied de-stripe '
                 f'corrections {STDDEVCorr:.3g}')
        LOG.info('clean_streak - Maximum of the last applied correction: '
                 f'{current_max_corr:.3g}.')

    # add (in quadratures) an error term associated with the accuracy of
    # bias stripe correction:
    truecorr_sig2 = ((sigcorr2 * (image.invflat ** 2).T).T).astype(image.err.dtype)  # noqa

    # update the ERR extension
    image.err[:, :] = np.sqrt(np.abs(imerr2 + truecorr_sig2))

    if warn_maxiter:
        LOG.warning(
            'clean_streak - Maximum number of clipping iterations '
            f'specified by the user ({maxiter}) has been reached.')

    # weighted mean, sample variance, and max value for
    # total (cummulative) corrections to the *RAW* image:
    trim_cnpix = cnpix[cnpix > 0]
    trim_cumcorr = cumcorr[cnpix > 0]
    tcnpix = np.sum(trim_cnpix)
    tcnpix2 = np.sum(trim_cnpix ** 2)
    cwmean = np.sum(trim_cnpix * trim_cumcorr) / tcnpix
    trim_cumcorr -= cwmean
    wvar = np.sum(trim_cnpix * trim_cumcorr ** 2) / tcnpix
    uwvar = wvar / (1.0 - float(tcnpix2) / float(tcnpix) ** 2)
    STDDEVCorr = np.sqrt(uwvar)
    MaxCorr = np.amax(np.abs(trim_cumcorr))

    NUpdRows = np.sum(updrows)

    return True, NUpdRows, NMaxIter, Bkgrnd, STDDEVCorr, MaxCorr, Nrpt-1

def _write_row_number(lineno, offset=1, pad=1):
    if lineno is None:
        return ''
    return (pad * ' ' + f'(row #{lineno + offset:d})')

def djs_iterstat(InputArr, MaxIter=10, SigRej=3.0,
                 Max=None, Min=None, Mask=None, lineno=None):
    """
    Iterative sigma-clipping.

    Parameters
    ----------
    InputArr : `numpy.ndarray`
        Input image array.

    MaxIter, SigRej : see `clean`

    Max, Min : float
        Max and min values for clipping.

    Mask : `numpy.ndarray`
        Mask array to indicate pixels to reject, in addition to clipping.
        Pixels where mask is zero will be rejected.
        If not given, all pixels will be used.

    lineno : int or None
        Line number to be used in log and/or warning messages.

    Returns
    -------
    FMean, FSig, FMedian, NPix : float
        Mean, sigma, and median of final result.

    NIter : int
        Number of performed clipping iterations

    BMask : `numpy.ndarray`
        Logical image mask from the final iteration.

    """
    NGood = InputArr.size
    ArrShape = InputArr.shape
    if NGood == 0:
        imrow = _write_row_number(lineno=lineno, offset=0, pad=1)
        LOG.warning(f'djs_iterstat - No data points given{imrow}')
        return 0, 0, 0, 0, 0, None
    if NGood == 1:
        imrow = _write_row_number(lineno=lineno, offset=0, pad=1)
        LOG.warning('djs_iterstat - Only one data point; '
                    f'cannot compute stats{imrow}')
        return 0, 0, 0, 0, 0, None
    if np.unique(InputArr).size == 1:
        imrow = _write_row_number(lineno=lineno, offset=0, pad=1)
        LOG.warning('djs_iterstat - Only one value in data; '
                    f'cannot compute stats{imrow}')
        return 0, 0, 0, 0, 0, None

    # Determine Max and Min
    if Max is None:
        Max = InputArr.max()
    if Min is None:
        Min = InputArr.min()

    # Use all pixels if no mask is provided
    if Mask is None:
        Mask = np.ones(ArrShape, dtype=np.byte)
    else:
        Mask = Mask.copy()

    # Reject those above Max and those below Min
    Mask[InputArr > Max] = 0
    Mask[InputArr < Min] = 0

    FMean = np.sum(1.0 * InputArr * Mask) / NGood
    FSig  = np.sqrt(np.sum((1.0 * InputArr - FMean) ** 2 * Mask) / (NGood - 1))

    NLast = -1
    Iter  = 0
    NGood = np.sum(Mask)
    if NGood < 2:
        imrow = _write_row_number(lineno=lineno, offset=0, pad=1)
        LOG.warning('djs_iterstat - No good data points; '
                    f'cannot compute stats{imrow}')
        return 0, 0, 0, 0, 0, None

    SaveMask = Mask.copy()
    if Iter >= MaxIter:  # to support MaxIter=0
        NLast = NGood

    while (Iter < MaxIter) and (NLast != NGood) and (NGood >= 2):
        LoVal = FMean - SigRej * FSig
        HiVal = FMean + SigRej * FSig

        Mask[InputArr < LoVal] = 0
        Mask[InputArr > HiVal] = 0
        NLast = NGood
        npix = np.sum(Mask)

        if npix >= 2:
            FMean = np.sum(1.0 * InputArr * Mask) / npix
            FSig = np.sqrt(np.sum(
                (1.0 * InputArr - FMean) ** 2 * Mask) / (npix - 1))
            SaveMask = Mask.copy()  # last mask used for computation of mean
            NGood = npix
            Iter += 1
        else:
            break

    logical_mask = SaveMask.astype(bool)

    if NLast > 1:
        FMedian = np.median(InputArr[logical_mask])
        NLast = NGood
    else:
        FMedian = FMean

    return FMean, FSig, FMedian, NLast, Iter, logical_mask

# --------------------------- #
# Interfaces for command line #
# --------------------------- #

def main():
    """Command line driver."""
    import argparse

    parser = argparse.ArgumentParser(
        prog=__taskname__,
        description='Remove horizontal stripes from ACS WFC post-SM4 data.')
    parser.add_argument(
        'arg0', metavar='input', type=str, help='Input file')
    parser.add_argument(
        'arg1', metavar='suffix', type=str, help='Output suffix')
    parser.add_argument(
        'maxiter', nargs='?', type=int, default=15, help='Max #iterations')
    parser.add_argument(
        'sigrej', nargs='?', type=float, default=2.0, help='Rejection sigma')
    parser.add_argument(
        '--stat', type=str, default='pmode1', help='Background statistics')
    parser.add_argument(
        '--lower', nargs='?', type=float, default=None,
        help='Lower limit for "good" pixels.')
    parser.add_argument(
        '--upper', nargs='?', type=float, default=None,
        help='Upper limit for "good" pixels.')
    parser.add_argument(
        '--binwidth', type=float, default=0.1,
        help='Bin width for distribution sampling.')
    parser.add_argument(
        '--mask1', nargs=1, type=str, default=None,
        help='Mask image for [SCI,1]')
    parser.add_argument(
        '--mask2', nargs=1, type=str, default=None,
        help='Mask image for [SCI,2]')
    parser.add_argument(
        '--dqbits', nargs='?', type=str, default=None,
        help='DQ bits to be considered "good".')
    parser.add_argument(
        '--rpt_clean', type=int, default=0,
        help='Number of *repeated* bias de-stripes to perform.')
    parser.add_argument(
        '--atol', nargs='?', type=float, default=0.01,
        help='Absolute tolerance to stop *repeated* bias de-stripes.')
    parser.add_argument(
        '-c', '--clobber', action="store_true", help='Clobber output')
    parser.add_argument(
        '-q', '--quiet', action='store_true',
        help='Do not print informational messages')
    parser.add_argument(
        '--version', action="version",
        version=f'{__taskname__} v{__version__} ({__vdate__})')
    args = parser.parse_args()

    if args.mask1:
        mask1 = args.mask1[0]
    else:
        mask1 = args.mask1

    if args.mask2:
        mask2 = args.mask2[0]
    else:
        mask2 = args.mask2

    clean(args.arg0, args.arg1, stat=args.stat, maxiter=args.maxiter,
          sigrej=args.sigrej, lower=args.lower, upper=args.upper,
          binwidth=args.binwidth, mask1=mask1, mask2=mask2, dqbits=args.dqbits,
          rpt_clean=args.rpt_clean, atol=args.atol,
          clobber=args.clobber, verbose=not args.quiet)

if __name__ == '__main__':
    main()


import copy
import json
import os
import zlib

import astropy
import numpy as np
from astropy import units as u
from astropy.coordinates import SkyCoord
from astropy.io import ascii
from pixell import utils
from sandbox import misc as sbmisc
from scipy import spatial

from . import misc

class HierarchicalCatalog(object):
    def __init__(self, output_dir):
        self.output_dir = output_dir
        sbmisc.create_dir(self.output_dir)
        self.structure_file = self.output_path('structure.json')
        self.structure = None
        self.higharch_map = {0: None, 1: {}}
        self.catalogs = {0: None, 1: {}, 2: {}}

    def output_path(self, file_name):
        return os.path.join(self.output_dir, file_name)

    def load(self, greedy=False):
        structure = self.__base_load(False)
        if greedy:
            # preload catalogs
            self.get_catalogs_per_freq()
        return structure

    def __base_load(self, reload=False):
        if self.structure is not None and not reload:
            pass
        else:
            try:
                self.structure = json.load(open(self.structure_file, 'r'))
                for freq in self.structure.keys():
                    target_file = self.output_path('{}_hcmap.csv'.format(freq))
                    misc.insert_to_dict(self.higharch_map, read_crush(target_file), [1, freq])
                    for season in self.structure[freq].keys():
                        target_file = self.output_path('{}_{}.csv'.format(freq, season))
                        misc.insert_to_dict(self.catalogs, read_crush(target_file), [2, freq, season])
            except:
                print("failed to load")
        return self.structure

    def get_catalogs_per_season(self):
        self.__base_load(reload=False)
        return self.catalogs[2]

    def get_catalogs_per_freq(self):
        self.__base_load(reload=False)
        if not self.catalogs[1]:
            for freq in self.structure.keys():
                target_file = self.output_path('{}.csv'.format(freq))
                if os.path.exists(target_file):
                    misc.insert_to_dict(self.catalogs, read_crush(target_file), [1, freq])
                else:
                    merged, _ = crossmatch(astropy.table.vstack(list(self.catalogs[2][freq].values())),
                                           0.5 * utils.arcmin)
                    misc.insert_to_dict(self.catalogs, merged, [1, freq])
                    astropy.io.ascii.write(merged, self.output_path('{}.csv'.format(freq)))
        else:
            pass
        return self.catalogs[1]

    def build(self, catalogs, seperation, cuts=[('snr', '>5')], verbose=True, overwrite=False):
        common_cols = None
        structure = {}
        for freq in catalogs.keys():
            structure[freq] = {}
            for season in catalogs[freq].keys():
                structure[freq][season] = ''
                target_file = self.output_path('{}_{}.csv'.format(freq, season))
                if os.path.exists(target_file) and not overwrite:
                    cat = read_crush(target_file)
                    catalogs[freq][season] = cat
                else:
                    if common_cols is None:
                        common_cols = get_common_colnames(misc.nested_dict_values(catalogs))
                    for patch in catalogs[freq][season].keys():
                        cat = catalogs[freq][season][patch]
                        cat = cat[list(common_cols)]
                        for cut in cuts:
                            cat = apply_cut(cat, cut[0], cut[1])
                        catalogs[freq][season][patch] = cat

                    catalogs[freq][season] = _postprocess_catalog(
                        astropy.table.vstack(list(catalogs[freq][season].values())))
                    if verbose: print("crossmatching {} {}".format(freq, season))
                    cat, childs = crossmatch(catalogs[freq][season], seperation, verbose=verbose)
                    for i, child in enumerate(childs):
                        nchild = len(child)
                        if nchild > 1:
                            seasons = cat['season'][i].split(',')
                            patches = cat['patch'][i].split(',')
                            freqs = cat['freq'][i].split(',')

                            note = []
                            for j in range(nchild):
                                note.append('_'.join([seasons[j], patches[j], freqs[j], str(child[j])]))
                            note = ','.join(note)
                            note = str(zlib.compress(note.encode('utf8')))
                            cat[i]['note'] = note
                        cat[i]['child_count'] = int(nchild)
                        cat[i]['season'] = season
                    cat.remove_columns(['patch', 'freq'])
                    idx_col = np.arange(len(cat)).astype(int).astype(object)
                    idx_col = astropy.table.Column(idx_col, name='index')
                    cat.add_column(idx_col)
                    catalogs[freq][season] = cat
                    astropy.io.ascii.write(catalogs[freq][season], self.output_path('{}_{}.csv'.format(freq, season)))
                misc.insert_to_dict(self.catalogs, cat, [2, freq, season])
            target_file = self.output_path('{}_hcmap.csv'.format(freq))

            if os.path.exists(target_file) and not overwrite:
                misc.insert_to_dict(self.higharch_map, read_crush(target_file), [1, freq])
            else:
                cat = astropy.table.vstack(list(catalogs[freq].values()))
                _, qbt, _ = build_kd_tree(cat, seperation)
                t = catalogs[freq][season]['parent', 'child', 'child_count'][:0].copy()
                for i, childs in enumerate(qbt):
                    t.add_row(np.empty(1, dtype=t.dtype)[0])
                    child_dict = {}
                    for child in childs:
                        season = cat[child]['season']
                        idx = cat[child]['index']
                        catalogs[freq][season][idx]['parent'] = i
                        misc.insert_to_dict(child_dict, '', [season, idx])
                    t[i]['child'] = str(child_dict)
                    t[i]['child_count'] = len(childs)

                freq_col = np.array([freq] * len(t)).astype(object)
                freq_col = astropy.table.Column(freq_col, name='freq')
                idx_col = np.arange(len(t)).astype(int).astype(object)
                idx_col = astropy.table.Column(idx_col, name='index')
                t.add_columns([idx_col, freq_col])

                astropy.io.ascii.write(t, self.output_path('{}_hcmap.csv'.format(freq)))
                misc.insert_to_dict(self.higharch_map, t, [1, freq])
                for season in catalogs[freq].keys():
                    astropy.io.ascii.write(catalogs[freq][season], self.output_path('{}_{}.csv'.format(freq, season)))
                    misc.insert_to_dict(self.catalogs, catalogs[freq][season], [2, freq, season])

        if os.path.exists(self.structure_file) and not overwrite:
            structure = json.load(open(self.structure_file, 'r'))
        else:
            json.dump(structure, open(self.structure_file, 'w'))
        self.structure = structure
        return structure

def crossmatch(catalog, seperation, method='kdtree', verbose=True):
    algs = {'kdtree': __kd_tree_matching, 'linear': __linear_matching}
    alg = algs[method]
    crossmatched = catalog[:0].copy()
    return alg(catalog, seperation, crossmatched, verbose)

def build_kd_tree(catalog, seperation):
    coords = np.zeros((len(catalog), 2), dtype=np.float64)
    coords[:, 0] = catalog['ra']
    coords[:, 1] = catalog['dec']
    tree = spatial.cKDTree(coords)
    qbt = tree.query_ball_tree(tree, seperation)
    return tree, qbt, coords

def __kd_tree_matching(catalog, seperation, crossmatched, verbose=True):
    _, qbt, coords = build_kd_tree(catalog, seperation)
    flags = np.zeros(coords.shape[0]).astype(bool)

    elmnt = len(flags)
    delta = int(elmnt / 25)
    delta = min(elmnt, delta)
    ctr = 0
    child = []
    for indxes in qbt:
        if flags[indxes[0]]: continue
        child.append(indxes.copy())
        ctr = np.sum(flags)
        if ctr % delta == 0 and verbose: misc.progress(ctr, elmnt, 'crossmatching')

        crossmatched = coadd_catalog(catalog[indxes], crossmatched)
        for idx in indxes:
            ctr += 1
            flags[idx] = True
    return crossmatched, child

def __linear_matching(catalog, seperation, crossmatched, verbose=True):
    assert False
    ref = catalog.copy()
    coords = SkyCoord(ra=catalog['ra'], dec=catalog['dec'], unit='radian')

    elmnt = len(coords)
    delta = int(elmnt / 25)
    delta = min(elmnt, delta)
    ctr = 0
    while len(coords) > 0:
        if ctr % delta == 0 and verbose: misc.progress(ctr, elmnt, 'crossmatching')

        coord = coords[0]
        loc = coord.separation(coords) < seperation * astropy.units.rad
        crossmatched = coadd_catalog(ref[loc], crossmatched)
        ref = ref[np.invert(loc)]
        coords = coords[np.invert(loc)]
        ctr += 1
    del ref, coords

    return crossmatched, None

def is_column(catalog, name):
    return name in catalog.colnames

def coadd_catalog(catalog, out_catalog=None):
    ivar = 1. / catalog['sigma_jy'] ** 2
    tot_weight = np.sum(ivar)

    def ivar_sum(vals, ivar=ivar, tot_weight=tot_weight):
        return np.sum((vals * ivar) / tot_weight)

    def char_sum(vals, compress=False):
        ret = ','.join(vals.astype(str).astype(object))
        if compress:
            ret = 'zlib_' + str(zlib.compress(ret.encode('utf8')))
        return ret

    t = catalog[:0].copy()
    t.add_row(np.empty(1, dtype=t.dtype)[0])
    t[0]['ra'] = ivar_sum(catalog['ra'])
    t[0]['dec'] = ivar_sum(catalog['dec'])
    t[0]['jy'] = ivar_sum(catalog['jy'])
    t[0]['sigma_jy'] = 1 / np.sqrt(tot_weight)
    t[0]['snr'] = np.sqrt(np.sum(catalog['snr'] ** 2))

    map_details = ['season', 'patch', 'freq']
    for idx in map_details:
        if is_column(t, idx): t[0][idx] = char_sum(catalog[idx])

    if out_catalog is None:
        out_catalog = t
    else:
        out_catalog = astropy.table.vstack([out_catalog, t])
    return out_catalog

def calibrate_flux(catalog, temp_cal=1.0, pol_cal=1.0, inplace=False):
    ret = catalog if inplace else catalog.copy()
    ret['jy'] *= temp_cal
    ret['sigma_jy'] *= temp_cal
    return ret

def get_common_colnames(catalogs):
    common_cols = None
    for cat in catalogs:
        cur_cols = set(cat.colnames)
        if common_cols is None:
            common_cols = cur_cols
        else:
            common_cols = common_cols & cur_cols
    return common_cols

def standarize_catalogs(catalogs, inplace=False):
    ret = copy.deepcopy(catalogs) if not inplace else catalogs
    common_cols = get_common_colnames(ret)
    for i, cat in enumerate(ret):
        ret[i] = cat[list(common_cols)]
    return ret

def apply_cut(cat, col, constrain, inplace=False):
    exc = "cat['%s']%s" % (col, constrain)
    ret = cat.copy() if not inplace else cat
    try:
        loc = np.where(eval(exc))
        ret = cat[loc]
    except:
        print("can't execute", exc)
    return ret

def _postprocess_catalog(catalog, dtype='astropy'):
    def _convert_table(catalog, dtype='astropy'):
        if dtype not in ['astropy', 'pandas', 'numpy']:
            print("can't recognize dtype={}. return the catalog as astropy table".format(dtype))

        if dtype == 'pandas':
            catalog = catalog.to_pandas()
        elif dtype == 'numpy':
            catalog = np.array(catalog.as_array(), dtype=catalog.dtype).view(np.recarray)
        else:
            pass

        return catalog

    def _add_units(catalog):
        if is_column(catalog, 'ra'): catalog['ra'].unit = u.rad
        if is_column(catalog, 'dec'): catalog['dec'].unit = u.rad
        if is_column(catalog, 'jy'): catalog['jy'].unit = u.Jy
        if is_column(catalog, 'sigma_jy'): catalog['sigma_jy'].unit = u.Jy

        return catalog

    for idx in catalog.colnames:
        if '<U' in str(catalog[idx].dtype):
            catalog[idx] = catalog[idx].astype('object')

    if 'name' in catalog.colnames: catalog['name'] = catalog['name'].astype(object)
    if 'child' in catalog.colnames: catalog['child'] = catalog['child'].astype(str).astype(object)
    if 'parent' in catalog.colnames: catalog['parent'] = catalog['parent'].astype(str).astype(object)
    if 'note' in catalog.colnames: catalog['note'] = catalog['note'].astype(str).astype(object)
    if 'idx' in catalog.colnames: catalog['idx'] = catalog['idx'].astype(str).astype(object)
    if dtype in ['pandas', 'numpy']:
        catalog = _convert_table(catalog, dtype=dtype)
    else:
        catalog = _add_units(catalog)
    return catalog

def read_crush(fname, dtype="astropy"):
    ret = ascii.read(fname)
    return _postprocess_catalog(ret, dtype=dtype)

def read_nemo(fname, dtype='astropy'):
    # # took it from pixell.pointsrcs and modified to be compatible with the output from the latest version of nemo (
    # Oct 14th, 2019)
    """Reads the nemo catalog
    Args:
        dtype:
        fname:
    """
    ret = ascii.read(fname)
    ret.rename_columns(['RADeg', 'decDeg', 'deltaT_c', 'err_deltaT_c', 'SNR', 'fluxJy', 'err_fluxJy'],
                       ['ra', 'dec', 'I', 'dI', 'snr', 'jy', 'sigma_jy'])
    ret['ra'] *= utils.degree
    ret['dec'] *= utils.degree
    del ret['template'], ret['tileName'], ret['numSigPix'], ret['galacticLatDeg'], ret['name'], ret['I'], ret['dI']
    return _postprocess_catalog(ret)

import numpy,from scipy,import astropy,from astropy


import re

import nltk
import numpy

nltk.download('punkt')

import warnings
warnings.filterwarnings("ignore")

from sklearn.feature_extraction.text import TfidfVectorizer
from tqdm import tqdm
from gensim import corpora, similarities
from gensim.models import TfidfModel
from gensim.corpora import Dictionary, MmCorpus
from gensim.similarities import Similarity
from nltk.stem import PorterStemmer

ps = PorterStemmer()

from sklearn.metrics import f1_score
from sklearn.metrics import precision_score
from sklearn.metrics import recall_score
from sklearn.metrics import accuracy_score
from sklearn.linear_model import LogisticRegression

from sklearn.naive_bayes import MultinomialNB
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.ensemble import BaggingClassifier
from sklearn.ensemble import VotingClassifier
from sklearn.ensemble import AdaBoostClassifier
from xgboost import XGBClassifier

import pandas as pd
import preprocessor as tp

pd.set_option('display.max_rows', None)
pd.set_option('display.max_columns', None)

tp.set_options(tp.OPT.URL,tp.OPT.MENTION)

#from imblearn.ensemble import BalancedBaggingClassifier
#https://scikit-learn.org/stable/modules/generated/sklearn.linear_model.LogisticRegression.html

import random
import copy
import argparse
import scipy.sparse as sp
import numpy as np
import time
import math
from random import randint
try:
    import queue
except ImportError:
    import Queue as queue

import pickle
import os
from multiprocessing import Pool as ProcessPool
import itertools
from functools import partial

# import user pythons file
from topic_description import TRECTopics
from systemReader import systemReader
from global_definition import *

rng = np.random.seed(5)
random.seed(5)

# actual active learning for TREC is happening here for a particular topicID
# here we run for either all documents in the collection
# or all documents in the official qrels
def active_learning_multi_processing(topicId, df, al_protocol, al_classifier, document_collection, topic_seed_info, topic_complete_qrels_address, train_per_centage, use_pooled_budget, per_topic_budget_from_trec_qrels, feature_type):
    train_index_list = topic_seed_info[topicId]

    #topic_complete_qrels = pickle.load(open(topic_complete_qrels_address + topicId + '.pickle', 'rb'))

    original_labels = {}
    for row_index, row in df.iterrows():
        original_labels[row_index] = row['majority_label']

    original_predicted_merged_dict = {}
    original_labels_list = []
    number_of_1 = 0
    for k, v in original_labels.items():
        original_predicted_merged_dict[k] = v
        if v == 1.0:
            number_of_1 = number_of_1 + 1
        original_labels_list.append(v)
    #exit(0)

    #print "tmp_l1:",original_labels_list.count(1)

    #print "tmp_l2:",predicted_labels_list.count(1)

    #print "sum", original_labels_list.count(1) + predicted_labels_list.count(1)

    original_predicted_merged_list = []
    for k in sorted(original_predicted_merged_dict.keys()):
        #print k, original_predicted_merged_dict[k]
        original_predicted_merged_list.append(original_predicted_merged_dict[k])

    #print "again sum", original_predicted_merged_list.count(1)

    # need to convert y to np.array the Y otherwise Y[train_index_list] does not work directly on a list
    y = np.array(original_predicted_merged_list)  # 2 is complete labels of all documents in document collection
    # type needed because y is an object need and throws error Unknown label type: 'unknown'
    y = y.astype('int')
    #print "numpy sum", np.count_nonzero(y)
    #print y

    #print y.shape
    #print train_index_list
    #print y[train_index_list]

    #exit(0)

    total_documents = len(y)
    total_document_set = set(np.arange(0, total_documents, 1))

    initial_X_test = []
    test_index_dictionary = {}
    test_index_counter = 0

    #print "Starting Test Set Generation:"
    #start = time.time()
    for train_index in range(0, total_documents):
        if train_index not in train_index_list:
            initial_X_test.append(document_collection[train_index])
            test_index_dictionary[test_index_counter] = train_index
            test_index_counter = test_index_counter + 1

    #print "Finshed Building Test Set:", time.time() - start

    predictableSize = len(initial_X_test)
    isPredictable = [1] * predictableSize  # initially we will predict all

    # initializing the train_size controller
    train_size_controller = len(train_index_list)
    loopCounter = 1  # loop starts from 1 because 0 is for seed_set
    topic_all_info = {}  # key is the loopCounter

    while True:
        #print "iteration:", loopCounter
        # here modeling is utilizing the document selected in previous
        # iteration for training
        # when loopCounter == 0
        # model is utilizing all the seed document collected at the begining
        if al_classifier == 'LR':
            #model = LogisticRegression(solver=large_data_solver, C=large_data_C_parameter, max_iter=200)
            model = LogisticRegression(C=1.0, class_weight=None, dual=False, fit_intercept=True,
                   intercept_scaling=1, l1_ratio=None, max_iter=100,
                   n_jobs=None, penalty='l2',
                   random_state=None, solver='saga', tol=0.0001, verbose=0,
                   warm_start=False)

        elif al_classifier == 'SVM':
            model = SVC(C=1.0, kernel='linear', degree=3, gamma='auto', probability = True)
        elif al_classifier == 'RF':
            model =  RandomForestClassifier(n_estimators=10, max_depth=10, random_state=0)
        elif al_classifier == 'RFN':
            model = RandomForestClassifier(n_estimators=10, max_depth=None, random_state=0)
        elif al_classifier == 'RFN100':
            model = RandomForestClassifier(n_estimators=100, max_depth=None, random_state=0)
        elif al_classifier == 'NB':
            model = MultinomialNB()
        elif al_classifier == 'Ada':
            # base model is decision tree
            # logistic regression will not help
            model = AdaBoostClassifier(n_estimators=50,
                                     learning_rate=1)
        elif al_classifier == 'Xgb':
            model = XGBClassifier(random_state=1, learning_rate=0.01)
        elif al_classifier == 'BagLR':
            LRmodel = LogisticRegression(solver=large_data_solver, C=large_data_C_parameter, max_iter=200)
            model = BaggingClassifier(LRmodel, n_estimators = 5, max_samples = 1) # If float, then draw max_samples * X.shape[0] samples. 1 means use all samples
        elif al_classifier == 'BagNB':
            model = BaggingClassifier(MultinomialNB(), n_estimators = 5, max_samples = 0.5) # If float, then draw max_samples * X.shape[0] samples. 1 means use all samples
        elif al_classifier == 'Vot':
            LRmodel = LogisticRegression(solver=large_data_solver, C=large_data_C_parameter, max_iter=200)
            NBmodel = MultinomialNB()
            model = VotingClassifier(estimators=[('lr', LRmodel), ('nb', NBmodel)], voting = 'soft')

        model.fit(document_collection[train_index_list], y[train_index_list])

        #print(model.coef_)
        #print(len(model.coef_[0]))

        test_index_list = list(total_document_set - set(train_index_list))
        pooled_document_count = len(set(train_index_list).intersection(set(original_labels_list)))

        y_actual = None
        y_pred = None
        y_pred_all = []

        if isPredictable.count(1) != 0:
            y_pred = []
            for test_index_elem in test_index_list:
                if feature_type == "tfidf":
                    y_pred.append(model.predict(document_collection[test_index_elem]))
                elif feature_type == "bert" or feature_type == "robert":
                    y_pred.append(model.predict([document_collection[test_index_elem]]))

            start = time.time()
            #print 'Statred y_pred_all'
            y_actual = np.concatenate((y[train_index_list], y[test_index_list]), axis=None)
            y_pred_all = np.concatenate((y[train_index_list], y_pred), axis=None)
            '''
            for doc_index in range(0,total_documents):
                if doc_index in train_index_list:
                    y_pred_all.append(y[doc_index])
                else:
                    # result_index in test_set
                    # test_index_list is a list of doc_index
                    # test_Index_list [25, 9, 12]
                    # test_index_list[0] = 25 and its prediction in y_pred[0] --one to one mapping
                    # so find the index of doc_index in test_index_list using
                    pred_index = test_index_list.index(doc_index)
                    y_pred_all.append(y_pred[pred_index])
            '''
            #print "Finsh y_pred_all", time.time() - start

        else: # everything in trainset
            y_pred = y
            y_actual = y
            y_pred_all = y
            test_index_list = train_index_list

        f1score = f1_score(y_actual, y_pred_all, average='binary')
        precision = precision_score(y_actual, y_pred_all, average='binary')
        recall = recall_score(y_actual, y_pred_all, average='binary')

        number_of_1_found_so_far = list(y[train_index_list]).count(1)
        prevalence = (number_of_1_found_so_far*1.0)/number_of_1
        #print f1score, precision, recall, len(train_index_list), len(test_index_list), len(y_pred_all)

        # save all info using (loopCounter - 1)
        # list should be deep_copy otherwise all will point to final referecne at final iterraion
        topic_all_info[loopCounter - 1] = (topicId, f1score, precision, recall, copy.deepcopy(train_index_list), test_index_list, y_pred, pooled_document_count, prevalence)

        #print(loopCounter - 1, number_of_1_found_so_far, number_of_1, prevalence)

        # it means everything in the train list and we do not need to predict
        # so we do not need any training of the model
        # so break here
        if isPredictable.count(1) == 0:
            break
        #print isPredictable.count(1)
        # suppose original budget is 5,
        # then when train_index_list is 5, we cannot just turn off Active learning
        # we need to use that AL with train_index_list of size 5 to train use that to predict the rest
        # so we cannot exit at 5, we should exit at 5 + 1
        # that is the reason we set per_topic_budget_from_trec_qrels[topicId] + 1 where 1 is the batch_size
        # it means everything of pooled_budget size in the train_list so we need not tany training of the model
        # so break here
        if use_pooled_budget == 1 and per_topic_budget_from_trec_qrels[topicId] == len(train_index_list):
            break

        queueSize = isPredictable.count(1)
        my_queue = queue.PriorityQueue(queueSize)

        # these are used for SPL
        randomArray = []

        for counter in range(0, predictableSize):
            if isPredictable[counter] == 1:
                # model.predict returns a list of values in so we need index [0] as we
                # have only one element in the list
                y_prob = None
                if feature_type == "tfidf":
                    y_prob = model.predict_proba(initial_X_test[counter])[0]
                elif feature_type == "bert" or feature_type == "robert":
                    y_prob = model.predict_proba([initial_X_test[counter]])[0]

                #print(y_prob)
                val = 0
                if al_protocol == 'CAL':
                    val = (-1)*y_prob[1] # -1 is needed priority do sorting increasing
                    my_queue.put((val, counter))
                elif al_protocol == 'SAL':
                    val = (-1)*calculate_entropy(y_prob[0], y_prob[1])
                    my_queue.put((val, counter))
                elif al_protocol == 'SPL':
                    randomArray.append(counter)

        if use_pooled_budget == 1:
            #print "use pooled budget"
            size_limit = math.ceil(train_per_centage[loopCounter] * per_topic_budget_from_trec_qrels[topicId])
            #print("size limit:", size_limit, "total_docs:", per_topic_budget_from_trec_qrels[topicId])

        else:
            size_limit = math.ceil(train_per_centage[loopCounter] * total_documents)
            #print("size limit:", size_limit, "total_docs:", total_documents)
        if al_protocol == 'SPL':
            random.shuffle(randomArray)
            #randomArray.reverse()
            batch_counter = 0
            # for batch_counter in range(0, batch_size):
            #    if batch_counter > len(randomArray) - 1:
            #        break
            while True:
                if train_size_controller == size_limit:
                    break

                itemIndex = randomArray[batch_counter]
                isPredictable[itemIndex] = 0
                train_index_list.append(test_index_dictionary[itemIndex])
                train_size_controller = train_size_controller + 1
                batch_counter = batch_counter + 1

        else:
            while not my_queue.empty():
                if train_size_controller == size_limit:
                    break
                item = my_queue.get()
                #print(item) # is a tuple where item[1] is the index , item[0] is the predit value
                isPredictable[item[1]] = 0  # not predictable

                train_index_list.append(test_index_dictionary[item[1]])
                train_size_controller = train_size_controller + 1

        loopCounter = loopCounter + 1
    return topic_all_info

def active_learning(topic_list, df, al_protocol, al_classifier, document_collection, topic_seed_info, topic_complete_qrels_address,train_per_centage, data_path, file_name, use_pooled_budget, per_topic_budget_from_trec_qrels, feature_type):
    num_workers = None
    workers = ProcessPool(processes = 1)
    with tqdm(total=len(topic_list)) as pbar:
        partial_active_learning_multi_processing = partial(active_learning_multi_processing, df=df, al_protocol=al_protocol, al_classifier = al_classifier, document_collection=document_collection,topic_seed_info=topic_seed_info,topic_complete_qrels_address=topic_complete_qrels_address,train_per_centage=train_per_centage, use_pooled_budget=use_pooled_budget, per_topic_budget_from_trec_qrels=per_topic_budget_from_trec_qrels, feature_type=feature_type)
        for topic_all_info in tqdm(workers.imap_unordered(partial_active_learning_multi_processing, topic_list)):
            topicId = topic_all_info[0][0] # 0 is the loopCounter Index and 0 is the first tuple
            file_complete_path = data_path + file_name + str(topicId) + ".pickle"
            pickle.dump(topic_all_info, open(file_complete_path, 'wb'))
            pbar.update()

def text_preprocessing(s):
    s = str(s)
    s = s.lower()
    # Change 't to 'not'
    s = re.sub(r"\'t", " not", s)
    # Remove @name
    s = re.sub(r'(@.*?)[\s]', ' ', s)
    # Isolate and remove punctuations except '?'
    s = re.sub(r'([\'\"\.\(\)\!\?\\\/\,])', r' \1 ', s)
    s = re.sub(r'[^\w\s\?]', ' ', s)
    # Remove some special characters
    s = re.sub(r'([\;\:\|\n])', ' ', s)
    # Remove stopwords except 'not' and 'can'
    #s = " ".join([word for word in s.split()
    #              if word not in stopwords.words('english')
    #              or word in ['not', 'can']])
    # Remove trailing whitespace
    s = s.replace("https", "")
    s = re.sub(r'\s+', ' ', s).strip()

    s = re.compile('RT @').sub('@', s, count=1)
    s = s.replace(":", "").strip()
    s = tp.clean(s)

    tokens = nltk.word_tokenize(s)
    stems = []
    for item in tokens:
        stems.append(ps.stem(item))

    #print(s)

    #print(stems)
    return stems

#########################################################################
# main script
# create all directories manually before running otherwise multi-processing will create lock condition
# for creating files

if __name__ == '__main__':

    # read the processed csv file in a dataframe

    input_file = "/home/nahid/Downloads/data_all/version2/mturk/processed/only_raw_manual_agrees.csv"
    hate_df = pd.read_csv(input_file)

    classifier_name = "LR"
    classifier = None
    if classifier_name == 'LR':
        classifier = LogisticRegression(solver=small_data_solver,C=small_data_C_parameter)

    # only IS and, RDS
    topicData = TRECTopics("Hate", 1, 2)
    seed_selection_type = "IS"
    topic_seed_info_file_name = "per_topic_seed_documents_" + seed_selection_type
    number_of_seeds = 10
    data_path = "/home/nahid/Downloads/data_all/version2/AL/"
    file_name = "seed_dcouments"
    seed_selection_type = "IS"
    topic_seed_info = topicData.get_seed_tweets(hate_df, number_of_seeds, data_path, file_name, seed_selection_type)

    print(topic_seed_info)

    '''
    # sanity check
    for topicId in sorted(topic_seed_info.iterkeys()):
       #print topicId, len(topic_seed_info[topicId]), topic_seed_info[topicId]
       print topicId, len(topic_seed_info[topicId])
    '''

    feature_type_list = ["tfidf", "robert"]
    al_protocol_list = ['CAL', 'SAL', 'SPL']

    #v = TfidfVectorizer(analyzer=text_preprocessing, min_df=1, max_features=4500)
    v = TfidfVectorizer(analyzer=text_preprocessing, ngram_range=(1, 3), binary=True, smooth_idf=False)

    document_collection = v.fit_transform(hate_df['tweet'])

    start = time.time()
    # topic_list = [str(topicID) for topicID in range(start_topic[datasource], end_topic[datasource])]
    start_top = 1
    end_top = 2
    use_pooled_budget = 0
    #al_classifier = "LR"
    al_classifier = classifier_name
    per_topic_budget_from_trec_qrels = None

    topic_list = [str(topicID) for topicID in range(start_top, end_top)]
    print("topic_list", topic_list)

    loaded_bert_features = pickle.load(open("/home/nahid/Downloads/data_all/mturk/processed/bert_feature.pickle", "rb"))

    loaded_robert_features = pickle.load(open("/home/nahid/Downloads/data_all/version2/mturk/processed/robert_feature.pickle", "rb"))
    loaded_robert_features = numpy.array(loaded_robert_features)

    print(loaded_robert_features.shape)
    print(len(loaded_robert_features))
    print(len(loaded_robert_features[0]))
    print(loaded_robert_features[0])

    for feature_type in feature_type_list:
        if feature_type == "bert":
            document_collection = loaded_bert_features
        elif feature_type == "robert":
            document_collection = loaded_robert_features
        for al_protocol in al_protocol_list:
            print(al_protocol)

            classifier_name = "LR"
            classifier = None
            if classifier_name == 'LR':
                classifier = LogisticRegression(solver=small_data_solver, C=small_data_C_parameter)

            topic_complete_qrels_address = data_path + "per_topic_complete_qrels_" + classifier_name + "_"
            topic_all_info_file_name = "per_topic_predictions_" + seed_selection_type + "_" + classifier_name + "_" + al_protocol + "_" + feature_type + "_"
            # print data_path + topic_all_info_file_name + topic_list[0] + '.pickle'
            topic_all_info = active_learning(topic_list, hate_df, al_protocol, al_classifier, document_collection,
                                             topic_seed_info, topic_complete_qrels_address,
                                             train_per_centage, data_path, topic_all_info_file_name,
                                             use_pooled_budget, per_topic_budget_from_trec_qrels, feature_type)

            #print(v.get_feature_names())
            #print(len(v.get_feature_names()))
            #print(len(document_collection))

            # sanity check
            for topicId in topic_list:
                print("topicId:", topicId)

                topic_all_info = pickle.load(open(data_path+ topic_all_info_file_name+ str(topicId) + ".pickle",'rb'))
                for k in sorted(topic_all_info.keys()):
                    #(0      ,       1,         2,      3,                              4,               5,   6
                    # topicId, f1score, precision, recall, copy.deepcopy(train_index_list), test_index_list, y_pred, pooled_document_count)

                    print (k*10, topic_all_info[k][0], topic_all_info[k][1], topic_all_info[k][3], len(topic_all_info[k][4]), len(topic_all_info[k][5]), len(topic_all_info[k][6]))

    '''
    datasource = sys.argv[1]  # can be 'TREC8','gov2', 'WT2013','WT2014'
    al_protocol = sys.argv[2]  # 'SAL', 'CAL', # SPL is not there yet
    seed_selection_type = sys.argv[3]  # 'IS' only
    classifier_name = sys.argv[4]  # "LR", "NR"--> means non-relevant all
    collection_size = sys.argv[5]  # 'all', 'qrels' qrels --> means consider documents inseide qrels only
    al_classifier = sys.argv[6]  # SVM, RF, NB and LR
    start_top = int(sys.argv[7])
    use_pooled_budget = int(sys.argv[8])  # 1 means use and 0 does not use that
    use_original_qrels = int(sys.argv[9])  # 1 means use original qrels, other value 0 means
    varied_qrels_directory_number = int(sys.argv[10])  # 1,2,3
    end_top = start_top + 1
    '''


import _init_paths

import os
import sys
from fast_rcnn.config import cfg, cfg_from_file, cfg_from_list
from fast_rcnn.test import im_detect
from fast_rcnn.nms_wrapper import nms
import caffe
import os
import cv2
from utils.timer import Timer
import numpy as np
import time

caffe.set_mode_gpu()
caffe.set_device(3)

#cfg_from_file("/tmp/test/submit_1019.yml")
cfg_from_file("/tmp/test/submit_0716.yml")
prototxt = "/tmp/test/weaponModel_test.prototxt"
caffemodel = "/tmp/test/weaponModel_iter_6000.caffemodel"
net = caffe.Net(prototxt, caffemodel, caffe.TEST)

im = cv2.imread("/tmp/test/test.jpg")

_t = {'im_preproc': Timer(), 'im_net': Timer(),
      'im_postproc': Timer(), 'misc': Timer()}

scores, boxes = im_detect(net, im, _t)

for i in range(10):
    _s = time.time()
    scores, boxes = im_detect(net, im, _t)
    _e = time.time()
    print "time: %s" % (_e-_s)
    time.sleep(1)


from os.path import exists
from pathlib import Path
import numpy as np
from . import util

# comparison of the correlations between CD4/CD8 expression levels and
# either the old cd8 score or this new one on two 10x datasets
# Looks like this new one is a better predictor (which is should be!)
#
# based on this I am switching this to be the default
#
#                   R    Pearson P  Spearman  Kendall  dataset
# old corr: CD4   -0.171  1.10e-19  5.25e-20  1.07e-19 hs_pbmc3
# new corr: CD4   -0.184  1.60e-22  1.03e-23  2.27e-23 hs_pbmc3
# old corr: CD8A   0.360  3.56e-86  1.42e-91  4.98e-85 hs_pbmc3
# new corr: CD8A   0.384  1.57e-98 7.91e-109 6.17e-100 hs_pbmc3
# old corr: CD8B   0.312  5.94e-64  3.23e-74  1.25e-69 hs_pbmc3
# new corr: CD8B   0.329  3.76e-71  1.55e-91  1.28e-84 hs_pbmc3
# old corr: CD4   -0.136  8.38e-08  8.22e-08  1.02e-07 hs_pbmc
# new corr: CD4   -0.154  1.37e-09  4.40e-09  5.52e-09 hs_pbmc
# old corr: CD8A   0.348  3.83e-45  2.75e-46  2.76e-43 hs_pbmc
# new corr: CD8A   0.400  2.49e-60  1.46e-63  1.40e-57 hs_pbmc
# old corr: CD8B   0.258  8.65e-25  2.64e-30  1.75e-28 hs_pbmc
# new corr: CD8B   0.303  5.09e-34  5.89e-45  1.38e-41 hs_pbmc

amino_acids = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', \
               'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']

# read the model parameters
all_models = {}

for ab in 'AB':
    model_params = {
        'window_size':None,
        'min_lenbin':None,
        'max_lenbin':None,
        'NV':None,
        'NJ':None,
        'NL':None,
        'NC':None,
        }
    model_file = Path.joinpath( Path(util.path_to_data),f'cd8_logreg_params_{ab}.txt')
    mtags, weights = [], []
    #print(f'loading cd8 model for chain {ab} {model_file}')
    with open(model_file, 'r') as data:
        for line in data:
            l = line.split()
            tag, value = l
            if tag in model_params:
                model_params[tag] = int(value)
            else:
                mtags.append(tag)
                weights.append(float(value))
    NV, NJ, NL, NC = [model_params[x] for x in 'NV NJ NL NC'.split()]
    NTOT = NV + NJ + NL + NC
    assert mtags[-1] == 'BIAS'
    bias = float(weights[-1])
    assert len(weights) == NTOT+1
    weights = np.array(weights)
    vgenes = mtags[:NV-1] # last is UNK gene
    jgenes = mtags[NV:NV+NJ-1]
    assert all(x.startswith(f'TR{ab}V') for x in vgenes)
    assert all(x.startswith(f'TR{ab}J') for x in jgenes)
    vgene_indexer = {x:i for i,x in enumerate(vgenes)}
    jgene_indexer = {x:i for i,x in enumerate(jgenes)}

    model_params['weights'] = weights
    model_params['vgene_indexer'] = vgene_indexer
    model_params['jgene_indexer'] = jgene_indexer

    all_models[ab] = model_params

def get_lenbin(L, min_lenbin, max_lenbin):
    if L <= min_lenbin:
        return min_lenbin
    elif L >= max_lenbin:
        return max_lenbin
    else:
        return L

def get_allele(g):
    assert g.count('*')==1
    return g[:g.index('*')]

def encode_single_chain_tcr(vgene, jgene, cdr3, model_params):
    ''' vgene and jgene still include '*01' (e.g.)
    '''
    window_size = model_params['window_size']
    min_lenbin = model_params['min_lenbin']
    max_lenbin = model_params['max_lenbin']
    vgene_indexer = model_params['vgene_indexer']
    jgene_indexer = model_params['jgene_indexer']
    NV = len(vgene_indexer)+1
    NJ = len(jgene_indexer)+1
    NL = max_lenbin-min_lenbin+1
    NC = 20 * (2*window_size+1)
    NTOT = NV + NJ + NL + NC

    # 1-hot encode the V/J genes and length
    x = np.zeros((NTOT+1,))
    iv = vgene_indexer.get(get_allele(vgene), NV-1)
    x[iv] = 1.0

    ij = jgene_indexer.get(get_allele(jgene), NJ-1)
    x[NV+ij] = 1.0

    il = get_lenbin(len(cdr3), min_lenbin, max_lenbin) - min_lenbin
    x[NV+NJ+il] = 1.0

    # 1-hot encode aas in the nterminal and cterminal cdr3 windows
    # k-hot encode the middle of the tcr
    nterm = min(window_size, len(cdr3)//2)
    cterm = min(window_size, len(cdr3)-nterm)
    #middle = len(cdr3)-nterm-cterm
    for i in range(window_size):
        if i<nterm:
            istart = NV+NJ+NL+20*i
            x[istart+amino_acids.index(cdr3[i])] += 1.0
        if i<cterm:
            istart = NV+NJ+NL+20*window_size+20*i
            x[istart+amino_acids.index(cdr3[-1-i])] += 1.0
    for aa in cdr3[nterm:-cterm]:
        assert len(cdr3)>2*window_size # sanity check
        istart = NV+NJ+NL+40*window_size
        x[istart+amino_acids.index(aa)] += 1.0

    x[NTOT] = 1. # for the bias term
    return x

def make_cd8_score_table_column(tcrs, use_sigmoid=False):
    ''' We fit a logistic regression model on single-chain bulk TCR data
    from sorted CD4/CD8 subsets. Here we just average the alpha-chain and beta-chain
    scores.
    '''
    global all_models
    score_totals = np.zeros((len(tcrs),))
    for iab, ab in enumerate('AB'):
        model = all_models[ab]
        tcr_vectors = [ encode_single_chain_tcr(x[iab][0], x[iab][1], x[iab][2], model)
                        for x in tcrs ]
        tcr_vectors = np.vstack(tcr_vectors)
        weights = model['weights']
        assert tcr_vectors.shape == (len(tcrs), weights.shape[0])
        ab_scores = np.dot( tcr_vectors, weights)
        if use_sigmoid:
            ab_scores = 1.0/(1.0+np.exp(-1*ab_scores))
        score_totals += 0.5 * ab_scores
    return score_totals


from datetime import datetime, timedelta

from ccxt import Exchange

from tradingkit.pubsub.core.event import Event
from tradingkit.pubsub.core.publisher import Publisher
from tradingkit.pubsub.event.book import Book
from tradingkit.pubsub.core.subscriber import Subscriber
from tradingkit.pubsub.event.candle import Candle
from tradingkit.pubsub.event.funding import Funding
from tradingkit.pubsub.event.open_order import OpenOrder
from tradingkit.pubsub.event.order import Order
from tradingkit.pubsub.event.plot import Plot
from tradingkit.pubsub.event.trade import Trade
from tradingkit.pubsub.event.liquidation import Liquidation
import numpy

class BridgeExchange(Publisher, Subscriber, Exchange):

    def __init__(self, exchange: Exchange):
        super().__init__()
        self.exchange = exchange
        self.closed_orders = {}
        self.orders_history = {}
        self.balance_history = []
        self.last_balance_check = None
        self.peak_balance = 0
        self.max_drawdown = 0
        self.last_price = None
        self.symbol = None
        self.has_position = True if "bitmex" in str(exchange.__class__) else False
        self.is_backtest = True if "TestEX" in str(exchange.__class__) or "BitmexBacktest" in str(
            exchange.__class__) else False
        exchange.seconds()

        self.candles = None
        self.last_candle = None
        self.timeframes = {
            "1m": 60,
            "5m": 300,
            "15m": 900,
            "30m": 1800,
            "1h": 3600,
            "4h": 14400,
            "6h": 21600,
            "12h": 43200,
            "1d": 86400,
            "2d": 172800,
            "4d": 345600,
            "1w": 604800,
            "2w": 1209600
        }

    def sec(self):
        return self.exchange.sec()

    def msec(self):
        return self.exchange.msec()

    def usec(self):
        return self.exchange.usec()

    def seconds(self):
        return self.exchange.seconds()

    def milliseconds(self):
        return self.exchange.milliseconds()

    def microseconds(self):
        return self.exchange.microseconds()

    def subscribed_events(self) -> list:
        return [Order, Trade, Book, Candle, Liquidation, Funding, OpenOrder]

    def get_order_asset(self, symbol):
        base, quote = symbol.split('/')
        return quote if self.is_bitmex() else base

    def is_bitmex(self):
        exchange_name = str(self.exchange.__class__.__name__).lower()
        return "bitmex" in exchange_name

    def is_inverse(self, symbol):
        return self.is_bitmex()

    def max_leverage(self, symbol=None):
        return 100 if self.is_bitmex() else 1

    def on_event(self, event: Event):
        if isinstance(event, Book):
            self.last_price = event.payload['bids'][0][0]
            if self.last_price is None:
                self.calculate_exchange_state(event.payload['timestamp'], event.payload['symbol'], event.payload['bids'][0][0])
                self.symbol = event.payload['symbol']
        if isinstance(event, Candle):
            self.plot_candle(event)
            if self.is_backtest:
                self.update_balance_hist(event)
        if isinstance(event, Order):
            order = event.payload.copy()
            if order['id'] in self.orders_history.keys():
                self.orders_history[order['id']].update(order)
                event.payload = self.orders_history[order['id']]
            self.plot_order(event)
            self.calculate_exchange_state(order['lastTradeTimestamp'], order['symbol'], self.last_price)
        if isinstance(event, OpenOrder):
            self.plot_order(event)
            order = event.payload.copy()
            self.calculate_exchange_state(order['timestamp'], order['symbol'], self.last_price)
        if isinstance(event, Liquidation):
            trade = event.payload
            self.calculate_exchange_state(trade['timestamp'], trade['symbol'], trade['price'])
        if isinstance(event, Trade):
            trade = event.payload
            self.last_price = event.payload['price']
            if self.symbol is None:
                self.symbol = event.payload['symbol']
            self.candle_dispatcher(trade)
        self.dispatch(event)

    def fetch_open_orders(self, symbol=None, since=None, limit=None, params={}):
        return self.exchange.fetch_open_orders(symbol, since, limit, params)

    def create_order(self, symbol, type, side=None, amount=0, price=None, params={}):
        if amount == 0:
            raise ValueError("Zero order amount is not allowed!!!")
        elif amount < 0 and side is not None:
            raise ValueError("Negative order amount is not allowed!!!")
        if side is None:
            side = "buy" if amount > 0 else "sell"
            if amount < 0:
                amount = -amount
        order = self.exchange.create_order(symbol, type, side, amount, price, params)
        self.orders_history[order['id']] = order
        return order

    def cancel_order(self, order_id, symbol=None, params={}):
        return self.exchange.cancel_order(order_id, symbol, params)

    def match_order(self, trade, order, price, base, quote):
        return self.exchange.match_order(trade, order, price, base, quote)

    def fetch_ticker(self, symbol, **kwargs):
        return self.exchange.fetch_ticker(symbol)

    def private_get_position(self):
        if self.has_position:
            position = self.exchange.private_get_position()
            if not position:
                return [{"currentQty": 0,
                         "homeNotional": 0,
                         "avgEntryPrice": 0,
                         "liquidationPrice": None,
                         "openOrderBuyQty": None,
                         "openOrderSellQty": None,
                         "lastPrice": None,
                         "markPrice": None
                         }]
            else:
                return position
        else:
            raise ValueError("Position not implemented only bitmex !!!")

    def fetch_balance(self):
        return self.exchange.fetch_balance()

    def fetch_free_balance(self):
        return self.exchange.fetch_free_balance()

    def fetch_used_balance(self):
        return self.exchange.fetch_used_balance()

    def fetch_total_balance(self):
        return self.exchange.fetch_total_balance()

    def fetchMarkets(self):
        return self.exchange.fetchMarkets()

    def get_max_draw_down(self):
        base, quote = self.symbol.split('/')
        balances = self.fetch_balance()
        self.calculate_max_drawdown(balances['total'][base], balances['total'][quote])
        return self.max_drawdown

    def getPairs(self, symbol):
        return symbol

    def plot_balances(self, exchange_state):
        quote = exchange_state['quote']
        base = exchange_state['base']

        ttip = "Total: %.2f %s
Quote: %.2f %s
Base: %.2f %s (in %s)" % (
            exchange_state['equity'],
            quote,
            exchange_state['quote_balance'],
            exchange_state['quote'],
            exchange_state['base_balance'] * exchange_state['price'],
            quote,
            base
        )
        self.dispatch(Plot({
            'name': 'Equity',
            'type': 'scatter',
            'mode': 'lines',
            'color': 'blue',
            'yaxis': 'balance',
            'quote': quote,
            'base': base,
            'data': {
                'x': exchange_state['exchange_date'],
                'y': exchange_state['equity'],
                'base_equity': exchange_state['base_equity'],
                'base_balance': exchange_state['all_balances']['total'][base],
                'quote_balance': exchange_state[
                    'all_balances']['total'][quote] if quote in exchange_state['all_balances'][ 'total'] else 0,
                'position_vol': exchange_state['position_vol'],
                'position_price': exchange_state['position_price'],
                'invested': exchange_state['base_balance'] * exchange_state['price'],
                'tooltip': ttip
            },
        }))

    def calculate_exchange_state(self, price, symbol, timestamp):
        exchange_date = datetime.fromtimestamp(timestamp / 1000.0).isoformat()
        base, quote = symbol.split('/')
        all_balances = self.fetch_balance()
        balances = all_balances['free'] if all_balances['free'][base] else all_balances['total']
        base_balance = balances[base] if base in balances else 0
        quote_balance = balances[quote] if quote in balances else 0
        equity = quote_balance + base_balance * price
        base_equity = base_balance + quote_balance / price
        if self.has_position:
            position = self.private_get_position()[0]
            position_vol = position['currentQty']
            position_price = position['avgEntryPrice']

        else:
            position_vol = base_balance * price
            position_price = 0

        self.plot_balances({
            'all_balances': all_balances,
            'base_balance': base_balance,
            'base_equity': base_equity,
            'equity': equity,
            'exchange_date': exchange_date,
            'position_price': position_price,
            'position_vol': position_vol,
            'quote': quote,
            'base': base,
            'price': price,
            'quote_balance': quote_balance
        })
        if self.is_backtest:
            self.calculate_max_drawdown(all_balances['total'][base], all_balances['total'][quote])

    def plot_order(self, event):
        order = event.payload
        name = None
        if order['status'].lower() == 'filled':
            fill_date = datetime.fromtimestamp(order['lastTradeTimestamp'] / 1000.0).isoformat()
            name = 'buy' if order['side'] == 'buy' else 'sell'
        else:
            fill_date = datetime.fromtimestamp(order['timestamp'] / 1000.0).isoformat()
            if order['status'].lower() == 'open':
                name = 'open_buy' if order['side'] == 'buy' else 'open_sell'
            elif order['status'].lower() == 'canceled':
                name = 'cancel_buy' if order['side'] == 'buy' else 'cancel_sell'
        if name:
            self.dispatch(Plot({
                'name': name,
                'type': 'scatter',
                'mode': 'markers',
                'color': 'green' if order['side'] == 'buy' else 'red',
                'yaxis': 'price',
                'data': {
                    'x': fill_date,
                    'y': order['price'],
                    'type': order['type'],
                    'id': order['id'],
                    'tooltip': "%.4f@%.2f" % (
                        order['amount'] if order['side'] == 'buy' else -order['amount'],
                        order['price']
                    )
                },
            }))

    def plot_candle(self, event):
        candle = event.payload
        candle['liquidationPrice'] = self.private_get_position()[0]['liquidationPrice'] if self.has_position else None
        self.dispatch(Plot({
            'name': 'Price',
            'type': 'candlestick',
            'yaxis': 'price',
            'data': candle,
        }))

    def calculate_max_drawdown(self, base_balance, quote_balance):
        if self.last_price:
            balance = quote_balance + base_balance * self.last_price if not self.has_position else base_balance

            if self.has_position:
                position = self.private_get_position()[0]
                if abs(position['currentQty']) > 0:
                    pnl = (self.last_price / position['avgEntryPrice'] * position['currentQty'] -
                           position['currentQty']) / self.last_price
                    balance += pnl / self.last_price

            if balance > self.peak_balance:
                self.peak_balance = balance

            drawdown = (balance - self.peak_balance) / self.peak_balance
            self.max_drawdown = min(self.max_drawdown, drawdown)

    def update_balance_hist(self, event):
        date = datetime.fromisoformat(event.payload['datetime'])
        if self.last_balance_check is None or date - self.last_balance_check > timedelta(hours=1):
            self.last_balance_check = date

            price = event.payload['close']
            base, quote = self.symbol.split('/')
            balances = self.fetch_balance()
            balances = balances['free'] if balances['free'][base] else balances['total']

            base_balance = balances[base] if base in balances else 0
            quote_balance = balances[quote] if quote in balances else 0

            # todo use quote equity for kraken (equity = quote_balance + base_balance * price)
            base_equity = base_balance + quote_balance / price

            self.balance_history.append(base_equity)

    def get_sharpe_ratio(self):
        standard_deviation = numpy.std(self.balance_history)
        sharpe_ratio = (self.balance_history[-1] / self.balance_history[0]) / standard_deviation
        return sharpe_ratio

    def candle_dispatcher(self, trade):

        if self.candles is None:
            self.candles = {x: {} for x in self.timeframes.keys()}

        if self.last_candle is None:
            self.last_candle = {x: None for x in self.timeframes.keys()}

        for tf in self.timeframes.keys():
            sec = self.timeframes[tf]
            key = trade['timestamp'] // (sec * 1000) * sec
            key = str(datetime.fromtimestamp(key))
            if self.last_candle[tf] is not None and key in self.candles[tf]:
                self.candles[tf][key]['high'] = max(self.candles[tf][key]['high'], trade['price'])
                self.candles[tf][key]['low'] = min(self.candles[tf][key]['low'], trade['price'])
                self.candles[tf][key]['close'] = trade['price']
                self.candles[tf][key]['vol'] += trade['amount']
                self.candles[tf][key]['cost'] += trade['cost']
                self.candles[tf][key]['trades'] += 1
            else:
                self.candles[tf][key] = {
                    'datetime': key,
                    'open': trade['price'],
                    'high': trade['price'],
                    'low': trade['price'],
                    'close': trade['price'],
                    'vol': trade['amount'],
                    'cost': trade['cost'],
                    'trades': 1,
                    'timeframe': tf
                }
                if self.last_candle[tf] is not None:
                    candle = Candle(self.last_candle[tf])
                    self.dispatch(candle)
                    self.plot_candle(candle)
                    if self.is_backtest:
                        self.update_balance_hist(candle)
            self.last_candle[tf] = self.candles[tf][key]