<reponame>erikhvatum/wautofocuser
# The MIT License (MIT)
#
# Copyright (c) 2014 WUSTL ZPLAB
#
# 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.
#
# Authors: <NAME> <<EMAIL>>

import collections
from misc.image_from_multi_image_statistics import generate_running_percentile_difference
import math
from misc.pca import pca_decompose
import numpy
from pathlib import Path
import scipy.ndimage
import scipy.ndimage.morphology
import skimage.exposure
import skimage.measure
import skimage.morphology
import skimage.io as skio
import sklearn.linear_model
import sklearn.neighbors
import sklearn.svm
for function in ('imread', 'imsave', 'imread_collection'):
    skio.use_plugin('freeimage', function)
import multiprocessing
import pickle
import sys

def generate__bf_bgs_masks__fluo_running_differences__bfs__composites(dpath, bgs_mask_alpha=0.3333, running_difference_alpha=0.3333, percentile=20, run_length=10, crop=None):
    if bgs_mask_alpha > 1 or bgs_mask_alpha < 0 \
       or running_difference_alpha > 1 or running_difference_alpha < 0 \
       or bgs_mask_alpha + running_difference_alpha > 1:
        raise ValueError('bgs_mask_alpha and running_difference_alpha must be in the range [0, 1], as must their sum.')
    dpath = Path(dpath)
    imfpaths = list((dpath / 'bestfmvs').glob('*.PNG'))
    indexes = sorted([int(imfpath.stem) for imfpath in imfpaths])
    fluo_running_percentile_difference_generator = generate_running_percentile_difference(percentile, run_length, crop)
    next(fluo_running_percentile_difference_generator)
    for index in indexes:
        try:
            mask = skio.imread(str(dpath / 'MixtureOfGaussianV2BGS' / '{}.png'.format(index)))
        except ValueError as ve:
            continue
        if (mask == 255).all():
            continue
        fluo = skio.imread(str(dpath / 'fluos' / '{}.PNG'.format(index)))
        if crop:
            mask = mask[:crop[0], :crop[1]]
            fluo = fluo[:crop[0], :crop[1]]
        fluo_difference = fluo_running_percentile_difference_generator.send(fluo)
        if fluo_difference is None:
            yield
            continue
        im = skio.imread(str(dpath / 'bestfmvs' / '{}.PNG'.format(index)))
        if crop:
            im = im[:crop[0], :crop[1]]
        yield ( \
               ( \
                ((im.astype(numpy.float32) - im.min()) / im.max()) * (1 - bgs_mask_alpha - running_difference_alpha) + \
                (mask > 0).astype(numpy.float32) * bgs_mask_alpha + \
                (fluo_difference > 2500).astype(numpy.float32) * running_difference_alpha \
               ) * 65535 \
              ).astype(numpy.uint16)

def fill_voids(mask, max_pixel_count_void_to_fill):
    mask = mask.copy()
    void_labels = skimage.measure.label(~mask)
    void_regions = skimage.measure.regionprops(void_labels)[1:]
    void_regions.sort(key=lambda region: region.area)
    for void_region in void_regions:
        if void_region.area > max_pixel_count_void_to_fill:
            break
        y, x = void_region.coords[0]
        if mask[y, x]:
            # Skip island
            continue
        # Fill void
        mask[void_labels == void_region.label] = True
    return mask

def generate_masks(dpath, percentile=0, run_length=10, d_threshold=3000, max_void_fill=3000, crop=(2160,2560)):
    def make_mask(im):
        mask = im > d_threshold
        mask = scipy.ndimage.binary_dilation(mask, iterations=5)
        mask = fill_voids(mask, max_void_fill)
        return scipy.ndimage.binary_erosion(mask, iterations=4)
    dpath = Path(dpath)
    imfpaths = list((dpath / 'bestfmvs').glob('*.PNG'))
    indexes = sorted([int(imfpath.stem) for imfpath in imfpaths])
    del imfpaths
    bf_rpdg = generate_running_percentile_difference(100 - percentile, run_length)
    fluo_rpdg = generate_running_percentile_difference(percentile, run_length)
    next(bf_rpdg)
    next(fluo_rpdg)
    for index in indexes:
        bf = skio.imread(str(dpath / 'bestfmvs' / '{}.PNG'.format(index)))
        fluo = skio.imread(str(dpath / 'fluos' / '{}.PNG'.format(index)))
        if crop:
            bf = bf[:crop[0], :crop[1]]
            fluo = fluo[:crop[0], :crop[1]]
        bf = bf_rpdg.send(bf)
        fluo = fluo_rpdg.send(fluo)
        if bf is None:
            yield None
            continue
        bf = make_mask(bf)
        fluo = make_mask(fluo)
        yield bf & fluo

def generate__bf_running_differences__fluo_running_differences__bfs__composites(dpath, bf_mask_alpha=0.3333, fluo_mask_alpha=0.3333, percentile=0, run_length=10, crop=None):
    if bf_mask_alpha > 1 or bf_mask_alpha < 0 \
       or fluo_mask_alpha > 1 or fluo_mask_alpha < 0 \
       or bf_mask_alpha + fluo_mask_alpha > 1:
        raise ValueError('bf_mask_alpha and fluo_mask_alpha must be in the range [0, 1], as must their sum.')
    dpath = Path(dpath)
    imfpaths = list((dpath / 'bestfmvs').glob('*.PNG'))
    indexes = sorted([int(imfpath.stem) for imfpath in imfpaths])
    del imfpaths
    bf_rpdg = generate_running_percentile_difference(100 - percentile, run_length)
    fluo_rpdg = generate_running_percentile_difference(percentile, run_length)
    next(bf_rpdg)
    next(fluo_rpdg)
    for index in indexes:
        try:
            bf = skio.imread(str(dpath / 'bestfmvs' / '{}.PNG'.format(index)))
        except ValueError as ve:
            continue
        fluo = skio.imread(str(dpath / 'fluos' / '{}.PNG'.format(index)))
        if crop:
            mask = mask[:crop[0], :crop[1]]
            fluo = fluo[:crop[0], :crop[1]]
        fluo_difference = fluo_rpdg.send(fluo)
        if fluo_difference is None:
            yield
            continue
        im = skio.imread(str(dpath / 'bestfmvs' / '{}.PNG'.format(index)))
        if crop:
            im = im[:crop[0], :crop[1]]
        yield ( \
               ( \
                ((im.astype(numpy.float32) - im.min()) / im.max()) * (1 - bgs_mask_alpha - running_difference_alpha) + \
                (mask > 0).astype(numpy.float32) * bgs_mask_alpha + \
                (fluo_difference > 2500).astype(numpy.float32) * running_difference_alpha \
               ) * 65535 \
              ).astype(numpy.uint16)

def overlay__masks__bfs__in_flipbook(dpath, rw, mask_alpha, mask_name='masks'):
    dpath = Path(dpath)
    imfpaths = list((dpath / 'bestfmvs').glob('*.PNG'))
    indexes = sorted([int(imfpath.stem) for imfpath in imfpaths])
    cs = []
    for index in indexes:
        try:
            mask = skio.imread(str(dpath / mask_name / '{}.png'.format(index)))[:2160,:2560]
        except ValueError as ve:
            try:
                mask = skio.imread(str(dpath / mask_name / '{:04}.png'.format(index)))[:2160,:2560]
            except ValueError as ve:
                continue
        if not (mask == 255).all():
            im = skio.imread(str(dpath / 'bestfmvs' / '{}.PNG'.format(index)))[:2160,:2560]
            cs.append((((im.astype(numpy.float32) / 65535) * (1 - mask_alpha) + (mask > 0).astype(numpy.float32) * mask_alpha)*65535).astype(numpy.uint16))
    rw.showImagesInNewFlipper(cs)

def make_multiclassifier_data_and_targets(im_fpath, mask_set_fpath, patch_width=9, background_sample_count=2000, worm_interior_sample_count=400, worm_wall_sample_count=100,
                                     pca_pcs=None, pca_means=None):
    mask_set_fpath_str = str(mask_set_fpath)
    imf = skimage.exposure.equalize_adapthist(skio.imread(str(im_fpath))).astype(numpy.float32)
    if imf.max() > 1:
        # For some reason, skimage.exposure.equalize_adapthist rescales to [0, 1] on OS X but not on Linux.
        # [0, 1] scaling is desired.
        imf -= imf.min()
        imf /= imf.max()
    masks = [
        ('_worm_interior.png', worm_interior_sample_count, 1),
        ('_worm_wall.png', worm_wall_sample_count, 2),
        ('_valid_exterior.png', background_sample_count, 0)]
    labels = []
    vectors = []
    for mask_fpath_suffix, sample_count, label in masks:
        coords = select_random_coords_in_mask(mask_set_fpath_str + mask_fpath_suffix, sample_count)
        for coord in coords:
            vector = make_patch_feature_vector(imf, patch_width, coord)
            labels.append(label)
            vectors.append(vector)
    if pca_pcs is not None:
        vectors = pca_decompose(vectors, pca_pcs, pca_means)
    else:
        vectors = numpy.array(vectors)
    return vectors, numpy.array(labels)

