from collections import namedtuple
import numpy as np
import scipy as sp
from scipy.sparse.csgraph import minimum_spanning_tree
from .. import logging as logg
from ..neighbors import Neighbors
from .. import utils
from .. import settings

def paga(
        adata,
        groups='louvain',
        use_rna_velocity=False,
        copy=False):
    """\
    Generate cellular maps of differentiation manifolds with complex
    topologies [Wolf17i]_.

    Partition-based graph abstraction (PAGA) quantifies the connectivities of
    partitions of a neighborhood graph of single cells, thereby generating a
    much simpler abstracted graph whose nodes label the partitions. Together
    with a random walk-based distance measure, this generates a partial
    coordinatization of data useful for exploring and explaining its variation.

    Parameters
    ----------
    adata : :class:`~scanpy.api.AnnData`
        Annotated data matrix.
    groups : categorical annotation of observations or 'louvain_groups', optional (default: 'louvain_groups')
        Criterion to determine the resulting partitions of the single-cell
        graph. 'louvain_groups' uses the Louvain algorithm and optimizes
        modularity of the graph. You can also pass your predefined groups by
        choosing any categorical annotation of observations (`adata.obs`).
    use_rna_velocity : `bool` (default: `False`)
        Use RNA velocity to orient edges in the abstracted graph and estimate transitions.
    copy : `bool`, optional (default: `False`)
        Copy `adata` before computation and return a copy. Otherwise, perform
        computation inplace and return `None`.

    Returns
    -------
    Returns or updates `adata` depending on `copy` with
    connectivities : np.ndarray (adata.uns['connectivities'])
        The full adjacency matrix of the abstracted graph, weights
        correspond to connectivities.
    confidence : np.ndarray (adata.uns['confidence'])
        The full adjacency matrix of the abstracted graph, weights
        correspond to confidence in the presence of an edge.
    confidence_tree : sc.sparse csr matrix (adata.uns['confidence_tree'])
        The adjacency matrix of the tree-like subgraph that best explains
        the topology.
    """
    if 'neighbors' not in adata.uns:
        raise ValueError(
            'You need to run `pp.neighbors` first to compute a neighborhood graph.')
    adata = adata.copy() if copy else adata
    utils.sanitize_anndata(adata)
    logg.info('running partition-based graph abstraction (PAGA)', reset=True)
    paga = PAGA(adata, groups, use_rna_velocity=use_rna_velocity)
    paga.compute()
    # only add if not present
    if 'paga' not in adata.uns:
        adata.uns['paga'] = {}
    if not use_rna_velocity:
        adata.uns['paga']['connectivities'] = paga.connectivities_coarse
        adata.uns['paga']['confidence'] = paga.confidence
        adata.uns['paga']['confidence_tree'] = paga.confidence_tree
        adata.uns[groups + '_sizes'] = np.array(paga.vc.sizes())
    else:
        adata.uns['paga']['transitions_confidence'] = paga.transitions_confidence
        adata.uns['paga']['transitions_ttest'] = paga.transitions_ttest
    adata.uns['paga']['groups'] = groups
    logg.info('    finished', time=True, end=' ' if settings.verbosity > 2 else '\n')
    if use_rna_velocity:
        logg.hint(
            'added\n'
            '    \'paga/transitions_confidence\', confidence adjacency (adata.uns)\n'
            '    \'paga/transitions_ttest\', confidence subtree (adata.uns)')
    else:
        logg.hint(
            'added\n'
            '    \'paga/connectivities\', connectivities adjacency (adata.uns)\n'
            '    \'paga/confidence\', confidence adjacency (adata.uns)\n'
            '    \'paga/confidence_tree\', confidence subtree (adata.uns)')
    return adata if copy else None

class PAGA(Neighbors):

    def __init__(self, adata, groups, use_rna_velocity=False,
                 tree_based_confidence=False):
        super(PAGA, self).__init__(adata)
        self._groups = groups
        self._tree_based_confidence = tree_based_confidence
        self._use_rna_velocity = use_rna_velocity

    def compute(self):
        if self._use_rna_velocity:
            self.compute_transitions_coarse()
        else:
            self.compute_connectivities_coarse()
            self.compute_confidence()

    def compute_connectivities_coarse(self):
        import igraph
        ones = self.connectivities.copy()
        # graph where edges carry weight 1
        ones.data = np.ones(len(ones.data))
        g = utils.get_igraph_from_adjacency(ones)
        self.vc = igraph.VertexClustering(
            g, membership=self._adata.obs[self._groups].cat.codes.values)
        cg = self.vc.cluster_graph(combine_edges='sum')
        self.connectivities_coarse = utils.get_sparse_from_igraph(cg, weight_attr='weight')/2

    def compute_confidence(self):
        """Translates the connectivities_coarse measure into a confidence measure.
        """
        pseudo_distance = self.connectivities_coarse.copy()
        pseudo_distance.data = 1./pseudo_distance.data
        connectivities_coarse_tree = minimum_spanning_tree(pseudo_distance)
        connectivities_coarse_tree.data = 1./connectivities_coarse_tree.data
        connectivities_coarse_tree_indices = [
            connectivities_coarse_tree[i].nonzero()[1]
            for i in range(connectivities_coarse_tree.shape[0])]
        # inter- and intra-cluster based confidence
        if not self._tree_based_confidence:
            total_n = self.n_neighbors * np.array(self.vc.sizes())
            maximum = self.connectivities_coarse.max()
            confidence = self.connectivities_coarse.copy()  # initializing
            for i in range(self.connectivities_coarse.shape[0]):
                for j in range(i+1, self.connectivities_coarse.shape[1]):
                    if self.connectivities_coarse[i, j] > 0:
                        geom_mean = np.sqrt(total_n[i] * total_n[j])
                        confidence[i, j] = self.connectivities_coarse[i, j] / geom_mean
                        confidence[j, i] = confidence[i, j]
        # tree-based confidence
        else:
            median_connectivities_coarse_tree = np.median(connectivities_coarse_tree.data)
            confidence = self.connectivities_coarse.copy()
            confidence.data[self.connectivities_coarse.data >= median_connectivities_coarse_tree] = 1
            connectivities_coarse_adjusted = self.connectivities_coarse.copy()
            connectivities_coarse_adjusted.data -= median_connectivities_coarse_tree
            connectivities_coarse_adjusted.data = np.exp(connectivities_coarse_adjusted.data)
            index = self.connectivities_coarse.data < median_connectivities_coarse_tree
            confidence.data[index] = connectivities_coarse_adjusted.data[index]
        confidence_tree = self.compute_confidence_tree(
            confidence, connectivities_coarse_tree_indices)
        self.confidence = confidence
        self.confidence_tree = confidence_tree

    def compute_confidence_tree(
            self, confidence, connectivities_coarse_tree_indices):
        confidence_tree = sp.sparse.lil_matrix(confidence.shape, dtype=float)
        for i, neighbors in enumerate(connectivities_coarse_tree_indices):
            if len(neighbors) > 0:
                confidence_tree[i, neighbors] = confidence[i, neighbors]
        return confidence_tree.tocsr()

    def compute_transitions_coarse(self):
        # analogous code using networkx
        # membership = adata.obs['clusters'].cat.codes.tolist()
        # partition = defaultdict(list)
        # for n, p in zip(list(range(len(G))), membership):
        #     partition[p].append(n)
        # partition = partition.values()
        # g_abstracted = nx.quotient_graph(g, partition, relabel=True)
        # for some reason, though, edges aren't oriented in the quotient
        # graph...
        import igraph
        g = utils.get_igraph_from_adjacency(
            self._adata.uns['velocyto_transitions'], directed=True)
        vc = igraph.VertexClustering(
            g, membership=self._adata.obs[self._groups].cat.codes.values)
        cg_full = vc.cluster_graph(combine_edges=False)

        g_bool = utils.get_igraph_from_adjacency(
            self._adata.uns['velocyto_transitions'].astype('bool'), directed=True)
        vc_bool = igraph.VertexClustering(
            g_bool, membership=self._adata.obs[self._groups].cat.codes.values)
        cg_bool = vc_bool.cluster_graph(combine_edges='sum')  # collapsed version
        transitions_coarse = utils.get_sparse_from_igraph(cg_bool, weight_attr='weight')
        # translate this into a confidence measure
        # the number of outgoing edges
        # total_n = np.zeros(len(vc.sizes()))
        # # (this is not the convention of standard stochastic matrices)
        # total_outgoing = transitions_coarse.sum(axis=1)
        # for i in range(len(total_n)):
        #     total_n[i] = vc.subgraph(i).ecount()
        #     total_n[i] += total_outgoing[i, 0]
        # use the topology based reference, the velocity one might have very small numbers
        total_n = self.n_neighbors * np.array(vc_bool.sizes())
        transitions_ttest = transitions_coarse.copy()
        transitions_confidence = transitions_coarse.copy()
        from scipy.stats import ttest_1samp
        for i in range(transitions_coarse.shape[0]):
            # no symmetry in transitions_coarse, hence we should not restrict to
            # upper triangle
            neighbors = transitions_coarse[i].nonzero()[1]
            for j in neighbors:
                forward = cg_full.es.select(_source=i, _target=j)['weight']
                backward = cg_full.es.select(_source=j, _target=i)['weight']
                # backward direction: add minus sign
                values = np.array(list(forward) + list(-np.array(backward)))
                # require some minimal number of observations
                if len(values) < 5:
                    transitions_ttest[i, j] = 0
                    transitions_ttest[j, i] = 0
                    transitions_confidence[i, j] = 0
                    transitions_confidence[j, i] = 0
                    continue
                t, prob = ttest_1samp(values, 0.0)
                if t > 0:
                    # number of outgoing edges greater than number of ingoing edges
                    # i.e., transition from i to j
                    transitions_ttest[i, j] = -np.log10(max(prob, 1e-10))
                    transitions_ttest[j, i] = 0
                else:
                    transitions_ttest[j, i] = -np.log10(max(prob, 1e-10))
                    transitions_ttest[i, j] = 0
                # geom_mean
                geom_mean = np.sqrt(total_n[i] * total_n[j])
                diff = (len(forward) - len(backward)) / geom_mean
                if diff > 0:
                    transitions_confidence[i, j] = diff
                    transitions_confidence[j, i] = 0
                else:
                    transitions_confidence[j, i] = -diff
                    transitions_confidence[i, j] = 0
        transitions_ttest.eliminate_zeros()
        transitions_confidence.eliminate_zeros()
        # transpose in order to match convention of stochastic matrices
        # entry ij means transition from j to i
        self.transitions_ttest = transitions_ttest.T
        self.transitions_confidence = transitions_confidence.T

