pyM2aia documentation master file, created by
sphinx-quickstart on Wed Jul 19 12:00:24 2023.
You can adapt this file completely to your liking, but it should at least
contain the root `toctree` directive.

BaseDataSet(images: ~typing.List[~m2aia.ImageIO.ImzMLReader], buffer_type: str)

images: ~typing.List[~m2aia.ImageIO.ImzMLReader], buffer_type: str

IonImageDataset(images: ~typing.List[~m2aia.ImageIO.ImzMLReader], centroids: ~typing.List[float], tolerance: float, tolerance_type: str = 'ppm', buffer_type='memory', transforms=None)

images: ~typing.List[~m2aia.ImageIO.ImzMLReader], centroids: ~typing.List[float], tolerance: float, tolerance_type: str = 'ppm', buffer_type='memory', transforms=None

SpectrumDataset(images: ~typing.List[~m2aia.ImageIO.ImzMLReader], neighborhood_size: int = 0, transforms=None, buffer_type='memory')

images: ~typing.List[~m2aia.ImageIO.ImzMLReader], neighborhood_size: int = 0, transforms=None, buffer_type='memory'

__len__ so that len(dataset) returns the size of the dataset, that is equal to the sum of the number of spectra (N) for each image.
__getitem__ to support the indexing such that dataset[i] can be used to get i'th sample. i is pointing to indices 0,...,p-1,p,...,q-1, ... N, with p=#SpectraImage1, q=#SpectraImage2 etc...
neighborhood_size so that 2*neighborhood_size+1 is the window size

__getitem__ to support the indexing such that dataset[i] can be used to get i’th sample. i is pointing to indices 0,…,p-1,p,…,q-1, … N, with p=#SpectraImage1, q=#SpectraImage2 etc…

.. automodule:: m2aia.Generators
   :members:
   :undoc-members:
   :show-inheritance:

BatchGenerator(dataset: ~m2aia.Dataset.BaseDataSet, batch_size: int, shuffle: bool = True)

dataset: ~m2aia.Dataset.BaseDataSet, batch_size: int, shuffle: bool = True

ImzMLReader(imzML_path, baseline_correction: ~typing.Literal['TopHat', 'Median', 'None'] = 'None', baseline_correction_half_window_size: int = 50, normalization: ~typing.Literal['TIC', 'Sum', 'Mean', 'Max', 'RMS', 'Internal', 'External', 'None'] = 'None', smoothing: ~typing.Literal['SavitzkyGolay', 'Gaussian', 'None'] = 'None', smoothing_half_window_size: int = 2, intensity_transformation: ~typing.Literal['Log2', 'Log10', 'SquareRoot', 'None'] = 'None', pooling: ~typing.Literal['Mean', 'Median', 'Maximum', 'Sum'] = 'Maximum')

imzML_path, baseline_correction: ~typing.Literal['TopHat', 'Median', 'None'] = 'None', baseline_correction_half_window_size: int = 50, normalization: ~typing.Literal['TIC', 'Sum', 'Mean', 'Max', 'RMS', 'Internal', 'External', 'None'] = 'None', smoothing: ~typing.Literal['SavitzkyGolay', 'Gaussian', 'None'] = 'None', smoothing_half_window_size: int = 2, intensity_transformation: ~typing.Literal['Log2', 'Log10', 'SquareRoot', 'None'] = 'None', pooling: ~typing.Literal['Mean', 'Median', 'Maximum', 'Sum'] = 'Maximum'

baseline_correction: Literal['TopHat', 'Median', 'None'] = 'None'

normalization: Literal['TIC', 'Sum', 'Mean', 'Max', 'RMS', 'Internal', 'External', 'None'] = 'None'

Literal['TIC', 'Sum', 'Mean', 'Max', 'RMS', 'Internal', 'External', 'None']

smoothing: Literal['SavitzkyGolay', 'Gaussian', 'None'] = 'None'

intensity_transformation: Literal['Log2', 'Log10', 'SquareRoot', 'None'] = 'None'

pooling: Literal['Mean', 'Median', 'Maximum', 'Sum'] = 'Maximum'

Wrapper class for M2aia's imzML reader https://github.com/m2aia/m2aia.

Complete processing examples with focus on deep learning can be found on https://github.com/m2aia/pym2aia-examples

Complete processing examples with focus on deep learning can be found on 

import m2aia as m2

I = m2.ImzMLReader("path/to/imzMl/file.imzML")
I.SetNormalization(m2.m2NormalizationTIC)
I.SetIntensityTransformation(m2.m2IntensityTransformationSquareRoot)
I.Execute()
ys_2 = I.GetMeanSpectrum()
i_2 = I.GetArray(imz, 75)

Check if the handle was initialized properly.
To prevent this check from throwing an exception you must call Execute() once.

ReferenceError: is invalid file name and or not yet called Execute().

ImzMLReader.GetArray(center, tol, dtype=<class 'numpy.float32'>, squeeze: bool = False) -> ~numpy.ndarray

(center, tol, dtype=<class 'numpy.float32'>, squeeze: bool = False)

Get the (ion) image data as numpy array.
The pixel values are the pooled intensities (ie. pooling strategies like
the 'Mean', 'Median', 'Maximum', or 'Sum') in the interval [center-tol, center+tol] of the spectra.

the ‘Mean’, ‘Median’, ‘Maximum’, or ‘Sum’) in the interval [center-tol, center+tol] of the spectra.

tol -- tolerance for query points on the x axis around the center.

TypeError: Image pixel type is not one of [np.float32, np.float64]

ImzMLReader.GetImage(center, tol, dtype=<class 'numpy.float32'>) -> ~SimpleITK.SimpleITK.Image

Get the index image data as numpy array.
The pixel values are the spectrum ids (pixel value >= 1) and background (pixel value == 0).

Get the index image data as parameterized SimpleITK.Image.

ys -- By passing a np.array (dtype=np.float32) from external, this np.array can
be reused and do not require an extra memory allocation. Otherwise
a new array is created.

By passing a np.array (dtype=np.float32) from external, this np.array can

A list of two np.array elements [xs,ys]. xs = x-values; ys = y-values.

IndexError: if index is not in the range of valid spectra indices [0,self.number_of_spectra-1].
TypeError: if the ImzML file format is not continuous profile/centroid!

Get the mask image data as numpy array.
The binary mask indicates valid spectra (pixel value >= 1) and background (pixel value == 0).

Get the mask image data as parameterized SimpleITK.Image.
The pixel values indicate valid spectra (pixel value >= 1) and background (pixel value == 0).

Returns a dictionary of all meta data information retrieved by m2aia.

ImzMLReader.GetNormalizationImage(type) -> ~SimpleITK.SimpleITK.Image

Get a normalization image data as parameterized SimpleITK.Image.

Get the number of valid spectra in the image.
This can be used to iterate over all spectra in the image using
a for loop::

a for loop:

for i in range(GetNumberOfSpectra()):
    xs, ys = reader.GetSpectrum(i)

The origin in world coordinates of the image as numpy array of size [3] and dtype=np.float64.

GetParametersAsFormattedString() (m2aia.ImageIO.ImzMLReader method)

Transform signal processing parameters into a fomatted string representation.

numpy array of size [3] for the x,y,z image dimensions (in number of pixels).

numpy array of size [3] and dtype=np.float64 for the pixel size in x,y,z dimension (in millimeter).

ImzMLReader.GetSpectra(indices: ~typing.List[int]) -> ~numpy.ndarray

Query a set of intensities by a list of indices.
Only continuous imzML files.

Query the x-axis (xs) and y-axis (ys) values for a given spectrum id.

Get the size of the x axis for a specific spectrum of the image.
This method is helpful for processed (centroid) imzML files. For continuous
(centroid/profile) imzML files use GetXAxisDepth(self).

id -- the id of a spectrum (may be different to the ids given in the imzML).

the id of a spectrum (may be different to the ids given in the imzML).

Position in index coordinates as numpy array of size [3] and dtype=np.int32.

Get the size of the x axis. For processed imzML files,
this value is the number of bins used to represent the x-axis.

ImzMLReader.SetBaselineCorrection(strategy: ~typing.Literal['TopHat', 'Median', 'None'], half_window_size=50)

strategy: ~typing.Literal['TopHat', 'Median', 'None'], half_window_size=50

strategy -- Set the basline correction strategy using one of the m2BaselineCorrection literals.

Set the basline correction strategy using one of the m2BaselineCorrection literals.

