'''
Deconvolution Pipeline
----------------------

The deconvolution process can be pipelined from start to finish
using the :class:`~.ScanProcessor` class. This includes precursor recalculation and
coisolation detection. :class:`~.ScanProcessor` is intended to be used either as a
replacement for :class:`~.ScanIterator` when deconvolution is desired, or that it's
:meth:`ScanProcessor.process` method will be used to handle individual :class:`~.ScanBunch`
objects/pairs of precursor :class:`~.Scan` objects and a :class:`list` of product :class:`~.Scan`
objects.

The free-function :func:`~.process` is a convenience wrapper around :class:`~.ScanProcessor`,
with fewer configurable parameters.

.. code:: python

    from ms_deisotope import ScanProcessor, glycopeptide, peptide
    from ms_deisotope.scoring import PenalizedMSDeconVFitter, MSDeconVFitter
    from ms_deisotope.test.common import datafile

    # Locate example dataset
    path = datafile("20150710_3um_AGP_001_29_30.mzML.gz")

    proc = processor.ScanProcessor(path, ms1_deconvolution_args={
        "averagine": glycopeptide,
        "scorer": PenalizedMSDeconVFitter(20., 2.),
        "truncate_after": 0.95
    }, msn_deconvolution_args={
        "averagine": peptide,
        "scorer": MSDeconVFitter(10.),
        "truncate_after": 0.8
    })

    bunch = next(proc)
    print(bunch)
    print(bunch.precursor.deconvoluted_peak_set)
'''
import logging

from six import string_types as basestring

import numpy as np

from ms_peak_picker import pick_peaks, PeakSet, PeakIndex
from ms_peak_picker.scan_filter import FTICRBaselineRemoval

from .averagine import AveragineCache, peptide
from .scoring import PenalizedMSDeconVFitter, MSDeconVFitter
from .deconvolution import deconvolute_peaks
from .data_source import MSFileLoader, ScanIterator
from .data_source.common import Scan, ScanBunch, ChargeNotProvided
from .utils import Base
from .peak_dependency_network import NoIsotopicClustersError
from .qc.isolation import PrecursorPurityEstimator
from .task import LogUtilsMixin

logger = logging.getLogger("deconvolution_scan_processor")
logger.addHandler(logging.NullHandler())

def _get_nearest_index(query_mz, peak_list):
    best_index = None
    best_error = float('inf')

    for i, peak in enumerate(peak_list):
        error = abs(peak.mz - query_mz)
        if error < best_error:
            best_error = error
            best_index = i
    return best_index

class PriorityTarget(Base):
    """Represent a targeted envelope deconvolution's parameters and constraints.

    This class is used to tell :func:`ms_deisotope.deconvolution.deconvolute_peaks`
    that the solution produced by this peak should be preferentially extracted.

    Attributes
    ----------
    info : :class:`~.PrecursorInformation`
        The associated precursor information block which contains
        the charge state hint.
    peak : :class:`~.FittedPeak`
        The peak from which to start the deconvolution
    trust_charge_hint : bool
        Whether or not to force the deconvoluter to only consider
        the charge specified in the hint.
    mz : float
        The m/z of :attr:`peak`
    charge : int
        The charge state hint from :attr:`info`
    isolation_window : :class:`~.IsolationWindow`
        The isolation window for this precursor ion. May be `None`
    """

    def __init__(self, peak, info, trust_charge_hint=True, precursor_scan_id=None,
                 product_scan_id=None, isolation_window=None):
        self.peak = peak
        self.info = info
        self.trust_charge_hint = trust_charge_hint
        self.precursor_scan_id = precursor_scan_id
        self.product_scan_id = product_scan_id
        self.isolation_window = isolation_window

    def __iter__(self):
        yield self.peak
        yield self.info

    @property
    def mz(self):
        '''
        The m/z of the matched peak

        Returns
        -------
        float
        '''
        return self.peak.mz

    @property
    def charge(self):
        '''
        The charge state of the precursor ion reported by the source.

        Returns
        -------
        int
        '''
        try:
            return int(self.info.charge)
        except TypeError:
            return 0

    def charge_range_hint(self, charge_range):
        """Create an updated charge range for a Deconvoluter to search.

        At the moment, this only amounts to either returning the charge
        range unchanged or returning a charge range that only contains
        the hinted charge state, depending upon whether :attr:`trust_charge_hint`
        is `False` or not.

        Parameters
        ----------
        charge_range : tuple
            The charge range to update

        Returns
        -------
        tuple
            The updated charge range
        """
        if self.trust_charge_hint and self.info.charge is not ChargeNotProvided:
            return (self.charge, self.charge)
        else:
            return charge_range

    def __repr__(self):
        return "PriorityTarget(mz=%0.4f, intensity=%0.4f, charge_hint=%d)" % (
            self.mz, self.peak.intensity, self.charge)

def _loader_creator(specification):
    if isinstance(specification, basestring):
        return MSFileLoader(specification)
    elif isinstance(specification, ScanIterator):
        return specification
    else:
        raise ValueError("Cannot determine how to get a ScanIterator from %r" % (specification,))

def _simplify_peak_set(peaks, bin_width=5.0):
    bin_edges = np.arange(0, peaks[-1].mz + bin_width, bin_width)
    bins = []
    for i, bin_edge in enumerate(bin_edges, 1):
        if i == len(bin_edges):
            next_edge = bin_edges[-1] + bin_width
        else:
            next_edge = bin_edges[i]
        subset = peaks.between(bin_edge, next_edge)
        bins.append(subset)

    thresholds = []
    reduced_subsets = {}
    k = 0
    for b in bins:
        if len(b) > 0:
            bin_intensities = np.array([p.intensity for p in b])
            thresholds.append(np.max(bin_intensities) / 3.)
            for p in b:
                if p.intensity > thresholds[-1]:
                    reduced_subsets[p.peak_count] = p
            k += (bin_intensities > thresholds[-1]).sum()
        else:
            thresholds.append(0.0)
    subset_peaks = PeakSet(
        sorted(reduced_subsets.values(), key=lambda x: x.mz)).clone()
    subset_peaks.reindex()
    return PeakIndex(np.array([]), np.array([]), subset_peaks)