def paga_degrees(adata):
    """Compute the degree of each node in the abstracted graph.

    Parameters
    ----------
    adata : AnnData
        Annotated data matrix.

    Returns
    -------
    degrees : list
        List of degrees for each node.
    """
    import networkx as nx
    g = nx.Graph(adata.uns['paga']['confidence'])
    degrees = [d for _, d in g.degree(weight='weight')]
    return degrees

def paga_expression_entropies(adata):
    """Compute the median expression entropy for each node-group.

    Parameters
    ----------
    adata : AnnData
        Annotated data matrix.

    Returns
    -------
    entropies : list
        Entropies of median expressions for each node.
    """
    from scipy.stats import entropy
    groups_order, groups_masks = utils.select_groups(
        adata, key=adata.uns['paga']['groups'])
    entropies = []
    for mask in groups_masks:
        X_mask = adata.X[mask]
        x_median = np.median(X_mask, axis=0)
        x_probs = (x_median - np.min(x_median)) / (np.max(x_median) - np.min(x_median))
        entropies.append(entropy(x_probs))
    return entropies

def paga_compare_paths(adata1, adata2,
                       adjacency_key='confidence', adjacency_key2=None):
    """Compare paths in abstracted graphs in two datasets.

    Compute the fraction of consistent paths between leafs, a measure for the
    topological similarity between graphs.

    By increasing the verbosity to level 4 and 5, the paths that do not agree
    and the paths that agree are written to the output, respectively.

    The PAGA "groups key" needs to be the same in both objects.

    Parameters
    ----------
    adata1, adata2 : AnnData
        Annotated data matrices to compare.
    adjacency_key : str
        Key for indexing the adjacency matrices in `.uns['paga']` to be used in
        adata1 and adata2.
    adjacency_key2 : str, None
        If provided, used for adata2.

    Returns
    -------
    OrderedTuple with attributes ``n_steps`` (total number of steps in paths)
    and ``frac_steps`` (fraction of consistent steps), ``n_paths`` and
    ``frac_paths``.
    """
    import networkx as nx
    g1 = nx.Graph(adata1.uns['paga'][adjacency_key])
    g2 = nx.Graph(adata2.uns['paga'][adjacency_key2 if adjacency_key2 is not None else adjacency_key])
    leaf_nodes1 = [str(x) for x in g1.nodes() if g1.degree(x) == 1]
    logg.msg('leaf nodes in graph 1: {}'.format(leaf_nodes1), v=5, no_indent=True)
    paga_groups = adata1.uns['paga']['groups']
    asso_groups1 = utils.identify_groups(adata1.obs[paga_groups].values,
                                         adata2.obs[paga_groups].values)
    asso_groups2 = utils.identify_groups(adata2.obs[paga_groups].values,
                                         adata1.obs[paga_groups].values)
    orig_names1 = adata1.obs[paga_groups].cat.categories
    orig_names2 = adata2.obs[paga_groups].cat.categories

    import itertools
    n_steps = 0
    n_agreeing_steps = 0
    n_paths = 0
    n_agreeing_paths = 0
    # loop over all pairs of leaf nodes in the reference adata1
    for (r, s) in itertools.combinations(leaf_nodes1, r=2):
        r2, s2 = asso_groups1[r][0], asso_groups1[s][0]
        orig_names = [orig_names1[int(i)] for i in [r, s]]
        orig_names += [orig_names2[int(i)] for i in [r2, s2]]
        logg.msg('compare shortest paths between leafs ({}, {}) in graph1 and ({}, {}) in graph2:'
               .format(*orig_names), v=4, no_indent=True)
        no_path1 = False
        try:
            path1 = [str(x) for x in nx.shortest_path(g1, int(r), int(s))]
        except nx.NetworkXNoPath:
            no_path1 = True
        no_path2 = False
        try:
            path2 = [str(x) for x in nx.shortest_path(g2, int(r2), int(s2))]
        except nx.NetworkXNoPath:
            no_path2 = True
        if no_path1 and no_path2:
            # consistent behavior
            n_paths += 1
            n_agreeing_paths += 1
            n_steps += 1
            n_agreeing_steps += 1
            logg.msg('there are no connecting paths in both graphs', v=5, no_indent=True)
            continue
        elif no_path1 or no_path2:
            # non-consistent result
            n_paths += 1
            n_steps += 1
            continue
        if len(path1) >= len(path2):
            path_mapped = [asso_groups1[l] for l in path1]
            path_compare = path2
            path_compare_id = 2
            path_compare_orig_names = [[orig_names2[int(s)] for s in l] for l in path_compare]
            path_mapped_orig_names = [[orig_names2[int(s)] for s in l] for l in path_mapped]
        else:
            path_mapped = [asso_groups2[l] for l in path2]
            path_compare = path1
            path_compare_id = 1
            path_compare_orig_names = [[orig_names1[int(s)] for s in l] for l in path_compare]
            path_mapped_orig_names = [[orig_names1[int(s)] for s in l] for l in path_mapped]
        n_agreeing_steps_path = 0
        ip_progress = 0
        for il, l in enumerate(path_compare[:-1]):
            for ip, p in enumerate(path_mapped):
                if ip >= ip_progress and l in p:
                    # check whether we can find the step forward of path_compare in path_mapped
                    if (ip + 1 < len(path_mapped)
                        and
                        path_compare[il + 1] in path_mapped[ip + 1]):
                        # make sure that a step backward leads us to the same value of l
                        # in case we "jumped"
                        logg.msg('found matching step ({} -> {}) at position {} in path{} and position {} in path_mapped'
                               .format(l, path_compare_orig_names[il + 1], il, path_compare_id, ip), v=6)
                        consistent_history = True
                        for iip in range(ip, ip_progress, -1):
                            if l not in path_mapped[iip - 1]:
                                consistent_history = False
                        if consistent_history:
                            # here, we take one step further back (ip_progress - 1); it's implied that this
                            # was ok in the previous step
                            logg.msg('    step(s) backward to position(s) {} in path_mapped are fine, too: valid step'
                                   .format(list(range(ip - 1, ip_progress - 2, -1))), v=6)
                            n_agreeing_steps_path += 1
                            ip_progress = ip + 1
                            break
        n_steps_path = len(path_compare) - 1
        n_agreeing_steps += n_agreeing_steps_path
        n_steps += n_steps_path
        n_paths += 1
        if n_agreeing_steps_path == n_steps_path: n_agreeing_paths += 1

        # only for the output, use original names
        path1_orig_names = [orig_names1[int(s)] for s in path1]
        path2_orig_names = [orig_names2[int(s)] for s in path2]
        logg.msg('      path1 = {},\n'
               'path_mapped = {},\n'
               '      path2 = {},\n'
               '-> n_agreeing_steps = {} / n_steps = {}.'
               .format(path1_orig_names,
                       [list(p) for p in path_mapped_orig_names],
                       path2_orig_names,
                       n_agreeing_steps_path, n_steps_path), v=5, no_indent=True)
    Result = namedtuple('paga_compare_paths_result',
                        ['frac_steps', 'n_steps', 'frac_paths', 'n_paths'])
    return Result(frac_steps=n_agreeing_steps/n_steps if n_steps > 0 else np.nan,
                  n_steps=n_steps if n_steps > 0 else np.nan,
                  frac_paths=n_agreeing_paths/n_paths if n_steps > 0 else np.nan,
                  n_paths=n_paths if n_steps > 0 else np.nan)

import numpy,import scipy,from scipy,import networkx


import glob, os, sys
import numpy as np
import scipy.integrate as integrate
import scipy.ndimage
import matplotlib.pylab as plt
from scipy.interpolate import griddata

class Dist:
    """
    This class load the data given a filename
    and gives the possibility to generate a plot with the uploaded data
    """
    def __init__(self, filename, is_avoid_zeros=True):
        # It is better to make general x,y arrays
        self.x, self.y = np.loadtxt(filename, comments="#", unpack=True)
        if is_avoid_zeros:
            s_len = len(self.x)
            self.x, self.y = self.avoid_zeros()
            print("%i lines deleted" % (s_len - len(self.x)))

    def avoid_zeros(self):
        is_not_zero = self.y != 0
        x = self.x[is_not_zero]
        y = self.y[is_not_zero]
        return x, y

    def plot(self, loglog=True):
        fig = plt.figure()
        ax = fig.add_subplot(111)
        if loglog:
            ax.loglog(self.x, self.y, 'o')
        else:
            ax.plot(self.x, self.y, 'o')