half_window_size -- 2*half_window_size + 1 spectrum points are used for BaselineCorrection.

2*half_window_size + 1 spectrum points are used for BaselineCorrection.

ImzMLReader.SetIntensityTransformation(strategy: ~typing.Literal['Log2', 'Log10', 'SquareRoot', 'None'])

strategy: ~typing.Literal['Log2', 'Log10', 'SquareRoot', 'None']

strategy -- m2IntensityTransformation
Set the intensity transformation strategy using one of the m2IntensityTransformation literals.

m2IntensityTransformation

ImzMLReader.SetNormalization(strategy: ~typing.Literal['TIC', 'Sum', 'Mean', 'Max', 'RMS', 'Internal', 'External', 'None'])

strategy: ~typing.Literal['TIC', 'Sum', 'Mean', 'Max', 'RMS', 'Internal', 'External', 'None']

strategy: Literal['TIC', 'Sum', 'Mean', 'Max', 'RMS', 'Internal', 'External', 'None']

strategy -- m2Normalization
Set the normalization strategy using one of the m2Normalization literals.

m2Normalization

ImzMLReader.SetPooling(strategy: ~typing.Literal['Mean', 'Median', 'Maximum', 'Sum'])

strategy -- m2Pooling
Set the pooling strategy using one of the m2Pooling literals.

m2Pooling

ImzMLReader.SetSmoothing(strategy: ~typing.Literal['SavitzkyGolay', 'Gaussian', 'None'], half_window_size=2)

strategy: ~typing.Literal['SavitzkyGolay', 'Gaussian', 'None'], half_window_size=2

strategy -- Set the smoothing strategy using one of the m2Smoothing literals.

Set the smoothing strategy using one of the m2Smoothing literals.

half_window_size -- 2*half_window_size + 1 spectrum points used for smoothing.

Create a spectrum iterator/generator, yielding all valid spectra.

a for loop:

SpectrumRandomBatchIterator() (m2aia.ImageIO.ImzMLReader method)

Create a spectrum batch iterator/generator, yielding a batch of ys-values of spectra in a random order with repetitions.

for ys_batch in reader.SpectrumRandomBatchIterator(batch_size = N):
    ...

a np.array as batch of intensities with shape [batch_size, self.depth]

WriteContinuousCentroidImzML() (m2aia.ImageIO.ImzMLReader method)

Given a list of centroids, write a continuous centroid imzML to the given path.
Use 'SetTolerance' to define the range query for each centroid (ppm).

Use ‘SetTolerance’ to define the range query for each centroid (ppm).

path -- Target file path the <path>.imzML and the <path>.ibd is written to.

Target file path the <path>.imzML and the <path>.ibd is written to.

I.SetTolerance(75)
I.WriteContinuousCentroidImzML("path/to/imzMl/file.imzML", [300, 400, 500])

load_library_dependencies_recursively() (in module m2aia.Library)

load_library_dependencies_recursively(search_path: ~pathlib.Path, library_name: str, dependencies: ~typing.List)

search_path: ~pathlib.Path, library_name: str, dependencies: ~typing.List

BatchSequence(dataset: ~m2aia.Dataset.BaseDataSet, batch_size: int, shuffle: bool = True)

on_epoch_end() (m2aia.keras.KerasGenerators.BatchSequence method)

bin_peaks(image: ~m2aia.ImageIO.ImzMLReader, mask: ~numpy.array | None = None)

image: ~m2aia.ImageIO.ImzMLReader, mask: ~numpy.array | None = None

xs: sorted list of m/z values of peaks
ys: sorted list of intensities according to xs sorting
ss: sorted list of source indices according to xs sorting
grouper: binning method
tolerance: TODO

bin_assignments:
assigns each entry in sorted list to the corresponding bin, so that
xs[bin_assignments==k] returns all xs values corresponding to the k'th bin

xs[bin_assignments==k] returns all xs values corresponding to the k’th bin

pick_peaks(I: ~m2aia.ImageIO.ImzMLReader, SNR: float = 8, sort_by_xs: bool = True, noise_estimator=<function noise_mad>)

(I: ~m2aia.ImageIO.ImzMLReader, SNR: float = 8, sort_by_xs: bool = True, noise_estimator=<function noise_mad>)

This function is iterating over all spectra in an image and picks peaks.
sort_by_xs: if True all values of the return arrays are sorted according to ascending peaks_xs values.

e.g.: peak_xs=[x1 x1 x1 x2 x2 x2 x2 x3 x3 x3 x3 x3 x3]
    3 peaks found in spectrum with id 1
    4 peaks found in spectrum with id 2
    6 peaks found in spectrum with id 3

3 peaks found in spectrum with id 1
4 peaks found in spectrum with id 2
6 peaks found in spectrum with id 3

peak_ys: list of intensities corresponding to peak_xs
e.g.: peak_ys=[y1 y1 y1 y2 y2 y2 y2 y3 y3 y3 y3 y3 y3]

peak_idx = list of array indices corresponding to a spectrum array
e.g.: peak_idx=[idx1 idx1 idx1 idx2 idx2 idx2 idx2 idx3 idx3 idx3 idx3 idx3 idx3]

peaks_source = indicators for each corresponding peak_xs, from which spectrum id it was picked
e.g.: peaks_source=[1 1 1 2 2 2 2 3 3 3 3 3 3]

volcano_data(image: ~m2aia.ImageIO.ImzMLReader, mask: ~numpy.array, x)

volcano_plot_data(image: ~m2aia.ImageIO.ImzMLReader, mask: ~numpy.array, xs, pThreshold=0.05, fcThreshold=1)

image: ~m2aia.ImageIO.ImzMLReader, mask: ~numpy.array, xs, pThreshold=0.05, fcThreshold=1

import tensorflow as tf
from .. import BatchGenerator
from .. import Dataset

class BatchSequence(tf.keras.utils.Sequence):
    def __init__(self, dataset: Dataset.BaseDataSet , batch_size: int, shuffle: bool=True):
        super().__init__()
        self.gen = BatchGenerator(dataset, batch_size, shuffle)
    
    def __len__(self):
        return self.gen.__len__()

    def on_epoch_end(self):
        self.gen.on_epoch_end()
    
    def __getitem__(self, index):
        return self.gen.__getitem__(index)

from operator import index
from os import times
from time import time
from typing import *
from . import ImageIO
import hashlib
import pathlib
import numpy as np

class BaseDataSet():
    def __init__(self, images : List[ImageIO.ImzMLReader], buffer_type : str) -> None:
        self.images = images
        self.elements = []

        self.buffer_type = buffer_type
        if 'disk' in self.buffer_type: # buffer images on disk; check path exists
            self.buffer_path = pathlib.Path(self.buffer_type.split(':')[1])
            if not self.buffer_path.exists():
                raise Exception(f"The path {str(self.buffer_path)} does not exist!")
    
    def __len__(self) -> int:
        pass

        pass

    def getitems(self, indexes: List[int]):
        pass