if __name__ == '__main__':
    def _worker_process_function(dpath, percentile, run_length, d_threshold, max_void_fill, crop):
        mask_dpath = dpath / 'masks'
        if not mask_dpath.exists():
            mask_dpath.mkdir()
        for idx, mask in enumerate(generate_masks(dpath, percentile, run_length, d_threshold, max_void_fill, crop)):
            if mask is None:
                continue
            skio.imsave(str(mask_dpath / '{:04}.png'.format(idx)), mask.astype(numpy.uint8) * 255)

    def _process_exception_callback(process_exception):
        print('warning: worker failed with exception:', process_exception)

    import argparse
    import os
    argparser = argparse.ArgumentParser(description='Experiment01_a data and target set generator.')
    argparser.add_argument('--wellDevelopmentalSuccessDb',
                           default=Path(os.path.expanduser('~')) / 'Data' / 'experiment01_a' / 'well_developmental_success_db.pickle',
                           type=Path)
    argparser.add_argument('--experiment01_a',
                           default=Path(os.path.expanduser('~')) / 'Data' / 'experiment01_a',
                           type=Path)
    argparser.add_argument('--percentile', default=0, type=float)
    argparser.add_argument('--run-length', default=10, type=int)
    argparser.add_argument('--d-threshold', default=3000, type=int)
    argparser.add_argument('--max-void-fill', default=3000, type=int)
    argparser.add_argument('--crop', default='None', type=str, help='Specify a 2 element tuple (y, x), or None to disable')
    args = argparser.parse_args()
    with open(str(args.wellDevelopmentalSuccessDb), 'rb') as f:
        well_developmental_success_db = pickle.load(f)
    with multiprocessing.Pool(multiprocessing.cpu_count() + 1) as pool:
        async_results = []
        for p, s in sorted(well_developmental_success_db.items(), key=lambda v: v[0]):
            if s != 'LittleOrNone':
                async_results.append(pool.apply_async(_worker_process_function,
                                                      (args.experiment01_a / p.parts[-1],
                                                       args.percentile,
                                                       args.run_length,
                                                       args.d_threshold,
                                                       args.max_void_fill,
                                                       eval(args.crop)),
                                                      error_callback=_process_exception_callback))
        pool.close()
        pool.join()


<gh_stars>1-10
"""
This class stores a dense list of sparse matrices. For external memory
usage, it writes each sparse matrix to a matrix market format file in a
directory. Then, that directory is gzipped to make the file a bit portable.
"""

import os
import shutil
import tarfile

import pyllars.utils as utils

import numpy as np
import scipy.io
import scipy.sparse

# tolerance value for checking equality
tolerance = 1e-10

class ExternalSparseMatrixList:
    def __init__(self, size, tarfile_contents=None, tarfile_names=None):
        self.sparse_matrix_list = [None] * size
        self.tarfile_contents = tarfile_contents
        self.tarfile_names = tarfile_names

    def __getitem__(self, key):
        # check if we should try to fetch this
        if (self.sparse_matrix_list[key] is None) and (self.tarfile_contents is not None):
            fn = '{}.mtx'.format(key)
            if fn in self.tarfile_names:
                f = self.tarfile_contents.extractfile(fn)
                m = scipy.io.mmread(f)
                self[key] = scipy.sparse.csr_matrix(m)
            else:
                self[key] = scipy.sparse.csr_matrix((1,0))

        return self.sparse_matrix_list[key]

    def __setitem__(self, key, item):
        self.sparse_matrix_list[key] = item

    def write(self,filename, overwrite=False):
        # first, strip the compression ending, if present
        filename = filename.replace('.tar.gz', '')
        filename = filename.replace('.tgz', '')

        # check that the path is safe
        if os.path.exists(filename) and not overwrite:
            raise OSError("Attempting to overwrite existing file: '{}'".format(filename))

        if os.path.exists(filename):
            if os.path.isfile(filename):
                os.remove(filename)
            else:
                shutil.rmtree(filename)

        # create the folder
        os.makedirs(filename)

        # first, write the metadata (just the size)
        fn = os.path.join(filename, "meta.txt")
        with open(fn, 'w') as f:
            f.write(str(len(self)))

        # write each matrix in MM format
        for i in range(len(self.sparse_matrix_list)):
            sm = self.sparse_matrix_list[i]
            if sm is None:
                continue
            fn = os.path.join(filename, str(i))
            scipy.io.mmwrite(fn, sm)

        # create the tgz file
        fn = '{}.tar.gz'.format(filename)
        tar = tarfile.open(fn, "w:gz")
        tar.add(filename, arcname='')
        tar.close()

        # remove the folder
        shutil.rmtree(filename)

        # and rename the tgz file
        os.rename(fn, filename)

    def __len__(self):
        return len(self.sparse_matrix_list)

    def __eq__(self, other):
        # make sure they are the same size
        if len(self) != len(other):
            return False

        # check that each element is equal
        for i in range(len(self)):
            diff = self[i] - other[i]
            if diff.nnz != 0:
                # check if the nonzero values exceed the tolerance
                nonzeros = diff[diff.nonzero()]
                nonzeros = np.squeeze(np.asarray(nonzeros))
                for nz in nonzeros:
                    if np.abs(nz) > tolerance:
                        return False

        # the all elements are equal
        return True

    def sum(self):
        """ This method calculates the sum of all matrices in the list.

        Returns:
            dtype: the sum of all matrices
        """
        s = sum(matrix.sum() for matrix in self.sparse_matrix_list if matrix is not None)
        return s

    def max_len(self, axis=1):
        """ This method finds the maximum size of any matrix along the given
            axis. Presumably, this is for use by the to_sparse_matrix method,
            but it could have other uses.

            Args:
                axis (int): 0 for rows, 1 for columns (i.e., matrix.shape[axis])

            Returns:
                int: the maximum length along the given axis of any matrix.
        """
        l = list(matrix.shape[axis] for matrix in self.sparse_matrix_list if matrix is not None)
        if len(l) == 0:
            return 0

        max_len = max(l)
        return max_len

    def to_sparse_matrix(self, min_cols=None):
        """ This method attempts to convert the sparse matrix list into a
            single sparse matrix. This operation iterates over the columns
            of the first row of each matrix. So it is only sensible if the
            matrix list consists of a set of column vectors.

            For some cases, the number of columns in the matrix should exceed
            the number of columns in the longest matrix. If this is the case,
            the number of columns can be specified.

            NOTE: This method does not do any type checking.

            Returns:
                scipy.sparse matrix: A sparse matrix in which each item in
                    the list is treated as a row in the matrix

                    OR

                    None, if the list contains no items
        """
        # create a 2d numpy array that contains all of the rows
        max_len = self.max_len()

        if (min_cols is not None) and (max_len < min_cols):
            max_len = min_cols

        # determine the dtype
        dtype = None
        for i in range(len(self)):
            if self[i] is not None:
                dtype = self[i].dtype
                break

        # make sure we found something
        if dtype is None:
            return None

        sparse_matrix = scipy.sparse.lil_matrix((len(self), max_len), dtype=dtype)

        # transfer all of the values over to the sparse array
        for i in range(len(self)):
            sv = self[i]
            if sv is None:
                continue

            sv = sv.tocoo()

            for j,k,v in zip(sv.row, sv.col, sv.data):
                sparse_matrix[i,k] = v

        return sparse_matrix

def lazy_read_external_sparse_matrix_list(filename):

    # open the gzipped tar file
    contents = tarfile.open(filename, "r:gz")
    names = contents.getnames()

    # read the meta file
    fn = 'meta.txt'
    f = contents.extractfile(fn)
    size = int(f.readline())
    f.close()

    # create the container
    esml = ExternalSparseMatrixList(size, tarfile_contents=contents, tarfile_names=names)
    return esml

def read_external_sparse_matrix_list(filename):
    # open the gzipped tar file
    contents = tarfile.open(filename, "r:gz")
    names = contents.getnames()

    # read the meta file
    fn = 'meta.txt'
    f = contents.extractfile(fn)
    size = int(f.readline())
    f.close()

    # create the container
    esml = ExternalSparseMatrixList(size)

    # read in each sparse matrix
    for i in range(size):
        fn = '{}.mtx'.format(i)

        if fn in names:
            f = contents.extractfile(fn)
            m = scipy.io.mmread(f)
            esml[i] = scipy.sparse.csr_matrix(m)
        else:
            esml[i] = scipy.sparse.csr_matrix((1,0))

    return esml