class DistCollector:
    """
    this is class to collects all the filenames
    in a dictionary of instances and allows to produce different plots
    selecting four main parameters.
    First the type of structure. Subsequently for the given nanostructure it is possible to select the quantity or type of distribution to plot,
    and restrict the analysis to a given diameter or thickness.
    Parameters:
    ===========
    mainDir: str
        Directory containing the files
    maxLex: int, opt
        max lenght of string describing the file types to consider
        such in dot_Hy_300_t20_bis.dat
    structure: string, opt
        the structure used in the experiment (dot,pillar,ring)
    """
    def __init__(self, mainDir, maxLen=4, structure="dot"):
        self._mainDir = mainDir
        # Check if the dist_type exists
        # How can we do it?
        self.plotTypes=dict()
        self.plotTypes= {'Hy': 'Hysteresis Loop','Hyst': 'Hysteresis Loop', 'Ener': 'Energies'}
        self.dis_types = self._get_distribution_types(maxLen)
        self.diameters = self._get_diameters(maxLen)
        self.thicknesses= self._get_thicknesses(maxLen)
        print(self.dis_types)
        print(self.thicknesses)
        self.distrs = dict()
        for dis_type in self.dis_types:
            self.distrs[dis_type] = dict()
            for diameter in self.diameters:
               pattern = "%s_%s_%s_t*_*.dat" % (structure, dis_type,diameter)
               pattern = os.path.join(self._mainDir, pattern)
               filenames = sorted(glob.glob(pattern))
               print('\n'.join(filenames))
               self.distrs[dis_type][diameter] = dict()
               for filename in filenames:
                    fname = os.path.join(self._mainDir, filename)
                    thick = self._get_thickness(fname)
                    self.distrs[dis_type][diameter][thick] = Dist(fname)

    def plot(self, dis_type,diameter="*",thickness="*", loglog=False):
        """
        plot all the distributions
        just giving the type ('S', 'T', 'E', etc)
        """
        if dis_type not in self.dis_types:
            print("Type %s does not exist, please check it" % dis_type)
            return
        if diameter != "*" and (diameter not in self.diameters):
            print("Diameter %s does not exist, please check it" % diameter)
            return
        if thickness != "*" and (thickness not in self.thicknesses):
            print("thickness %s does not exist, please check it" % thickness)
            return
        fig = plt.figure()
        ax = fig.add_subplot(111)
        ax.set_title('%s' % self.plotTypes[dis_type])
        if diameter != "*":
            if thickness != "*":
                ax.set_title('%s , diameter = %s nm, thickness = %s nm' % (self.plotTypes[dis_type],diameter,thickness))
            else:
                ax.set_title('%s , diameter = %s nm' % (self.plotTypes[dis_type],diameter))

        if (thickness != "*" and diameter == "*"):
            ax.set_title('%s , thickness = %s nm' % (self.plotTypes[dis_type],thickness))

        for diam in sorted(self.distrs[dis_type]):
            if (diam==diameter and diameter!="*") or diameter=="*":
                for thick in sorted(self.distrs[dis_type][diam]):
                    if (thick==thickness and thickness!="*") or thickness=="*":
                        d = self.distrs[dis_type][diam][thick]
                        if thickness=="*" and diameter=="*":
                            lb = " d= %s nm, t= %s nm" % (diam,thick)
                        else:
                            if diameter=="*":
                                lb = "d= %s nm" % (diam)
                            else:
                                lb = "t= %s nm" % (thick)
                        ax.plot(d.x, d.y, label=lb)

        ax.legend(numpoints=1,loc=4)
        ax.grid(True)
        # Here we need to explicity say to show the plot
        plt.show()

    def _get_distribution_types(self, maxLen=4):
        """
        find the type of distributions in the given directory, reading the 2nd position in the files name
        and returns all the availble diameters as in dot_Hyst_100_00_s20.dat
        Parameters:
        ===========
            maxLen: int, opt
            max length of the string to be searched
        """
        filenames = glob.glob(os.path.join(self._mainDir, "*.dat"))
        filenames = [os.path.splitext(filename)[0] for filename in filenames]
        filenames = [os.path.split(filename)[1] for filename in filenames]
        filenames = [filename.split("_", 2)[1] for filename in filenames]
        dis_types = [filename for filename in filenames if len(filename) <= maxLen]
        dis_types = set(dis_types)
        return dis_types

    def _get_diameters(self, maxLen=3):
        """
        find the diameter or maxdimension of the object (denoted by dimension in nanometers)
        in the given directory, reading the 3rd position in the file names and returns all the availble diameters
        as in dot_Hyst_100_00_s20.dat
        Parameters:
        ===========
            maxLen: int, opt
            max length of the string to be searched
        """
        filenames = glob.glob(os.path.join(self._mainDir, "*.dat"))
        filenames = [os.path.splitext(filename)[0] for filename in filenames]
        filenames = [os.path.split(filename)[1] for filename in filenames]
        print('\n'.join(filenames))
        filenames = [filename.split("_",3)[2] for filename in filenames]
        diameters = [filename for filename in filenames if len(filename) <= maxLen]
        diameters = set(diameters)
        return diameters
    def _get_thicknesses(self, maxLen=4):
        """
        find the diameter or maxdimension of the objecr (denoted by dimension in nanometers)
        looking at the last character of the filenames
        as in dot_Hyst_100_00_s20.dat
        Parameters:
        ===========
            maxLen: int, opt
            max length of the string to be searched
        """
        filenames = glob.glob(os.path.join(self._mainDir, "*.dat"))
        filenames = [os.path.splitext(filename)[0] for filename in filenames]
        filenames = [os.path.split(filename)[1] for filename in filenames]
        filenames = [filename.split("_t",1)[1] for filename in filenames]
        filenames = [filename.split("_",1)[0] for filename in filenames]
        for filename in filenames:
           if "v" in filename:
              filename = ['%s.%s' %(filename.split("v",1)[0],filename.split("v",1)[1])]
        thicknesses = [filename for filename in filenames if len(filename) <= maxLen]
        thicknesses = set(thicknesses)
        return thicknesses
    def _get_diameter(self,filename,maxLen=3):
        """
        find the diameter or maxdimension of the objecr (denoted by dimension in nanometers)
        looking at the last character of the filenames
        as in dot_Hyst_100_00_s20.dat
        Parameters:
        ===========
            filename
            maxLen: int, opt
            max length of the string to be searched
        """
        filename = os.path.splitext(filename)[0]
        filename = os.path.split(filename)[1]
        filename = filename.split("_",3)[2]
        diameter = filename
        return diameter

    def _get_thickness(self,filename, maxLen=3):
        """
        find the diameter or maxdimension of the objecr (denoted by dimension in nanometers)
        looking at the last character of the filenames
        as in dot_Hyst_100_00_s20.dat
        Parameters:
        ===========
            filename
            maxLen: int, opt
            max length of the string to be searched
        """
        filename = os.path.splitext(filename)[0]
        filename = os.path.split(filename)[1]
        filename = filename.split("_t")[-1]
        filename = filename.split("_")[0]
        if "v" in filename:
              part1=filename.split("v",1)[0]
              part2=filename.split("v",1)[1]
              filename = ''.join((filename.split("v",1)[0],'.',filename.split("v",1)[1]))#['%s.%s' %(filename.split("v",1)[0],filename.split("v",1)[1])]
        print(filename)
        thickness = filename
        return thickness

class integral:
    """
    This class load the data given a filename and integrates the curve
    """
    def __init__(self, filename, mainDir, is_avoid_zeros=True):
        # It is better to make general x,y arrays
        self._mainDir = mainDir
        fname = os.path.join(self._mainDir, filename)
        self.x, self.y = np.loadtxt(fname , comments="#", unpack=True)
        if is_avoid_zeros:
            s_len = len(self.x)
            self.x, self.y = self.avoid_zeros()
            print("%i lines deleted" % (s_len - len(self.x)))
        self.fullHyst=self.x[-1]-self.x[0]
        value=self.integra()
        self.energy=2*4*np.pi*1.e-7*value

    def avoid_zeros(self):
        is_not_zero = self.y != 0
        x = self.x[is_not_zero]
        y = self.y[is_not_zero]
        return x, y

    def integra(self):

        if self.fullHyst==0:
           middle=int(np.round(self.x.size/4))
           top=int(np.round(self.x.size/2))
           self._branchup=integrate.simps(self.y[0:middle],self.x[0:middle])
           self._branchdown=integrate.simps(self.y[middle:top],self.x[middle:top])
           self.result=-self._branchdown-self._branchup
        else:
           middle=int(np.round(self.x.size/2))
           self._branchdown=integrate.simps(self.y,self.x)
           self._branchup=integrate.simps(-np.flipud(self.y),-np.flipud(self.x))
           self.result=(-self._branchdown+self._branchup)/2
        return self.result

class mapsHystEnergy:
    """
    this is class to collect all the filenames
    in a dictionary of instances and plot the desired map givien the parameters
    Parameters:
    ===========
    mainDir: str
        Directory containing the files
    maxLex: int, opt
        max lenght of string describing the file types to consider
        such in dot_Hyst_500_00_s30.dat
    structure: string, opt
        the structure used in the experiment (dot,pillar,thorus)
    """
    def __init__(self, mainDir,structure="dot"):
        self.dist=DistCollector(mainDir)
    def integra(self,x,y):
        self.fullHyst=x[-1]-x[0]
        if self.fullHyst==0:
           middle=int(np.round(x.size/4))
           top=int(np.round(x.size/2))
           self._branchup=integrate.simps(y[0:middle],x[0:middle])
           self._branchdown=integrate.simps(y[middle:top],x[middle:top])
           self.result=-self._branchdown-self._branchup
        else:
           middle=int(np.round(x.size/2))
           self._branchdown=integrate.simps(y,x)
           self._branchup=integrate.simps(-np.flipud(y),-np.flipud(x))
           self.result=(-self._branchdown+self._branchup)/2
        return self.result

    def setData(self,outName="mapdata",structure="dot",dis_type="Hy"):
        points=np.array([])
        values = np.array([])
        mappatxt = np.array([])
        for diam in sorted(self.dist.distrs[dis_type]):
            for thick in sorted(self.dist.distrs[dis_type][diam]):
                points=np.append(points,(int(diam),float(thick)))
                value=self.integra(self.dist.distrs[dis_type][diam][thick].x,self.dist.distrs[dis_type][diam][thick].y)
                self.energy=2*4*np.pi*1.e-7*value
                values=np.append(values,self.energy)
                print(self.energy,diam,thick)
                mappatxt=np.append(mappatxt,(int(diam),float(thick),self.energy))
        mappatxt=np.reshape(mappatxt,(-1,3))
        points=np.reshape(points,(-1,2))
        values=np.reshape(values,(-1,1))
        np.savetxt(outName,mappatxt[:],"%4d  %4.2f %12.8e")
        return (points,values)

    def plotMap(self):
        points,values=self.setData()
        #xmin=np.min(points[:,0])
        #xmax=np.max(points[:,0])
        #ymin=np.min(points[:,1])
        #ymax=np.max(points[:,1])

        grid_x, grid_y = np.mgrid[np.min(points[:,0]):np.max(points[:,0]):100j, np.min(points[:,1]):np.max(points[:,1]):100j]
        #print(grid_x, grid_y)

        grid_z0 = griddata((points[:,0],points[:,1]), values, (grid_x, grid_y), method='linear',fill_value=0)
        print(np.min(values)/2)
        #
        origin = 'lower'
        CS=plt.contourf(grid_x,grid_y,grid_z0[:,:,0],100)
        plt.rcParams['contour.negative_linestyle'] = 'solid'
        CS2 = plt.contour(CS, levels=CS.levels[::10],
                          colors='k',
                          origin=origin,
                          hold='on')
        plt.clabel(CS2, fontsize=9, inline=1)
        #plt.axis([200, 650, 10, 30])
        plt.colorbar(CS)
        plt.show()

if __name__ == "__main__":
    mainDir = "W:\\Micro\\Riccardo\\Dot\\Single\\Results\\Hyst_new\\Bis"
    #mainDir = "D:\\git\\Python-In-The-Lab_Project\\Python-In-The-Lab_Project\\Hyst"

    dcoll = DistCollector(mainDir)
    dcoll.plot("Hy", diameter="300")
    #integ=integral("dot_Hy_650_t30_bis.dat",mainDir)
    #print(integ.result)
    #integ=integral("dot_Hy_650_t25_bis.dat",mainDir)
    #print(integ.result)
    #integ=integral("dot_Hy_500_t30_bis.dat",mainDir)
    #print(integ.result)
    #integ=integral("dot_Hy_500_t25_bis.dat",mainDir)
    #print(integ.result)
    maps=mapsHystEnergy(mainDir)
    maps.plotMap()
    #point,values=maps.setData()

    #print(dcoll.distrs["Hyst"]["300"]["30"].x)
    #print(integ.energy)