class ScanProcessor(Base, LogUtilsMixin):
    """Orchestrates the deconvolution of a :class:`~.ScanIterator` scan by scan. This process will
    apply different rules for MS1 scans and MSn scans. This type itself mimics a :class:`~.ScanIterator`,
    consuming (raw) mass spectral data and producing deisotoped and charge deconvolved spectra.

    The algorithms used for each task are independent and can be specified in the appropriate
    attribute dictionary, however there is information sharing between each MS1 scan and its
    MSn scans as the precursor monoisotopic mass is recalibrated according to the MS1 processing
    arguments, and the selected charge state is used to limit the charge range used in the matching
    MSn scan. These are described by :class:`PriorityTarget` objects.

    If an averagine-based deconvoluter is used, the averagine cache will be pre-populated.

    At the moment, MSn assumes only MS2. Until MS3 data become available for testing, this limit
    will remain.

    Attributes
    ----------
    data_source : :class:`str`, :class:`~.ScanIterator` or file-like
        Any valid object to be passed to the `loader_type` callable to produce
        a :class:`~.ScanIterator` instance. A path to a mass spectrometry data file,
        a file-like object, or an instance of :class:`~.ScanIterator`. Used to populate :attr:`reader`
    loader_type : callable
        A callable, which when passed :attr:`data_source` returns an instance of :class:`~.ScanIterator`.
        By default, this is :func:`~.MSFileLoader`. Used to populate :attr:`reader`
    reader: ScanIterator
        Any object implementing the :class:`~.ScanIterator` interface, produced by calling
        :attr:`loader_type` on :attr:`data_source`.
    ms1_deconvolution_args : :class:`dict`
        The arguments passed to :func:`~ms_deisotope.deconvolution.deconvolute_peaks` for MS1
        scans.
    ms1_peak_picking_args : :class:`dict`
        The arguments passed to :func:`ms_peak_picker.pick_peaks` for MS1 scans.
    msn_deconvolution_args : :class:`dict`
        The arguments passed to :func:`~ms_deisotope.deconvolution.deconvolute_peaks` for MSn
        scans.
    msn_peak_picking_args : :class:`dict`
        The arguments passed to :func:`ms_peak_picker.pick_peaks` for MSn scans.
    pick_only_tandem_envelopes : :class:`bool`
        Whether or not to process whole MS1 scans or just the regions around those peaks
        chosen for MSn
    default_precursor_ion_selection_window : :class:`float`
        Size of the selection window to use when :attr:`pick_only_tandem_envelopes` is `True`
        and the information is not available in the scan.
    trust_charge_hint : :class:`bool`
        Whether or not to trust the charge provided by the data source when determining
        the charge state of precursor isotopic patterns. Defaults to `True`
    respect_isolation_window: :class:`bool`
        Whether to use the bounds of the isolation window to reject a monoisotopic peak
        solution
    terminate_on_error: :class:`bool`
        Whether or not  to stop processing on an error. Defaults to `True`
    ms1_averaging: :class:`int`
        The number of adjacent MS1 scans to average prior to picking peaks.
    """

    def __init__(self, data_source, ms1_peak_picking_args=None,
                 msn_peak_picking_args=None,
                 ms1_deconvolution_args=None,
                 msn_deconvolution_args=None,
                 pick_only_tandem_envelopes=False,
                 default_precursor_ion_selection_window=1.5,
                 trust_charge_hint=True,
                 loader_type=None,
                 envelope_selector=None,
                 terminate_on_error=True,
                 ms1_averaging=0,
                 respect_isolation_window=False,
                 too_many_peaks_threshold=7000):
        if loader_type is None:
            loader_type = _loader_creator

        self.data_source = data_source
        self.ms1_peak_picking_args = ms1_peak_picking_args or {}
        self.msn_peak_picking_args = msn_peak_picking_args or ms1_peak_picking_args or {}
        self.ms1_deconvolution_args = ms1_deconvolution_args or {}
        self.ms1_deconvolution_args.setdefault("charge_range", (1, 8))
        self.msn_deconvolution_args = msn_deconvolution_args or {}
        self.msn_deconvolution_args.setdefault("charge_range", (1, 8))
        self.pick_only_tandem_envelopes = pick_only_tandem_envelopes

        self.too_many_peaks_threshold = too_many_peaks_threshold

        self.default_precursor_ion_selection_window = default_precursor_ion_selection_window
        self.respect_isolation_window = respect_isolation_window

        self.trust_charge_hint = trust_charge_hint
        self.ms1_averaging = int(ms1_averaging) if ms1_averaging else 0

        self.loader_type = loader_type

        self._signal_source = self.loader_type(data_source)
        self.envelope_selector = envelope_selector
        self.terminate_on_error = terminate_on_error
        self._prepopulate_averagine_cache()

    def _prepopulate_averagine_cache(self):
        if 'averagine' in self.ms1_deconvolution_args:
            averagine = self.ms1_deconvolution_args['averagine']
            if isinstance(averagine, (list, tuple)):
                averagine = [AveragineCache(a).populate() for a in averagine]
            else:
                averagine = AveragineCache(averagine).populate()
            self.ms1_deconvolution_args['averagine'] = averagine
        if 'averagine' in self.msn_deconvolution_args:
            averagine = self.msn_deconvolution_args['averagine']
            if isinstance(averagine, (list, tuple)):
                averagine = [AveragineCache(a).populate() for a in averagine]
            else:
                averagine = AveragineCache(averagine).populate()
            self.msn_deconvolution_args['averagine'] = averagine

    def _reject_candidate_precursor_peak(self, peak, product_scan):
        isolation = product_scan.isolation_window
        if isolation is None or isolation.is_empty():
            pinfo = product_scan.precursor_information
            err = peak.mz - pinfo.mz
            return abs(err) > self.default_precursor_ion_selection_window
        else:
            return peak.mz not in isolation and self.respect_isolation_window

    @property
    def reader(self):
        '''The :class:`~.ScanIterator` which generates the raw scans that will
        be processed.

        Returns
        -------
        :class:`~.ScanIterator`
        '''
        return self._signal_source

    def _get_envelopes(self, precursor_scan):
        """Get the m/z intervals to pick peaks from for the
        given MS1 scan

        Parameters
        ----------
        precursor_scan: Scan

        Returns
        -------
        list or None
        """
        if not self.pick_only_tandem_envelopes and self.envelope_selector is None:
            return None
        elif self.envelope_selector is None:
            chosen_envelopes = [s.precursor_information for s in precursor_scan.product_scans]
            chosen_envelopes = sorted([(p.mz - 5, p.mz + 10) for p in chosen_envelopes])
        else:
            chosen_envelopes = self.envelope_selector(precursor_scan)
        return chosen_envelopes

    def _pick_precursor_scan_peaks(self, precursor_scan, chosen_envelopes=None):
        """Pick peaks from the given precursor scan

        Parameters
        ----------
        precursor_scan: Scan
            Scan to pick peaks from
        chosen_envelopes: list, optional
            list of m/z intervals to pick peaks for

        Returns
        -------
        PeakSet
        """
        if precursor_scan.is_profile:
            peak_mode = 'profile'
        else:
            peak_mode = 'centroid'
        prec_mz, prec_intensity = precursor_scan.arrays
        if not self.pick_only_tandem_envelopes and self.envelope_selector is None:
            prec_peaks = pick_peaks(prec_mz, prec_intensity, peak_mode=peak_mode, **self.ms1_peak_picking_args)
        else:
            if chosen_envelopes is None:
                chosen_envelopes = self._get_envelopes(precursor_scan)
            prec_peaks = pick_peaks(prec_mz, prec_intensity, peak_mode=peak_mode,
                                    target_envelopes=chosen_envelopes,
                                    **self.ms1_peak_picking_args)
        return prec_peaks

    def _average_ms1(self, precursor_scan):
        """Average signal from :attr:`self.ms1_averaging` scans from
        before and after ``precursor_scan`` and pick peaks from the
        averaged arrays.

        Parameters
        ----------
        precursor_scan: Scan
            The scan to use as a point of reference

        Returns
        -------
        PeakSet
        """
        # averaged scans are always profile mode
        new_scan = precursor_scan.average(self.ms1_averaging)
        prec_peaks = pick_peaks(*new_scan.arrays,
                                target_envelopes=self._get_envelopes(precursor_scan),
                                **self.ms1_peak_picking_args)
        return prec_peaks

    def pick_precursor_scan_peaks(self, precursor_scan):
        """Picks peaks for the given ``precursor_scan`` using the
        appropriate strategy.

        If :attr:`ms1_averaging` > 0, then the signal averaging strategy
        is used, otherwise peaks are picked directly.

        Parameters
        ----------
        precursor_scan: Scan

        Returns
        -------
        PeakSet
        """
        self.log("Picking Precursor Scan Peaks: %r" % (precursor_scan, ))
        if self.ms1_averaging > 0:
            prec_peaks = self._average_ms1(precursor_scan)
        else:
            prec_peaks = self._pick_precursor_scan_peaks(precursor_scan)
        n_peaks = len(prec_peaks)
        if n_peaks > self.too_many_peaks_threshold:
            self.log("%d peaks found for %r, applying local intensity threshold." % (n_peaks, precursor_scan))
            prec_peaks = _simplify_peak_set(prec_peaks)
        precursor_scan.peak_set = prec_peaks
        return prec_peaks

    def pick_product_scan_peaks(self, product_scan):
        """Pick the peaks of product scan

        Parameters
        ----------
        product_scan: :class:`~.Scan`
            The scan to pick peaks from.

        Returns
        -------
        PeakSet
        """
        if product_scan.is_profile:
            peak_mode = 'profile'
        else:
            peak_mode = 'centroid'
        product_mz, product_intensity = product_scan.arrays
        peaks = pick_peaks(product_mz, product_intensity, peak_mode=peak_mode, **self.msn_peak_picking_args)

        if peaks is None:
            raise EmptyScanError(
                "Could not pick peaks for empty product scan", self)

        product_scan.peak_set = peaks
        return peaks

    def get_precursor_peak_for_product_scans(self, precursor_scan):  # pragma: no cover
        """A utility method to obtain :class:`PriorityTarget` objects for
        each product scan of `precursor_scan`.

        Parameters
        ----------
        precursor_scan: :class:`~.Scan`
            The scan to extract.

        Returns
        -------
        :class:`list` of :class:`PriorityTarget`
        """
        priorities = []
        peaks = precursor_scan.peak_set
        for scan in precursor_scan.product_scans:
            precursor_ion = scan.precursor_information
            if peaks is None:
                peaks = self.pick_precursor_scan_peaks(precursor_scan)
            peak, _ = peaks.get_nearest_peak(precursor_ion.mz)
            precursor_ion.peak = peak
            target = PriorityTarget(
                peak,
                precursor_ion,
                self.trust_charge_hint,
                isolation_window=scan.isolation_window)
            if self._reject_candidate_precursor_peak(peak, scan):
                self.log(
                    "Unable to locate a peak for precursor ion %r for tandem scan %s of precursor scan %s" % (
                        precursor_ion, scan.title,
                        precursor_scan.title))
            else:
                priorities.append(target)
        return priorities

    def process_scan_group(self, precursor_scan, product_scans):
        """Performs the initial extraction of information relating
        `precursor_scan` to `product_scans` and picks peaks for ``precursor_scan``.
        Called by :meth:`process`. May be used separately if doing the process step
        by step.

        Parameters
        ----------
        precursor_scan : :class:`~.Scan`
            An MS1 Scan
        product_scans : :class:`list` of :class:`~.Scan`
            A :class:`list` of MSn Scans related to `precursor_scan`

        Returns
        -------
        precursor_scan: :class:`~.Scan`
            As Parameter
        prioritiies: :class:`list` of :class:`~PriorityTarget`
            :class:`list` of the peak target windows in `precursor_scan` which
            are related to `product_scans`
        product_scans: :class:`list` of :class:`~.Scan`
            As Parameter
        """
        prec_peaks = self.pick_precursor_scan_peaks(precursor_scan)
        priorities = []

        if prec_peaks is None:
            raise EmptyScanError(
                "Could not pick peaks for empty precursor scan", self)

        for scan in product_scans:
            precursor_ion = scan.precursor_information
            peak = prec_peaks.has_peak(precursor_ion.mz)
            if peak is not None:
                err = abs(peak.mz - precursor_ion.mz)
            else:
                peak, err = prec_peaks.get_nearest_peak(precursor_ion.mz)
            self.debug("For Precursor at %0.4f, found Peak at %0.4f with error %0.4f" % (
                precursor_ion.mz, peak.mz, err))
            precursor_ion.peak = peak
            target = PriorityTarget(
                peak, precursor_ion, self.trust_charge_hint,
                scan.precursor_information.precursor_scan_id,
                scan.precursor_information.product_scan_id,
                isolation_window=scan.isolation_window)
            if self._reject_candidate_precursor_peak(peak, scan):
                self.log(
                    "Unable to locate a peak for precursor ion %r for tandem scan %s of precursor scan %s" % (
                        precursor_ion, scan.id,
                        precursor_scan.id))
            else:
                priorities.append(target)

        return precursor_scan, priorities, product_scans

    def _default_all_precursor_information(self, scans):
        for scan in scans:
            if scan.ms_level > 1:
                scan.precursor_information.default(orphan=True)

    def deconvolute_precursor_scan(self, precursor_scan, priorities=None):
        """Deconvolute the given precursor scan, giving priority to its product ions,
        correcting the :attr:`precursor_information` attributes of priority targets,
        as well as calculating the degree of precursor purity and coisolating ions.

        Parameters
        ----------
        precursor_scan : :class:`~.Scan`
            The precursor scan to deconvolute
        priorities : :class:`list` of :class:`PriorityTarget`, optional
            The priority targets for the product ions derived from `precursor_scan`

        Returns
        -------
        :class:`~DeconvolutedPeakSet`
            The deconvoluted peaks of ``precursor_scan``
        :class:`list` of :class:`PriorityTarget`
            The precursor ions selected, with updated mass and charge information

        Raises
        ------
        Exception
            Any errors which are thrown during the deconvolution process may be thrown
            if :attr:`terminate_on_error` is `True`.
        """
        if priorities is None:
            priorities = []

        self.log("Deconvoluting Precursor Scan %r" % precursor_scan)
        self.log("Priorities: %r" % priorities)

        ms1_deconvolution_args = self.ms1_deconvolution_args.copy()

        if precursor_scan.polarity in (1, -1):
            polarity = precursor_scan.polarity
            ms1_deconvolution_args['charge_range'] = tuple(
                polarity * abs(c) for c in ms1_deconvolution_args['charge_range'])
        try:
            decon_result = deconvolute_peaks(
                precursor_scan.peak_set, priority_list=priorities,
                **ms1_deconvolution_args)
        except NoIsotopicClustersError as e:
            e.scan_id = precursor_scan.id
            if self.terminate_on_error:
                raise e
            else:
                self.log("No isotopic clusters found in %r" % precursor_scan.id)

        dec_peaks, priority_results = decon_result
        if decon_result.errors:
            self.error("Errors occurred during deconvolution of %s, %r" % (
                precursor_scan.id, decon_result.errors))
        precursor_scan.deconvoluted_peak_set = dec_peaks
        for pr in priority_results:
            if pr is None:
                continue
            else:
                pr.chosen_for_msms = True

        # `priorities` and `priority_results` are parallel lists. The
        # ith position in `priorities` corresponds to the ith deconvoluted
        # priority result in `priority_results`. The entry in `priority_results`
        # may be `None` if the deconvolution failed, but elements of `priorities`
        # should always be FittedPeak or Peak-like instances

        coisolation_detection = PrecursorPurityEstimator(default_width=self.default_precursor_ion_selection_window)
        self.debug("Priority Targets for %s: %r" % (
            precursor_scan.id, [
                (p.mz, p.charge) if p is not None else None for p in priorities
            ]))
        for product_scan in precursor_scan.product_scans:
            precursor_information = product_scan.precursor_information

            # unknown precursor purity
            product_scan.annotations['precursor purity'] = 0.0
            i = _get_nearest_index(precursor_information.mz, priorities)

            # If no peak is found in the priority list, it means the priority list is empty.
            # This should never happen in the current implementation. If it did, then we forgot
            # to pass the priority list to this function.
            if i is None:
                self.log(
                    "Could not find deconvolution for %r (No nearby peak in the priority list)" %
                    precursor_information)
                precursor_information.default(orphan=True)
                continue

            peak = priority_results[i]
            # If the deconvolution result is None, then we have no answer
            if peak is None:
                self.log(
                    "Could not find deconvolution for %r (No solution was found for this region)" %
                    precursor_information)
                precursor_information.default(orphan=True)

                continue
            elif peak.charge == 1 or (peak.charge != precursor_information.charge and self.trust_charge_hint):
                if precursor_information.charge != ChargeNotProvided:
                    self.log(
                        "Could not find deconvolution for %r (Unacceptable solution was proposed: %r)" %
                        (precursor_information, peak))
                    precursor_information.default()
                    continue

            precursor_purity = -1.0
            if peak is not None:
                precursor_purity, coisolation = coisolation_detection(
                    precursor_scan,
                    peak,
                    product_scan.isolation_window)
                precursor_information.coisolation = coisolation
                self.debug(
                    "Precursor m/z %f\nExperimental = %r\nTheoretical = %r" % (
                        peak.mz,
                        ', '.join(["(%0.4f, %0.1f)" % (p.mz, p.intensity) for p in peak.envelope]),
                        ', '.join(["(%0.4f, %0.1f)" % (p.mz, p.intensity) for p in peak.fit.theoretical]))
                )

            product_scan.annotations['precursor purity'] = precursor_purity
            precursor_information.extract(peak)
        return dec_peaks, priority_results

    def deconvolute_product_scan(self, product_scan):
        """Deconvolute the peaks of `product_scan`.

        This method will override the upper limit "charge_range" of
        :attr:`msn_deconvolution_args` to the charge information of
        the precursor ion.

        This method sets the :attr:`~.Scan.deconvoluted_peak_set` of
        `product_scan`.

        Parameters
        ----------
        product_scan : :class:`~.Scan`
            The scan to deconvolute.

        Returns
        -------
        :class:`~.DeconvolutedPeakSet`

        Raises
        ------
        Exception
            Any errors which are thrown during the deconvolution process may be thrown
            if :attr:`terminate_on_error` is `True`.
        """
        self.log("Deconvoluting Product Scan %r" % (product_scan, ))
        precursor_ion = product_scan.precursor_information
        top_charge_state = precursor_ion.extracted_charge
        if not top_charge_state:
            top_charge_state = precursor_ion.charge
        deconargs = dict(self.msn_deconvolution_args)
        charge_range = list(deconargs.get("charge_range", [1, top_charge_state]))
        if top_charge_state is not None and top_charge_state is not ChargeNotProvided and\
           top_charge_state != 0 and abs(top_charge_state) < abs(charge_range[1]):
            charge_range[1] = top_charge_state

        deconargs["charge_range"] = charge_range

        if product_scan.polarity in (-1, 1):
            polarity = product_scan.polarity
            deconargs["charge_range"] = [
                polarity * abs(c) for c in deconargs["charge_range"]]

        try:
            dec_peaks, _ = deconvolute_peaks(product_scan.peak_set, **deconargs)
        except NoIsotopicClustersError as e:
            self.log("No Isotopic Clusters found in %r" % product_scan.id)
            e.scan_id = product_scan.id
            if self.terminate_on_error:
                raise e

        product_scan.deconvoluted_peak_set = dec_peaks
        return dec_peaks

    def _get_next_scans(self):
        bunch = next(self.reader)
        try:
            precursor, products = bunch
        except ValueError:
            if isinstance(bunch, Scan):
                if bunch.ms_level == 1:
                    precursor = bunch
                    products = []
                else:
                    precursor = None
                    products = [bunch]

        if self.pick_only_tandem_envelopes:
            while len(products) == 0:
                precursor, products = next(self.reader)

        return precursor, products

    def process(self, precursor, products):
        """Fully preprocesses the `precursor` and `products` scans, performing
        any necessary information sharing.

        This method may be used to process scans from other sources not from the
        wrapped :class:`~.ScanIterator`.

        Parameters
        ----------
        precursor : :class:`~.Scan`
            An MS1 Scan
        products : :class:`list` of :class:`~.Scan`
            A list of MSn Scans related to `precursor`

        Returns
        -------
        :class:`~.ScanBunch`
            The fully processed version of `precursor` and `products`
        """
        precursor_scan, priorities, product_scans = self.process_scan_group(precursor, products)
        if precursor_scan is not None:
            self.deconvolute_precursor_scan(precursor_scan, priorities)
        else:
            self._default_all_precursor_information(product_scans)

        for product_scan in product_scans:
            self.pick_product_scan_peaks(product_scan)
            self.deconvolute_product_scan(product_scan)

        return ScanBunch(precursor_scan, product_scans)

    def next(self):
        """Fetches the next bunch of scans from :attr:`reader` and
        invokes :meth:`process` on them, picking peaks and deconvoluting them.

        Returns
        -------
        ScanBunch
        """
        precursor, products = self._get_next_scans()
        bunch = self.process(precursor, products)
        return bunch

    def __next__(self):
        """Fetches the next bunch of scans from :attr:`reader` and
        invokes :meth:`process` on them, picking peaks and deconvoluting them.

        Returns
        -------
        ScanBunch
        """
        return self.next()

    def __iter__(self):
        return self

    def pack_next(self):
        """As :meth:`next`, except instead of producing :class:`ScanBunch` of
        :class:`Scan` instances, instead it uses :class:`ProcessedScan` to strip away
        much of the heavy information like the raw data arrays.

        Returns
        -------
        ScanBunch
        """
        precursor, products = self._get_next_scans()
        precursor_scan, product_scans = self.process(precursor, products)
        return ScanBunch(precursor_scan.pack() if precursor_scan else None, [p.pack() for p in product_scans])

    def start_from_scan(self, *args, **kwargs):
        """A wrapper around :meth:`~.RandomAccessScanSource.start_from_scan` provided by
        :attr:`reader`, if available.

        Returns
        -------
        self

        See Also
        --------
        :meth:`~.RandomAccessScanSource.start_from_scan`
        """
        self.reader.start_from_scan(*args, **kwargs)
        return self