def concatenate(lists):
    total_size = 0
    for l in lists:
        total_size += len(l)

    joined_list = ExternalSparseMatrixList(total_size)
    index = 0
    for i in range(len(lists)):
        l = lists[i]
        for j in range(len(l)):
            # convert just to make absolutely sure this is a sparse matrix
            joined_list[index] = scipy.sparse.lil_matrix(l[j])
            index += 1

    return joined_list

def to_sparse_matrix(list_of_sparse_row_vectors):
    """ This function converts a list of sparse row vectors, i.e., sparse
        matrices with a shape like (1, X), into a single sparse matrix. Each
        row of the resulting matrix corresponds to the respective sparse row
        vector in the input list.

        Args:
            list_of_sparse_row_matrices (list of scipy.sparse): a list of scipy
                sparse matrices with shape (1,X).

        Returns:
            scipy.sparse.lil_matrix: a sparse matrix in which the rows
                correspond to the vectors in the list. The dimension of the
                matrix is (N, Y), where N is the number of items in the list
                and Y is the largest size of any input matrix.

        Imports:
            logging

        Raises:
            ValueError: if any of the input matrices do not have a shape of
                the form (1, X)
    """
    import logging

    esml = ExternalSparseMatrixList(len(list_of_sparse_row_vectors))

    msg = "Copying sparse row vectors to sparse matrix list"
    logging.info(msg)

    for i, sparse_row_vector in enumerate(list_of_sparse_row_vectors):
        if sparse_row_vector.shape[0] != 1:
            msg = ("list_of_sparse_row_vectors[{}] did not have the correct shape. "
                "Expected something of the form (1,X); found: ({},{})".format(i,
                sparse_row_vector.shape[0], sparse_row_vector.shape[1]))
            raise ValueError(msg)

        esml[i] = sparse_row_vector

    msg = "Converting matrix list to single sparse matrix"
    logging.info(msg)

    sparse_matrix = esml.to_sparse_matrix()

    return sparse_matrix

def main():
    # create a random sparse matrix list
    esml = ExternalSparseMatrixList(3)
    for i in range(3):
        esml[i] = scipy.sparse.rand(8, 4, density=0.4)

    # write it to disk
    esml.write('test_esml', True)

    # read it back
    esml_read = read_external_sparse_matrix_list('test_esml')

    # make sure they are equal
    assert(esml_read == esml)

    print("TEST SUCCEEDED: The created sparse matrix list and the one read "
        "from disk are equal.")

if __name__ == '__main__':
    main()


<gh_stars>0
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
import math

# def butter_highpass_filter(f_s, f_pass, f_stop, fs = 0.5, td = 1, g_pass = 1, g_stop = 50, wc = None):
#     """
#     Return explanation:
#     N: number of orders
#     b: numerator of the Filter
#     a: denominator of the filter
#     """
#     wp = f_pass / (f_s / 2)
#     ws = f_stop / (f_s / 2)
#     omega_p = (2 / td) * np.tan(wp / 2)
#     omega_s = (2 / td) * np.tan(ws / 2)
#     # Find the order and natural frequency of the highpass filter
#     N, wn = signal.buttord(omega_p, omega_s, g_pass, g_stop, analog = True)
#     # Find the Gain of the highpass filter
#     if wc:
#         b, a = signal.butter(N, wc, btype = 'high', analog = True)
#         wn = wc
#     else:
#         b, a = signal.butter(N, wn, btype = 'high', analog = True)
#     return N, b, a, wn

def butter_highpass_filter(fs, fc, order = 5):
    """
    Function explanation:
    fs: sampling rate
    fc: cutoff frequency
    order: Number of orders
    Return: denominator and numerator of the filter's transfer function
    """
    nyq = 0.5 * fs
    normal_fc = fc / nyq
    b, a = signal.butter(order, normal_fc, btype = 'high', analog = False)
    return b, a

def butter_bandpass_filter(fs, lowcut, highcut, order = 5):
    """
    Function explanation:
    fs: sampling rate
    lowcut: lowcut Frequency
    highcut:highcut Frequency
    order: Number of orders
    Return: denominator and numerator of the filter's transfer function
    """
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype = 'band', analog = False)
    return b, a


from __future__ import with_statement
from __future__ import division
from __future__ import absolute_import
import numpy as np
import tensorflow as tf
import scipy.signal
from gym.spaces import Box, Discrete

EPS = 1e-8

def combined_shape(length, shape=None):
    if shape is None:
        return (length,)
    if(np.isscalar(shape)):
        return (length, shape)
    else:
        l = list(shape)
        l.insert(0, length)
        l = tuple(l)
        return l
#        tp = (length, shape)
#        return convert([element for tupl in tp for element in tupl])

def placeholder(dim=None):
    return tf.placeholder(dtype=tf.float32, shape=combined_shape(None,dim))

def placeholders(*args):
    return [placeholder(dim) for dim in args]

def placeholder_from_space(space):
    if isinstance(space, Box):
        return placeholder(space.shape)
    elif isinstance(space, Discrete):
        return tf.placeholder(dtype=tf.int32, shape=(None,))
    raise NotImplementedError

def placeholders_from_spaces(*args):
    return [placeholder_from_space(space) for space in args]

def mlp(x, hidden_sizes=(32,), activation=tf.tanh, output_activation=None):
    for h in hidden_sizes[:-1]:
        x = tf.layers.dense(x, units=h, activation=activation)
    x = tf.layers.dense(x, units=hidden_sizes[-1], activation=output_activation)
    return x

def get_vars(scope=u''):
    return [x for x in tf.trainable_variables() if scope in x.name]

def count_vars(scope=u''):
    v = get_vars(scope)
    return sum([np.prod(var.shape.as_list()) for var in v])

def gaussian_likelihood(x, mu, log_std):
    pre_sum = -0.5 * (((x-mu)/(tf.exp(log_std)+EPS))**2 + 2*log_std + np.log(2*np.pi))
    return tf.reduce_sum(pre_sum, axis=1)

def discount_cumsum(x, discount):
    u"""
    magic from rllab for computing discounted cumulative sums of vectors.

    input:
        vector x,
        [x0,
         x1,
         x2]

    output:
        [x0 + discount * x1 + discount^2 * x2,
         x1 + discount * x2,
         x2]
    """
    return scipy.signal.lfilter([1], [1, float(-discount)], x[::-1], axis=0)[::-1]

def nextOrderPolynomial(x, length, lastOrderPolynomial, incrementList):
    #empty list for next order
    nextOrderPolynomial = []

    for i in xrange(length):
        for j in xrange(incrementList[i]):
            nextOrderPolynomial.append(tf.multiply(x[:,i],lastOrderPolynomial[j]))

    return nextOrderPolynomial

def nextIncrementList(incrementList):
    nextIncrementList = [1]
    for i in xrange(len(incrementList)-1):
        nextIncrementList.append(nextIncrementList[i]+incrementList[i+1])
    return nextIncrementList

def polynomial(x, order, u_dim):
    u""" Computes u, a polynomial function of x (for each row of x).
    Polynomial will be composed of monomials of degree up to and including
    the specified 'order'.
    """
    NUM_STATE_DIMS = x.get_shape().as_list()[-1]

    print u"Generating monomials..."
    FullPolynomial = []
    z = tf.ones_like(x[:,0])
    FullPolynomial.append(z)
    incrementList = []
    lastPolynomial = []
    for i in xrange(NUM_STATE_DIMS):
        incrementList.append(i+1)
        lastPolynomial.append(x[:,i])

    #generate full polynomial
    FullPolynomial.extend(lastPolynomial)
    for i in xrange(order - 1):
        nextPolynomial = nextOrderPolynomial(x,NUM_STATE_DIMS,lastPolynomial,incrementList)
        FullPolynomial.extend(nextPolynomial)
        incrementList = nextIncrementList(incrementList)
        lastPolynomial = nextPolynomial
    print u"number of monomials: ", len(FullPolynomial)
    print u"...done!"

    # Now declare tf variables for the coefficients of all these monomials.
    coeffs = tf.get_variable(u"polynomial_coefficients",
                             initializer=tf.zeros((len(FullPolynomial), u_dim)),dtype=tf.float32)

    FullPolynomial = tf.transpose(tf.stack(FullPolynomial))

    # Compute polynomial output for each state.
    action = tf.matmul(FullPolynomial,coeffs)
    return action

u"""
Policies
"""

def mlp_categorical_policy(x, a, hidden_sizes, activation, output_activation, action_space):
    act_dim = action_space.n
    logits = mlp(x, list(hidden_sizes)+[act_dim], activation, None)
    logp_all = tf.nn.log_softmax(logits)
    pi = tf.squeeze(tf.multinomial(logits,1), axis=1)
    logp = tf.reduce_sum(tf.one_hot(a, depth=act_dim) * logp_all, axis=1)
    logp_pi = tf.reduce_sum(tf.one_hot(pi, depth=act_dim) * logp_all, axis=1)
    return pi, logp, logp_pi