""" Simple Solar abundance calculations.
"""
from __future__ import (print_function, absolute_import, division,
                        unicode_literals)

# Python 2 & 3 compatibility
try:
    basestring
except NameError:
    basestring = str

import numpy as np
import numbers
import imp

from astropy.io import ascii
from astropy.utils.misc import isiterable

#from xastropy.xutils import xdebug as xdb
l_path = imp.find_module('linetools')[1]

#
class SolarAbund(object):
    """Class to handle simple Solar Abundance calculations

    Parameters
    ----------
    ref: str, optional
       'Asplund2009' :: Asplund et al. 2009, ARA&A, 47, 481 meteoritic
       table (several photospheric though)
    """
    # Init
    def __init__(self, ref='Asplund2009', verbose=False):

        # Error catching
        if not isinstance(ref, basestring):
            raise TypeError('SolarAbund__init__: Wrong ref type for '
                            'SolarAbund input')
        self.ref = ref

        # Load Data
        print('Loading abundances from {:s}'.format(self.ref))
        self.load_data()
        print('Abundances are relative by number on a '
              'logarithmic scale with H=12')

    def load_data(self):
        """Grab the Solar Abundance data (in linetools/abund)
        """
        # Data file
        if self.ref == 'Asplund2009':
            dat_file = l_path + '/data/abund/solar_Asplund2009.dat'
            # Read table
            names = ('Elm', 'Abund', 'Z')
            table = ascii.read(dat_file, format='no_header', names=names)
            #
        else:
            raise ValueError('Unrecognized reference for SolarAbund: {:s}'.format(self.ref))
        # Save
        self._data = table

    def get_ratio(self, rtio):
        """ Return abundance ratio

        Parameters
        ----------
        rtio : str
          Element ratio (e.g. 'Si/Fe')
        """
        # Elements
        elm1, elm2 = rtio.split('/')
        # Abundances
        ab1 = self[elm1]
        ab2 = self[elm2]
        # ratio
        return ab1 - ab2

    def __getitem__(self, k):
        """ Return abundance given an element

        Parameters
        ----------
        k : int or str or list/tuple
          * int -- Atomic number (6)
          * str -- Element name (e.g. 'C')

        Returns
        -------
        Abund : float
        """
        # Iterate?
        if isiterable(k) and not isinstance(k, basestring):
            out_abnd = []
            for ik in k:
                out_abnd.append(self[ik])
            out_abnd = np.array(out_abnd)
            return out_abnd

        if isinstance(k, numbers.Integral): # Atomic number
            mt = np.where(self._data['Z'] == k)[0]
            if len(mt) != 1:
                raise ValueError('Atomic Number not in Table: {:d}'.format(k))
        elif isinstance(k, basestring): # Name
            mt = np.where(self._data['Elm'] == k)[0]
            if len(mt) != 1:
                raise ValueError('Element not in Table: {:s}'.format(k))
        else:
            raise IndexError('Not prepared for this type of input', k)

        # Return
        return self._data['Abund'][mt][0]

    # Printing
    def __repr__(self):
        # Generate sets string
        return '<SolarAbund: {:s}>'.format(self.ref)


#!/usr/bin/env python
"""
Populates 2D or 3D arrays with data from an Insight3 analysis file.

Hazen 7/09

Modified (incompletely) to support lazy loading, which helps
when dealing with super huge insight3 files.

Hazen 11/11
"""

import numpy
import os
import sys

import storm_analysis.sa_library.datareader as datareader
import storm_analysis.sa_library.grid_c as grid_c
import storm_analysis.sa_library.i3dtype as i3dtype
import storm_analysis.sa_library.readinsight3 as readinsight3
import storm_analysis.sa_library.regfilereader as regfilereader

def getFilmSize(filename, i3_data):
    """
    Determine the (analyzed) film size.
    """

    # First try to load meta data.
    metadata = readinsight3.loadI3Metadata(filename, verbose = False)
    if metadata is not None:
        movie_data = metadata.find("movie")
        movie_x = int(movie_data.find("movie_x").text)
        movie_y = int(movie_data.find("movie_y").text)
        movie_l = int(movie_data.find("movie_l").text)

        # Check if analysis stopped before the end of the movie.
        settings = metadata.find("settings")
        max_frame = int(settings.find("max_frame").text)
        if (max_frame > 0):
            movie_l = max_frame
        return [movie_x, movie_y, movie_l]

    # Next try and load the corresponding movie file.
    names = [filename[:-9], filename[:-10]]
    extensions = [".dax", ".spe", ".tif"]
    for name in names:
        for ext in extensions:
            if os.path.exists(name + ext):
                movie_file = datareader.inferReader(name + ext)
                return movie_file.filmSize()

    # Finally, just guess / hope. Assume that
    # the image size is a power of 2..
    film_l = int(numpy.max(i3_data['fr']))+1

    max_x = numpy.max(i3_data['x'])
    x_size = 2
    while(x_size < max_x):
        x_size = x_size * 2

    max_y = numpy.max(i3_data['y'])
    y_size = 2
    while(y_size < max_y):
        y_size = y_size * 2

    print("Could not find movie file for", filename, "assuming", x_size, "x", y_size, "by", str(film_l))
    return [x_size, y_size, film_l]

class I3GGeneric(object):
    """
    Generic Insight3 grid class.
    """
    def __init__(self, filename, scale = 4, verbose = True):

        # Setup names.
        self.filename = filename
        self.fullname = filename
        self.dirname = os.path.dirname(filename)
        if (len(self.dirname) > 0):
            self.dirname = self.dirname + "/"
            self.filename = os.path.splitext(os.path.basename(filename))[0]

        # Other class variables
        self.scale = int(scale)
        self.z_range = 1000.0

class I3GData(I3GGeneric):
    """
    The I3 grid class.

    This class will attempt to load the entire localization list
    into memory. This is fine for smaller data sets but can
    be problematic for large data sets.
    """
    def __init__(self, filename, scale = 4, verbose = True):
        I3GGeneric.__init__(self,
                            filename,
                            scale = scale,
                            verbose = verbose)

        self.i3data = readinsight3.loadI3GoodOnly(filename, verbose = verbose)
        self.i3data['fr'] -= 1

        # Determine film size.
        [image_x, image_y, self.film_l] = getFilmSize(filename, self.i3data)
        self.im_size = [image_x, image_y]

        # Determine what channels the image has.
        self.channels = []
        for i in range(10):
            mask = (self.i3data['c'] == i)
            if mask.sum() > 0:
                self.channels.append(i)

    # Utility
    def applyXYDriftCorrection(self, dx, dy):
        if isinstance(dx, numpy.ndarray):
            f = self.i3data['fr']
            i = numpy.arange(f.size, dtype=int)
            dx = dx.astype(numpy.float32)
            dy = dy.astype(numpy.float32)
            self.i3data['xc'][i] = self.i3data['x'] + dx[f]
            self.i3data['yc'][i] = self.i3data['y'] + dy[f]
        else:
            self.i3data['xc'] = self.i3data['x'] + dx
            self.i3data['yc'] = self.i3data['y'] + dy

    def applyZDriftCorrection(self, dz):
        if isinstance(dz, numpy.ndarray):
            f = self.i3data['fr']
            i = numpy.arange(f.size, dtype=int)
            self.i3data['zc'][i] = self.i3data['z'] + dz[f]
        else:
            self.i3data['zc'] = self.i3data['z'] + dz

    def get2D(self, zmin = -1000.0, zmax = 1000.0):
        return self.i3To2DGridAllChannelsMerged(verbose = 0, zmin = zmin, zmax = zmax)

    def get3D(self, zbins, zmin = -1000.0, zmax = 1000.0):
        return self.i3To3DGridAllChannelsMerged(zbins, zmin = zmin, zmax = zmax, verbose = 0)

    def getData(self):
        return self.i3data

    def getDirname(self):
        return self.dirname

    def getFilename(self):
        return self.filename

    def getFilmLength(self):
        return self.film_l

    def getFullname(self):
        return self.fullname

    def getImageSize(self):
        return self.im_size

    def getNumberMolecules(self):
        return self.i3data['x'].size

    def getScale(self):
        return self.scale

    def getXY(self):
        return [self.i3data['xc'],
                self.i3data['yc']]

    def getXYZ(self):
        return [self.i3data['xc'],
                self.i3data['yc'],
                self.i3data['zc']]

    def getXYZCat(self):
        return [self.i3data['xc'],
                self.i3data['yc'],
                self.i3data['zc'],
                self.i3data['c']]

    def getXYZICat(self):
        return [self.i3data['xc'],
                self.i3data['yc'],
                self.i3data['zc'],
                self.i3data['i'],
                self.i3data['c']]

    def getZRange(self):
        return self.z_range

    def offsetX(self, dx):
        self.i3data['xc'] += dx

    def offsetY(self, dy):
        self.i3data['yc'] += dy

    def offsetZ(self, dz):
        self.i3data['zc'] += dz

    def setScale(self, scale):
        self.scale = scale

    # Gridding data
    def i3To2DGrid(self, fmin = 0, fmax = 500000, zmin = -1000.0, zmax = 1000.0, uncorrected = False, matrix = False, translate = False, verbose = True):

        if uncorrected:
            [x, y, z] = [self.i3data['x'],
                         self.i3data['y'],
                         self.i3data['z']]
        else:
            [x, y, z] = [self.i3data['xc'],
                         self.i3data['yc'],
                         self.i3data['zc']]

        cat = self.i3data['c']
        f = self.i3data['fr']

        [image_x, image_y] = self.im_size
        scale = int(self.scale)

        if isinstance(matrix, numpy.ndarray):
            if translate:
                [x, y] = regfilereader.applyTransform(matrix, x, y)
            else:
                [x, y] = regfilereader.applyTransformNoTranslation(matrix, x, y)

        max_max = 0.0
        max_counts = []
        image_data = []
        for channel in self.channels:
            mask = (cat == channel) & (f >= fmin) & (f < fmax) & (z > zmin) & (z < zmax)
            #print numpy.sum(mask)
            i_x = numpy.floor(x[mask] * scale).astype(int)
            i_y = numpy.floor(y[mask] * scale).astype(int)
            if (i_x.shape[0] > 0):
                channel_data = grid_c.grid2D(i_x,i_y,(image_x*scale,image_y*scale))
            else:
                channel_data = numpy.zeros((image_x*scale, image_y*scale))
            image_data.append(channel_data.astype(numpy.float32))
            max_count = numpy.max(channel_data)
            if max_count > max_max:
                max_max = max_count
            max_counts.append(max_count)

        return [image_data, self.channels, max_counts, max_max]

    def i3To2DGridAllChannelsMerged(self, fmin = 0, fmax = 500000, zmin = -1000.0, zmax = 1000.0, uncorrected = False, verbose = True):
        [image_data, channels, max_counts, max_max] = self.i3To2DGrid(fmin = fmin,
                                                                      zmin = zmin,
                                                                      zmax = zmax,
                                                                      fmax = fmax,
                                                                      uncorrected = uncorrected,
                                                                      verbose = verbose)
        if (len(image_data) > 0):
            merged_image = image_data[0]
            for i in range(len(image_data)-1):
                merged_image += image_data[i+1]
            return merged_image

    def i3To3DGrid(self, z_bins, fmin = 0, fmax = 500000, zmin = -1000.0, zmax = 1000.0, uncorrected = False, verbose = True):

        if uncorrected:
            [x, y, z] = [self.i3data['x'],
                         self.i3data['y'],
                         self.i3data['z']]
        else:
            [x, y, z] = [self.i3data['xc'],
                         self.i3data['yc'],
                         self.i3data['zc']]

        cat = self.i3data['c']
        f = self.i3data['fr']

        [image_x, image_y] = self.im_size
        xy_scale = int(self.scale)
        z_bins = int(z_bins)
        z_range = zmax - zmin

        max_max = 0.0
        max_counts = []
        image_data = []
        for channel in self.channels:
            mask = (cat == channel) & (f >= fmin) & (f < fmax) & (z > zmin) & (z < zmax)
            i_x = numpy.floor(x[mask] * xy_scale).astype(int)
            i_y = numpy.floor(y[mask] * xy_scale).astype(int)
            i_z = numpy.floor((z[mask] - zmin) * float(z_bins)/z_range).astype(int)
            if (i_x.shape[0] > 0):
                channel_data = grid_c.grid3D(i_x, i_y, i_z, (image_x*xy_scale, image_y*xy_scale, z_bins))
            else:
                channel_data = numpy.zeros((image_x*xy_scale, image_y*xy_scale, z_bins))
            image_data.append(channel_data.astype(numpy.float32))
            max_count = numpy.max(channel_data)
            if max_count > max_max:
                max_max = max_count
            max_counts.append(max_count)

        return [image_data, self.channels, max_counts, max_max]

    def i3To3DGridAllChannelsMerged(self, z_bins, fmin = 0, fmax = 500000, zmin = -1000.0, zmax = 1000.0, uncorrected = False, verbose = True):
        [image_data, channels, max_counts, max_max] = self.i3To3DGrid(z_bins,
                                                                      fmin = fmin,
                                                                      fmax = fmax,
                                                                      zmin = zmin,
                                                                      zmax = zmax,
                                                                      uncorrected = uncorrected,
                                                                      verbose = verbose)
        merged_image = image_data[0]
        for i in range(len(image_data)-1):
            merged_image += image_data[i+1]
        return merged_image