class SpectrumDataset(BaseDataSet):
    """Dataset for accession of individual spectra. 
    
    __len__ so that len(dataset) returns the size of the dataset, that is equal to the sum of the number of spectra (N) for each image.
    __getitem__ to support the indexing such that dataset[i] can be used to get i'th sample. i is pointing to indices 0,...,p-1,p,...,q-1, ... N, with p=#SpectraImage1, q=#SpectraImage2 etc...
    neighborhood_size so that 2*neighborhood_size+1 is the window size

    Parameters
    ----------
    images : List[ImageIO.ImzMLReader]
       
    """

    def __init__(self,images : List[ImageIO.ImzMLReader], neighborhood_size: int = 0, transforms = None, buffer_type='memory')-> None:
        super().__init__(images, buffer_type)
        
        self.spectrum_depth = self.images[0].GetXAxisDepth()
        if "memory" == self.buffer_type:
            self.buffer = [(np.array([None] * self.images[k].GetNumberOfSpectra()), np.zeros((self.images[k].GetNumberOfSpectra(), self.spectrum_depth), dtype = np.float32)) for k in range(len(self.images))]
        
        self.neighborhood_size = neighborhood_size
        self.index_images = []
        self.footprint = np.ones([2*neighborhood_size+1, 2*neighborhood_size+1, 1])
        self.transforms = transforms

        self.hit_counter = [0]*len(self.images)

        for imageID, handle in enumerate(self.images):
            assert(self.spectrum_depth == handle.GetXAxisDepth())
            imageElements = []
            # get array returns a 2D image
            index_image = handle.GetArray(handle.GetXAxis()[self.spectrum_depth//2], 1, squeeze=False).astype(np.int32)
            index_image.fill(-1)

            for spectrumID in range(handle.GetNumberOfSpectra()):
                (x,y,z) = handle.GetSpectrumPosition(spectrumID)
                imageElements.append((imageID, spectrumID, (x,y,z)))
                index_image[z,y,x] = spectrumID

            self.elements.extend(imageElements)
            self.index_images.append(index_image)

        return len(self.elements)

        return self.getitems([index])

        ids_split_to_images = {}

        #sort by images and create list of image related ids         
        for index in indexes:
            imageID, spectrumID, (x,y,z) = self.elements[index]
            shape = self.images[imageID].GetShape()
            if imageID not in ids_split_to_images:
                ids_split_to_images[imageID] = []
            
            if self.neighborhood_size > 0: 
                # neighborhood do not violate border regions
                left = np.clip(x-self.neighborhood_size, 0, shape[0])
                right = np.clip(x+self.neighborhood_size+1, 0, shape[0])
                bottom = np.clip(y-self.neighborhood_size,0, shape[1])
                top = np.clip(y+self.neighborhood_size+1,0, shape[1])

                # get all spectra indices for the requested spectrum position and neighbors
                indices = self.index_images[imageID][0,bottom:top,left:right]
                indices = indices.flatten()

                # handle invalid spectra positions (indicated by -1 in index images)
                if np.any(indices < 0): 
                    choices = np.random.choice(indices[indices>=0], len(indices[indices < 0]))
                    indices[indices < 0] = choices

                # handling missing values if at border
                expected_indices = (self.neighborhood_size*2+1)**2
                if len(indices) < expected_indices:
                    missing_indices = expected_indices - len(indices)
                    indices = np.concatenate([indices,np.random.choice(indices, missing_indices)])

                indices = indices.tolist()
            else:
                indices = [spectrumID]

            ids_split_to_images[imageID].extend(indices)

        result = None
        interim = None
        BUFFER_QUERY = 0
        BUFFER_DATA = 1
        for imageID, indices in ids_split_to_images.items():
            
            # use buffering
            if self.buffer_type is not None:

                indices = np.array(indices)
                interim = np.zeros((len(indices), self.spectrum_depth))

                # mask buffer entries which are not already filled with data
                query_mask = self.buffer[imageID][0][indices] == None
                
                # load data from imzML
                if query_mask.any():
                    self.hit_counter[imageID] = self.hit_counter[imageID] + 1
                    print(self.hit_counter, end='\r')
                    # get data for those indices from the image
                    interim[query_mask] = self.images[imageID].GetSpectra(indices[query_mask])
                    # mark buffer query structure to prevent double queries
                    self.buffer[imageID][BUFFER_QUERY][indices[query_mask]] = True

                    # copy data to buffer
                    self.buffer[imageID][BUFFER_DATA][indices[query_mask]] = interim[query_mask]
                    
                # for all entries in the BUFFER_QUERY structure which already have data
                # copy buffered data to interim data structure
                if ~query_mask.any():
                    interim[~query_mask] = self.buffer[imageID][BUFFER_DATA][indices[~query_mask]]
            else:
                interim = self.images[imageID].GetSpectra(indices)
            
            if result is None:
                result = interim
            else:
                result = np.concatenate([result, interim])
        
        if self.transforms is not None:
            result = self.transforms(result)
        
        return result

class IonImageDataset(BaseDataSet):
    def __init__(self, images : List[ImageIO.ImzMLReader], 
                       centroids:List[float], 
                       tolerance:float, 
                       tolerance_type:str='ppm', 
                       buffer_type='memory', 
                       transforms=None)-> None:
        self.elements = centroids
        self.tolerance = tolerance
        self.tolerance_type = tolerance_type
        self.buffer = [{} for _ in images]
        
    def get_tolerance(self, c): 
        if self.tolerance_type == "ppm":
            return c * self.tolerance * 10e-6
        else:
            return self.tolerance

    def make_buffered_image(self, index):
        c = self.elements[index%len(self.elements)]
        image_id = index//len(self.elements)
        if c not in self.buffer[image_id] or self.buffer_type == 'none':
            # create ion image
            ii = self.images[image_id].GetArray(c, self.get_tolerance(c), squeeze=False)
            
            if self.buffer_type == 'memory': # buffer image in memory
                self.buffer[image_id][c] = ii
            else: # buffer image on disk
                hash = hashlib.sha224("{:.8f}".format(c).encode())
                hash.update("{:.8f}".format(self.get_tolerance(c)).encode())
                path = self.buffer_path.joinpath(self.images[image_id].GetImageName())
                path.mkdir(exist_ok=True)
                path = path.joinpath(hash.hexdigest())
                np.save(path ,ii)
                self.buffer[image_id][c] = path

        else: # load buffered entry
            if self.buffer_type == 'memory':
                ii = self.buffer[image_id][c]
            else:
                ii = np.load(str(self.buffer[image_id][c])+".npy")
        
            ii = self.transforms(ii)
        return ii

        return len(self.elements) * len(self.images)

        return self.make_buffered_image(index)

        return np.stack([self.__getitem__(index) for index in indexes])

import random
from . import Dataset

class BatchGenerator():

    def __init__(self, dataset: Dataset.BaseDataSet, batch_size: int, shuffle: bool = True):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        
        self.elements = [i for i in range(len(self.dataset))]
        missing_values = self.batch_size-(len(self.elements) % self.batch_size)
        self.elements.extend(random.sample(self.elements, missing_values))
        self.on_epoch_end()

        if self.shuffle:
            random.shuffle(self.elements)

        'Denotes the number of batches per epoch'
        return len(self.elements) // self.batch_size

        return self.dataset.getitems(self.elements[index*self.batch_size:(index+1)*self.batch_size])

# class IonImageBatchGenerator():

#     def __init__(self, dataset: Dataset.IonImageDataset, batch_size: int, shuffle: bool = True):
        
#         self.dataset = dataset
        
#         if tolerance_type == "ppm":
#             self.tolerance = self.tolerance * 10e-6

#         if elements is not None:
#             self.elements = elements
#         else:
#             self.elements = image_handle.GetXAxis().tolist()
#             print("WARNING!", "All m/z bins are used to generate ion images", f"[{len(self.elements)}]")

#         self.buffer_type = buffer_type
#         self.buffer_path = None
#         self.buffer = {}
#         if 'disk' in self.buffer_type: # buffer images on disk; check path exists
#             self.buffer_path = pathlib.Path(self.buffer_type.split(':')[1])
#             if not self.buffer_path.exists():
#                 raise Exception(f"The path {str(self.buffer_path)} does not exist!")

#         missing_values = self.batch_size-(len(self.elements) % self.batch_size)
#         self.elements.extend(random.sample(self.elements, missing_values))
        
#         self.on_epoch_end()

#     def on_epoch_end(self):
#         if self.shuffle:
#             random.shuffle(self.elements)

#     def __len__(self):
#         'Denotes the number of batches per epoch'
#         return len(self.elements) // self.batch_size

#     def get_tolerance(self, c): 
#         if self.tolerance_type == "ppm":
#             return c * self.tolerance * 10e-6
#         else:
#             return self.tolerance

#     def make_buffered_image(self, c):
#         if c not in self.buffer or self.buffer_type == 'none':
#             # create ion image
#             ii = self.image_handle.GetArray(c, self.get_tolerance(c), self.dtype)
            
#             if self.buffer_type == 'memory': # buffer image in memory
#                 self.buffer[c] = ii
#             else: # buffer image on disk
#                 hash = hashlib.sha224("{:.8f}".format(c).encode())
#                 hash.update("{:.8f}".format(self.get_tolerance(c)))
#                 path = self.buffer_path.joinpath(self.image_handle.GetImageName())
#                 path.mkdir(exist_ok=True)
#                 path = path.joinpath(hash.hexdigest())
#                 np.save(path ,ii)
#                 self.buffer[c] = path
            
#             return ii

#         else: # load buffered entry
#             if self.buffer_type == 'memory':
#                 return self.buffer[c]
#             else:
#                 return np.load(str(self.buffer[c])+".npy")
        
#     def __getitem__(self, index):
#         X = np.array([self.make_buffered_image(c) for c in self.elements[index*self.batch_size:(index+1)*self.batch_size]])
        
#         return X

from typing import Literal
from typing import List, Dict
from ctypes import create_string_buffer, c_void_p, c_uint32, c_char_p, c_double, c_float, c_ushort, POINTER
import SimpleITK as sitk

from .Library import get_library

m2NormalizationNone: str = 'None'
m2NormalizationTIC: str = 'TIC'
m2NormalizationSum: str = 'Sum'
m2NormalizationMean: str = 'Mean'
m2NormalizationMax: str = 'Max'
m2NormalizationRMS: str = 'RMS'
m2NormalizationInternal: str = 'Internal'
m2NormalizationExternal: str = 'External'
m2Normalization = Literal[f"{m2NormalizationTIC}", f"{m2NormalizationSum}", f"{m2NormalizationMean}",
                          f"{m2NormalizationMax}", f"{m2NormalizationRMS}", f"{m2NormalizationInternal}",
                          f"{m2NormalizationExternal}", f"{m2NormalizationNone}"]

m2SmoothingSavitzkyGolay: str = "SavitzkyGolay"
m2SmoothingGaussian: str = "Gaussian"
m2Smoothing = Literal[f"{m2SmoothingSavitzkyGolay}",
                      f"{m2SmoothingGaussian}", "None"]

m2PoolingMean: str = "Mean"
m2PoolingMedian: str = "Median"
m2PoolingMaximum: str = "Maximum"
m2PoolingSum: str = "Sum"
m2Pooling = Literal[f"{m2PoolingMean}", f"{m2PoolingMedian}",
                    f"{m2PoolingMaximum}", f"{m2PoolingSum}"]

m2BaselineCorrectionTopHat: str = "TopHat"
m2BaselineCorrectionMedian: str = "Median"
m2BaselineCorrection = Literal[f"{m2BaselineCorrectionTopHat}",
                               f"{m2BaselineCorrectionMedian}", "None"]

m2IntensityTransformationLog2: str = "Log2"
m2IntensityTransformationLog10: str = "Log10"
m2IntensityTransformationSquareRoot: str = "SquareRoot"
m2IntensityTransformation = Literal[f"{m2IntensityTransformationLog2}",
                                    f"{m2IntensityTransformationLog10}", f"{m2IntensityTransformationSquareRoot}", "None"]

class ImzMLReader(object):
    ''' Wrapper class for M2aia's imzML reader https://github.com/m2aia/m2aia.
    
    Complete processing examples with focus on deep learning can be found on https://github.com/m2aia/pym2aia-examples 
    
    Example usage::

        import m2aia as m2

        I = m2.ImzMLReader("path/to/imzMl/file.imzML")
        I.SetNormalization(m2.m2NormalizationTIC)
        I.SetIntensityTransformation(m2.m2IntensityTransformationSquareRoot)
        I.Execute()
        ys_2 = I.GetMeanSpectrum()
        i_2 = I.GetArray(imz, 75)

    '''
    def __init__(self, imzML_path, baseline_correction: m2BaselineCorrection = "None",
                 baseline_correction_half_window_size: int = 50,
                 normalization: m2Normalization = "None",
                 smoothing: m2Smoothing = "None",
                 smoothing_half_window_size: int = 2,
                 intensity_transformation: m2IntensityTransformation = "None",
                 pooling: m2Pooling = m2PoolingMaximum):

        self.lib = get_library()

        HANDLE_PTR = c_void_p

        self.lib.CreateImageHandle.argtypes = [c_char_p]
        self.lib.CreateImageHandle.restype = HANDLE_PTR

        self.lib.DestroyImageHandle.argtypes = [HANDLE_PTR]
        self.lib.DestroyImageHandle.restype = None

        self.lib.GetSize.argtypes = [HANDLE_PTR, POINTER(c_uint32)]
        self.lib.GetSize.restype = None

        self.lib.GetSpacing.argtypes = [
            HANDLE_PTR, POINTER(c_double)]
        self.lib.GetSpacing.restype = None

        self.lib.GetOrigin.argtypes = [
        self.lib.GetOrigin.restype = None

        self.lib.GetXAxis.argtypes = [
        self.lib.GetXAxis.restype = None

        self.lib.GetXAxisDepth.argtypes = [HANDLE_PTR]
        self.lib.GetXAxisDepth.restype = c_uint32

        self.lib.GetImageArrayFloat64.argtypes = [
            HANDLE_PTR, c_double, c_double, POINTER(c_double)]
        self.lib.GetImageArrayFloat64.restype = None

        self.lib.GetImageArrayFloat32.argtypes = [
            HANDLE_PTR, c_double, c_double, POINTER(c_float)]
        self.lib.GetImageArrayFloat32.restype = None

        self.lib.GetMaskArray.argtypes = [
            HANDLE_PTR, POINTER(c_ushort)]
        self.lib.GetMaskArray.restype = None

        self.lib.GetIndexArray.argtypes = [
            HANDLE_PTR, POINTER(c_uint32)]
        self.lib.GetIndexArray.restype = None

        self.lib.GetNormalizationArray.argtypes =[
            HANDLE_PTR, c_char_p, POINTER(c_double)]
        self.lib.GetNormalizationArray.restype = None

        self.lib.GetSpectrumType.argtypes = [HANDLE_PTR]
        self.lib.GetSpectrumType.restype = c_uint32

        self.lib.GetSpectrumDepth.argtypes = [HANDLE_PTR, c_uint32]
        self.lib.GetSpectrumDepth.restype = c_uint32

        self.lib.GetSizeInBytesOfYAxisType.argtypes = [HANDLE_PTR]
        self.lib.GetSizeInBytesOfYAxisType.restype = c_uint32

        self.lib.GetMeanSpectrum.argtypes = [
        self.lib.GetMeanSpectrum.restype = None

        self.lib.GetSpectrumPosition.argtypes = [
            HANDLE_PTR, c_uint32, POINTER(c_uint32)]
        self.lib.GetSpectrumPosition.restype = None

        self.lib.GetMaxSpectrum.argtypes = [
        self.lib.GetMaxSpectrum.restype = None

        self.lib.GetYDataTypeSizeInBytes.argtypes = [HANDLE_PTR]
        self.lib.GetYDataTypeSizeInBytes.restype = c_uint32

        self.lib.GetNumberOfSpectra.argtypes = [HANDLE_PTR]
        self.lib.GetNumberOfSpectra.restype = c_uint32

        self.lib.GetMetaDataDictionary.argtypes = [HANDLE_PTR]
        self.lib.GetMetaDataDictionary.restype = c_char_p

        self.lib.DestroyCharBuffer.argtypes = [HANDLE_PTR]
        self.lib.DestroyCharBuffer.restype = None

        self.lib.GetSpectrum.argtypes = [
            HANDLE_PTR, c_uint32, POINTER(c_float), POINTER(c_float)]
        self.lib.GetSpectrum.restype = None

        self.lib.GetSpectra.argtypes = [
            HANDLE_PTR, c_void_p, c_uint32, POINTER(c_float)]
        self.lib.GetSpectra.restype = None

        self.lib.GetIntensities.argtypes = [
        self.lib.GetIntensities.restype = None

        self.lib.WriteContinuousCentroidImzML.argtypes = [
            HANDLE_PTR, c_char_p, POINTER(c_double), c_uint32]
        self.lib.WriteContinuousCentroidImzML.restype = None

        self.lib.SetSmoothing.argtypes = [HANDLE_PTR, c_char_p, c_uint32]
        self.lib.SetSmoothing.restype = None

        self.lib.SetBaselineCorrection.argtypes = [HANDLE_PTR, c_char_p, c_uint32]
        self.lib.SetBaselineCorrection.restype = None
        
        self.lib.SetNormalization.argtypes = [HANDLE_PTR, c_char_p]
        self.lib.SetNormalization.restype = None

        self.lib.SetIntensityTransformation.argtypes = [HANDLE_PTR, c_char_p]
        self.lib.SetIntensityTransformation.restype = None

        self.lib.SetPooling.argtypes = [HANDLE_PTR, c_char_p]
        self.lib.SetPooling.restype = None

        self.lib.SetTolerance.argtypes = [HANDLE_PTR, c_float]
        self.lib.SetTolerance.restype = None

        self.lib.GetTolerance.argtypes = [HANDLE_PTR]
        self.lib.GetTolerance.restype = c_float
        
        self.lib.Update.argtypes = [HANDLE_PTR]
        self.lib.Update.restype = None

        self.x_axis = None

        self.imzML_path = imzML_path
        self.handle = None
        self.spectrum_type_id = None
        self.spectrum_types = {
                0 :'None' ,
                1: 'Profile',
                2: 'Centroid',
                4: 'Continuous',
                8: 'Processed',
                21: 'ContinuousProfile',
                41:'ProcessedProfile',
                70: 'ContinuousCentroid',
                138:'ProcessedCentroid'
        }

        self.baseline_correction = baseline_correction
        self.baseline_correction_hws = baseline_correction_half_window_size
        self.smoothing = smoothing
        self.smoothing_hws = smoothing_half_window_size
        self.normalization = normalization
        self.intensity_transformation = intensity_transformation
        self.pooling = pooling

        # Read and initialize the image by creating a handle
        self.Load()

    def __delete__(self):
        self.lib.DestroyImageHandle(self.handle)

    def Load(self):
        
        if self.handle is not None:
            self.lib.DestroyImageHandle(self.handle)

        cPath = create_string_buffer(self.imzML_path.encode())
        self.handle = self.lib.CreateImageHandle(cPath)

        self.SetBaselineCorrection(self.baseline_correction, self.baseline_correction_hws)
        self.SetSmoothing(self.smoothing, self.smoothing_hws)
        self.SetNormalization(self.normalization)
        self.SetIntensityTransformation(self.intensity_transformation)
        self.SetPooling(self.pooling)

        self.lib.Update(self.handle)

        self.depth = self.lib.GetXAxisDepth(self.handle)

        # mean overview spectrum
        self.mean_spectrum = np.zeros(self.depth, dtype=np.float64)
        self.lib.GetMeanSpectrum(self.handle, self.mean_spectrum.ctypes.data_as(
            POINTER(c_double)))

        # max overview spectrum
        self.max_spectrum = np.zeros(self.depth, dtype=np.float64)
        self.lib.GetMaxSpectrum(self.handle, self.max_spectrum.ctypes.data_as(

        self.number_of_spectra = self.lib.GetNumberOfSpectra(self.handle)
        self.shape = self.GetShape()

        self.spectrum_type_id = self.lib.GetSpectrumType(self.handle)
        
        # XAxis
        self.x_axis = np.zeros(self.depth, dtype=np.float64)
        self.lib.GetXAxis(self.handle, self.x_axis.ctypes.data_as(

    def path(self) -> pathlib.Path:
        '''Absolute path to the referenced imzML'''
        return pathlib.Path(self.imzML_path)

    def dir(self) -> pathlib.Path:
        '''Absolute path to directory containing the referenced imzML '''
        return self.path().parent

    def name(self) -> str:
        '''Name (including file ending) of the given imzML'''
        return self.path().name

    def WriteContinuousCentroidImzML(self, path : str, centroids):
        ''' Given a list of centroids, write a continuous centroid imzML to the given path.
            Use 'SetTolerance' to define the range query for each centroid (ppm).

            :param path: Target file path the <path>.imzML and the <path>.ibd is written to.
            :param centroids: a list of centroids.

            Example usage::

                import m2aia as m2

                I = m2.ImzMLReader("path/to/imzMl/file.imzML")
                I.Execute()
                I.SetTolerance(75)
                I.WriteContinuousCentroidImzML("path/to/imzMl/file.imzML", [300, 400, 500])

        '''
        cPath = create_string_buffer(path.encode())
        
        centroids = np.array(centroids, dtype=np.double)
        self.lib.WriteContinuousCentroidImzML(self.handle, cPath, centroids.ctypes.data_as(
            POINTER(c_double)), len(centroids))

    def CheckHandle(self):
        ''' Check if the handle was initialized properly.
        To prevent this check from throwing an exception you must call Execute() once.

        :raises: ReferenceError: is invalid file name and or not yet called Execute().
        '''
        if self.handle is None:
            raise ReferenceError(
                "Please initialize image handle by providing a valid file name and run the Execute() function of the reader!")

    def GetParametersAsFormattedString(self):
        '''Transform signal processing parameters into a fomatted string representation.'''
        s = str()
        s += f"(baseline-correction {self.baseline_correction})\n"
        s += f"(baseline-correction-hw {self.baseline_correction_hws})\n"
        s += f"(smoothing {self.smoothing})\n"
        s += f"(smoothing-hw {self.smoothing_hws})\n"
        s += f"(normalization {self.normalization})\n"
        s += f"(pooling {self.pooling})\n"
        s += f"(transform {self.intensity_transformation})\n"
        return s

    def SetSmoothing(self, strategy: m2Smoothing, half_window_size=2):
        '''Set the spectrum smoothing strategy.

        :param strategy: Set the smoothing strategy using one of the m2Smoothing literals.
        :param half_window_size: 2*half_window_size + 1 spectrum points used for smoothing.
        '''
        
        self.smoothing = strategy
        self.smoothing_hws = half_window_size
        arg = create_string_buffer(strategy.encode())
        self.lib.SetSmoothing(self.handle, arg, half_window_size)

    def SetIntensityTransformation(self, strategy: m2IntensityTransformation):
        '''
        Set the intensity transformation strategy.

        :param strategy: m2IntensityTransformation
            Set the intensity transformation strategy using one of the m2IntensityTransformation literals.
        '''
        self.intensity_transformation = strategy
        self.lib.SetIntensityTransformation(self.handle, arg)

    def SetBaselineCorrection(self, strategy: m2BaselineCorrection, half_window_size=50):
        '''Set the baseline correction strategy.

        :param strategy: Set the basline correction strategy using one of the m2BaselineCorrection literals.
        :param half_window_size: 2*half_window_size + 1 spectrum points are used for BaselineCorrection.
        '''
        self.baseline_correction = strategy
        self.baseline_correction_hws = half_window_size
        self.lib.SetBaselineCorrection(self.handle, arg, half_window_size)

    def SetNormalization(self, strategy: m2Normalization):
        '''
        Set the normalization strategy.

        :param strategy: m2Normalization
            Set the normalization strategy using one of the m2Normalization literals.
        '''
        self.normalization = strategy
        self.lib.SetNormalization(self.handle, arg)

    def SetPooling(self, strategy: m2Pooling):
        '''
        Set the pooling strategy.

        :param strategy: m2Pooling
            Set the pooling strategy using one of the m2Pooling literals.
        '''
        self.pooling = strategy
        self.lib.SetPooling(self.handle, arg)

    def SetTolerance(self, tol: np.float32):
        '''
        Set the tolerance value.

        :param tol: np.float32
            The tolerance value to be set.
        '''
        self.lib.SetTolerance(self.handle, tol)

    def GetTolerance(self) -> np.float32:
        '''
        Get the current tolerance value.

        :return: np.float32
            The current tolerance value.
        '''
        return self.lib.GetTolerance(self.handle)

    def GetYDataType(self):
        '''Return the intensity data type defined in the imzML file.

        :return: np.float32 or np.float64; if not defined return None
        '''
        self.CheckHandle()
        size_in_bytes = self.lib.GetYDataTypeSizeInBytes(self.handle)
        if size_in_bytes == 4:
            return np.float32
        elif size_in_bytes == 8:
            return np.float64
        else:
            return None

    def GetShape(self) -> np.array:
        '''Get the shape of the image.

        :return: numpy array of size [3] for the x,y,z image dimensions (in number of pixels).
        '''
        shape = np.zeros((3), dtype=np.int32)
        self.lib.GetSize(self.handle, shape.ctypes.data_as(
            POINTER(c_uint32)))
        return shape

    def GetSpacing(self) -> np.array:
        '''Get the pixel spacing of the image

        :return: numpy array of size [3] and dtype=np.float64 for the pixel size in x,y,z dimension (in millimeter).
        '''
        spacing = np.zeros((3), dtype=np.float64)
        self.lib.GetSpacing(self.handle, spacing.ctypes.data_as(
        return spacing

    def GetSpectrumPosition(self, id) -> np.array:
        '''Get the image index position for a given spectrum id.

        :param id: the id of a spectrum (may be different to the ids given in the imzML).

        :return: Position in index coordinates as numpy array of size [3] and dtype=np.int32.
        '''

        pos = np.zeros((3), dtype=np.int32)
        self.lib.GetSpectrumPosition(self.handle, id, pos.ctypes.data_as(
        return pos

    def GetMetaData(self) -> Dict[str,str]:
        '''Returns a dictionary of all meta data information retrieved by m2aia.

        :return: List of strings of meta data.
        '''
        separator = '\t'
        data = self.lib.GetMetaDataDictionary(self.handle)
        lines = [f.strip() for f in data.decode("utf-8").split('\n') if len(f.strip()) > 0]
        return { line.split(separator)[0]:line.split(separator)[1]  if len(line.split(separator)) > 0  else "true" for line in lines}

    def GetOrigin(self) -> np.array:
        '''Get the image origin.

        :return: The origin in world coordinates of the image as numpy array of size [3] and dtype=np.float64.
        '''
        origin = np.zeros((3), dtype=np.float64)
        self.lib.GetOrigin(self.handle, origin.ctypes.data_as(
        return origin

    def GetMaskArray(self) -> np.ndarray:
        ''' Get the mask image data as numpy array.
        The binary mask indicates valid spectra (pixel value >= 1) and background (pixel value == 0).

        :return: Numpy array of size [x,y,z] with dtype=np.ushort.
        '''
        slice = np.zeros(self.GetShape()[::-1], dtype=np.ushort)
        self.lib.GetMaskArray(
            self.handle, slice.ctypes.data_as(POINTER(c_ushort)))
        return slice

    def GetMaskImage(self) -> sitk.Image:
        ''' Get the mask image data as parameterized SimpleITK.Image.
        The pixel values indicate valid spectra (pixel value >= 1) and background (pixel value == 0).

        :return: sitk.Image of size [x,y,z] with dtype=np.ushort.
        '''
        slice = self.GetMaskArray()
        spacing = self.GetSpacing()
        origin = self.GetOrigin()
        I = sitk.GetImageFromArray(slice)
        I.SetSpacing(spacing)
        I.SetOrigin(origin)
        return I

    def GetIndexArray(self) -> np.ndarray:
        ''' Get the index image data as numpy array.
        The pixel values are the spectrum ids (pixel value >= 1) and background (pixel value == 0).

        :return: Numpy array of size [x,y,z] with dtype=np.uint32.
        '''
        slice = np.zeros(self.GetShape()[::-1], dtype=np.uint32)
        self.lib.GetIndexArray(
            self.handle, slice.ctypes.data_as(POINTER(c_uint32)))

    def GetIndexImage(self):
        ''' Get the index image data as parameterized SimpleITK.Image.

        :return: sitk.Image of size [x,y,z] with dtype=np.uint32.
        '''
        slice = self.GetIndexArray()
        return I
    
    def GetNormalizationArray(self, type) -> np.ndarray:
        ''' Get a normalization image data as numpy array.

        :return: Numpy array of size [x,y,z] with dtype=np.float64.
        '''
        arg = create_string_buffer(type.encode())
        
        image = np.zeros(self.GetShape()[::-1], dtype=np.float64)
        self.lib.GetNormalizationArray(
            self.handle, arg, image.ctypes.data_as(POINTER(c_double)))
        return image

    def GetNormalizationImage(self, type) -> sitk.Image:
        ''' Get a normalization image data as parameterized SimpleITK.Image.
        
        :return: sitk.Image of size [x,y,z] with dtype=np.float64.
        '''
        slice = self.GetNormalizationArray(type)
        return I

    def GetArray(self, center, tol, dtype=np.float32, squeeze: bool = False) -> np.ndarray:
        ''' Get the (ion) image data as numpy array.
        The pixel values are the pooled intensities (ie. pooling strategies like 
        the 'Mean', 'Median', 'Maximum', or 'Sum') in the interval [center-tol, center+tol] of the spectra.

        :return: Numpy array of size [x,y,z] with dtype=dtype.

        :param center: value on the x axis.
        :param tol: tolerance for query points on the x axis around the center.
        :param dtype: array element type [np.float32, np.float64].
        :param squeeze: Remove all dimensions if any is smaller or equals 1.

        :raises:
            TypeError: Image pixel type is not one of [np.float32, np.float64]
        '''
        xs = self.GetXAxis()

        if center < np.min(xs) or center > np.max(xs):
            raise ValueError("Center is out of x-axis range!", center, tol, dtype, squeeze)

        slice = np.zeros(self.GetShape()[::-1], dtype=dtype)
        if dtype == np.float32:
            self.lib.GetImageArrayFloat32(
                self.handle, center, tol, slice.ctypes.data_as(POINTER(c_float)))
        elif dtype == np.float64:
            self.lib.GetImageArrayFloat64(
                self.handle, center, tol, slice.ctypes.data_as(POINTER(c_double)))
        else:
            raise TypeError(
                "Image pixel type is not one of [np.float32, np.float64].")

        if squeeze:
            return np.squeeze(slice)

    def GetImage(self, center, tol, dtype=np.float32) -> sitk.Image:
        ''' Get the (ion) image data as parameterized SimpleITK.Image.
        
        :meth:`m2aia.ImzMLReader.GetArray`

        :return: sitk.Image of size [x,y,z] with dtype=dtype.

        :raises:
        '''
        array = self.GetArray(center, tol, dtype)
        I = sitk.GetImageFromArray(array)
        return I

    def GetMeanSpectrum(self) -> np.array:
        ''' Get the overview spectrum (mean over all spectra).

        :return: np.array with mean intensity values of all spectra.
        '''
        return self.mean_spectrum

    def GetMaxSpectrum(self) -> np.array:
        ''' Get the overview spectrum (max over all spectra).

        :return: np.array with maximum intensity values of all spectra.
        '''
        return self.max_spectrum

    def GetXAxis(self) -> np.array:
        ''' Get the x axis values (i.e. m/z values on the x axis).

        :return: np.array
        '''
        return self.x_axis

    def GetXAxisDepth(self) -> int:
        ''' Get the size of the x axis. For processed imzML files, 
        this value is the number of bins used to represent the x-axis.

        :return: Number of x values.
        '''
        return self.depth

    def GetSpectrumDepth(self, id) -> int:
        ''' Get the size of the x axis for a specific spectrum of the image. 
        This method is helpful for processed (centroid) imzML files. For continuous
        (centroid/profile) imzML files use GetXAxisDepth(self).

        :param id: Id of a spectrum in the image.

        '''
        depth = self.lib.GetSpectrumDepth(self.handle, id)
        if depth <= 0:
            raise RuntimeError("Spectrum depth can not be 0!")
        return depth

    def GetSpectrumType(self) -> str:
        '''Get the imzML type, i.e. continuous/processed profile/centroid.
        
        :return: The type of the imzML as string.
        '''
        return self.spectrum_types[self.spectrum_type_id]

    def GetSizeInBytesOfYAxisType(self) -> int:
        '''Get number of bytes used to store the intensity values.
        
        :return: The number of bytes.
        '''
        return self.lib.GetSizeInBytesOfYAxisType(self.handle)

    def GetSpectrum(self, index) -> List[np.array]:
        ''' Query the x-axis (xs) and y-axis (ys) values for a given spectrum id.

        :param index: Id of a spectrum in the image.
        
        :return: A list of two np.array elements [xs,ys]. xs = x-values; ys = y-values.

        :raises: IndexError: if index is not in the range of valid spectra indices [0,self.number_of_spectra-1].
        '''
        if index < 0 or index >= self.number_of_spectra:
            raise IndexError(
                "Index " + str(index) + " out of range of valid spectrum indices [0," + str(self.number_of_spectra - 1) + "] ")

        depth = self.GetSpectrumDepth(index)
        xs = np.zeros(depth, dtype=np.float32)
        ys = np.zeros(depth, dtype=np.float32)

        self.lib.GetSpectrum(self.handle, index, xs.ctypes.data_as(
            POINTER(c_float)), ys.ctypes.data_as(
            POINTER(c_float)))

        return [xs, ys]

    def GetIntensities(self, index, ys = None) -> np.array:
        ''' Query the y-axis (ys) values for a given spectrum id.

        :param ys: By passing a np.array (dtype=np.float32) from external, this np.array can 
            be reused and do not require an extra memory allocation. Otherwise
            a new array is created.
        
        :raises: 
            IndexError: if index is not in the range of valid spectra indices [0,self.number_of_spectra-1].
            TypeError: if the ImzML file format is not continuous profile/centroid!

        '''

        if not 'Continuous' in self.GetSpectrumType():
                f"The ImzML format has to be in Continuous format! Current format ist {self.GetSpectrumType()}!\nUse GetSpectrum(...) instead.")

        if ys is None:
            ys = np.zeros(self.depth, dtype=np.float32)
        
        if ys.shape[0] != self.depth:
            np.resize(ys, (self.depth))
        
        self.lib.GetIntensities(
            self.handle, index, ys.ctypes.data_as(POINTER(c_float)))
        return ys

    def GetSpectra(self, indices: List[int]) -> np.ndarray:
        ''' Query a set of intensities by a list of indices.
        Only continuous imzML files.

        :param indices: List of Ids of spectra in the image.

        :return: a np.ndarray of shape [len(indices), self.depth].
        
        :raises: 
        '''
        for index in indices:
            if index < 0 or index >= self.number_of_spectra:
                raise IndexError(
                    "Index " + str(index) + " out of range of valid spectrum indices [0," + str(self.number_of_spectra - 1) + "] ")

        batch_size = len(indices)
        ys_batch = np.zeros([batch_size, self.depth], dtype=np.float32)
        idx = np.array(indices, dtype=np.uint32)

        self.lib.GetSpectra(self.handle, idx.ctypes.data_as(
            POINTER(c_uint32)), batch_size, ys_batch.ctypes.data_as(

        return ys_batch

    def GetNumberOfSpectra(self) -> int:
        '''Get the number of valid spectra in the image.
        This can be used to iterate over all spectra in the image using
        a for loop:: 
        
            for i in range(GetNumberOfSpectra()):
                xs, ys = reader.GetSpectrum(i)
                
        :return: Number of valid spectra.
        '''
        return self.number_of_spectra

    def SpectrumIterator(self):
        '''Create a spectrum iterator/generator, yielding all valid spectra.
        
            for i,xs,ys in reader.SpectrumIterator():
                ...
                
        :return: a triplet with (i=spectrum-id, xs=x-values, ys=y-values)
        '''
        for i in range(self.number_of_spectra):
            xs, ys = self.GetSpectrum(i)
            yield i, xs, ys

    def SpectrumRandomBatchIterator(self, batch_size):
        '''Create a spectrum batch iterator/generator, yielding a batch of ys-values of spectra in a random order with repetitions.
        
        Example:: 
        
            for ys_batch in reader.SpectrumRandomBatchIterator(batch_size = N):
                ...
                
        :return: a np.array as batch of intensities with shape [batch_size, self.depth]
        '''
        while (True):
            ids = np.random.random_integers(
                0, self.number_of_spectra-1, batch_size)
            data = np.array([self.GetSpectrum(int(i))[1] for i in ids])
            yield data

    # def _NeighborsStructureElement(self,neighbors, shape):
    #     N = np.prod(shape)
    #     d = np.array(N*[0])
    #     d[:neighbors] = 1
    #     random.shuffle(d)
    #     while d[N//2] == 1:
    #         random.shuffle(d)
    #     d[N//2] = 1
    #     return np.reshape(d,shape)

    # def BatchIteratorWithNNeighbors(self, batch_size, shuffle=True, neighbors=3):
    #     """
    #         1 random entry + subsequently n neighbors
    #         [random sample x]
    #         [neighbor n-1]
    #         [neighbor n-2]
    #         ...
    #         [neighbor n-n]
    #     """
    #     assert(batch_size%(neighbors+1) == 0)

    #     ids = [i+1 for i in range(self.GetNumberOfSpectra())]
    #     pos = np.zeros((len(ids),3),np.int32)
    #     for i in ids:
    #         pos[i-1] = self.GetSpectrumPosition(i-1)

    #     spatial_representation = np.zeros(self.GetShape()[:2])
    #     spatial_representation[pos[:,0], pos[:,1]] = ids

    #     missing_values = len(ids) % batch_size
    #     # ensure full batches by padding the id list
    #     ids = np.pad(ids, (missing_values//2 + missing_values%2, missing_values//2), mode='reflect')
    #     assert((len(ids) % batch_size) != 0)

    #     structure_element = self._NeighborsStructureElement(neighbors, [3,3])

    #     if shuffle:
    #         random.shuffle(ids)

    #     number_of_seeds = batch_size // (neighbors+1)
    #     # generate batches until all ids are touched once
    #     for p in range(len(ids) // batch_size - 1):
    #         neighbor_ids = []

    #         for seed in ids[p*number_of_seeds:(p+1)*number_of_seeds]:
    #             seed_id = seed
    #             while True:
    #                 mask = ndimage.binary_dilation(spatial_representation == seed_id, structure=structure_element).astype(bool)
    #                 masking_result = spatial_representation[mask]
    #                 if len(masking_result) != (neighbors + 1):
    #                     seed_id = ids[random.randint(0,len(ids)-1)]
    #                 else:
    #                     break
    #             masking_result[masking_result == 0] = masking_result[masking_result != 0][0]
    #             neighbor_ids.append(masking_result)

    #         data = np.array([self.GetSpectrum(int(i-1))[1] for i in np.reshape(neighbor_ids,(-1))])

    #         yield data

    #     yield None

import ctypes
import ctypes.util
import platform
import os
import subprocess
import sys
from typing import List, Set

def get_shared_lib_dependencies(so_file_path):
    try:
        ldd_output = subprocess.check_output(['ldd', so_file_path], stderr=subprocess.STDOUT, universal_newlines=True)
        # print("ldd_output",ldd_output)
        dependencies = {line.split(' => ')[0].strip() for line in ldd_output.splitlines() if "not found" in line}
        return dependencies
    except subprocess.CalledProcessError:
        return {}

def load_library_dependencies_recursively(search_path : pathlib.Path, library_name: str, dependencies: List):
    """ Load required M2aia libraries recursively
    """
    lib_path = str(search_path.joinpath(library_name))
    lib_missing_dependencies = get_shared_lib_dependencies(lib_path)
    
    try:
        ctypes.cdll.LoadLibrary(lib_path)
        if lib_path not in dependencies:
            dependencies.append(lib_path)
        return
    except:
        while lib_missing_dependencies:
            lib_working = lib_missing_dependencies.pop()
            load_library_dependencies_recursively(search_path, lib_working, dependencies)
        
    ctypes.cdll.LoadLibrary(lib_path)
    if lib_path not in dependencies:
        dependencies.append(lib_path)
        
def load_m2aia_library():
    search_path = pathlib.Path(os.environ["M2AIA_PATH"])
    target_library_path_parts = ["libM2aiaCore.so"]
    
    if "Windows" in platform.platform():
        os.add_dll_directory(search_path)
        # os.add_dll_directory(search_path.joinpath("MitkCore"))
        return ctypes.cdll.LoadLibrary("M2aiaCore.dll")
    
    else: #"Linux" in platform.platform():

        if "Darwin" in platform.platform():
            raise ImportError("macOS/Darwin based systems are currently not tested.")

        dependencies = []
        for lib_name in target_library_path_parts:
            load_library_dependencies_recursively(search_path, lib_name, dependencies)
        
        os.environ["M2AIA_LIBRARIES"] = ';'.join(dependencies)       
        for dep in dependencies:
            ctypes.cdll.LoadLibrary(dep)
        
        return ctypes.cdll.LoadLibrary((search_path /  "libM2aiaCore.so").absolute())
    
def get_library():
    try:
        return load_m2aia_library()
    except SystemExit:
        pass
    except ImportError as e:
        print(e)
        raise ImportError(
"""Could not find the required M2aia libraries.
pyM2aia requires a valid M2aia installation/build. 
Go to https://m2aia.github.io/m2aia and download the latest version of M2aia. 
Then, follow the setup procedure for pyM2aia on https://github.com/m2aia/pym2aia.
"""
                        ,name="m2aia")

from .ImageIO import *
from .Generators import *
from .Dataset import *
from .utils import *

import os
import logging

try:
    import tensorflow
    from .keras import KerasGenerators
except:
    pass

def validate_environment():
    search_path = pathlib.Path(os.environ['M2AIA_PATH'])
    if search_path.is_dir() and (len([p for p in search_path.glob("**/*M2aiaCore*")])):
        logging.debug("os.environ['M2AIA_PATH'] = " + os.environ["M2AIA_PATH"])
    else:
        logging.debug("Variable: M2AIA_PATH; Description: Binary search path for M2aia's libraries.")
        if not search_path.exists():
            logging.debug("\t- does not exist!")
        
        if search_path.is_file():
            logging.debug("\t- is not a directory!")
        
        # if not search_path.joinpath("MitkCore").exists():
        #     logging.debug("\t- does not contain a folder called MitkCore!")

        if not len([p for p in search_path.glob("**/*M2aiaCore*")]):
            logging.debug("\t- missing library libM2aiaCore.so or M2aiaCore.dll!")
        
        logging.debug("You can fix this problem by adding 'M2AIA_PATH' to your system environment variables. To do so, download the latest M2aia binaries from https://m2aia.github.io/m2aia.")
        raise ImportError("Loading M2aia was not possible!")

def prepare_environment():
    # os.environ["M2AIA_DEBUG"] = ""
    # default search path is pointing to packaged binaries
    if not "M2AIA_PATH" in os.environ:        
        os.environ["M2AIA_PATH"] = str(pathlib.Path(os.path.abspath(__file__)).parent.joinpath("bin"))
        logging.debug("Default library search path: ", os.environ["M2AIA_PATH"])
    else:
        logging.debug("Manually defined library search path: ", os.environ["M2AIA_PATH"])

prepare_environment()
validate_environment()

# dry-load M2aia binary libraries
get_library()

from collections import Counter

## .grouperStrict
##  strict grouping function
##  Don't allow peaks of one sample in the same bin.
##
## params:
##  mass: double, sorted mass
##  intensities: double, corresponding intensities
##  samples: double, corresponding sample id numbers
##  tolerance: double, maximal deviation of a peak position to be
##             considered as same peak
##
## returns:
##  NA if further splitting is needed
##  meanMass (double) if all criteria are matched
def grouper_strict(xs, ys, sources, tolerance):

    counts = Counter(sources)
    if any([key for key, count in counts.items() if count >1]):
        return None

    mean_xs = np.mean(xs)

    if any(np.abs(xs - mean_xs) / mean_xs > tolerance):

    return mean_xs, counts

def group_binning(xs, ys, ss, grouper, tolerance=0.002):
    """
    xs: sorted list of m/z values of peaks
    ys: sorted list of intensities according to xs sorting
    ss: sorted list of source indices according to xs sorting
    grouper: binning method
    tolerance: TODO

    Returns

    bin_assignments:
        assigns each entry in sorted list to the corresponding bin, so that
        xs[bin_assignments==k] returns all xs values corresponding to the k'th bin
    bin_xs:
        new xs center values of the bins
    counts:
        the count of pixels associated with each peak

    """

    d = xs[1:] - xs[:-1]
    bin_assignments = np.zeros_like(xs,dtype=int)

    n = len(xs)
    boundary = []
    boundary.append((0,n))
    current_id = 0
    bin_counts = []
    bin_xs = []
    
    while len(boundary):
        left, right = boundary.pop()
        gapIdx = np.argmax(d[left:right-1]) + left + 1

        # check left interval
        l = grouper(xs[left: gapIdx], ys[left: gapIdx], ss[left: gapIdx], tolerance)
        if l == None:
            boundary.append((left,gapIdx))
        else:
            bin_assignments[left: gapIdx] = current_id
            x, c = l
            bin_counts.append(c)
            bin_xs.append(x)
            current_id+=1

        # check right interval
        r = grouper(xs[gapIdx: right], ys[gapIdx: right], ss[gapIdx:right], tolerance)
        if r == None:
            boundary.append((gapIdx, right))
        else:
            bin_assignments[gapIdx:right] = current_id
            x, c = r

    bin_assign = np.zeros_like(bin_assignments)
    for i, k in enumerate(list( dict.fromkeys(bin_assignments.tolist()))):
        bin_assign[bin_assignments==k] = i #bin id starts by 1

    bin_xs = np.array(bin_xs)
    sort_indices = np.argsort(bin_xs)

    return bin_assign, bin_xs[sort_indices], [bin_counts[c] for c in sort_indices]

def bin_peaks(image: m2.ImzMLReader, mask: np.array = None):
    
    print("Bin peaks for :", image.GetSpectrumType())
    if "ProcessedCentroid" != image.GetSpectrumType():
        raise TypeError("Image has to be of type ProcessedCentroid")

    if mask is not None:
        indices = image.GetIndexArray()[mask>0]
    else:
        indices = image.GetIndexArray()[image.GetMaskArray()>0]

    # BIN PEAKS
    xs_list = [] # list of all mz values
    ys_list = [] # list of all intensities
    ss_list = [] # list of source indices

    for i in indices:
        xs,ys = image.GetSpectrum(i)
        xs_list.extend(xs)
        ys_list.extend(ys)
        ss_list.extend(len(xs) * [i])

    # sort lists
    sort_indices = np.argsort(xs_list)
    xs_list = np.array(xs_list)[sort_indices]
    ys_list = np.array(ys_list)[sort_indices]
    ss_list = np.array(ss_list)[sort_indices]

    print("Start binning on ", len(xs_list), "individual m/z values found in", np.sum(mask>0), "pixels.")

    return group_binning(xs_list, ys_list, ss_list, grouper_strict)

from ..ImageIO import ImzMLReader
from scipy.signal import find_peaks
from scipy import stats

def noise_mad(xs, ys):
    return stats.median_abs_deviation(ys, scale=1/1.4826)

def pick_peaks(I: ImzMLReader, 
               SNR: float=8,
               sort_by_xs: bool=True,
               noise_estimator=noise_mad):
    """Pick peaks for a given ImzMLReader object. 
    
    This function is iterating over all spectra in an image and picks peaks.
    sort_by_xs: if True all values of the return arrays are sorted according to ascending peaks_xs values.

    returns:
    peak_xs: list of m/z positions 
        e.g.: peak_xs=[x1 x1 x1 x2 x2 x2 x2 x3 x3 x3 x3 x3 x3]
            3 peaks found in spectrum with id 1
            4 peaks found in spectrum with id 2
            6 peaks found in spectrum with id 3
            
    peak_ys: list of intensities corresponding to peak_xs
        e.g.: peak_ys=[y1 y1 y1 y2 y2 y2 y2 y3 y3 y3 y3 y3 y3]

    peak_idx = list of array indices corresponding to a spectrum array
        e.g.: peak_idx=[idx1 idx1 idx1 idx2 idx2 idx2 idx2 idx3 idx3 idx3 idx3 idx3 idx3]

    peaks_source = indicators for each corresponding peak_xs, from which spectrum id it was picked
        e.g.: peaks_source=[1 1 1 2 2 2 2 3 3 3 3 3 3]

    """
    
    peak_xs = []
    peak_ys = []
    peak_idx = []
    peaks_source = []

    for i, xs, ys in I.SpectrumIterator():

        noise = noise_estimator(ys)
        p = find_peaks(ys, height=noise*SNR)
        
        # masses
        peak_xs.extend(xs[p[0]])
        
        # xs-index for peak
        peak_idx.extend(p[0])

        # intensities
        peak_ys.extend(ys[p[0]])
        
        # spectrum id for each peak
        peaks_source.extend(len(p[0]) * [i])
    
    if sort_by_xs == True:
        sorted_indices = np.argsort(peak_xs)
        peak_xs = np.array(peak_xs)[sorted_indices]
        peak_ys = np.array(peak_ys)[sorted_indices]
        peak_idx = np.array(peak_idx)[sorted_indices]
        peaks_source = np.array(peaks_source)[sorted_indices]
    
    return peak_xs, peak_ys, peak_idx, peaks_source

from scipy.stats import ttest_ind
from mne.stats import fdr_correction

def volcano_data(image: m2.ImzMLReader, mask: np.array, x):
    data = image.GetArray(x, x*100*10e-6)
    A = data[mask==1]
    B = data[mask==2]
    return ttest_ind(A,B)[1], np.log2(np.mean(A)/np.mean(B))
    
def volcano_plot_data(image: m2.ImzMLReader, mask: np.array, xs, pThreshold = 0.05, fcThreshold = 1):
    print("Start volcano")
    p_fc_sig = np.zeros((len(xs),3))
    for i, x in enumerate(xs):
        p_fc_sig[i,:2] = volcano_data(image, mask, x)
    p_fc_sig[:,0][p_fc_sig[:,0] == 0] = np.finfo(np.float64).tiny
    p_fc_sig[:,0] = fdr_correction(p_fc_sig[:,0])[1]
    p_fc_sig[:,2] = np.logical_and(p_fc_sig[:,0] < pThreshold, abs(p_fc_sig[:,1]) > fcThreshold)

    return p_fc_sig