def mlp_gaussian_policy(x, a, hidden_sizes, activation, output_activation, action_space):
    act_dim = a.shape.as_list()[-1]
    mu = mlp(x, list(hidden_sizes)+[act_dim], activation, output_activation)
    log_std = tf.get_variable(name=u'log_std', initializer=-0.5*np.ones(act_dim, dtype=np.float32))
    std = tf.exp(log_std)
    pi = mu + tf.random_normal(tf.shape(mu)) * std
    logp = gaussian_likelihood(a, mu, log_std)
    logp_pi = gaussian_likelihood(pi, mu, log_std)
    return pi, logp, logp_pi

def polynomial_gaussian_policy(x, a, order, action_space):
    act_dim = a.shape.as_list()[-1]
    mu = polynomial(x, order, act_dim)
    log_std = tf.get_variable(name=u'log_std', initializer=-1.5*np.ones(act_dim, dtype=np.float32))
    std = tf.exp(log_std)
    pi = mu + tf.random_normal(tf.shape(mu)) * std
    logp = gaussian_likelihood(a, mu, log_std)
    logp_pi = gaussian_likelihood(pi, mu, log_std)
    return pi, logp, logp_pi

u"""
Actor-Critics
"""
def mlp_actor_critic(x, a, hidden_sizes=(64,64), activation=tf.tanh,
                     output_activation=None, policy=None, action_space=None):

    # default policy builder depends on action space
    if policy is None and isinstance(action_space, Box):
        policy = mlp_gaussian_policy
    elif policy is None and isinstance(action_space, Discrete):
        policy = mlp_categorical_policy

    with tf.variable_scope(u'pi'):
        pi, logp, logp_pi = policy(x, a, hidden_sizes, activation, output_activation, action_space)
    with tf.variable_scope(u'v'):
        # DFK modified: want unbiased gradient estimate, so replacing MLP with
        # zero (for all states).
        # TODO(@eric): figure out how to do this right and not just multiply by zero.
        #v = tf.zeros_like(x)
        v = 0.0 * tf.squeeze(mlp(x, list(hidden_sizes)+[1], activation, None), axis=1)
    return pi, logp, logp_pi, v

def polynomial_actor_critic(x, a, order, policy=None, action_space=None):
    # default policy builder depends on action space
    if policy is None and isinstance(action_space, Box):
        policy = polynomial_gaussian_policy
    elif policy is None and isinstance(action_space, Discrete):
        assert(False)

    with tf.variable_scope(u'pi',reuse=tf.AUTO_REUSE):
        pi, logp, logp_pi = policy(x, a, order, action_space)
    with tf.variable_scope(u'v',reuse=tf.AUTO_REUSE):
        # DFK modified: want unbiased gradient estimate, so replacing MLP with
        # zero (for all states).
        # TODO(@eric): figure out how to do this right and not just multiply by zero.
        v = 0.0 * tf.squeeze(mlp(x, list((1,1))+[1], tf.tanh, None), axis=1)
    return pi, logp, logp_pi, v


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import math
import os
import pickle

import matplotlib
import matplotlib.gridspec as gridspec
import numpy as np
import scipy
from mat4py import loadmat
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

import plotly.graph_objects as go
from plotly.subplots import make_subplots

import time
import multiprocessing as mp

import plotly.offline as py
import plotly.graph_objs as go
from matplotlib import cm

from dependencies import ROOTDIR
from utils.HelperFunctions import Output, Configuration, Imaging

class PlotRoutines: #  656 lines so far
    """Plots results from the electrode reconstruction/model creation in order to validate the results"""

    def __init__(self, subject, inputfolder=''):
        self.cfg = Configuration.load_config(ROOTDIR)
        self.debug=False
        GetData()
        # TODO: new order with a) get 'static data' such as background, trajectory, lead_model, b) plot all results.
        # TODO: c) include some variable data (markers, rotation) and add some callback functions in form of arrows

        # Get static data, that is lead data, backgrounds and trajectories
        filename_leadmodel = os.path.join(inputfolder, 'elecModels_' + subject + '.pkl')
        lead_data_raw, intensityProfiles, skelSkalms = self.load_leadModel(inputfolder, filename=filename_leadmodel)
        lead_data, sides = self.estimate_hemisphere(lead_data_raw) # determines which side is used in lead_data
        lead_model, default_positions, default_coordinates = GetData.get_default_lead(lead_data_raw[0])  # all in [mm]

        # Get initial data for both leads
        marker, coordinates, trajectory, resize, emp_dist = [{} for _ in range(5)]
        for hemisphere in sides:
            marker[hemisphere], coordinates[hemisphere], trajectory[hemisphere], resize[hemisphere] = \
                GetData.get_leadModel(self, lead_data[hemisphere], default_positions, default_coordinates, side=hemisphere)
            _, emp_dist[hemisphere] = GetData.resize_coordinates(self, coordinates[hemisphere], lead_data[hemisphere])
            _, lead_data[hemisphere]['trajectory'], _, marker_temp = \
                GetData.resolve_coordinates(self, marker[hemisphere], default_coordinates, default_positions,
                                            lead_data[hemisphere], resize_bool=resize[hemisphere],
                                            rszfactor=emp_dist[hemisphere]) # This part doesn't do anything meaningful!

        # Get data for xy-plane estimation/plot
        intensity_matrix, bounding_box, fitvolume = GetData.multiprocessing_xyplanes(self, lead_data, trajectory)

        # Start plotting the 'fixed parts', that is lead model (left) and CT intensitiy planes (middle)
        # fig, grid = self.create_figure()
        # fig, layout = self.create_figure_plotly()
        layout = self.create_figure_plotly()
        data = []
        data = self.plotCTintensitiesPlotLy(intensity_matrix['left'], fitvolume['left'], data)

        fig = go.Figure(
            data=data,
            layout=layout)
        fig.show()
        # https://community.plotly.com/t/moving-the-location-of-a-graph-point-interactively/7161/2 # TODO Next steps
        # CTintensitiesLeft = fig.add_subplot(grid[0:3,0], facecolor='None', projection='3d')
        # CTintensitiesLeft = self.plotCTintensitiesNEW(CTintensitiesLeft, intensity_matrix['left'],
        #                                              bounding_box['left'], fitvolume['left'], trajectory['left'])

        # CTintensitiesRight = fig.add_subplot(grid[0:,2:3], facecolor='None', projection='3d')

        self.interactive_plot(lead_data['right'], intensityProfiles, skelSkalms)

    @staticmethod
    def estimate_hemisphere(lead_data):
        """estimates the available sides and returns a list of all available leads; as all data is in LPS (Left,
        Posterior, Superior) so that side can be deduced from trajectory"""

        sides = []
        renamed_lead_data = dict()
        for info in lead_data:
            side_temp = 'right' if not info['trajectory'][0, 0] > info['trajectory'][-1, 0] else 'left'
            renamed_lead_data[side_temp] = info
            sides.append(side_temp)
        return renamed_lead_data, sides

    def create_figure(self, num_cols=2, num_rows=2):
        """creates a figure which is filled with content in later steps"""
        fig = plt.figure(facecolor=np.append(np.array([.64] * 3), .25), frameon=False)
        col_width = [5, 1]
        row_height = [5, 1]
        grid = gridspec.GridSpec(ncols=int(num_cols), nrows=int(num_rows), figure=fig, width_ratios=col_width,
                                 height_ratios=row_height, hspace=0.08, wspace=0.1)
        return fig, grid

    def create_figure_plotly(self, num_cols=3, num_rows=1):
        """creates a figure which is filled with content in later steps"""
        #fig = make_subplots(
        #    rows=num_rows, cols=num_cols,
        #    column_widths=[0.6, 0.4],
        #    row_heights=[0.4, 0.6],
        #    specs=[[{'is_3d': True}, {'is_3d': True}, {'is_3d': False}]],
        #    subplot_titles=['model', 'data', 'xplane'])

        axis = dict(showbackground=True,
                    backgroundcolor="rgb(230, 230,230)",
                    gridcolor="rgb(255, 255, 255)",
                    zerolinecolor="rgb(255, 255, 255)",
                    )

        layout = go.Layout(
            width=700,
            height=700,
            # scene=go.layout.Scene(xaxis=go.layout.XAxis(axis),
            #               yaxis=go.layout.YAxis(axis),
            #               zaxis=go.layout.ZAxis(axis, range=[-2, 2]),
            #               aspectratio=dict(x=1, y=1, z=1),
            #               ),
        )