class I3GDataLL(I3GData):
    """
    The I3 grid lazy-load class.

    This class will only load the localizations as needed, making
    it quite a bit less memory intensive.

    FIXME: Ugh, this inherits I3GData, but initializes using I3GGeneric.
    """
    def __init__(self, filename, scale = 4, verbose = True):
        I3GGeneric.__init__(self,
                            filename,
                            scale = scale,
                            verbose = verbose)

        self.i3_in = readinsight3.I3Reader(filename)
        self.i3data = self.i3_in.nextBlock()
        self.resetFp()

        # Determine film size.
        [image_x, image_y, self.film_l] = getFilmSize(filename, self.i3data)
        self.im_size = [image_x, image_y]

        # Determine what channels the image has.
        self.channels = []
        if (self.getNumberMolecules() > 0):
            for i in range(10):
                mask = (self.i3data['c'] == i)
                if mask.sum() > 0:
                    self.channels.append(i)

    def close(self):
        self.i3_in.close()

    def dataIsGood(self):
        if(type(self.i3data)==type(numpy.array([]))):
            return True
        else:
            return False

    def getCurrentFrameRange(self):
        return [self.i3data['fr'][0], self.i3data['fr'][-1]]

    def getNumberMolecules(self):
        return self.i3_in.getNumberMolecules()

    def loadDataInFrames(self, fmin = 0, fmax = 500000):
        self.i3data = self.i3_in.getMoleculesInFrameRange(fmin+1, fmax+1)
        self.i3data['fr'] -= 1

    def i3ToXDGridAllChannelsMergedLL(self, grid_fn, verbose):
        self.i3_in.resetFp()
        self.i3data = self.i3_in.nextBlock()
        merged = grid_fn()

        self.i3data = self.i3_in.nextBlock()
        while(isinstance(self.i3data, numpy.ndarray)):
            if verbose:
                sys.stdout.write(".")
                sys.stdout.flush()
            merged += grid_fn()
            self.i3data = self.i3_in.nextBlock()
        if verbose:
            print("")

        return merged

    def i3To2DGridAllChannelsMergedLL(self, zmin = -1000.0, zmax = 1000.0, uncorrected = False, verbose = True):
        def gridFn():
            return self.i3To2DGridAllChannelsMerged(zmin = zmin,
                                                    zmax = zmax,
                                                    uncorrected = uncorrected,
                                                    verbose = verbose)
        return self.i3ToXDGridAllChannelsMergedLL(gridFn, verbose)

    def i3To2DGridSpecificChannelLL(self, channel, zmin = -1000.0, zmax = 1000.0, uncorrected = False, verbose = True):
        def gridFn():
            mask = (self.i3data['c'] == channel)
            self.i3data = i3dtype.maskData(self.i3data, mask)
            return self.i3To2DGridAllChannelsMerged(zmin = zmin,
                                                    zmax = zmax,
                                                    uncorrected = uncorrected,
                                                    verbose = verbose)
        return self.i3ToXDGridAllChannelsMergedLL(gridFn, verbose)

    def i3To3DGridAllChannelsMergedLL(self, z_bins, zmin = -1000.0, zmax = 1000.0, uncorrected = False, verbose = True):
        def gridFn():
            return self.i3To3DGridAllChannelsMerged(z_bins,
                                                    zmin = zmin,
                                                    zmax = zmax,
                                                    uncorrected = uncorrected,
                                                    verbose = verbose)
        return self.i3ToXDGridAllChannelsMergedLL(gridFn, verbose)

    def i3To3DGridSpecificChannelLL(self, channel, z_bins, zmin = -1000.0, zmax = 1000.0, uncorrected = False, verbose = True):
        def gridFn():
            mask = (self.i3data['c'] == channel)
            self.i3data = i3dtype.maskData(self.i3data, mask)
            return self.i3To3DGridAllChannelsMerged(z_bins,
                                                    zmin = zmin,
                                                    zmax = zmax,
                                                    uncorrected = uncorrected,
                                                    verbose = verbose)
        return self.i3ToXDGridAllChannelsMergedLL(gridFn, verbose)

    def nextBlock(self, block_size = 500000):
        self.i3data = self.i3_in.nextBlock(block_size = block_size)
        if(type(self.i3data)==type(numpy.array([]))):
            self.i3data['fr'] -= 1
            return True
        else:
            return False

    def resetFp(self):
        self.i3_in.resetFp()

#
# The MIT License
#
# Copyright (c) 2012 Zhuang Lab, Harvard University
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#


import os
from typing import TYPE_CHECKING

import numpy as np
from typing_extensions import Literal

from qubekit.molecules import CloudPen, Dipole, Quadrupole
from qubekit.utils.constants import ANGS_TO_NM

if TYPE_CHECKING:
    from qubekit.molecules import Ligand

class ExtractChargeData:
    """
    Choose between extracting (the more extensive) data from Chargemol, or from ONETEP.
    Store all info back into the molecule object
    """

    @classmethod
    def extract_charge_data_chargemol(
        cls,
        molecule: "Ligand",
        dir_path: str,
        ddec_version: Literal[3, 6],
    ) -> "Ligand":
        """
        From Chargemol output files, extract the necessary parameters for calculation of L-J.
        Args:
            molecule: Usual molecule object.
            ddec_version: Which ddec version was used to run chargemol? 3 or 6.
            dir_path: Directory containing the ddec output files.
        Return:
            Updated molecule object.
        """

        if ddec_version == 6:
            net_charge_file_path = os.path.join(
                dir_path, "DDEC6_even_tempered_net_atomic_charges.xyz"
            )

        elif ddec_version == 3:
            net_charge_file_path = os.path.join(
                dir_path, "DDEC3_net_atomic_charges.xyz"
            )

        else:
            raise ValueError("Unsupported DDEC version; please use version 3 or 6.")

        try:
            with open(net_charge_file_path, "r+") as charge_file:
                lines = charge_file.readlines()
        except FileNotFoundError:
            raise FileNotFoundError(
                "Cannot find the DDEC output file.\nThis could be indicative of several issues.\n"
                "Please check Chargemol is installed in the correct location and that the configs "
                "point to that location."
            )

        # Find number of atoms
        atom_total = int(lines[0])

        # Find data markers:
        ddec_start_pos, cloud_pen_pos = 0, 0
        for pos, row in enumerate(lines):
            if "The following XYZ" in row:
                ddec_start_pos = pos + 2

            # [sic]
            elif "The sperically averaged" in row:
                cloud_pen_pos = pos + 2

            if ddec_start_pos and cloud_pen_pos:
                break
        else:
            raise EOFError(
                f"Cannot find charge or cloud penetration data in {net_charge_file_path}."
            )

        chargemol_xyz = []
        for line in lines[ddec_start_pos : ddec_start_pos + atom_total]:
            # _'s are the atomic symbol, xyz coords, then the quadrupole moment tensor eigenvalues.
            (
                atom_count,
                _,
                x,
                y,
                z,
                charge,
                x_dipole,
                y_dipole,
                z_dipole,
                _,
                q_xy,
                q_xz,
                q_yz,
                q_x2_y2,
                q_3z2_r2,
                *_,
            ) = line.split()

            # File counts from 1 not 0; thereby requiring -1 to get the index.
            atom_index = int(atom_count) - 1

            molecule.atoms[atom_index].aim.charge = float(charge)

            molecule.atoms[atom_index].dipole = Dipole(
                x=float(x_dipole),
                y=float(y_dipole),
                z=float(z_dipole),
            )

            molecule.atoms[atom_index].quadrupole = Quadrupole(
                q_xy=float(q_xy),
                q_xz=float(q_xz),
                q_yz=float(q_yz),
                q_xx=(float(q_x2_y2) / 2 - float(q_3z2_r2) / 6),
                q_yy=(-float(q_x2_y2) / 2 - float(q_3z2_r2) / 6),
                q_zz=float(q_3z2_r2) / 3,
            )

            # save the chargemol reorientated positions
            chargemol_xyz.append([x, y, z])

        molecule.chargemol_coords = chargemol_xyz

        for line in lines[cloud_pen_pos : cloud_pen_pos + atom_total]:
            # _'s are the xyz coords and the r_squared.
            atom_count, _, _, _, _, a, b, _ = line.split()
            atom_index = int(atom_count) - 1
            molecule.atoms[atom_index].cloud_pen = CloudPen(
                a=float(a),
                b=float(b),
            )

        r_cubed_file_name = os.path.join(dir_path, "DDEC_atomic_Rcubed_moments.xyz")

        with open(r_cubed_file_name, "r+") as vol_file:
            lines = vol_file.readlines()

        volumes = [float(line.split()[-1]) for line in lines[2 : atom_total + 2]]

        for atom_index in range(atom_total):
            molecule.atoms[atom_index].aim.volume = volumes[atom_index]

        return molecule

    @classmethod
    def extract_charge_data_onetep(cls, molecule: "Ligand", dir_path: str) -> "Ligand":
        """
        From ONETEP output files, extract the necessary parameters for calculation of L-J.
        Args:
            molecule: Usual molecule object.
            dir_path: Directory containing the ddec output files.
        Return:
            Updated molecule object.
        """

        # Second file contains the rest (charges, dipoles and volumes):
        ddec_output_file = (
            "ddec.onetep" if os.path.exists("ddec.onetep") else "iter_1/ddec.onetep"
        )

        ddec_file_path = os.path.join(dir_path, ddec_output_file)

        with open(ddec_file_path, "r") as file:
            lines = file.readlines()

        charge_pos, vol_pos = None, None
        for pos, line in enumerate(lines):

            # Charges marker in file:
            if "DDEC density" in line:
                charge_pos = pos + 7

            # Volumes marker in file:
            if "DDEC Radial" in line:
                vol_pos = pos + 4

        if any(position is None for position in [charge_pos, vol_pos]):
            raise EOFError(
                "Cannot locate charges and / or volumes in ddec.onetep file."
            )

        charges = [
            float(line.split()[-1])
            for line in lines[charge_pos : charge_pos + molecule.n_atoms]
        ]

        # Add the AIM-Valence and the AIM-Core to get V^AIM
        volumes = [
            float(line.split()[2]) + float(line.split()[3])
            for line in lines[vol_pos : vol_pos + molecule.n_atoms]
        ]

        for atom_index in range(molecule.n_atoms):
            molecule.atoms[atom_index].aim.volume = volumes[atom_index]
            molecule.atoms[atom_index].aim.charge = charges[atom_index]

        return molecule