ScanProcessor.log_with_logger(logger)

def process(data_source, ms1_averagine=peptide, msn_averagine=peptide,
            ms1_score_threshold=20, msn_score_threshold=5, denoise=False,
            ms1_max_missed_peaks=1, pick_only_tandem_envelopes=False,
            trust_charge_hint=True, envelope_selector=None, terminate_on_error=True,
            ms1_averaging=0, respect_isolation_window=False, use_quick_charge=True):
    """Construct a deconvolution pipeline for common applications.

    Parameters
    ----------
    data_source : :class:`str` or :class:`~.ScanIterator` or file-like object
        The scan data source to read raw spectra from
    ms1_averagine : :class:`~.Averagine` or :class:`~.AveragineCache`, optional
        The :class:`~.Averagine` model to use for MS1 scans. Defaults to
        :data:`ms_deisotope.averagine.peptide`.
    msn_averagine : :class:`~.Averagine` or :class:`~.AveragineCache`, optional
        The :class:`~.Averagine` model to use for MSn scans. Defaults to
        :data:`ms_deisotope.averagine.peptide`.
    ms1_score_threshold : float, optional
        The score threshold to use to reject isotopic pattern fits for MS1 scans.
        The default is 20.0.
    msn_score_threshold : float, optional
        The score threshold to use to reject isotopic pattern fits for MS1 scans.
        The default is 5.0.
    denoise : :class:`bool` or :class:`float`, optional
        Whether to denoise MS1 scans. If the value is not false-y, it may either be
        a float to set the scale of the denoising process, or 5.0 if the value is
        :const:`True`.
    ms1_max_missed_peaks : :class:`int`, optional
        The maximum number of missed peaks to permit for MS1 scans. The default is 1.
    pick_only_tandem_envelopes : :class:`bool`
        Whether or not to process whole MS1 scans or just the regions around those peaks
        chosen for MSn
    default_precursor_ion_selection_window : :class:`float`
        Size of the selection window to use when :attr:`pick_only_tandem_envelopes` is `True`
        and the information is not available in the scan.
    trust_charge_hint : :class:`bool`
        Whether or not to trust the charge provided by the data source when determining
        the charge state of precursor isotopic patterns. Defaults to `True`
    terminate_on_error: :class:`bool`
        Whether or not  to stop processing on an error. Defaults to `True`
    ms1_averaging: :class:`int`
        The number of adjacent MS1 scans to average prior to picking peaks.
    respect_isolation_window: :class:`bool`
        Whether to use the bounds of the isolation window to reject a monoisotopic peak
        solution
    use_quick_charge : :class:`bool`, optional
        Whether or not to used the QuickCharge algorithm for expediting charge calculation.

    Returns
    -------
    :class:`ScanProcessor`
    """
    if denoise:
        ms1_peak_picking_args = {
            "filters": [
                FTICRBaselineRemoval(
                    scale=denoise if denoise is not True else 5., window_length=2)
            ]
        }
    else:
        ms1_peak_picking_args = None

    ms1_deconvolution_args = {
        "averagine": ms1_averagine,
        "scorer": PenalizedMSDeconVFitter(ms1_score_threshold, 2.0),
        "use_quick_charge": use_quick_charge,
        "max_missed_peaks": ms1_max_missed_peaks,
        "truncate_after": 0.95,
    }
    msn_deconvolution_args = {
        "averagine": msn_averagine,
        "scorer": MSDeconVFitter(msn_score_threshold),
        "use_quick_charge": use_quick_charge,
        "truncate_after": 0.8,
    }
    processor = ScanProcessor(
        data_source, ms1_peak_picking_args,
        None, ms1_deconvolution_args,
        msn_deconvolution_args, pick_only_tandem_envelopes,
        trust_charge_hint=trust_charge_hint,
        envelope_selector=envelope_selector,
        terminate_on_error=terminate_on_error,
        ms1_averaging=ms1_averaging,
        respect_isolation_window=respect_isolation_window)
    return processor

class EmptyScanError(ValueError):
    """A sub-type of :class:`ValueError` which is used to indicate
    that a spectrum is empty and could not be manipulated.
    """
    def __init__(self, msg, scan_id=None):
        ValueError.__init__(self, msg)
        self.scan_id = scan_id


import numpy as np
import typing
from collections import namedtuple
from .materials import _IMMaterial
from ._elastic_sub_surface_stresses import normal_conv_kernels, tangential_conv_kernels

__all__ = ['Elastic', 'elastic_influence_matrix_spatial', 'elastic_influence_matrix_frequency', 'get_angular_velocity']

ElasticProps = namedtuple('ElasticProperties', 'K E v Lam M G', defaults=(None,) * 6)