#        return fig, layout
        return layout

    def interactive_plot(self, lead_data, intensityProfiles, skelSkalms):
        """ Start plotting routine according to Lead-DBS implementation [ea_autocoord and ea_manualreconstruction]"""

        lead_model, default_positions, default_coordinates = GetData.get_default_lead(lead_data)  # all in [mm]

        # Get estimated positions, coordinates, rotation information and markers for leads
        marker, coordinates, trajectory, resize = GetData.get_leadModel(self, lead_data, default_positions,
                                                                        default_coordinates, 'right')

        #TODO: elecModels should be updated somehow

        coords_temp, emp_dist = GetData.resize_coordinates(self, coordinates, lead_data)
        mean_empdist = np.mean(emp_dist)

        _, lead_data['trajectory'], _, marker_temp = \
            GetData.resolve_coordinates(self, marker, default_coordinates, default_positions, lead_data,
                                        resize_bool=resize, rszfactor=mean_empdist)
        lead_data = GetData.leadInformation_update(marker_temp, lead_data) # TODO doesn't make sense as markers donÄt change at this point

        # Start plotting
        # fig = plt.figure(facecolor=self.getbgsidecolor(side=0), tight_layout=True) # TODO: is it really the side, what is this color actually
        fig = plt.figure(facecolor=np.append(np.array([.64] * 3), .25), tight_layout=True)  # TODO: is it really the side, what is this color actually

        #fig.set_tight_layout(True)
        width = [1, 10, 10, 1]
        height = [1, 25, 1]
        grid = gridspec.GridSpec(ncols=4, nrows=3, figure=fig, width_ratios=width, height_ratios=height,
                                 hspace=0.08, wspace=0.1)

        grid_indices = [(0,-1), (0,-1)]
        self.plotCTintensities(lead_data, coordinates, trajectory, fig, grid, grid_indices)

        grid_indices = [(-1), (-1)]
        self.plot_leadInformation(lead_data, mean_empdist, fig, grid, grid_indices)

        self.plotCTaxial(lead_data, fig, grid, cmap='gist_gray')
        self.plot_leadModel(lead_model, fig, grid)

    @staticmethod
    def get_the_slice(x, y, z, surfacecolor, colorscale='gray', showscale=False):
        return go.Surface(x=x,  # https://plot.ly/python/reference/#surface
                          y=y,
                          z=z,
                          surfacecolor=surfacecolor,
                          colorscale=colorscale,
                          showscale=showscale,
                          name="")

    def plotCTintensitiesPlotLy(self, intensity_matrix, fitvolume, data):
        """function plotting perpendicular images of the intensities obtained from trajectory coordinates; separate
        function needed here as changes occur in trejectories/markers"""

        for idx, (direction, intensity) in enumerate(intensity_matrix.items()):
            matrix2plot = (intensity - np.min(intensity)) / (np.max(intensity) - np.min(intensity))  # Prepare data!
            X, Y, Z = fitvolume[direction][0, :, :], fitvolume[direction][1, :, :], fitvolume[direction][2, :, :]
            slice_temp = self.get_the_slice(X,Y,Z, matrix2plot)
            data.append(slice_temp)

        return data

        # TODO this part goes somewhere else as it is the dynamic part to be changed by using arrows
        # for c in coordinates:
        #     self.plot_coordinates(c, mainax1, marker='p', s=250)

        # Start plotting the markers for the 'head' and 'tail' of the lead
        # marker_plot = dict([(k, r) for k, r in lead_model.items() if k.startswith('marker')])
        # self.plot_coordinates(marker_plot["markers_head"], mainax1, marker='o', facecolor='g', edgecolor='g',
        #                      linewidth=1.5)
        # self.plot_coordinates(marker_plot["markers_tail"], mainax1, marker='o', facecolor='r', edgecolor='r',
        #                      linewidth=1.5)
        # mainax1.set_zlim(np.multiply(mainax1.get_zlim(), [1, .98]))

    @staticmethod
    def color_plotsurface(imat2plot, cmap='gist_gray', alpha=1):
        """defines a colormap which is used for plotting data later """
        import matplotlib.cm as cm
        from matplotlib import colors as colors

        color_dimension = imat2plot  # change to desired fourth dimension
        minn, maxx = color_dimension.min(), color_dimension.max()
        norm = matplotlib.colors.Normalize(minn, maxx)
        m = cm.ScalarMappable(norm=norm, cmap=cmap)
        m.set_array([])
        fcolors = m.to_rgba(color_dimension)
        fcolors[:,:,-1] = np.ones((fcolors[:,:,-1].shape))*alpha

        return fcolors

    @staticmethod
    def plot_coordinates(coordinates, axis, marker='', s=25, facecolor='k', linewidth=1.5, edgecolor=[.9]*3):
        """plots all coordinates of the electrodes estimated from the lead model"""

        axis.scatter(coordinates[0], coordinates[1], coordinates[2], s=s, c=facecolor, edgecolor=edgecolor,
                     marker=marker, linewidth=linewidth)

    @staticmethod
    def plot_leadInformation(lead_data, mean_empdist, fig, grid, grid_indices):
        """plots lead information in the lower right corner"""

        axis = fig.add_subplot(grid[grid_indices[0], grid_indices[1]], facecolor='None')
        fig.tight_layout()
        axis.set_axis_off()
        axis.set_facecolor('None')

        stn_direction = 'right' if not lead_data['trajectory'][0, 0] > lead_data['trajectory'][-1, 0] else 'left'
        text2plot = 'Lead: {} STN \nLead spacing: {:.2f} mm\nRotation:  {} ' \
                    'deg'.format(stn_direction, mean_empdist, lead_data['rotation'])
        axis.text(.6, .25, text2plot, horizontalalignment='center',fontsize=12, ha='center', va='center',
                  bbox=dict(boxstyle='circle', facecolor='#D8D8D8', ec="0.5", pad=0.5, alpha=1), fontweight='bold')

    def plotCTaxial(self, lead_data, fig, grid, cmap='gist_gray'):
        """function plotting axial sclices at the level of head and tail markers"""

        print("\t...extracting axial slices for corresponding markers of CTimaging\t...")
        marker_of_interest = ['markers_head', 'markers_tail']
        marker_plot = dict([(k, r) for k, r in lead_data.items() if k.startswith('marker') and
                            any(z in k for z in marker_of_interest)])
        color_specs = ['g', 'r']
        item = -1
        for marker, coordinates in marker_plot.items():
            item += 1
            intensity_matrix, bounding_box, _ = self.get_axialplanes(coordinates, lead_data=lead_data,
                                                                     window_size=15, resolution=.5)
            transversal_axis = fig.add_subplot(grid[item, -1], facecolor='None')
            transversal_axis.imshow(intensity_matrix, cmap=cmap, extent=[np.min(bounding_box, axis=1)[0],
                                                                         np.max(bounding_box, axis=1)[0],
                                                                         np.min(bounding_box, axis=1)[1],
                                                                         np.max(bounding_box, axis=1)[1]],
                                    interpolation='bicubic')
            transversal_axis.set_axis_off()
            transversal_axis.set_facecolor('None')

            transversal_axis.scatter(coordinates[0], coordinates[1], s=200, c=color_specs[item],
                                     edgecolor=color_specs[item], marker='x', linewidth=1.5)

        grid.tight_layout(fig)

    @staticmethod
    def plot_leadModel(lead_model, fig, grid, lead_data, colors=[.3, .5], items2plot=['insulation', 'contacts']):
        """Plots schematic lead model at the left side in order to visualise what it should look like """

        leadModelSubplot = fig.add_subplot(grid[0:, 0], facecolor='None', projection='3d')
        lead_plot = dict([(k, r) for k, r in lead_model.items() if any(z in k for z in items2plot)])
        max_coords, min_coords = [[] for _ in range(2)]

        for idx, item in enumerate(items2plot):
            for idx_vertices in (range(0, len(lead_plot[item]['vertices']))):
                verts_temp = np.array(lead_plot[item]['vertices'][idx_vertices]) - 1
                faces_temp = np.array(lead_plot[item]['faces'][idx_vertices]) - 1
                mesh = Poly3DCollection(verts_temp[faces_temp], facecolors=[colors[idx]] * 3, edgecolor='none',
                                        rasterized=True)
                leadModelSubplot.add_collection(mesh)

                if item == 'contacts':
                    max_coords.append(np.max(verts_temp[faces_temp], axis=1))
                    min_coords.append(np.min(verts_temp[faces_temp], axis=1))

        marker = {k: [] for k in ['head', 'tail']}
        if lead_data['left']['model'] == 'Boston Vercise Directional':
            for n in range(3):
                marker['tail'].append(lead_model['tail_position'][n])
                marker['head'].append(lead_model['tail_position'][n])
            marker['tail'][2] = np.min(min_coords[0])
        else:
            for n in range(3):
                marker['tail'].append(lead_model['tail_position'][n])
                marker['head'].append(lead_model['head_position'][n])

        color = ['r', 'g']
        for idx, loc in enumerate(['head', 'tail']):
            leadModelSubplot.scatter(marker[loc][0] - 1, marker[loc][1] - 2, marker[loc][2],c=color[idx], s=25)

        leadModelSubplot.set_zlim([0, 15])
        leadModelSubplot.set_ylim(-4, 4)
        leadModelSubplot.set_xlim(-4, 4)
        leadModelSubplot.set_facecolor('None')
        leadModelSubplot.view_init(elev=0, azim=-120)
        leadModelSubplot.set_axis_off()
        # grid.tight_layout(fig) # TODO is this necessary?

    def get_xyplanes(self, trajectory, lead_model, limits=[(-4,4),(-4,4),(-10,20)], sample_width=10,
                     direction=['sag', 'cor']):
        # TODO: this was moved into the GetData class and MUST be removed
        hd_trajectories = self.interpolate_trajectory(trajectory, resolution=10) # TODO: rename traj to trajectory after assigning in function

        slices = {k: [] for k in direction}
        imat = {k: [] for k in direction}
        fitvolume = {k: [] for k in direction}
        bounding_box = {k: [] for k in direction}

        for idx, plane in enumerate(direction):
            slices[plane] = list(range(limits[idx][0], limits[idx][1]+1, 1))
            imat[plane], fitvolume[plane] = self.resample_CTplanes(hd_trajectories, plane, lead_model, resolution=.35)

            span_vector = [sample_width, 0, 0] if plane != 'sag' else [0, sample_width, 0]

            idx = [0, -1]
            bounding_box_coords = []
            for k in idx:
                bounding_box_coords.append(hd_trajectories[k,:]-span_vector)
                bounding_box_coords.append(hd_trajectories[k, :] + span_vector)
            bounding_box_coords = np.array(bounding_box_coords)

            axes_name = ['xx', 'yy', 'zz']
            box = {k: [] for k in axes_name}
            for i, dim in enumerate(axes_name):
                box[dim] = bounding_box_coords.T[i,:].tolist()

            bounding_box[plane] = box

        return imat, bounding_box, fitvolume

    def get_axialplanes(self, marker_coordinates, lead_data, window_size=15, resolution=.5):
        """returns a plane at a specific window with a certain direction"""

        if lead_data['transformation_matrix'].shape[0] == 3:
            lead_data['transformation_matrix'] = np.eye(4)*lead_data['transformation_matrix'][0,0]
            lead_data['transformation_matrix'][-1,-1] = 1
        transformation_matrix = lead_data['transformation_matrix']
        transformation_matrix = np.eye(4)

        bounding_box_coords = []
        for k in range(2):
            bounding_box_coords.append(np.arange(start=marker_coordinates[k]-window_size,
                                                 stop=marker_coordinates[k]+window_size, step=resolution))
        bounding_box_coords.append(np.repeat(marker_coordinates[-1], len(bounding_box_coords[1])))
        bounding_box = np.array(bounding_box_coords)

        meshX, meshY = np.meshgrid(bounding_box[0,:], bounding_box[1,:])
        meshZ = np.repeat(bounding_box[-1,0], len(meshX.flatten()))
        fitvolume_orig = np.array([meshX.flatten(), meshY.flatten(), meshZ.flatten(), np.ones(meshX.flatten().shape)])
        fitvolume = np.linalg.solve(transformation_matrix, fitvolume_orig)
        resampled_points = PlotRoutines.interpolate_CTintensities(lead_data, fitvolume)
        imat = np.reshape(resampled_points, (meshX.shape[0], -1), order='F')

        return imat, bounding_box, fitvolume

    @staticmethod
    def resample_CTplanes(hd_trajectories, direction, lead_data, resolution=.2, sample_width=10, use_transformation_matrix=False):
        """Function resampling intesities of the source imaging to a grid which is later used to visualise the
        leads. [ea_mancor_updatescene lines 264f]"""

        direction = ''.join(direction) if type(direction) == list else direction  # in case direction is entered as list

        if use_transformation_matrix:  #  not necessary as all data in cDBS stay within the LPS coordinate system
            if lead_data['transformation_matrix'].shape[0] == 3:
                lead_data['transformation_matrix'] = np.eye(4)*lead_data['transformation_matrix'][0,0]
                lead_data['transformation_matrix'][-1,-1] = 1
            transformation_matrix = lead_data['transformation_matrix']
        else:
            transformation_matrix = np.eye(4)

        xvec = np.arange(start=-sample_width, stop=sample_width+resolution, step=resolution)
        meanfitline = np.vstack((hd_trajectories.T, np.ones(shape=(1, hd_trajectories.T.shape[1])))) # needed for transformation
        addvolume = np.tile(xvec,(len(meanfitline.T),1))

        fitvolume = []
        for t in range(4):
            fitvolume.append(np.tile(meanfitline[t,:], xvec.shape).reshape(xvec.shape[0], meanfitline.shape[1]).T)
        fitvolume_orig = np.stack(fitvolume)

        if direction == 'cor':
            fitvolume_orig[0,:,:] += addvolume
        elif direction == 'sag':
            fitvolume_orig[1, :, :] += addvolume
        elif direction == 'tra':
            fitvolume_orig[2, :, :] += addvolume

        fitvolume = np.linalg.solve(transformation_matrix, np.reshape(fitvolume_orig, (4, -1), order='F'))
        resampled_points = PlotRoutines.interpolate_CTintensities(lead_data, fitvolume)
        imat = np.reshape(resampled_points, (meanfitline.shape[1], -1), order='F')

        return imat, fitvolume_orig

    # ========================================    Interpolations   ========================================
    @staticmethod
    def interpolate_trajectory(orig_trajectory, resolution=20):
        """interpolates between trajectory points thus creating a „high resolution“ version of it"""

        hd_trajectory = []
        for idx in range(np.array(orig_trajectory).shape[1]):
            f = scipy.interpolate.interp1d(np.linspace(start=1, stop=50), np.array(orig_trajectory)[:, idx])
            hd_trajectory.append(f(np.linspace(start=1, stop=50, num=(-1 + len(orig_trajectory[:, idx]))
                                                                     * resolution + 1)))
        return np.stack(hd_trajectory).T

    def interpolate_CTintensities(lead_model, fitvolume):
        import SimpleITK as sitk

        img = sitk.ReadImage(os.path.join(*lead_model['filenameCTimaging']))
        physical_points = list(map(tuple, fitvolume[0:3,:].T))
        #physical_points = physical_points[0:5]
        num_samples = len(physical_points)
        physical_points = [img.TransformContinuousIndexToPhysicalPoint(pnt) for pnt in physical_points]

        #interp_grid_img = sitk.Image((len(physical_points) *([1] * (img.GetDimension() - 1))), sitk.sitkUInt8)
        interp_grid_img = sitk.Image([num_samples] + [1] * (img.GetDimension() - 1), sitk.sitkUInt8)
        displacement_img = sitk.Image([num_samples] + [1] * (img.GetDimension() - 1), sitk.sitkVectorFloat64,
                                      img.GetDimension())

        for i, pnt in enumerate(physical_points):
            displacement_img[[i] + [0] * (img.GetDimension() - 1)] = np.array(pnt) - np.array(
                interp_grid_img.TransformIndexToPhysicalPoint([i] + [0] * (img.GetDimension() - 1)))

        interpolator_enum = sitk.sitkLinear
        default_output_pixel_value = 0.0
        output_pixel_type = sitk.sitkFloat32 if img.GetNumberOfComponentsPerPixel() == 1 else sitk.sitkVectorFloat32
        resampled_temp = sitk.Resample(img, interp_grid_img, sitk.DisplacementFieldTransform(displacement_img),
                                       interpolator_enum, default_output_pixel_value, output_pixel_type)

        resampled_points = [resampled_temp[x,0,0] for x in range(resampled_temp.GetWidth())]
        debug = False
        if debug:
            for i in range(resampled_temp.GetWidth()):
                print(str(img.TransformPhysicalPointToContinuousIndex(physical_points[i])) + ': ' + str(resampled_temp[[i] + [0] * (img.GetDimension() - 1)]) + '\n')

        return np.array(resampled_points)

    # ====================    General Helper Functions for manual correction   ====================
    @staticmethod
    def load_leadModel(inputdir, filename):
        """Function loading results from [preprocLeadCT.py] which emulates the PaCER toolbox"""

        if not inputdir:
            Output.msg_box(text="No input folder provided, please double-check!", title="Missing input folder")
            return
        elif not os.path.isfile(filename):
            Output.msg_box(text="Models for electrode unavailable, please run detection first!",
                           title="Models not available")
            return
        else:
            with open(filename, "rb") as model: # roughly ea_loadreconstruction in the LeadDBS script
                lead_models = pickle.load(model)
                intensityProfiles = pickle.load(model)
                skelSkalms = pickle.load(model)

        return lead_models, intensityProfiles, skelSkalms

    @staticmethod
    def save_leadModel(lead_models, intensityProfiles, skelSkalms, filename=''):

        if not filename:
            Output.msg_box(text='No filename for saving lead model provided', title='No filename provided')
            return

        with open(filename, "wb") as f:
            pickle.dump(lead_models, f)
            pickle.dump(intensityProfiles, f)
            pickle.dump(skelSkalms, f)

    # ====================    Helper Functions fpr plotting Data   ====================
    @staticmethod
    def getbgsidecolor(side, xray=False):
        """ """
        from matplotlib import colors

        line_cols = matplotlib.cm.get_cmap('Set1', 64) # TODO: maybe a cmap would make sense
        line_cols = colors.rgb_to_hsv(line_cols(np.linspace(0, 1, 64))[:,0:3])
        line_cols[:,-1] = line_cols[:,-1]/3
        if xray:
            line_cols[:, 1] = line_cols[:, 1] / 1.5
            line_cols[:, 2] = line_cols[:, 2] * 1.5

        line_cols = colors.hsv_to_rgb(line_cols)
        col = line_cols[side,:] # TODO: why on earth is this so complicated to get these colors

        return col