def extract_c8_params(molecule: "Ligand", dir_path: str) -> "Ligand":
    """
    Extract the C8 dispersion coefficients from the MCLF calculation's output file.
    Args:
        molecule: Usual molecule object.
        dir_path: Directory containing the C8 xyz file from Chargemol.
    returns: Updated molecule object.
    """

    with open(os.path.join(dir_path, "MCLF_C8_dispersion_coefficients.xyz")) as c8_file:
        lines = c8_file.readlines()
        for i, line in enumerate(lines):
            if line.startswith(" The following "):
                lines = lines[i + 2 : -2]
                break
        else:
            raise EOFError("Cannot locate c8 parameters in file.")

        # c8 params IN ATOMIC UNITS
        c8_params = [float(line.split()[-1].strip()) for line in lines]
        for atom_index in range(molecule.n_atoms):
            molecule.atoms[atom_index].aim.c8 = c8_params[atom_index]

    return molecule

def extract_extra_sites_onetep(molecule: "Ligand"):
    """
    Gather the extra sites from the xyz file and insert them into the molecule object.
    * Find parent and 2 reference atoms
    * Calculate the local coords site
    Args:
        molecule: Usual molecule object.
    """

    with open("xyz_with_extra_point_charges.xyz") as xyz_sites:
        lines = xyz_sites.readlines()

    parent = 0
    site_number = 0

    for i, line in enumerate(lines[2:]):
        if line.split()[0] != "X":
            parent += 1
            # Search the following entries for sites connected to this atom
            for virtual_site in lines[i + 3 :]:
                site_data = {}
                element, *site_coords, site_charge = virtual_site.split()
                # Not a virtual site:
                if element != "X":
                    break
                else:
                    site_data["charge"] = site_charge
                    site_data["parent_index"] = parent
                    site_coords = np.array([float(coord) for coord in site_coords])

                    closest_atoms = list(molecule.to_topology().neighbors(parent))
                    if (len(closest_atoms) < 2) or (
                        len(molecule.atoms[parent].bonds) > 3
                    ):
                        for atom in list(
                            molecule.to_topology().neighbors(closest_atoms[0])
                        ):
                            if atom not in closest_atoms and atom != parent:
                                closest_atoms.append(atom)
                                break

                    # Get the xyz coordinates of the reference atoms
                    coords = molecule.coordinates
                    parent_coords = coords[parent]
                    close_a_coords = coords[closest_atoms[0]]
                    close_b_coords = coords[closest_atoms[1]]

                    site_data["closest_a_index"] = closest_atoms[0]
                    site_data["closest_b_index"] = closest_atoms[1]

                    parent_atom = molecule.atoms[parent]
                    if parent_atom.atomic_symbol == "N" and len(parent_atom.bonds) == 3:
                        close_c_coords = coords[closest_atoms[2]]
                        site_data["closest_c_index"] = closest_atoms[2]

                        x_dir = (
                            (close_a_coords + close_b_coords + close_c_coords) / 3
                        ) - parent_coords
                        x_dir /= np.linalg.norm(x_dir)

                        site_data["p2"] = 0
                        site_data["p3"] = 0

                        site_data["o_weights"] = [1.0, 0.0, 0.0, 0.0]
                        site_data["x_weights"] = [
                            -1.0,
                            0.33333333,
                            0.33333333,
                            0.33333333,
                        ]
                        site_data["y_weights"] = [1.0, -1.0, 0.0, 0.0]

                    else:
                        x_dir = close_a_coords - parent_coords
                        x_dir /= np.linalg.norm(x_dir)

                        z_dir = np.cross(
                            (close_a_coords - parent_coords),
                            (close_b_coords - parent_coords),
                        )
                        z_dir /= np.linalg.norm(z_dir)

                        y_dir = np.cross(z_dir, x_dir)

                        p2 = float(
                            np.dot((site_coords - parent_coords), y_dir.reshape(3, 1))
                            * ANGS_TO_NM
                        )
                        site_data["p2"] = round(p2, 4)
                        p3 = float(
                            np.dot((site_coords - parent_coords), z_dir.reshape(3, 1))
                            * ANGS_TO_NM
                        )
                        site_data["p3"] = round(p3, 4)

                        site_data["o_weights"] = [1.0, 0.0, 0.0]
                        site_data["x_weights"] = [-1.0, 1.0, 0.0]
                        site_data["y_weights"] = [-1.0, 0.0, 1.0]

                    p1 = float(
                        np.dot((site_coords - parent_coords), x_dir.reshape(3, 1))
                        * ANGS_TO_NM
                    )
                    site_data["p1"] = round(p1, 4)

                    molecule.extra_sites.create_site(**site_data)

                    site_number += 1


import copy
import itertools

import numpy as np
from sortedcontainers import SortedListWithKey

from .tabled_trie import Trie, make_trie