# noinspection PyPep8Naming
class Elastic(_IMMaterial):
    """ A Class for defining elastic materials

    Parameters
    ----------
    name: str
        The name of the material
    properties: dict
        dict of properties, dicts must have exactly 2 items.
        Allowed keys are : 'E', 'v', 'G', 'K', 'M', 'Lam'
        See notes for definitions
    max_load: float, optional (float('inf'))
        The maximum load on the surface, loads above this will be cropped during analysis, if this is specified a
        plastic deformation sub model should be added to the end of each model step to make the deformation permanent
    use_frequency_domain: bool, optional (True)
        If True the frequency domain definition of the influence matrix is used, otherwise the spatial domain definition
        is used.
    periodic_im_repeats: tuple, optional (1,1)
        The number of times the influence matrix should be wrapped in each dimension, used with spatially defined
        influence matrices and to set the zero frequency value for frequency domain influence matrices. Should not
        be used with non periodic contacts, for periodic contacts the total size should match the physical size. This
        is necessary to ensure truly periodic behaviour, no physical limit exists:
        (an infinitely long contact with any load per unit length will cause infinite displacement).
    zero_frequency_value: float, optional (None)
        If the frequency domain influence matrix is used the zero frequency value can be set, this defaults to the sum
        of the spatial influence matrix of the correct size. Should be set to 0 for fully periodic contacts.

    Methods
    -------
    speed_of_sound

    See Also
    --------

    Notes
    -----

    Keys refer to:
        - E   - Young's modulus
        - v   - Poission's ratio
        - K   - Bulk Modulus
        - Lam - Lame's first parameter
        - G   - Shear modulus
        - M   - P wave modulus

    Examples
    --------
    >>> # Make a material model for elastic steel
    >>> steel = Elastic('steel', {'E': 200e9, 'v': 0.3})
    >>> # Find it's p-wave modulus:
    >>> pwm = steel.M
    >>> # Find the speeds of sound:
    >>> sos = steel.speed_of_sound(7890)
    """

    material_type = 'Elastic'

    _properties = {'E': None,
                   'v': None,
                   'G': None,
                   'K': None,
                   'Lam': None,
                   'M': None, }

    _last_set = []
    density = None

    def __init__(self, name: str, properties: dict, max_load: float = np.inf,
                 use_frequency_domain: bool = True, periodic_im_repeats: tuple = (1, 1),
                 zero_frequency_value: float = None):
        super().__init__(name, use_frequency_domain, max_load,
                         periodic_im_repeats, zero_frequency_value)

        if len(properties) > 2:
            raise ValueError("Too many properties supplied, must be 1 or 2")

        for item in properties.items():
            self._set_props(*item)

    def _influence_matrix_spatial(self, components: typing.Union[typing.Sequence[str], str],
                                  grid_spacing: {typing.Sequence[float], float}, span: typing.Sequence[int]):
        """
        Influence matrix for an elastic material

        Parameters
        ----------
        grid_spacing: tuple
            The spacing between grid points in the x and y directions
        span: tuple
            The span required in the x and y directions in number of grid points
        components: str or Sequence {'xx','xy','xz','yx','yy','yz','zx','zy','zz','all'}
            The required components eg the 'xy' component represents the x
            deflection caused by loads in the y direction

        Returns
        -------
        dict
            dict of the requested influence matrix or matrices

        See Also
        --------
        elastic_loading
        elastic_deflection

        Notes
        -----

        K^{ij i'j'}_zz=(1-v)/(2*pi*G)*Czz

        Czz=(hx*(k*log((m+sqrt(k**2+m**2))/(n+sqrt(k**2+n**2)))+
                 l*log((n+sqrt(l**2+n**2))/(m+sqrt(l**2+m**2))))+
             hy*(m*log((k+sqrt(k**2+m**2))/(l+sqrt(l**2+m**2)))+
                 n*log((l+sqrt(l**2+n**2))/(k+sqrt(k**2+n**2)))))

        In which:

        k=i'-i+0.5
        l=i'-i-0.5
        m=j'-j+0.5
        n=j'-j-0.5
        hx=grid_spacing[0]
        hy=grid_spacing[1]

        If both shear_modulus_2 and v_2 are supplied and are not None the combined IM is returned for the surface
        pair

        Examples
        --------

        References
        ----------
        Complete boundary element method formulation for normal and tangential
        contact problems

        """
        shear_modulus_2 = None
        v_2 = None

        shear_modulus = self.G
        v = self.v

        components = {comp: elastic_influence_matrix_spatial(comp, span, grid_spacing, shear_modulus, v,
                                                             shear_mod_2=shear_modulus_2, v_2=v_2) for comp in
                      components}

        return components

    def _influence_matrix_frequency(self, components: typing.Sequence[str], grid_spacing: typing.Sequence[float],
                                    span: typing.Sequence[int]):
        x_omega, y_omega = get_angular_velocity(span, grid_spacing)
        norm_2 = x_omega ** 2 + y_omega ** 2
        norm = np.sqrt(norm_2)
        rtn_dict = {key: elastic_influence_matrix_frequency(y_omega, x_omega, norm, norm_2, key, self.E, self.v)
                    for key in components}
        return rtn_dict

    def _del_props(self, prop):
        # delete any of the material properties
        keys = list(self._properties.keys())
        if self._last_set == prop:
            self._properties = {key: None for key in keys}
            self._last_set = None
        else:
            self._properties = {key: None for key in keys
                                if not key == self._last_set}

    def _set_props(self, prop, value):
        allowed_props = ['E', 'v', 'G', 'K', 'Lam', 'M']
        if prop not in allowed_props:
            msg = (f'property {prop} not recognised allowed propertied are: ' +
                   ' '.join(allowed_props))
            raise ValueError(msg)

        self._properties[prop] = np.float64(value)

        if len(self._last_set) == 0:
            self._last_set.append(prop)  # if none ever set just set it
        elif self._last_set[-1] != prop:
            self._last_set.append(prop)  # if the last set is different replace it

        if len(self._last_set) > 1:  # if 2 props have been set update all
            set_props = {prop: np.float64(value),
                         self._last_set[-2]: self._properties[self._last_set[-2]]}
            self._properties = _get_properties(set_props)
        return

    @property
    def E(self):
        """The Young's modulus of the material"""
        return self._properties['E']

    @E.deleter
    def E(self):
        self._del_props('E')

    @E.setter
    def E(self, value):
        self._set_props('E', value)

    @property
    def v(self):
        """The Poissions's ratio of the material"""
        return self._properties['v']

    @v.deleter
    def v(self):
        self._del_props('v')

    @v.setter
    def v(self, value):
        self._set_props('v', value)

    @property
    def G(self):
        """The shear modulus of the material"""
        return self._properties['G']

    @G.deleter
    def G(self):
        self._del_props('G')

    @G.setter
    def G(self, value):
        self._set_props('G', value)

    @property
    def K(self):
        """The bulk modulus of the material"""
        return self._properties['K']

    @K.deleter
    def K(self):
        self._del_props('K')

    @K.setter
    def K(self, value):
        self._set_props('K', value)

    @property
    def Lam(self):
        """Lame's first parameter for the material"""
        return self._properties['Lam']

    @Lam.deleter
    def Lam(self):
        self._del_props('Lam')

    @Lam.setter
    def Lam(self, value):
        self._set_props('Lam', value)

    @property
    def M(self):
        """The p wave modulus of the material"""
        return self._properties['M']

    @M.deleter
    def M(self):
        self._del_props('M')

    @M.setter
    def M(self, value):
        self._set_props('M', value)

    def speed_of_sound(self, density: float = None):
        """find the speed of sound in the material

        Parameters
        ----------
        density : float optional (None)
            The density of the material

        Returns
        -------

        speeds : dict
            With keys 's' and 'p' giving the s and p wave speeds

        Notes
        -----

        Finds speeds according to the following equations:

        Vs=sqrt(G/rho)
        Vp=sqrt(M/rho)

        Where rho is the density, G is the shear modulus and M is the p wave
        modulus

        Examples
        --------
        >>> # Find the speed of sound in steel
        >>> my_material = Elastic({'E': 200e9, 'v': 0.3})
        >>> my_material.speed_of_sound(7850)

        """
        if density is not None:
            self.density = density
        elif self.density is None:
            raise ValueError("Density not given or set")

        speeds = {'s': np.sqrt(self.G / self.density),
                  'p': np.sqrt(self.M / self.density)}

        return speeds

    def sss_influence_matrices_normal(self, components: typing.Sequence[str], grid_spacing: typing.Sequence[float],
                                      span: typing.Sequence[int], z: typing.Sequence[float] = None, cuda: bool = False):
        if z is None:
            z_len = min(span)
            print(span)
            gs = grid_spacing[span.index(z_len)]
            print(gs)
            z = gs * np.arange(z_len // 2)
            z[0] = z[1] * 1e-4
        all_matrices = normal_conv_kernels(span, z, grid_spacing, self.E, self.v, cuda=cuda)
        return {comp: all_matrices[comp] for comp in components}

    def sss_influence_matrices_tangential_x(self, components: typing.Sequence[str],
                                            grid_spacing: typing.Sequence[float], span: typing.Sequence[int],
                                            z: typing.Sequence[float] = None, cuda: bool = False):
        if z is None:
            z_len = min(span)
            gs = grid_spacing[span.index(z_len)]
            z = gs * np.arange(z_len // 2)
            z[0] = z[1] * 1e-4
        all_matrices = tangential_conv_kernels(span, z, grid_spacing, self.v, cuda=cuda)
        return {comp: all_matrices[comp] for comp in components}

    def __repr__(self):
        return "Elastic(name = '" + self.name + f"', properties = {{ 'E':{self.E}, 'v':{self.v} }}"

def _get_properties(set_props: dict):
    """Get all elastic properties from any pair

    Parameters
    ----------
    set_props : dict
        dict of properties must have exactly 2 members valid keys are: 'K',
        'E', 'v', 'Lam', 'M', 'G'

    Returns
    -------
    out : dict
        dict of all material properties keys are: 'K', 'E', 'v', 'Lam', 'M', 'G'

    Notes
    -----

    Keys refer to:
        - E - Young's modulus
        - v - Poission's ratio
        - K - Bulk Modulus
        - Lam - Lame's first parameter
        - G - Shear modulus
        - M - P wave modulus

    """
    if len(set_props) != 2:
        raise ValueError("Exactly 2 properties must be set,"
                         " {} found".format(len(set_props)))

    valid_keys = ['K', 'E', 'v', 'G', 'Lam', 'M']

    set_params = [key for key in list(set_props.keys()) if key in valid_keys]

    if len(set_params) != 2:
        msg = ("Invalid keys in set_props keys found are: " +
               "{}".format(set_props.keys()) +
               ". Valid keys are: " + " ".join(valid_keys))
        raise ValueError(msg)

    out = set_props.copy()

    set_params = list(set_props.keys())
    set_params.sort()
    # p is properties this saves a lot of space
    p = ElasticProps(**set_props)

    if set_params[0] == 'E':
        if set_params[1] == 'G':
            out['K'] = p.E * p.G / (3 * (3 * p.G - p.E))
            out['Lam'] = p.G * (p.E - 2 * p.G) / (3 * p.G - p.E)
            out['M'] = p.G * (4 * p.G - p.E) / (3 * p.G - p.E)
            out['v'] = p.E / (2 * p.G) - 1
        elif set_params[1] == 'K':
            out['G'] = 3 * p.K * p.E / (9 * p.K - p.E)
            out['Lam'] = 3 * p.K * (3 * p.K - p.E) / (9 * p.K - p.E)
            out['M'] = 3 * p.K * (3 * p.K + p.E) / (9 * p.K - p.E)
            out['v'] = (3 * p.K - p.E) / (6 * p.K)
        elif set_params[1] == 'Lam':
            R = np.sqrt(p.E ** 2 + 9 * p.Lam ** 2 + 2 * p.E * p.Lam)
            out['G'] = (p.E - 3 * p.Lam + R) / 4
            out['K'] = (p.E + 3 * p.Lam + R) / 6
            out['M'] = (p.E - p.Lam + R) / 2
            out['v'] = 2 * p.Lam / (p.E + p.Lam + R)
        elif set_params[1] == 'M':
            S = np.sqrt(p.E ** 2 + 9 * p.M ** 2 - 10 * p.E * p.M)
            out['G'] = (3 * p.M + p.E - S) / 8
            out['K'] = (3 * p.M - p.E + S) / 6
            out['Lam'] = (p.M - p.E + S) / 4
            out['v'] = (p.E - p.M + S) / (4 * p.M)
        else:  # set_params[1]=='v'
            out['G'] = p.E / (2 * (1 + p.v))
            out['K'] = p.E / (3 * (1 - 2 * p.v))
            out['Lam'] = p.E * p.v / ((1 + p.v) * (1 - 2 * p.v))
            out['M'] = p.E * (1 - p.v) / ((1 + p.v) * (1 - 2 * p.v))
    elif set_params[0] == 'G':
        if set_params[1] == 'K':
            out['E'] = 9 * p.K * p.G / (3 * p.K + p.G)
            out['Lam'] = p.K - 2 * p.G / 3
            out['M'] = p.K + 4 * p.G / 3
            out['v'] = (3 * p.K - 2 * p.G) / (2 * (3 * p.K + p.G))
        elif set_params[1] == 'Lam':
            out['E'] = p.G * (3 * p.Lam + 2 * p.G) / (p.Lam + p.G)
            out['K'] = p.Lam + 2 * p.G / 3
            out['M'] = p.Lam + 2 * p.G
            out['v'] = p.Lam / (2 * (p.Lam + p.G))
        elif set_params[1] == 'M':
            out['E'] = p.G * (3 * p.M - 4 * p.G) / (p.M - p.G)
            out['K'] = p.M - 4 * p.G / 3
            out['Lam'] = p.M - 2 * p.G
            out['v'] = (p.M - 2 * p.G) / (2 * p.M - 2 * p.G)
        else:  # set_params[1]=='v'
            out['E'] = 2 * p.G * (1 + p.v)
            out['K'] = 2 * p.G * (1 + p.v) / (3 * (1 - 2 * p.v))
            out['Lam'] = 2 * p.G * p.v / (1 - 2 * p.v)
            out['M'] = 2 * p.G * (1 - p.v) / (1 - 2 * p.v)
    elif set_params[0] == 'K':
        if set_params[1] == 'Lam':
            out['E'] = 9 * p.K * (p.K - p.Lam) / (3 * p.K - p.Lam)
            out['G'] = 3 * (p.K - p.Lam) / 2
            out['M'] = 3 * p.K - 2 * p.Lam
            out['v'] = p.Lam / (3 * p.K - p.Lam)
        elif set_params[1] == 'M':
            out['E'] = 9 * p.K * (p.M - p.K) / (3 * p.K + p.M)
            out['G'] = 3 * (p.M - p.K) / 4
            out['Lam'] = (3 * p.K - p.M) / 2
            out['v'] = (3 * p.K - p.M) / (3 * p.K + p.M)
        else:  # set_params[1]=='v'
            out['E'] = 3 * p.K * (1 - 2 * p.v)
            out['G'] = (3 * p.K * (1 - 2 * p.v)) / (2 * (1 + p.v))
            out['Lam'] = 3 * p.K * p.v / (1 + p.v)
            out['M'] = 3 * p.K * (1 - p.v) / (1 + p.v)
    elif set_params[0] == 'Lam':
        if set_params[1] == 'M':
            out['E'] = (p.M - p.Lam) * (p.M + 2 * p.Lam) / (p.M + p.Lam)
            out['G'] = (p.M - p.Lam) / 2
            out['K'] = (p.M + 2 * p.Lam) / 3
            out['v'] = p.Lam / (p.M + p.Lam)
        else:
            out['E'] = p.Lam * (1 + p.v) * (1 - 2 * p.v) / p.v
            out['G'] = p.Lam(1 - 2 * p.v) / (2 * p.v)
            out['K'] = p.Lam * (1 + p.v) / (3 * p.v)
            out['M'] = p.Lam * (1 - p.v) / p.v
    else:
        out['E'] = p.M * (1 + p.v) * (1 - 2 * p.v) / (1 - p.v)
        out['G'] = p.M * (1 - 2 * p.v) / (2 * (1 - p.v))
        out['K'] = p.M * (1 + p.v) / (3 * (1 - p.v))
        out['Lam'] = p.M * p.v / (1 - p.v)

    return out

def get_angular_velocity(span, gs):
    r_omega = np.fft.fftfreq(span[1], d=gs[1]) * (2 * np.pi)
    c_omega = np.fft.fftfreq(span[0], d=gs[0]) * (2 * np.pi)
    return np.meshgrid(r_omega, c_omega)

def elastic_influence_matrix_frequency(y_omega, x_omega, norm, norm_2, comp: str, e: float, v: float):
    """

    Parameters
    ----------
    y_omega: np.ndarray
        rotational velocity components in the y direction
    x_omega: np.ndarray
        rotational velocity components in the x direction
    norm: np.ndarray
        The norm of the velocity vectors
    norm_2: np.ndarray
        The squared norm of the velocity vectors
    comp: str
        The component to find for example 'xy' gives the deformations in the y direction caused by a load in the x
        direction
    e: float
        The Young's modulus of the material
    v: float
        The Poission's ratio of the material

    Returns
    -------
    The influence matrix component in the frequency domain

    """
    with np.errstate(divide='ignore', invalid='ignore'):
        if comp == 'zz':
            fact = 2 * (1 - v ** 2)
        elif comp == 'yy':
            fact = 2 * (1 + v) * (1 - v * y_omega ** 2 / norm_2)
        elif comp == 'xx':
            fact = 2 * (1 + v) * (1 - v * x_omega ** 2 / norm_2)
        elif comp in ('xy', 'yx'):
            fact = y_omega * x_omega * 2 * v * (1 + v) / norm_2
        elif comp in ('yz', 'zy'):
            fact = 1j * y_omega * (1 + v) * (1 - 2 * v) / norm * (-1 if comp == 'zy' else 1)
        elif comp in ('xz', 'zx'):
            fact = 1j * x_omega * (1 + v) * (1 - 2 * v) / norm * (-1 if comp == 'zx' else 1)
        else:
            raise ValueError('component name not recognised: ' + comp + ', components must be lower case')
        rtn = fact * (1 / (e * norm))
    rtn[0, 0] = 0.0
    return rtn

# noinspection PyTypeChecker
def elastic_influence_matrix_spatial(comp: str, span: typing.Sequence[int], grid_spacing: typing.Sequence[float],
                                     shear_mod: float, v: float,
                                     shear_mod_2: typing.Optional[float] = None,
                                     v_2: typing.Optional[float] = None) -> np.array:
    """Find influence matrix components for an elastic contact problem

    Parameters
    ----------
    comp : str {'xx','xy','xz','yx','yy','yz','zx','zy','zz'}
        The component to be returned
    span: Sequence[int]
        The span of the influence matrix in the x and y directions
    grid_spacing: Sequence[float]
        The grid spacings in the x and y directions
    shear_mod : float
        The shear modulus of the surface material
    v : float
        The Poission's ratio of the surface material
    shear_mod_2: float (optional) None
        The shear modulus of the second surface for a combined stiffness matrix
    v_2: float (optional) None
        The Poisson's ratio of the second surface for a combined stiffness matrix
    fft: bool
        If true the fft of the influence matrix will be returned

    Returns
    -------
    C : array
        The influence matrix component requested

    See Also
    --------
    elastic_im

    Notes
    -----

    Don't use this function, used by: elastic_im

    References
    ----------
    Complete boundary element method formulation for normal and tangential
    contact problems

    """
    span = tuple(span)
    try:
        # lets just see how this changes
        # i'-i and j'-j
        idmi = (np.arange(span[1]) - span[1] // 2 + (1 - span[1] % 2))
        jdmj = (np.arange(span[0]) - span[0] // 2 + (1 - span[0] % 2))
        mesh_idmi = np.tile(idmi, (span[0], 1))
        mesh_jdmj = np.tile(np.expand_dims(jdmj, -1), (1, span[1]))

    except TypeError:
        raise TypeError("Span should be a tuple of integers")

    k = mesh_idmi + 0.5
    el = mesh_idmi - 0.5
    m = mesh_jdmj + 0.5
    n = mesh_jdmj - 0.5

    hy = grid_spacing[1]
    hx = grid_spacing[0]

    second_surface = (shear_mod_2 is not None) and (v_2 is not None)
    if not second_surface:
        v_2 = 1
        shear_mod_2 = 1

    if (shear_mod_2 is not None) != (v_2 is not None):
        raise ValueError('Either both or neither of the second surface parameters must be set')

    if comp == 'zz':
        c_zz = (hx * (k * np.log((m + np.sqrt(k ** 2 + m ** 2)) / (n + np.sqrt(k ** 2 + n ** 2))) +
                      el * np.log((n + np.sqrt(el ** 2 + n ** 2)) / (m + np.sqrt(el ** 2 + m ** 2)))) +
                hy * (m * np.log((k + np.sqrt(k ** 2 + m ** 2)) / (el + np.sqrt(el ** 2 + m ** 2))) +
                      n * np.log((el + np.sqrt(el ** 2 + n ** 2)) / (k + np.sqrt(k ** 2 + n ** 2)))))

        const = (1 - v) / (2 * np.pi * shear_mod) + second_surface * ((1 - v_2) / (2 * np.pi * shear_mod_2))
        ret = const * c_zz
    elif comp == 'xx':
        c_xx = (hx * (1 - v) * (k * np.log((m + np.sqrt(k ** 2 + m ** 2)) / (n + np.sqrt(k ** 2 + n ** 2))) +
                                el * np.log(
                (n + np.sqrt(el ** 2 + n ** 2)) / (m + np.sqrt(el ** 2 + m ** 2)))) +
                hy * (m * np.log((k + np.sqrt(k ** 2 + m ** 2)) / (el + np.sqrt(el ** 2 + m ** 2))) +
                      n * np.log((el + np.sqrt(el ** 2 + n ** 2)) / (k + np.sqrt(k ** 2 + n ** 2)))))
        const = 1 / (2 * np.pi * shear_mod) + second_surface * (1 / (2 * np.pi * shear_mod_2))
        ret = const * c_xx
    elif comp == 'yy':
        c_yy = (hx * (k * np.log((m + np.sqrt(k ** 2 + m ** 2)) / (n + np.sqrt(k ** 2 + n ** 2))) +
                      el * np.log((n + np.sqrt(el ** 2 + n ** 2)) / (m + np.sqrt(el ** 2 + m ** 2)))) +
                hy * (1 - v) * (m * np.log((k + np.sqrt(k ** 2 + m ** 2)) / (el + np.sqrt(el ** 2 + m ** 2))) +
                                n * np.log((el + np.sqrt(el ** 2 + n ** 2)) / (k + np.sqrt(k ** 2 + n ** 2)))))
        const = 1 / (2 * np.pi * shear_mod) + second_surface * (1 / (2 * np.pi * shear_mod_2))
        ret = const * c_yy
    elif comp in ['xz', 'zx']:
        c_xz = (hy / 2 * (m * np.log((k ** 2 + m ** 2) / (el ** 2 + m ** 2)) +
                          n * np.log((el ** 2 + n ** 2) / (k ** 2 + n ** 2))) +
                hx * (k * (np.arctan(m / k) - np.arctan(n / k)) +
                      el * (np.arctan(n / el) - np.arctan(m / el))))
        const = ((2 * v - 1) / (4 * np.pi * shear_mod) + second_surface * (
            (2 * v_2 - 1) / (4 * np.pi * shear_mod_2))) * (-1 if comp == 'zx' else 1)
        ret = const * c_xz
    elif comp in ['yx', 'xy']:
        c_yx = (np.sqrt(hy ** 2 * n ** 2 + hx ** 2 * k ** 2) -
                np.sqrt(hy ** 2 * m ** 2 + hx ** 2 * k ** 2) +
                np.sqrt(hy ** 2 * m ** 2 + hx ** 2 * el ** 2) -
                np.sqrt(hy ** 2 * n ** 2 + hx ** 2 * el ** 2))
        const = v / (2 * np.pi * shear_mod) + second_surface * (v_2 / (2 * np.pi * shear_mod_2))
        ret = const * c_yx
    elif comp in ['zy', 'yz']:
        c_zy = (hx / 2 * (k * np.log((k ** 2 + m ** 2) / (n ** 2 + k ** 2)) +
                          el * np.log((el ** 2 + n ** 2) / (m ** 2 + el ** 2))) +
                hy * (m * (np.arctan(k / m) - np.arctan(el / m)) +
                      n * (np.arctan(el / n) - np.arctan(k / n))))
        const = ((1 - 2 * v) / (4 * np.pi * shear_mod) + second_surface * (1 - 2 * v_2) / (
            4 * np.pi * shear_mod_2)) * (-1 if comp == 'zy' else 1)
        ret = const * c_zy
    else:
        ValueError('component name not recognised: ' + comp + ', components must be lower case')
    return ret


"""
Module for dealing with LSFs of various astronomical instruments.
"""
from __future__ import print_function, absolute_import, division, unicode_literals

import numpy as np
from scipy.interpolate import interp1d
from astropy.io import fits, ascii
from astropy.units import Quantity
import astropy.units as u
from astropy.table import Table, QTable, Column
import glob, imp
from linetools.analysis.interp import interp_Akima
import warnings

lt_path = imp.find_module('linetools')[1]

class LSF(object):
    """Class to deal with line-spread-functions (LSFs) from
    various different astronomical spectrographs.

    Note: only implemented for HST/COS and HST/STIS at the moment.

    Parameters
    ----------
    instr_config : dict
        A dictionary with the instrument configuration details relevant
        to the required LSF. Mandatory keywords of the dict are: ['name'],
        all of which must be either string or None.
        Note: There must be extra relevant keywords specific to each instrument.

    """

    def __init__(self, instr_config):
        # mandatory keys for characterizing a spectrograph mode
        self.mandatory_dict_keys = ['name']

        # Check correct format
        if not isinstance(instr_config, dict):
            raise TypeError('`instr_config` must be a dictionary.')
        elif not all([key in instr_config.keys() for key in self.mandatory_dict_keys]):
            raise SyntaxError('`instr_config` must have the following mandatory keys {}:'.format(self.mandatory_dict_keys))

        # Initialize basics
        self.instr_config = instr_config
        self.name = instr_config['name']
        if self.name not in ['COS', 'STIS']:
            raise NotImplementedError('Not ready for this instrument: {}'.format(self.name))

        # initialize specific to given instrument name
        # only implemented for HST/COS  and HST/STIS so far
        if self.name == 'COS':
            self.pixel_scale, self._data = self.load_COS_data()
        elif self.name == 'STIS':
            self.pixel_scale, self._data = self.load_STIS_data()
        # IMPORTANT: make sure that LSFs are given in linear wavelength scales !!!

        #reformat self._data
        self.check_and_reformat_data()

        #other relevant values to initialize?

    def get_lsf(self, wv_array, kind='Akima'):
        """ Given a wavelength array `wv_array`, it returns
        the LSF kernel at the central wavelength of the array,
        using the same pixel scale and extent of `wv_array`.

        Method for non-Gaussian kernels: First, tabulated LSFs
        are linearly interpolated to the center of `wv_array`
        (see LSF.interpolate_to_wv0() for details); then, the
        LSF is interpolated to match the `wv_array`
        scale and extent using Akima (or cubic) interpolation (see
        LSF.interpolate_to_wv_array() for details).

        Parameters
        ----------
        wv_array : Quantity numpy.ndarray, shape(N,)
            Wavelength array for which the LSF kernel is defined. The
            central wavelength value of `wv_array` define the wavelength
            at which the LSF is defined, while the limits of `wv_array`
            define the extent of the kernel.
        kind : str, optional
            Specifies the kind of interpolation as a string either
            ('cubic', 'Akima')

        Returns
        -------
        lsf_array : numpy.ndarray, shape(N,)
            The lsf kernel.

        """
        lsf_array = self.interpolate_to_wv_array(wv_array, kind=kind)
        return lsf_array['kernel'].data

    def check_and_reformat_data(self):
        """Any re-formating of self._data should happen here.

        At the moment this function does the following:
         - Make sure that the number of relative pixels of the LSF is odd integer
         - Impose the middle value to define the 0 relative pixel
         - Make sure tables with 'rel_pix' given in fraction of pixels (e.g. COS NUV) work properly
         - Normalize tabulated LSFs
        """

        self._data['rel_pix'] = self.check_and_reformat_relpix(self._data['rel_pix'].data)

        #normalize given LSFs
        for col_name in self._data.keys()[1:]:
            self._data[col_name] /= np.sum(self._data[col_name])

    def check_and_reformat_relpix(self, relpix):
        """Performs checks and reformating on a relative pixel array.

        At the moment this function does the following:
         - Make sure that the number of relative pixels of the LSF is an odd integer
         - Impose the middle value to define the 0 relative pixel

        Parameters
        ----------
        relpix : np.array
            Relative pixel array

        Returns
        -------
        new_relpix array that conforms to the LSF class standards.
        """

        # odd integer for total number of relative pixels
        n_pix = len(relpix)
        assert n_pix % 2 != 0, ValueError('LSF tables must be given as odd integers!')

        # make sure relpix == 0 is in the middle of the array
        n_half = int( (n_pix - 1) / 2)
        mid_value = relpix[n_half]
        new_relpix = relpix - mid_value  # 0 in the middle

        return new_relpix

    def load_COS_data(self):
        """Load the right data according to `instr_config` for HST/COS
        instrument"""

        # define pixel scales; values obtained from STScI
        # these values must be consistent with the given LSFs
        pixel_scale_dict = {'G130M': 9.97 / 1000. * u.AA,
                    'G160M': 12.23 / 1000. * u.AA,
                    'G140L': 80.3 / 1000. * u.AA,
                    'G230L': 390. / 1000. * u.AA,
                    'G185M': 37. / 1000. * u.AA,
                    'G225M': 33. / 1000. * u.AA,
                    'G285M': 40. / 1000. * u.AA}
        # define channel based on grating name
        channel_dict = {'G130M':  'FUV',
                    'G160M': 'FUV',
                    'G140L': 'FUV',
                    'G230L': 'NUV',
                    'G185M': 'NUV',
                    'G225M': 'NUV',
                    'G285M': 'NUV'}

        try:
            grating = self.instr_config['grating']
        except:
            raise SyntaxError('`grating` keyword missing in `instr_config` dictionary.')

        if grating not in channel_dict.keys():
            raise NotImplementedError('Not ready for this HST/COS grating: {}'.format(grating))

        if channel_dict[grating] == 'NUV': #there is only 1 LSF file for NUV data
            file_name = 'nuv_all_lp1.txt'
        # COS
        elif channel_dict[grating] == 'FUV':
            # Use the ones corrected by scattering when possible
            # (currently, these are only available for lifetime-position 1)
            # check: http://www.stsci.edu/hst/cos/performance/spectral_resolution
            try:
                life_position = self.instr_config['life_position']
            except:
                raise SyntaxError('`life_position` keyword missing in `instr_config` dictionary.')

            if life_position not in ['1','2','3']:
                raise ValueError('HST/COS `life_position` should be either `1` or `2` or `3` (strings)')

            if life_position == '1':
                if grating == 'G140L': #use theoretical values
                    file_name = 'fuv_G140L_lp1.txt'

                elif grating == 'G130M': #use empirical values corrected by scattering
                    file_name = 'fuv_G130M_lp1_empir.txt'

                elif grating == 'G160M': #use empirical values corrected by scattering
                    file_name = 'fuv_G160M_lp1_empir.txt'

            elif life_position in ['2','3']:
                try:
                    cen_wave = self.instr_config['cen_wave']
                except:
                    raise SyntaxError('`cen_wave` keyword missing in `instr_config` dictionary. This should provide the central wavelength of the grating in Angstroms as a string.')
                #adjust format in cases where cen_wave is of the form: str(1230A)
                if cen_wave.endswith('A'): #adjust format
                    cen_wave = cen_wave[:-1]

                #filenames in this case have a well defined naming convention, and strict format.
                if life_position == '2':
                    file_name = 'fuv_{}_{}_lp2.txt'.format(grating,cen_wave)
                elif life_position == '3':
                    file_name = 'fuv_{}_{}_lp3.txt'.format(grating,cen_wave)
                else: # this should never happen
                    raise NotImplementedError('Unexpected error: please contact linetools developers!')

        else: # Wrong COS channel
            raise NotImplementedError('Not ready for the given HST/COS channel; only `NUV` and `FUV` channels allowed.')

        # point to the right file
        file_name = lt_path + '/data/lsf/{}/{}'.format(self.name,file_name)

        # get column names
        f = open(file_name,'r')
        line = f.readline()  # first line of file
        f.close()
        # get rid of '\n' in first line
        line = line.split('\n')[0]
        # by construction first column should be separated by `,`
        col_names = line.split(',')
        col_names[0] = 'rel_pix'

        pixel_scale = pixel_scale_dict[grating]  # read from dictionary defined above
        # read data
        data = ascii.read(file_name, data_start=1, names=col_names)

        return pixel_scale, data

    def load_STIS_data(self):
        """Load the right data according to `instr_config` for HST/STIS
        instrument"""

        # define pixel scales; values obtained from STScI
        # these values must be consistent with the given LSFs
        pixel_scale_dict = {
            'G140L': 0.60 * u.AA,
            'G140M': 0.05 * u.AA,
            'G230L': 1.58 * u.AA,
            'G230M': 0.09 * u.AA,
            'E140M': 'lambda/91700',
            'E140H': 'lambda/228000',
            'E230M': 'lambda/60000',
            'E230H': 'lambda/228000',
            'G230LB': 1.35 * u.AA,
            'G230MB': 0.15 * u.AA,
            'G430L': 2.73 * u.AA,
            'G430M': 0.28 * u.AA,
            'G750L': 4.92 * u.AA,
            'G750M': 0.56 * u.AA
        }
        # define channel based on grating name
        channel_dict = {
            'G140L': 'FUV-MAMA',
            'G140M': 'FUV-MAMA',
            'G230L': 'NUV-MAMA',
            'G230M': 'NUV_MAMA',
            'E140M': 'FUV-MAMA',
            'E140H': 'FUV-MAMA',
            'E230M': 'NUV-MAMA',
            'E230H': 'NUV-MAMA',
            'G230LB': 'CCD',
            'G230MB': 'CCD',
            'G430L': 'CCD',
            'G430M': 'CCD',
            'G750L': 'CCD',
            'G750M': 'CCD'
        }
        # also need slits
        available_slits = {
            'G140L': ['52x0.1', '52x0.2', '52x0.5', '52x2.0'],
            'G140M': ['52x0.1', '52x0.2', '52x0.5', '52x2.0'],
            'G230L': ['52x0.1', '52x0.2', '52x0.5', '52x2.0'],
            'G230M': ['52x0.1', '52x0.2', '52x0.5', '52x2.0'],
            'E140H': ['0.1x0.03', '0.2x0.09', '0.2x0.2', '6x0.2'],
            'E140M': ['0.1x0.03', '0.2x0.06', '0.2x0.2', '6x0.2'],
            'E230H': ['0.1x0.03', '0.1x0.09', '0.1x0.2', '6x0.2'],
            'E230M': ['0.1x0.03', '0.2x0.06', '0.2x0.2', '6x0.2'],
            'G430L': ['52x0.1', '52x0.2', '52x0.5', '52x2.0'],
            'G750L': ['52x0.1', '52x0.2', '52x0.5', '52x2.0']
        }

        try:
            grating = self.instr_config['grating']
        except:
            raise SyntaxError('`grating` keyword missing in `instr_config` dictionary.')
        if grating not in channel_dict.keys():
            raise NotImplementedError('Not ready for this HST/STIS grating: {}. '
                                      'Available gratings for HST/STIS are: {}'.format(grating, channel_dict.keys()))
        if grating in ['G750M', 'G430M', 'G430L']:
            raise NotImplementedError('{} not implemented yet; coming soon...'.format(grating))

        # We also need to know the slit width
        try:
            slit = self.instr_config['slit']
        except:
            raise SyntaxError('`slit` keyword missing in `instr_config` dictionary.')
        if slit not in available_slits[grating]:
            raise NotImplementedError('Not ready for this HST/STIS slit: {}. '
                                      'Available slits for HST/STIS grating {} are: {}'.format(slit, grating, available_slits[grating]))

        # now we need to read the right files (new ones for echelle modes)
        if grating[0]=='E':
            lsf_files = glob.glob(lt_path + '/data/lsf/STIS/stis_LSF_{}_????_LTmod.txt'.format(grating))
            # figure relevant wavelengths from file names
            wa_names = [fname.split('/')[-1].split('_')[-2].split('.')[0] for fname in lsf_files]
        else:
            lsf_files = glob.glob(lt_path + '/data/lsf/STIS/stis_LSF_{}_????.txt'.format(grating))
            # figure relevant wavelengths from file names
            wa_names = [fname.split('/')[-1].split('_')[-1].split('.')[0] for fname in lsf_files]
        #TODO: Remove following lines upon testing
        '''
        # figure relevant wavelengths from file names
        #wa_names = [fname.split('/')[-1].split('_')[-1].split('.')[0] for fname in lsf_files]
        '''
        # sort them
        sorted_inds = np.argsort(wa_names)
        lsf_files = np.array(lsf_files)[sorted_inds]
        wa_names = np.array(wa_names)[sorted_inds]

        # read the relevant kernels; they may have different rel_pix values depending on wave
        kernels_dict = dict()
        for ii, file_name in enumerate(lsf_files):
            # get the column names
            f = open(file_name,'r')
            lines = f.readlines()  # get all lines
            f.close()
            col_names = lines[1] # column names are in second line
            # get rid of '\n'
            col_names = col_names.split('\n')[0]
            # split by blank space(s) and remove the first element
            col_names = col_names.split()[1:]
            # rename new first column
            col_names[0] = 'rel_pix'

            # get original data
            data_aux = ascii.read(file_name, data_start=2, names=col_names)
            # reformat rel_pix
            data_aux['rel_pix'] = self.check_and_reformat_relpix(data_aux['rel_pix'])

            #TODO: remove the following block upon testing
            '''
            # handle asymmetric STIS LSF
            if data_aux['rel_pix'][len(data_aux['rel_pix']) // 2] == 0.:
                pass
            elif data_aux['rel_pix'][(len(data_aux['rel_pix']) // 2) - 1 ] ==0.:
                data_aux.insert_row(0,data_aux[-1])
                data_aux['rel_pix'][0]=-data_aux['rel_pix'][-1]
            elif data_aux['rel_pix'][(len(data_aux['rel_pix']) // 2) + 1 ] ==0.:
                data_aux.add_row(data_aux[0])
                data_aux['rel_pix'][-1]=-data_aux['rel_pix'][0]
            '''

            # create column with absolute wavelength based on pixel scales
            if isinstance(pixel_scale_dict[grating],np.unicode):
                # deal with wavelength-dependent pixel scale (i.e., STIS echelle)
                scalefac = 1./float(pixel_scale_dict[grating].split('/')[-1])
                wave_aux = float(wa_names[ii])*u.AA * (1. + data_aux['rel_pix']*scalefac)
            else:
                wave_aux = float(wa_names[ii])*u.AA + data_aux['rel_pix']*pixel_scale_dict[grating]
            data_aux['wv'] = wave_aux  # not used for now, but may be useful with a different approach
            # create column with normalized kernel for relevant slit
            kernel = data_aux[slit] / np.sum(data_aux[slit])
            data_aux['kernel'] = kernel

            # store only rel_pix, wv, kernel as Table
            kernels_dict[wa_names[ii]] = data_aux['rel_pix', 'wv', 'kernel']

        # at this point the kernels_dict have the original information as given by the STScI
        # tables, in their various formats...

        # project to a single rel_pix scale using interpolation; for simplicity use a
        # custom rel_pix because STScI sometimes provide them as non-constant pixel fractions.
        # This new rel_pix grid will apply to all the wavelengths

        rel_pix = kernels_dict[wa_names[0]]['rel_pix']  # use the first one as reference
        rel_pix = np.linspace(-1*np.max(rel_pix), np.max(rel_pix), len(rel_pix))  # impose them linear and symmetric
        rel_pix = self.check_and_reformat_relpix(rel_pix)  # again just in case something odd happened
        data_table = Table()
        data_table['rel_pix'] = rel_pix
        for wa_name in wa_names:
            kernel_aux = interp_Akima(data_table['rel_pix'],
                                      kernels_dict[wa_name]['rel_pix'], kernels_dict[wa_name]['kernel'])
            data_table['{}A'.format(wa_name)] = kernel_aux

        pixel_scale = pixel_scale_dict[grating]  # read from dictionary defined above
        # import pdb; pdb.set_trace()
        # todo: work out a cleverer approach to this whole issue of having different rel_pix, pixel_scales, etc
        return pixel_scale, data_table

    def interpolate_to_wv0(self, wv0):
        """Retrieves a unique LSF valid at wavelength wv0

        This is done by linearly interpolating from tabulated values
        at different wavelengths. These tabulated values (stored
        internally in self._data) are usually given as calibration
        products by instrument developers and should be loaded by
        self.load_XX_data() in the initialization stage of LSF(),
        where XX is the name of the instrument)

        Parameters
        ----------
        wv0 : Quantity
            Wavelength at which an LSF solution is required

        Returns
        -------
        lsf_table : Table
            The interpolated lsf at wv0. This table has two
            columns: 'wv' and 'kernel'
        """
        # get wa0 to Angstroms
        wv0 = wv0.to('AA').value

        # transform to wavelength in float() form assuming Angstroms
        col_names = self._data.keys()
        col_waves = np.array([float(name.split('A')[0]) for name in col_names[1:]])

        # find out the closest 2 columns in self._data to wv0 (on each side); these kernels will be used for interpolation
        seps =  np.fabs(col_waves - wv0)
        closest_ind = np.argsort(seps)[0]  # the 1 closest
        if closest_ind == 0:  # lower edge
            ind_blue = closest_ind
        elif closest_ind == len(col_waves) - 1:  # upper edge
            ind_blue = closest_ind - 1
        elif (col_waves[closest_ind] - wv0) > 0:  # middle case 1
            ind_blue = closest_ind - 1
        else:  # middle case 2
            ind_blue = closest_ind
        ind_red = ind_blue + 1

        # create a smaller version of self._data with the 2 most relevant columns
        good_keys = col_names[1:]  # get rid of the first name, i.e. 'rel_pix'
        good_keys = good_keys[ind_blue : ind_red + 1]
        data_aux = self._data[good_keys]
        col_waves_aux = col_waves[ind_blue : ind_red + 1]
        lsf_vals = []
        for row in data_aux:
            aux_val = []
            for i in range(0, len(row)):
                aux_val += [row[i]]

            # we don't want to extrapolate wildly, but allow LSF instantiations for wv0 outside range of 'col_waves'
            if (wv0 >= col_waves_aux[0]) & (wv0 <= col_waves_aux[-1]):
                f = interp1d(col_waves_aux,aux_val,bounds_error=True,kind='linear')  # no need to extrapolate
                lsf_vals += [f(wv0)]

            elif (wv0 < col_waves[0]) & ((col_waves[0] - wv0) < np.abs(col_waves[1] - col_waves[0])):
                f = interp1d(col_waves_aux, aux_val, bounds_error=False,
                             fill_value=aux_val[0], kind='linear')  # assign shortest wv LSF definition
                lsf_vals += [f(wv0)]

                # warning
                if (col_waves[0] - wv0) > (np.abs(col_waves[1] - col_waves[0])/2.):
                    warnings.warn(
                        "LSF may result from extrapolation outside wavelength range characterized for current grating.")

            elif (wv0 > col_waves[-1]) & ((wv0 - col_waves[-1]) < np.abs(col_waves[-1] - col_waves[-2])):
                f = interp1d(col_waves_aux, aux_val, bounds_error=False,
                             fill_value=aux_val[-1], kind='linear') #assign longest wv LSF definition
                lsf_vals += [f(wv0)]

                # warning
                if (wv0 - col_waves[-1]) > (np.abs(col_waves[-1] - col_waves[-2])/2.):
                    warnings.warn(
                        "LSF may result from extrapolation outside wavelength range characterized for current grating.")
            else:
                raise ValueError("wv0 too far outside range of defined LSFs. Perhaps you've chosen the wrong grating?")
        lsf_vals = np.array(lsf_vals)

        # normalize
        lsf_vals /= np.sum(lsf_vals)

        # create Column to store the interpolated LSF
        # lsf_vals = Column(name='{:.0f}A'.format(wv0.value),data=lsf_vals)
        lsf_vals = Column(name='kernel', data=lsf_vals)

        # create column of relative pixel in absolute wavelength
        if isinstance(self.pixel_scale,np.unicode):
            # deal with wavelength-dependent pixel scale (i.e., STIS echelle)
            scalefac = 1./float(self.pixel_scale.split('/')[-1])
            wv_array = [(wv0*u.AA * (1. + scalefac * self._data['rel_pix'][i])).value
                        for i in range(len(self._data))]
        else:
            wv_array = [(self.pixel_scale * self._data['rel_pix'][i] + wv0*u.AA).value for i in range(len(self._data))]
        wv = Column(name='wv',data=wv_array, unit=u.AA)

        # create lsf Table
        lsf = Table()
        lsf.add_column(wv)
        lsf.add_column(lsf_vals)

        # return lsf Table()
        return lsf

    def interpolate_to_wv_array(self, wv_array, kind='Akima', debug=False):
        """ Interpolate an LSF to a wavelength array.

        Given `wv_array` this function interpolates an LSF
        to match both scale and extent of `wv_array` using the
        Akima or cubic-spline interpolators (default is Akima).
        Some checks are performed too.

        Parameters
        ----------
        wv_array : Quantity numpy.ndarray, shape(N,)
            Wavelength array for which the LSF kernel is defined. The
            central wavelength value of `wv_array` define the wavelength
            at which the LSF is defined, while the limits of `wv_array`
            define the extent of the kernel.
        kind : str, optional
            Specifies the kind of interpolation as a string either
            ('cubic', 'Akima'); default is `Akima`.

        Returns
        -------
        lsf_table : Table
            The interpolated lsf using at the central wavelength of
            `wv_array`, using the same pixel scale as `wv_array`.
            This table has two columns: 'wv' and 'kernel'. (lst_table['wv']
            is equal to `wv_array` by construction.)

        """
        # Check correct format
        if not ((isinstance(wv_array, np.ndarray)) or (isinstance(wv_array, Quantity))):
            raise SyntaxError('`wv_array` must be Quantity numpy.ndarray')
        elif len(wv_array.shape) != 1:
            raise SyntaxError('`wv_array` must be of shape(N,), i.e. 1-dimensional array')
        if kind not in ['cubic','Akima','akima']:
            raise ValueError('Only `cubic` or `Akima` interpolation available.')

        # define useful quantities
        wv_min = np.min(wv_array)
        wv_max = np.max(wv_array)
        wv0 = 0.5 * (wv_max + wv_min)

        lsf_tab = self.interpolate_to_wv0(wv0)

        # make sure the wv_array is dense enough to sample the LSF kernel
        kernel_wvmin = np.min(lsf_tab['wv']) * u.AA
        kernel_wvmax = np.max(lsf_tab['wv']) * u.AA
        cond = (wv_array >= kernel_wvmin) & (wv_array <= kernel_wvmax)
        if np.sum(cond) < 10:  # this number is somewhat arbitrary but reasonable
            raise ValueError('The input `wv_array` is undersampling the LSF kernel! Try a finer grid.')

        # convert to Angstroms
        wv_array_AA = np.array([wv.to('AA').value for wv in wv_array])

        # interpolate to wv_array
        if kind == 'cubic':
            f = interp1d(lsf_tab['wv'], lsf_tab['kernel'], kind='cubic', bounds_error= False, fill_value=0)
            lsf_vals =  f(wv_array_AA)
        elif kind in ('Akima','akima'):
            # f = Akima1DInterpolator(lsf_tab['wv'],lsf_tab['kernel'])
            # NT: I tried Akima interpolator from scipy.interpolate
            # and is not robust in extreme situations where the
            # wv_array is large compared to the kernel FWHM.
            # Let's try linetools.analysis.interp Akima version
            lsf_vals = interp_Akima(wv_array_AA,lsf_tab['wv'],lsf_tab['kernel'])

        # make sure the kernel is never negative
        cond = lsf_vals < 0
        if np.sum(cond) > 0:
            warnings.warn('The interpolated kernel has negative values; imposing them to be 0.')
            if debug:
                import matplotlib.pyplot as plt
                plt.plot(wv_array_AA, lsf_vals, 'k-')
                # import pdb; pdb.set_trace()
            lsf_vals = np.where(lsf_vals < 0, 0., lsf_vals)

        # normalize
        lsf_vals /= np.sum(lsf_vals)

        # re-define Table
        lsf_tab = Table()
        lsf_tab.add_column(Column(name='wv', data=wv_array))
        lsf_tab.add_column(Column(name='kernel', data=lsf_vals))

        return lsf_tab

import numpy,from scipy,import astropy,from astropy


'''
Commonly used global constants. **New users should change
`ANCILLARY_DATA_PATHS` if it doesn't match their file system.**

Source for EASE-Grid 2.0 parameters:

    https://nsidc.org/ease/ease-grid-projection-gt

NOTE: EASE-Grid 2.0 parameters do not match exactly those specified at NSIDC
because, in fact, they must be 1000 or 9000, depending on the grid size, in
order to get the right number of rows and columns in the output image.
Grid resolution is only a whole number for the polar (north or south
hemisphere) grids.
'''

import csv
import numpy as np
from collections import OrderedDict

ANCILLARY_DATA_PATHS = {
    'smap_l4c_ancillary_data_file_path': 'SPL4C_Vv4040_SMAP_L4_C.Ancillary.h5',
    'smap_l4c_1km_ancillary_data_lc_path': 'MCD12Q1_M01_lc_dom_uint8',
    'smap_l4c_9km_ancillary_data_lc_path': 'MOD12Q1_M09_lc_dom_uint8',
    'smap_l4c_1km_ancillary_data_x_coord_path': 'SMAP_L4_C_LON_14616_x_34704_M01_flt32',
    'smap_l4c_1km_ancillary_data_y_coord_path': 'SMAP_L4_C_LAT_14616_x_34704_M01_flt32',
    'smap_l4c_9km_ancillary_data_x_coord_path': 'SMAP_L4_C_LON_1624_x_3856_M09_flt32',
    'smap_l4c_9km_ancillary_data_y_coord_path': 'SMAP_L4_C_LAT_1624_x_3856_M09_flt32',
    'smap_l4c_9km_pft_subgrid_counts_CONUS': 'SMAP_L4C_Vv4040_1km_subgrid_PFT_counts_CONUS.h5',
    'smap_l4c_9km_sparse_col_index': 'MCD12Q1_M09land_col.uint16',
    'smap_l4c_9km_sparse_row_index': 'MCD12Q1_M09land_row.uint16',
    'transcom_netcdf_path': 'CarbonTracker_TransCom_and_other_regions.nc'
}

# Version 4020 BPLUT parameters in order of PFT numeric code (PFT 0 through 9)
BPLUT = OrderedDict({
    # Only reason for OrderedDict (instead of dict) is for consistency with
    #   pyl4c.data.fixtures.BPLUT()
    '_version': '4',
    'LUE': np.array([[ # gC per MJ
        np.nan, 1.71, 1.38, 1.71, 1.19, 1.95, 1.68, 2.53, 3.6, np.nan
    ]]),
    'CUE': np.array([[
        np.nan, 0.7, 0.57, 0.79, 0.78, 0.65, 0.6, 0.72, 0.71, np.nan
    ]]),
    'tmin': np.array([ # degrees K
        [np.nan, 235, 230, 259, 260, 259, 251, 246, 266, np.nan],
        [np.nan, 309, 303, 301, 284, 304, 302, 314, 319, np.nan]
    ]),
    'vpd': np.array([ # Pascals
        [np.nan, 0, 15, 869, 1500, 4, 0, 228, 1500, np.nan],
        [np.nan, 4169, 7000, 3452, 5401, 4282, 4229, 4516, 7000, np.nan]
    ]),
    'smrz': np.array([ # Percent saturation
        [np.nan, -30, 18, -30, -26, 22, -30, -15, 10, np.nan],
        [np.nan, 74, 26, 49, 87, 72, 76, 30, 68, np.nan]
    ]),
    'smsf': np.array([ # Percent saturation
        [np.nan, -50, -46, -9, -50, -15, -3, -42, -49, np.nan],
        [np.nan, 46, 59, 43, 53, 48, 51, 41, 30, np.nan]
    ]),
    'ft': np.array([ # Frozen = 0, Thawed = 1
        [np.nan, 0.58, 0.36, 0.55, 0.9, 0.91, 0.85, 0.77, 1, np.nan],
        [np.nan, 1, 1, 1, 1, 1, 1, 1, 1, np.nan]
    ]),
    'tsoil': np.array([[ # degrees K
        np.nan, 265.04, 477.83, 238.81, 267.26, 292.82, 232.24, 263.97, 329.63, np.nan
    ]]),
    # NOTE: The medium/structural and slow/recalcitrant decay constants
    #   (2nd and 3rd rows) are true decay constants, unlike in Ops
    #   BPLUT which presents only dimensionless scalars.
    'decay_rates': np.array([ # days^-1
        [np.nan, 0.027, 0.028, 0.028, 0.03, 0.015, 0.025, 0.019, 0.035, np.nan],
        [np.nan, 1.08e-2, 1.12e-2, 1.12e-2, 1.2e-2, 0.6e-2, 1e-2, 0.76e-2, 1.4e-2, np.nan],
        [np.nan, 2.51e-4, 2.6e-4, 2.6e-4, 2.79e-4, 1.4e-4, 2.33e-4, 1.77e-4, 3.26e-4, np.nan]
    ]),
    'f_metabolic': np.array([[ # Fraction of daily litterfall entering "fast" or metabolic pool
        np.nan, 0.49, 0.71, 0.67, 0.67, 0.62, 0.76, 0.78, 0.78, np.nan
    ]]), # Also known as "f_met"
    'f_structural': np.array([[ # Fraction of structural (str) pool transferred in "humification"
        np.nan, 0.3, 0.3, 0.7, 0.3, 0.35, 0.55, 0.5, 0.8, np.nan
    ]]) # Also known as "f_str" -- See Jones et al. (2017, IEEE TGARS, p.5)
})

EASE2_GRID_PARAMS = {
    # A GeoTransform for a north-up raster is:
    #   (x_min, pixel_width, 0, y_max, 0, -pixel_height)
    'M01': {
        'epsg': 6933,
        'geotransform': (-17367530.45, 1000, 0, 7314540.83, 0, -1000),
        'resolution': 1000.89502334956, # From Brodzik et al.
        'shape': (14616, 34704),
        'size': 14616*34704
    },
    'M09': {
        'epsg': 6933,
        'geotransform': (-17367530.45, 9000, 0, 7314540.83, 0, -9000),
        'resolution': 9008.055210146, # From Brodzik et al.
        'shape': (1624, 3856),
        'size': 1624*3856
    },
    'N09': {
        'epsg': 6931,
        'geotransform': (-9000000.0, 9000, 0, 9000000.0, 0, -9000),
        'resolution': 9000,
        'shape': (2000, 2000),
        'size': 2000*2000
    },
    'M25': {
        'epsg': 6933,
        'geotransform': (-17367530.45, 25000, 0, 7307375.92, 0, -25000),
        'resolution': 25000,
        'shape': (584, 1388),
        'size': 584*1388
    },
    'M36': {
        'epsg': 6933,
        'geotransform': (-17367530.45, 36000, 0, 7314540.83, 0, -36000),
        'resolution': 36032.22,
        'shape': (406, 964),
        'size': 406*964
    }
}

HDF_PATHS = { # Where in the HDF hierarchy certain variables live
    'SPL4CMDL': { # By Earthdata Dataset ID
        '4': { # By Version number
            'longitude': 'GEO/longitude',
            'latitude': 'GEO/latitude',
            'SOC': 'SOC/soc_mean',
            'SOC*': 'SOC/soc_pft%d_mean',
            'GPP': 'GPP/gpp_mean',
            'GPP*': 'GPP/gpp_pft%d_mean',
            'NEE': 'NEE/nee_mean',
            'NEE*': 'NEE/nee_pft%d_mean',
            'RH': 'RH/rh_mean',
            'RH*': 'RH/rh_pft%d_mean',
        }
    },
    'SPL4SMGP': {
        '4': {'longitude': 'cell_lon', 'latitude': 'cell_lat',}
    }
}

PFT = OrderedDict({
    0: ('Water', 'WET'),
    1: ('Evergreen Needleleaf', 'ENF'),
    2: ('Evergreen Broadleaf', 'EBF'),
    3: ('Deciduous Needleleaf', 'DNF'),
    4: ('Deciduous Broadleaf', 'DBF'),
    5: ('Shrub', 'SHB'),
    6: ('Grass', 'GRS'),
    7: ('Cereal Crop', 'CCR'),
    8: ('Broadleaf Crop', 'BCR'),
    9: ('Urban and Built-Up', 'URB')
})

# BBOX extents, in decimal degrees, based on gdal_rasterize -te option:
#   https://gdal.org/programs/gdal_rasterize.html#cmdoption-gdal-translate-te
SUBSETS_BBOX = { # <xmin>, <ymin>, <xmax>, <ymax>
    'CONUS': [-124.5, 24.4, -66.7, 50.0],
    'WesternHemisphere': [-180, -90, 0, 90],
    'WesternHemisphere2': [-180, -90, -25, 90],
    'NorthernHemisphere': [-180, 0, 180, 90],
    'NorthernHemisphere45': [-180, 45, 180, 90],
    'NorthernHemisphere40': [-180, 40, 180, 90],
    'Nigeria': [2.3, 4, 14.8, 14],
    # AOI for: Chiodi and Harrison (2013) Journal of Climate 26 (3):822–837
    'ChiodiHarrison2013': [-160, -5, -110, 5],
    'Kang2014EMS': [-125, 47, -113, 49],
    'NorthernPlains': [-116.1, 40.9, -96.4, 49],
    'Montana': [-116.1, 44.4, -104.01, 49],
    'Iowa': [-96.7, 40.58, -90.1, 43.5],
    'Bangladesh1kmSubset': [78.9, 23.2, 96.1, 29.3]
}

def parameter_mapped(name, pft_array, bplut = BPLUT):
    '''
    Given a BPLUT parameter and a PFT array, returns an array with the
    corresponding parameter values for each PFT code.

    Parameters
    ----------
    name : str
        The name of the BPLUT parameter
    pft_array : numpy.ndarray
        Array of any size or shape with numeric elements corresponding to
        PFT codes
    bplut : dict
        (Optional) A collection of BPLUT parameters

    Returns
    -------
    numpy.ndarray
        Array of same size, shape as pft_array but with parameter values
        in place of PFT codes
    '''
    param = bplut[name]
    # Basically, index the <name> array, in PFT order, by PFT numeric codes
    return np.asarray(param)[np.ravel(pft_array)].reshape(pft_array.shape)

def restore_bplut(csv_file_path, version_id = None):
    '''
    Translates a BPLUT CSV file to a Python internal representation
    (OrderedDict instance).

    Parameters
    ----------
    csv_file_path : str
        File path to the CSV representation of the BPLUT
    version_id : str
        (Optional) Version identifier for the BPLUT

    Returns
    -------
    OrderedDict
    '''
    header = ('LC_index', 'LC_Label', 'model_code', 'NDVItoFPAR_scale',
        'NDVItoFPAR_offset', 'LUEmax', 'Tmin_min_K', 'Tmin_max_K',
        'VPD_min_Pa', 'VPD_max_Pa', 'SMrz_min', 'SMrz_max', 'FT_min',
        'FT_max', 'SMtop_min', 'SMtop_max', 'Tsoil_beta0', 'Tsoil_beta1',
        'Tsoil_beta2', 'fraut', 'fmet', 'fstr', 'kopt', 'kstr', 'kslw',
        'Nee_QA_Rank_min', 'Nee_QA_Rank_max', 'Nee_QA_Error_min',
        'Nee_QA_Error_max', 'Fpar_QA_Rank_min', 'Fpar_QA_Rank_max',
        'Fpar_QA_Error_min', 'Fpar_QA_Error_max', 'FtMethod_QA_mult',
        'FtAge_QA_Rank_min', 'FtAge_QA_Rank_max', 'FtAge_QA_Error_min',
        'FtAge_QA_Error_max', 'Par_QA_Error', 'Tmin_QA_Error',
        'Vpd_QA_Error', 'Smrz_QA_Error', 'Tsoil_QA_Error', 'Smtop_QA_Error')
    contents = []
    with open(csv_file_path, 'r') as stream:
        reader = csv.DictReader(
            filter(lambda row: row[0] != '#', stream), fieldnames = header)
        for row in reader:
            contents.append(row)

    result = BPLUT.copy()
    result['_version'] = 'UNKNOWN' if version_id is None else version_id
    result['LUE'] = np.array([[
        contents[p-1]['LUEmax'] if p in range(1, 9) else np.nan
        for p in range(0, 10)
    ]], dtype = np.float32).round(4)
    result['CUE'] = np.array([[
        (1 - float(contents[p-1]['fraut'])) if p in range(1, 9) else np.nan
        for p in range(0, 10)
    ]], dtype = np.float32).round(4)
    result['tmin'] = np.array([
        [contents[p-1]['Tmin_min_K'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
        [contents[p-1]['Tmin_max_K'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
    ], dtype = np.float32).round(1)
    result['vpd'] = np.array([
        [contents[p-1]['VPD_min_Pa'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
        [contents[p-1]['VPD_max_Pa'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
    ], dtype = np.float32).round(2)
    result['smrz'] = np.array([
        [contents[p-1]['SMrz_min'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
        [contents[p-1]['SMrz_max'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
    ], dtype = np.float32).round(1)
    result['smsf'] = np.array([
        [contents[p-1]['SMtop_min'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
        [contents[p-1]['SMtop_max'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
    ], dtype = np.float32).round(1)
    result['ft'] = np.array([
        [contents[p-1]['FT_min'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
        [contents[p-1]['FT_max'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
    ], dtype = np.float32).round(3)
    result['tsoil'] = np.array([[
        contents[p-1]['Tsoil_beta0'] if p in range(1, 9) else np.nan
        for p in range(0, 10)
    ]], dtype = np.float32).round(1)
    result['f_metabolic'] = np.array([[
        contents[p-1]['fmet'] if p in range(1, 9) else np.nan
        for p in range(0, 10)
    ]], dtype = np.float32).round(3)
    result['f_structural'] = np.array([[
        contents[p-1]['fstr'] if p in range(1, 9) else np.nan
        for p in range(0, 10)
    ]], dtype = np.float32).round(1)
    result['decay_rates'] = np.array([
        [contents[p-1]['kopt'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
        [contents[p-1]['kstr'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
        [contents[p-1]['kslw'] if p in range(1, 9) else np.nan
            for p in range(0, 10)],
    ], dtype = np.float32).round(3)
    # The "kstr" and "kslw" values are really the fraction of kopt
    #   assigned to the second and third pools
    result['decay_rates'][1,:] = np.multiply(
        result['decay_rates'][0,:], result['decay_rates'][1,:])
    result['decay_rates'][2,:] = np.multiply(
        result['decay_rates'][0,:], result['decay_rates'][2,:])
    return result