class GetData:
    def __init__(self, parent=PlotRoutines):
        self.parent = parent

    @staticmethod
    def leadInformation_update(information2update, lead_data):
        """replaces values of lead_models with updated values; information2update can be marker, rotation, etc."""

        for key_name, val in information2update.items():
            lead_data[key_name] = val
        return lead_data

    # ===================================    Functions related with coordinates   ===================================
    def resize_coordinates(self, lead_coords, lead_data):
        """function which enables resizing cooridnates (e.g. 8 contacts to 4 contacts if needed; additional
        functionality contains """

        if lead_data['model'] == 'Boston Vercise Directional' or 'St Jude 6172' or 'St Jude 6173':
            coordinates = np.zeros((4, 3))
            coordinates[0, :] = lead_coords[0, :]
            coordinates[1, :] = np.mean(lead_coords[1: 4, :], axis=0)
            coordinates[2, :] = np.mean(lead_coords[4: 7, :], axis=0)
            coordinates[3, :] = lead_coords[7, :]

            emp_dist = GetData.lead_dist(coords=coordinates)
        else:
            coordinates = lead_coords
            emp_dist = GetData.lead_dist(coords=coordinates, factor=lead_data['numel'])

        return coordinates, emp_dist

    @staticmethod
    def lead_dist(coords, factor=3):
        """calculate lead distances according to its coordinates"""

        spatial_distance = scipy.spatial.distance.cdist(coords, coords, 'euclidean')
        emp_dist = np.sum(np.sum(np.tril(np.triu(spatial_distance, 1), 1))) / factor

        return emp_dist

    def resolve_coordinates(self, marker, lead_coords_mm, lead_positions, lead_data, resize_bool=False, rszfactor=0):
        """emulates the function from Lead-DBS ea_resolvecoords; unlike in Lead DBS this is done one at a time cf.
        https://github.com/netstim/leaddbs/blob/master/templates/electrode_models/ea_resolvecoords.m"""

        if resize_bool:
            can_dist = np.linalg.norm(lead_positions["head_position"] - lead_positions["tail_position"])
            coords_temp, can_eldist = GetData.resize_coordinates(self, lead_coords_mm, lead_data)

            stretch = can_dist * (rszfactor / can_eldist) if rszfactor != 0 else can_dist
            vec = np.divide((marker["markers_tail"] - marker["markers_head"]),
                            np.linalg.norm(marker["markers_tail"] - marker["markers_head"]))
            marker["markers_tail"] = marker["markers_head"] + vec * stretch

        coords, traj_vector, trajectory, can_eldist = [[] for _ in range(4)]
        if not marker["markers_head"].size==0:
            M = np.stack((np.append(marker["markers_head"], 1), np.append(marker["markers_tail"], 1),
                          np.append(marker["markers_x"], 1), np.append(marker["markers_y"], 1)))
            E = np.stack((np.append(lead_positions["head_position"], 1), np.append(lead_positions["tail_position"], 1),
                          np.append(lead_positions["x_position"], 1), np.append(lead_positions["y_position"], 1)))
            X = np.linalg.lstsq(E, M, rcond=None)

            coords_mm = np.concatenate([lead_coords_mm, np.ones(shape=(lead_coords_mm.shape[0],1))], axis=1)
            coords = (coords_mm @ X[0]).T
            coords = coords[0: 3,:].T

            traj_vector = (marker["markers_tail"] - marker["markers_head"]) / \
                          np.linalg.norm(marker["markers_tail"] - marker["markers_head"])

            trajectory = np.stack((marker["markers_head"] - traj_vector*5, marker["markers_head"] + traj_vector*25))
            trajectory = np.array((np.linspace(trajectory[0, 0], trajectory[1, 0], num=50),
                                   np.linspace(trajectory[0, 1], trajectory[1, 1], num=50),
                                   np.linspace(trajectory[0, 2], trajectory[1, 2], num=50))).T

        return coords, trajectory, can_eldist, marker

    # ===================================    Functions in combination with leads   ===================================
    @staticmethod
    def get_default_lead(lead_data):
        """obtains default lead properties according to the model proposed in the PaCER algorithm @ ./template"""

        if lead_data['model'] == 'Boston Vercise Directional': # load mat-file to proceed
            mat_filename = 'boston_vercise_directed.mat'
            lead_model = loadmat(os.path.join(ROOTDIR, 'ext', 'LeadDBS', mat_filename), 'r')['electrode']
            default_positions = {x: np.hstack(vals) for x, vals in lead_model.items() if x.endswith('position')}
            default_coordinates = np.array(lead_model['coords_mm'])  # in [mm]
        else:
            Output.msg_box(text="Lead type not yet implemented.", title="Lead type not implemented")
            return

        return lead_model, default_positions, default_coordinates

    def get_leadModel(self, lead_data, default_positions, default_coordinates, side, resize=False):
        """ reads and estimates all necessary markers and trajectories for the corresponding lead model """
        print("\t... reading lead data properties for {} side and estimating rotation".format(side))

        marker_unprocessed = dict([(k, r) for k, r in lead_data.items() if k.startswith('marker')])

        if not (lead_data['first_run'] and lead_data['manual_correction']):
            resize = True  # TODO: Not sure if this is doing the job; some warning/information should be displayed that > first run
            lead_data['first_run'] = False

        _, lead_data['trajectory'], _, marker_temp = \
            GetData.resolve_coordinates(self, marker_unprocessed, default_coordinates, default_positions, lead_data,
                                        resize_bool=resize)

        lead_data['rotation'] = GetData.initialise_rotation(lead_data, marker_temp)
        xvec, yvec, lead_data['rotation'], marker_rotation = GetData.estimate_rotation(lead_data, marker_temp)
        lead_data = GetData.leadInformation_update(marker_rotation, lead_data)

        if xvec.size == 0 or yvec.size == 0:
            xvec, yvec, lead_data['rotation'], marker_rotation = GetData.estimate_rotation(lead_data, marker_rotation)
            lead_data = GetData.leadInformation_update(marker_rotation, lead_data)

        options2process = {'xvec': [1, 0, 0], 'yvec': [0, 1, 0]}
        unrot = {k: [] for k in options2process.keys()}
        for key in options2process:
            vec_temp = np.cross(lead_data["normtraj_vector"], options2process[key])
            unrot[key] = np.divide(vec_temp, np.linalg.norm(vec_temp))

        marker = dict([(k, r) for k, r in lead_data.items() if k.startswith('marker')])
        coordinates, trajectory, _, _ = GetData.resolve_coordinates(self, marker, default_coordinates,
                                                                    default_positions, lead_data, resize_bool=False)  # ea_mancor_updatescene line 144
        return marker, coordinates, trajectory, resize

    def worker_xy_planes(self, traj, lead, dimension, side, IM_dict, BB_dict, FV_dict, queue):
        print("\t...extracting intensities for CTimages @ {} side\t...".format(side))
        queue.put('.')
        intensity_matrix, bounding_box, fitvolume = GetData.get_xyplanes(self, trajectory=traj, lead_model=lead,
                                                                         direction=dimension)
        IM_dict[side] = intensity_matrix
        BB_dict[side] = bounding_box
        FV_dict[side] = fitvolume

    def multiprocessing_xyplanes(self, lead_data, trajectory, dimension=['sag', 'cor']):
        """extracts the intensities corresponding to the CTimaging, which is used as 'background' in the figure"""

        queue = mp.Queue()
        start_multi = time.time()
        manager = mp.Manager()
        IM_dict, BB_dict, FV_dict = manager.dict(), manager.dict(), manager.dict()
        processes = [mp.Process(target=GetData.worker_xy_planes,
                                args=(self, trajectory[side], lead_data[side], dimension, side,
                                      IM_dict, BB_dict, FV_dict, queue))
                     for side in trajectory.keys()]
        _ = [p.start() for p in processes]

        #while any([p.is_alive() for p in processes]):
        #    while not queue.empty():
        #        status = queue.get()
        #        #print("{}".format(status))
        #    time.sleep(0.1)

        _ = [p.join(timeout=1) for p in processes]
        _ = [queue.get() for p in processes]

        print("\n\tIntensities successfully extracted for {} side(s) in "
              "{:.2f} secs".format(len(list(trajectory.keys())), time.time() - start_multi))

        return IM_dict, BB_dict, FV_dict

    def multiprocessing_xyplanesPOOL(self, lead_data, trajectory, dimension=['sag', 'cor']):
        """extracts the intensities corresponding to the CTimaging, which is used as 'background' in the figure"""

        start_multi = time.time()
        manager = mp.Manager()
        queue = mp.Queue()
        IM_dict, BB_dict, FV_dict = manager.dict(), manager.dict(), manager.dict()
        pool = mp.Pool(processes=4)
        processes = [pool.apply_async(func=GetData.worker_xy_planes,
                                      args=(self, trajectory[side], lead_data[side], dimension, side,
                                            IM_dict, BB_dict, FV_dict, queue, ))
                     for side in trajectory.keys()]
        pool.close()
        pool.join()
        final_result = [process.get() for process in processes]

        # prevent adding anything more to the queue and wait for queue to empty
        #_ = [p.start() for p in processes]
        #_ = [p.join(timeout=1) for p in processes]

        print("\n\tIntensities successfully extracted for {} side(s) in "
              "{:.2f} secs".format(len(list(trajectory.keys())), time.time() - start_multi))

        return IM_dict, BB_dict, FV_dict

    def get_xyplanes(self, trajectory, lead_model, limits=[(-4,4),(-4,4),(-10,20)], sample_width=10,
                     direction=['sag', 'cor']):
        # TODO: this was moved into the GetData class and MUST be removed
        hd_trajectories = Imaging.interpolate_trajectory(trajectory, resolution=10) # TODO: rename traj to trajectory after assigning in function

        slices = {k: [] for k in direction}
        imat = {k: [] for k in direction}
        fitvolume = {k: [] for k in direction}
        bounding_box = {k: [] for k in direction}

        for idx, plane in enumerate(direction):
            slices[plane] = list(range(limits[idx][0], limits[idx][1]+1, 1))
            imat[plane], fitvolume[plane] = self.resample_CTplanes(hd_trajectories, plane, lead_model, resolution=.35)

            span_vector = [sample_width, 0, 0] if plane != 'sag' else [0, sample_width, 0]

            idx = [0, -1]
            bounding_box_coords = []
            for k in idx:
                bounding_box_coords.append(hd_trajectories[k,:]-span_vector)
                bounding_box_coords.append(hd_trajectories[k, :] + span_vector)
            bounding_box_coords = np.array(bounding_box_coords)

            axes_name = ['xx', 'yy', 'zz']
            box = {k: [] for k in axes_name}
            for i, dim in enumerate(axes_name):
                box[dim] = bounding_box_coords.T[i,:].tolist()

            bounding_box[plane] = box

        return imat, bounding_box, fitvolume

    # ==============================    Functions related to estimation of rotation   ==============================
    @staticmethod
    def initialise_rotation(lead_model, marker):
        """script iniitalising the estimation of rotation angles; necessary as at the beginning there is no
        information available; This function is followed by estimate_rotation.py (see below)"""

        if lead_model['manual_correction'] and not lead_model['rotation']:
            vec_temp = marker['markers_y'] - marker['markers_head']
            vec_temp[2] = 0
            vec_temp = np.divide(vec_temp, np.linalg.norm(vec_temp))
            initial_rotation = np.degrees(math.atan2(np.linalg.norm(np.cross([0,1,0], vec_temp)),
                                                     np.dot([0,1,0], vec_temp)))
            if marker['markers_y'][0] > marker['markers_head'][0]:
                initial_rotation = - initial_rotation
            rotation = initial_rotation
        elif not lead_model['manual_correction'] and not lead_model['rotation']:
            rotation = 0

        return rotation

    @staticmethod
    def estimate_rotation(lead_models, marker):
        """determination of rotation according to markers provided; follows steps in ea_mancor_updatescene of
         Lead-DBS package (cf. https://github.com/ningfei/lead/blob/develop/ea_mancor_updatescene.m) """

        rotation, normtrajvector = lead_models['rotation'], lead_models["normtraj_vector"]

        yvec = np.zeros((3,1))
        yvec[0] = -np.cos(0) * np.sin(np.deg2rad(rotation))
        yvec[1] = (np.cos(0) * np.cos(np.deg2rad(rotation))) + (np.sin(0) * np.sin(np.deg2rad(rotation)) * np.sin(0))
        yvec[2] = (-np.sin(0) * np.cos(np.deg2rad(rotation))) + (np.cos(0) * np.sin(np.deg2rad(rotation)) * np.sin(0))

        xvec = np.cross(yvec.T, [0,0,1])
        xvec = xvec - (np.dot(xvec, normtrajvector) / np.linalg.norm(normtrajvector)**2) * normtrajvector
        xvec = np.divide(xvec, np.linalg.norm(xvec))
        yvec = -np.cross(xvec, normtrajvector)

        marker['markers_x'] = marker['markers_head'] + (xvec * lead_models['lead_diameter'] / 2)
        marker['markers_y'] = marker['markers_head'] + (yvec * lead_models['lead_diameter'] / 2)

        return xvec, yvec, rotation, marker