class LevenshteinSearcher:
    """
    Класс для поиска близких слов
    в соответствии с расстоянием Левенштейна

    """
    def __init__(self, alphabet, dictionary, operation_costs=None,
                 allow_spaces=False, euristics='none'):
        self.alphabet = alphabet
        self.allow_spaces = allow_spaces
        if isinstance(euristics, int):
            if euristics < 0:
                raise ValueError("Euristics should be non-negative integer or None")
            else:
                self.euristics = euristics if euristics != 0 else None
        elif euristics in ["none", "None", None]:
            self.euristics = None
        else:
            raise ValueError("Euristics should be non-negative integer or None")
        if isinstance(dictionary, Trie):
            # словарь передан уже в виде бора
            self.dictionary = dictionary
        else:
            self.dictionary = make_trie(alphabet, dictionary, make_cashed=True,
                                        precompute_symbols=self.euristics,
                                        allow_spaces=self.allow_spaces)
        self.transducer = SegmentTransducer(
            alphabet, operation_costs=operation_costs, allow_spaces=allow_spaces)
        self._precompute_euristics()
        self._define_h_function()

    def __contains__(self, word):
        return word in self.dictionary

    def search(self, word, d, allow_spaces=True, return_cost=True):
        """
        Finds all dictionary words in d-window from word
        """
        if not all((c in self.alphabet
                    or (c == " " and self.allow_spaces)) for c in word):
            return []
            # raise ValueError("{0} contains an incorrect symbol".format(word))
        return self._trie_search(
            word, d, allow_spaces=allow_spaces, return_cost=return_cost)

    def _trie_search(self, word, d, transducer=None,
                     allow_spaces=True, return_cost=True):
        """
        Находит все слова в префиксном боре, расстояние до которых
        в соответствии с заданным преобразователем не превышает d
        """
        if transducer is None:
            # разобраться с пробелами
            transducer = self.transducer.inverse()
        allow_spaces &= self.allow_spaces
        trie = self.dictionary
        #  инициализация переменных
        used_agenda_keys = set()
        agenda = SortedListWithKey(key=(lambda x:x[1]))
        h = self.h_func(word, trie.root)
        # agenda[self.agenda_key("", 0, trie.root)] = (0.0, 0.0, h)
        key, value = ("", 0, trie.root), (0.0, 0.0, h)
        agenda.add((key, value))
        answer = dict()
        k = 0
        # очередь с приоритетом с промежуточными результатами
        while len(agenda) > 0:
            key, value = agenda.pop(0)
            if key in used_agenda_keys:
                continue
            used_agenda_keys.add(key)
            low, pos, index = key
            cost, g, h = value
            # g --- текущая стоимость, h --- нижняя оценка будущей стоимости
            # cost = g + h --- нижняя оценка суммарной стоимости
            k += 1
            max_upperside_length = min(len(word) - pos, transducer.max_up_length)
            for upperside_length in range(max_upperside_length + 1):
                new_pos = pos + upperside_length
                curr_up = word[pos: new_pos]
                if curr_up not in transducer.operation_costs:
                    continue
                for curr_low, curr_cost in transducer.operation_costs[curr_up].items():
                    new_g = g + curr_cost
                    if new_g > d:  #если g > d, то h можно не вычислять
                        continue
                    if curr_low == " ":
                        if allow_spaces and trie.is_final(index):
                            new_index = trie.root
                        else:
                            new_index = Trie.NO_NODE
                    else:
                        new_index = trie.descend(index, curr_low)
                    if new_index is Trie.NO_NODE:
                        continue
                    new_low = low + curr_low
                    new_h = self.h_func(word[new_pos: ], new_index)
                    new_cost = new_g + new_h
                    if new_cost > d:
                        continue
                    new_key = (new_low, new_pos, new_index)
                    new_value = (new_cost, new_g, new_h)
                    if new_pos == len(word) and trie.is_final(new_index):
                        old_g = answer.get(new_low, None)
                        if old_g is None or new_g < old_g:
                            answer[new_low] = new_g
                    agenda.add((new_key, new_value))
        answer = sorted(answer.items(), key=(lambda x: x[1]))
        if return_cost:
            return answer
        else:
            return [elem[0] for elem in answer]

    def _precompute_euristics(self):
        """
        Предвычисляет будущие символы и стоимости операций с ними
        для h-эвристики
        """
        if self.euristics is None:
            return
        # вычисление минимальной стоимости операции,
        # приводящей к появлению ('+') или исчезновению ('-') данного символа
        removal_costs = {a : np.inf for a in self.alphabet}
        insertion_costs = {a : np.inf for a in self.alphabet}
        if self.allow_spaces:
            removal_costs[' '] = np.inf
            insertion_costs[' '] = np.inf
        for up, costs in self.transducer.operation_costs.items():
            for low, cost in costs.items():
                if up == low:
                    continue
                if up != '':
                    removal_cost = cost / len(up)
                    for a in up:
                        removal_costs[a] = min(removal_costs[a], removal_cost)
                if low != '':
                    insertion_cost = cost / len(low)
                    for a in low:
                        insertion_costs[a] = min(insertion_costs[a], insertion_cost)
        # предвычисление возможных будущих символов в узлах дерева
        # precompute_future_symbols(self.dictionary, self.euristics, self.allow_spaces)
        # предвычисление стоимостей потери символа в узлах дерева
        self._absense_costs_by_node = _precompute_absense_costs(
            self.dictionary, removal_costs, insertion_costs,
            self.euristics, self.allow_spaces)
        # массив для сохранения эвристик
        self._temporary_euristics = [dict() for i in range(len(self.dictionary))]

    def _define_h_function(self):
        if self.euristics in [None, 0]:
            self.h_func = (lambda *x: 0.0)
        else:
            self.h_func = self._euristic_h_function

    def _euristic_h_function(self, suffix, index):
        """
        Вычисление h-эвристики из работы Hulden,2009 для текущей вершины словаря

        Аргументы:
        ----------
        suffix : string
            непрочитанный суффикс входного слова
        index : int
            индекс текущего узла в словаре

        Возвращает:
        -----------
        cost : float
            оценка снизу для стоимости замены,
            приводящей к входному слову с суффиксом suffix,
            если прочитанный префикс слова без опечатки
            привёл в вершину с номером index
        """
        if self.euristics > 0:
            suffix = suffix[:self.euristics]
        # кэширование результатов
        index_temporary_euristics = self._temporary_euristics[index]
        cost = index_temporary_euristics.get(suffix, None)
        if cost is not None:
            return cost
        # извлечение нужных данных из массивов
        absense_costs = self._absense_costs_by_node[index]
        data = self.dictionary.data[index]
        costs = np.zeros(dtype=np.float64, shape=(self.euristics,))
        # costs[j] --- оценка штрафа при предпросмотре вперёд на j символов
        for i, a in enumerate(suffix):
            costs[i:] += absense_costs[a][i:]
        cost = max(costs)
        index_temporary_euristics[suffix] = cost
        return cost

    def _minimal_replacement_cost(self, first, second):
        first_symbols, second_symbols = set(), set()
        removal_cost, insertion_cost = 0, 0
        for a, b in itertools.zip_longest(first, second, fillvalue=None):
            if a is not None:
                first_symbols.add(a)
            if b is not None:
                second_symbols.add(b)
            removal_cost = max(removal_cost, len(first_symbols - second_symbols))
            insertion_cost = max(insertion_cost, len(second_symbols - first_symbols))
        return min(removal_cost, insertion_cost)

def _precompute_absense_costs(dictionary, removal_costs, insertion_costs, n,
                              allow_spaces=False):
    """
    Вычисляет минимальную стоимость появления нового символа в узлах словаря
    в соответствии со штрафами из costs

    Аргументы:
    ---------------
    dictionary : Trie
        словарь, хранящийся в виде ациклического автомата

    removal_costs : dict
        штрафы за удаление символов

    insertion_costs : dict
        штрафы за вставку символов

    n : int
        глубина ``заглядывания вперёд'' в словаре

    Возвращает
    ---------------
    answer : list of dicts, len(answer)=len(dictionary)
        answer[i][a][j] равно минимальному штрафу за появление символа a
        в j-ой позиции в вершине с номером i
    """
    answer = [dict() for node in dictionary.data]
    if n == 0:
        return answer
    curr_alphabet = copy.copy(dictionary.alphabet)
    if allow_spaces:
        curr_alphabet += [' ']
    for l, (costs_in_node, node) in enumerate(zip(answer, dictionary.data)):
        # определение минимальной стоимости удаления символов
        curr_node_removal_costs = np.empty(dtype=np.float64, shape=(n,))
        if len(node[0]) > 0:
            curr_node_removal_costs[0] = min(removal_costs[symbol] for symbol in node[0])
            for j, symbols in enumerate(node[1:], 1):
                if len(symbols) == 0:
                    curr_node_removal_costs[j:] = curr_node_removal_costs[j-1]
                    break
                curr_cost = min(removal_costs[symbol] for symbol in symbols)
                curr_node_removal_costs[j] = min(curr_node_removal_costs[j-1], curr_cost)
        else:
            curr_node_removal_costs[:] = np.inf
        # определение минимальной стоимости вставки
        for a in curr_alphabet:
            curr_symbol_costs = np.empty(dtype=np.float64, shape=(n,))
            curr_symbol_costs.fill(insertion_costs[a])
            for j, symbols in enumerate(node):
                if a in symbols:
                    curr_symbol_costs[j:] = 0.0
                    break
                curr_symbol_costs[j] = min(curr_symbol_costs[j], curr_node_removal_costs[j])
            costs_in_node[a] = curr_symbol_costs
    return answer

class SegmentTransducer:
    """
    Класс, реализующий взвешенный конечный преобразователь,
    осуществляющий замены из заданного списка операций

    Аргументы:
    ----------
    alphabet : list
        алфавит

    operation_costs : dict or None(optional, default=None)
        словарь вида {(up,low) : cost}

    allow_spaces : bool(optional, default=False)
        разрешены ли элементы трансдукции, содержащие пробел
        (используется только если явно не заданы operation costs
        и они равны значению по умолчанию)

    """
    def __init__(self, alphabet, operation_costs=None, allow_spaces=False):
        self.alphabet = alphabet
        if operation_costs is None:
            self._make_default_operation_costs(allow_spaces=allow_spaces)
        elif not isinstance(operation_costs, dict):
            raise TypeError("Operation costs must be a dictionary")
        else:
            self.operation_costs = operation_costs
        self._make_reversed_operation_costs()
        self._make_maximal_key_lengths()
        # self.maximal_value_lengths = {}
        # for up, probs in self.operation_costs.items():
            # СЛИШКОМ МНОГО ВЫЗОВОВ, НАДО КАК-ТО ЗАПОМНИТЬ
            # МАКСИМАЛЬНЫЕ ДЛИНЫ КЛЮЧЕЙ ПРИ ОБРАЩЕНИИ
            # max_low_length = max(len(low) for low in probs) if (len(probs) > 0) else -1
            # self.maximal_value_lengths[up] = self.maximal_key_length

    def get_operation_cost(self, up, low):
        """
        Возвращает стоимость элементарной трансдукции up->low
        или np.inf, если такой элементарной трансдукции нет

        Аргументы:
        ----------
        up, low : string
            элементы элементарной трансдукции

        Возвращает:
        -----------
        cost : float
            стоимость элементарной трансдукции up->low
            (np.inf, если такая трансдукция отсутствует)
        """
        up_costs = self.operation_costs.get(up, None)
        if up_costs is None:
            return np.inf
        cost = up_costs.get(low, np.inf)
        return cost

    def inverse(self):
        """
        Строит пробразователь, задающий обратное конечное преобразование
        """
        # УПРОСТИТЬ ОБРАЩЕНИЕ!!!
        inversed_transducer = SegmentTransducer(self.alphabet, operation_costs=dict())
        inversed_transducer.operation_costs = self._reversed_operation_costs
        inversed_transducer._reversed_operation_costs = self.operation_costs
        inversed_transducer.max_low_length = self.max_up_length
        inversed_transducer.max_up_length = self.max_low_length
        inversed_transducer.max_low_lengths_by_up = self.max_up_lengths_by_low
        inversed_transducer.max_up_lengths_by_low = self.max_low_lengths_by_up
        return inversed_transducer

    def distance(self, first, second, return_transduction = False):
        """
        Вычисляет трансдукцию минимальной стоимости,
        отображающую first в second

        Аргументы:
        -----------
        first : string
        second : string
            Верхний и нижний элементы трансдукции

        return_transduction : bool (optional, default=False)
            следует ли возвращать трансдукцию минимального веса
            (см. возвращаемое значение)

        Возвращает:
        -----------
        (final_cost, transductions) : tuple(float, list)
            если return_transduction=True, то возвращает
            минимальную стоимость трансдукции, переводящей first в second
            и список трансдукций с данной стоимостью

        final_cost : float
            если return_transduction=False, то возвращает
            минимальную стоимость трансдукции, переводящей first в second
        """
        if return_transduction:
            add_pred = (lambda x, y: (y == np.inf or x < y))
        else:
            add_pred = (lambda x, y: (y == np.inf or x <= y))
        clear_pred = (lambda x, y: (y < np.inf and x < y))
        update_func = lambda x, y: min(x, y)
        costs, backtraces = self._fill_levenshtein_table(first, second,
                                                        update_func, add_pred, clear_pred)
        final_cost = costs[-1][-1]
        if final_cost == np.inf:
            transductions = [None]
        elif return_transduction:
            transductions = self._backtraces_to_transductions(first, second, backtraces,
                                                              final_cost, return_cost=False)
        if return_transduction:
            return final_cost, transductions
        else:
            return final_cost

    def transduce(self, first, second, threshold):
        """
        Возвращает все трансдукции, переводящие first в second,
        чья стоимость не превышает threshold

        Возвращает:
        ----------
        result : list
            список вида [(трансдукция, стоимость)]
        """
        add_pred = (lambda x, y: x <= threshold)
        clear_pred =(lambda x, y: False)
        update_func = (lambda x, y: min(x, y))
        costs, backtraces = self._fill_levenshtein_table(first, second,
                                                        update_func, add_pred, clear_pred,
                                                        threshold=threshold)
        result = self._backtraces_to_transductions(first, second,
                                                   backtraces, threshold, return_cost=True)
        return result

    def lower_transductions(self, word, max_cost, return_cost=True):
        """
        Возвращает все трансдукции с верхним элементом word,
        чья стоимость не превышает max_cost

    `   Возвращает:
        ----------
        result : list
            список вида [(трансдукция, стоимость)], если return_cost=True
            список трансдукций, если return_cost=False
            список отсортирован в порядке возрастания стоимости трансдукции
        """
        prefixes = [[] for i in range(len(word) + 1)]
        prefixes[0].append(((), 0.0))
        for pos in range(len(prefixes)):
            # вставки
            prefixes[pos] = self._perform_insertions(prefixes[pos], max_cost)
            max_upperside_length = min(len(word) - pos, self.max_up_length)
            for upperside_length in range(1, max_upperside_length + 1):
                up = word[pos: pos + upperside_length]
                for low, low_cost in self.operation_costs.get(up, dict()).items():
                    for transduction, cost in prefixes[pos]:
                        new_cost = cost + low_cost
                        if new_cost <= max_cost:
                            new_transduction = transduction +(up, low)
                            prefixes[pos + upperside_length].append((new_transduction, new_cost))
        answer = sorted(prefixes[-1], key=(lambda x: x[0]))
        if return_cost:
            return answer
        else:
            return [elem[0] for elem in answer]

    def lower(self, word, max_cost, return_cost=True):
        transductions = self.lower_transductions(word, max_cost, return_cost=True)
        answer = dict()
        for transduction, cost in transductions:
            low = "".join(elem[1] for elem in transductions)
            curr_cost = answer.get(low, None)
            if curr_cost is None or cost < curr_cost:
                answer[low] = cost
        answer = sorted(answer.items(), key=(lambda x: x[1]))
        if return_cost:
            return answer
        else:
            return [elem[0] for elem in answer]

    def upper(self, word, max_cost, return_cost=True):
        inversed_transducer = self.inverse()
        return inversed_transducer.lower(word, max_cost, return_cost)

    def upper_transductions(self, word, max_cost, return_cost=True):
        inversed_transducer = self.inverse()
        return inversed_transducer.lower_transductions(word, max_cost, return_cost)

    def _fill_levenshtein_table(self, first, second, update_func, add_pred, clear_pred,
                               threshold=None):
        """
        Функция, динамически заполняющая таблицу costs стоимости трансдукций,
        costs[i][j] --- минимальная стоимость трансдукции,
        переводящей first[:i] в second[:j]

        Аргументы:
        ----------
        first, second : string
            Верхний и нижний элементы трансдукции
        update_func : callable, float*float -> bool
            update_func(x, y) возвращает новое значение в ячейке таблицы costs,
            если старое значение --- y, а потенциально новое значение --- x
            везде update_func = min
        add_pred : callable : float*float -> bool
            add_pred(x, y) возвращает, производится ли добавление
            нового элемента p стоимости x в ячейку backtraces[i][j]
            в зависимости от значения costs[i][j]=y и текущей стоимости x
        clear_pred : callable : float*float -> bool
            clear_pred(x, y) возвращает, производится ли очистка
            ячейки backtraces[i][j] в зависимости от значения costs[i][j]=y
            и текущей стоимости x элемента p, добавляемого в эту ячейку

        Возвращает:
        -----------
        costs : array, dtype=float, shape=(len(first)+1, len(second)+1)
            массив, в ячейке с индексами i, j которого хранится
            минимальная стоимость трансдукции, переводящей first[:i] в second[:j]
        backtraces : array, dtype=list, shape=(len(first)+1, len(second)+1)
            массив, в ячейке с индексами i, j которого хранятся
            обратные ссылки на предыдущую ячейку в оптимальной трансдукции,
            приводящей в ячейку backtraces[i][j]
        """
        m, n = len(first), len(second)
        # если threshold=None, то в качестве порога берётся удвоенная стоимость
        # трансдукции, отображающей символы на одинаковых позициях друг в друга
        if threshold is None:
            threshold = 0.0
            for a, b in zip(first, second):
                threshold += self.get_operation_cost(a, b)
            if m > n:
                for a in first[n: ]:
                    threshold += self.get_operation_cost(a, '')
            elif m < n:
                for b in second[m: ]:
                    threshold += self.get_operation_cost('', b)
            threshold *= 2
        # инициализация возвращаемых массивов
        costs = np.zeros(shape=(m + 1, n + 1), dtype=np.float64)
        costs[:] = np.inf
        backtraces = [None] * (m + 1)
        for i in range(m + 1):
            backtraces[i] = [[] for j in range(n + 1)]
        costs[0][0] = 0.0
        for i in range(m + 1):
            for i_right in range(i, min(i + self.max_up_length, m) + 1):
                up = first[i: i_right]
                max_low_length = self.max_low_lengths_by_up.get(up, -1)
                if max_low_length == -1: # no up key in transduction
                    continue
                up_costs = self.operation_costs[up]
                for j in range(n + 1):
                    if costs[i][j] > threshold:
                        continue
                    if len(backtraces[i][j]) == 0 and i + j > 0:
                        continue # не нашлось обратных ссылок
                    for j_right in range((j if i_right > i else j + 1),
                                         min(j + max_low_length, n) + 1):
                        low = second[j: j_right]
                        curr_cost = up_costs.get(low, np.inf)
                        old_cost = costs[i_right][j_right]
                        new_cost = costs[i][j] + curr_cost
                        if new_cost > threshold:
                            continue
                        if add_pred(new_cost, old_cost):
                            if clear_pred(new_cost, old_cost):
                                backtraces[i_right][j_right] = []
                            costs[i_right][j_right] = update_func(new_cost, old_cost)
                            backtraces[i_right][j_right].append((i, j))
        return costs, backtraces

    def _make_reversed_operation_costs(self):
        """
        Заполняет массив _reversed_operation_costs
        на основе имеющегося массива operation_costs
        """
        _reversed_operation_costs = dict()
        for up, costs in self.operation_costs.items():
            for low, cost in costs.items():
                if low not in _reversed_operation_costs:
                    _reversed_operation_costs[low] = dict()
                _reversed_operation_costs[low][up] = cost
        self._reversed_operation_costs = _reversed_operation_costs

    def _make_maximal_key_lengths(self):
        """
        Вычисляет максимальную длину элемента low
        в элементарной трансдукции (up, low) для каждого up
        и максимальную длину элемента up
        в элементарной трансдукции (up, low) для каждого low
        """
        self.max_up_length =\
            (max(len(up) for up in self.operation_costs)
             if len(self.operation_costs) > 0 else -1)
        self.max_low_length =\
            (max(len(low) for low in self._reversed_operation_costs)
             if len(self._reversed_operation_costs) > 0 else -1)
        self.max_low_lengths_by_up, self.max_up_lengths_by_low = dict(), dict()
        for up, costs in self.operation_costs.items():
            self.max_low_lengths_by_up[up] =\
                max(len(low) for low in costs) if len(costs) > 0 else -1
        for low, costs in self._reversed_operation_costs.items():
            self.max_up_lengths_by_low[low] =\
                max(len(up) for up in costs) if len(costs) > 0 else -1

    def _backtraces_to_transductions(self, first, second, backtraces, threshold, return_cost=False):
        """
        Восстанавливает трансдукции по таблице обратных ссылок

        Аргументы:
        ----------
        first, second : string
            верхние и нижние элементы трансдукции
        backtraces : array-like, dtype=list, shape=(len(first)+1, len(second)+1)
            таблица обратных ссылок
        threshold : float
            порог для отсева трансдукций,
            возвращаются только трансдукции стоимостью <= threshold
        return_cost : bool (optional, default=False)
            если True, то вместе с трансдукциями возвращается их стоимость

        Возвращает:
        -----------
        result : list
            список вида [(трансдукция, стоимость)], если return_cost=True
            и вида [трансдукция], если return_cost=False,
            содержащий все трансдукции, переводящие first в second,
            чья стоимость не превышает threshold
        """
        m, n = len(first), len(second)
        agenda = [None] * (m + 1)
        for i in range(m + 1):
            agenda[i] = [[] for j in range(n+1)]
        agenda[m][n] = [((), 0.0)]
        for i_right in range(m, -1, -1):
            for j_right in range(n, -1, -1):
                current_agenda = agenda[i_right][j_right]
                if len(current_agenda) == 0:
                    continue
                for (i, j) in backtraces[i_right][j_right]:
                    up, low = first[i:i_right], second[j:j_right]
                    add_cost = self.operation_costs[up][low]
                    for elem, cost in current_agenda:
                        new_cost = cost + add_cost
                        if new_cost <= threshold: # удаление трансдукций большой стоимости
                            agenda[i][j].append((((up, low),) + elem, new_cost))
        if return_cost:
            return agenda[0][0]
        else:
            return [elem[0] for elem in agenda[0][0]]

    def _perform_insertions(self, initial, max_cost):
        """
        возвращает все трансдукции стоимости <= max_cost,
        которые можно получить из элементов initial

        Аргументы:
        ----------
        initial : list of tuples
            список исходных трансдукций вида [(трансдукция, стоимость)]
        max_cost : float
            максимальная стоимость трансдукции

        Возвращает:
        -----------
        final : list of tuples
            финальный список трансдукций вида [(трансдукция, стоимость)]
        """
        queue = list(initial)
        final = initial
        while len(queue) > 0:
            transduction, cost = queue[0]
            queue = queue[1:]
            for string, string_cost in self.operation_costs[""].items():
                new_cost = cost + string_cost
                if new_cost <= max_cost:
                    new_transduction = transduction + ("", string)
                    final.append((new_transduction, new_cost))
                    queue.append((new_transduction, new_cost))
        return final

    def _make_default_operation_costs(self, allow_spaces=False):
        """
        sets 1.0 cost for every replacement, insertion, deletion and transposition
        """
        self.operation_costs = dict()
        self.operation_costs[""] = {c: 1.0 for c in list(self.alphabet) + [' ']}
        for a in self.alphabet:
            current_costs = {c: 1.0 for c in self.alphabet}
            current_costs[a] = 0.0
            current_costs[""] = 1.0
            if allow_spaces:
                current_costs[" "] = 1.0
            self.operation_costs[a] = current_costs
        # транспозиции
        for a, b in itertools.permutations(self.alphabet, 2):
            self.operation_costs[a + b] = {b + a: 1.0}
        # пробелы
        if allow_spaces:
            self.operation_costs[" "] = {c: 1.0 for c in self.alphabet}
            self.operation_costs[" "][""] = 1.0