<reponame>MayerJohannes/Klimamodelle
# 	Validation of the analog-method results using spartacus RR
#	<NAME>, <NAME>
#	Klimamodelle SS19
#	version:	2019-07-23
#
#	Execute:
#		python3 04-validation_iterate.py --save

import numpy as np
import xarray as xr
import dask
import matplotlib.pylab as plt
import iris
from eofs.xarray import Eof
from scipy.signal import detrend
from eofs.multivariate.iris import MultivariateEof
import datetime
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import argparse
import csv
import scipy as sp
import pandas as pd
import numpy.ma as ma

# Parse target day as argument
parser = argparse.ArgumentParser(description='man')
parser.add_argument('--save', action='store_true', dest='ssave', help='save output')
args = parser.parse_args()
ssave = args.ssave

# Spartacus first and last year
syear = 1961
eyear = 2017

# Load date list (= EOF.py output)
with open('datelist-v4.csv', 'r') as f:
	reader = csv.reader(f)
	your_list = list(reader)

if ssave: file1 = csv.writer(open("corrcoef-2.csv","w"))

# Iterate over all entries
max_cc = 0
for istr in your_list:
	#if istr[0] == TD: break

	str_td = istr[0] #'1981-01-01'
	str_a1 = istr[1] #'2005-01-02'
	str_a2 = istr[2] #'2012-01-04'

	# Target day
	year_td = str_td[:4]

	# Analogon 1
	year_a1 = str_a1[:4]

	# Analogon 2
	year_a2 = str_a2[:4]

	if int(year_td) < syear or  int(year_a1) < syear or int(year_a2) < syear: continue
	if int(year_td) > eyear or  int(year_a1) > eyear or int(year_a2) > eyear: continue

	ifile_td = '/home/johannes/klimamodelle/data/spartacus_rr/RR'+year_td+'.nc'
	ifile_a1 = '/home/johannes/klimamodelle/data/spartacus_rr/RR'+year_a1+'.nc'
	ifile_a2 = '/home/johannes/klimamodelle/data/spartacus_rr/RR'+year_a2+'.nc'

	idata_td = xr.open_mfdataset(ifile_td,chunks={'time': 366})
	idata_a1 = xr.open_mfdataset(ifile_a1,chunks={'time': 366})
	idata_a2 = xr.open_mfdataset(ifile_a2,chunks={'time': 366})

	td_field = idata_td.sel(time=slice(str_td,str_td))
	a1_field = idata_a1.sel(time=slice(str_a1,str_a1))
	a2_field = idata_a2.sel(time=slice(str_a2,str_a2))

	td_values = td_field.RR.values[:,:][0]
	a1_values = a1_field.RR.values[:,:][0]
	a2_values = a2_field.RR.values[:,:][0]

	# np.ravel : array to vector transformation
	# masked_invalid : Mask an array where invalid values occur
	# ma.corrcoef : compute corr coeff. matrix, considers missing values

	# Correlation TD and best analogon
	cc_a1 = ma.corrcoef(ma.masked_invalid(np.ravel(td_values)), ma.masked_invalid(np.ravel(a1_values)))[0,1]

	# Correlation TD and second best analogon
	cc_a2 = ma.corrcoef(ma.masked_invalid(np.ravel(td_values)), ma.masked_invalid(np.ravel(a2_values)))[0,1]
	print(str_td, cc_a1, cc_a2)

	# Write output
	if ssave: file1.writerow([str_td, cc_a1, cc_a2])

	# Compute largest correlation coefficient
	if cc_a1 > max_cc:
		max_cc = cc_a1
		max_td = str_td[:]

	if cc_a2 > max_cc:
		max_cc = cc_a2
		max_td = str_td

print(max_cc,max_td)