#! /usr/bin/env python

"""
Module with local/smart PCA (annulus or patch-wise) model PSF subtraction for
ADI, ADI+SDI (IFS) and ADI+RDI datasets. This implementation make use of
Python multiprocessing capabilities.
"""

from __future__ import division, print_function

__author__ = 'Carlos Alberto Gomez Gonzalez'
__all__ = ['pca_annular',
           'pca_rdi_annular']

import numpy as np
from scipy import stats
from multiprocessing import cpu_count
from ..preproc import (cube_derotate, cube_collapse, check_pa_vector,
                       check_scal_vector)
from ..preproc import cube_rescaling_wavelengths as scwave
from ..preproc.derotation import _find_indices_adi, _define_annuli
from ..preproc.rescaling import _find_indices_sdi
from ..conf import time_ini, timing
from ..conf.utils_conf import pool_map, fixed
from ..var import get_annulus_segments, matrix_scaling
from ..stats import descriptive_stats
from .svd import get_eigenvectors

def pca_annular(cube, angle_list, scale_list=None, radius_int=0, fwhm=4,
                asize=4, n_segments=1, delta_rot=1, delta_sep=(0.1, 1), ncomp=1,
                ncomp2=1, svd_mode='lapack', nproc=1, min_frames_lib=2,
                max_frames_lib=200, tol=1e-1, scaling=None, imlib='opencv',
                interpolation='lanczos4', collapse='median', full_output=False,
                verbose=True):
    """ PCA model PSF subtraction for ADI and ADI + mSDI (IFS) data. The PCA
    model is computed locally in each annulus (or annular sectors according to
    ``n_segments``). For each sector we discard reference frames taking into
    account a parallactic angle threshold (``delta_rot``) and a radial movement
    threshold (``delta_sep``).

    Parameters
    ----------
    cube : array_like, 3d or 4d
        Input cube.
    angle_list : array_like, 1d
        Corresponding parallactic angle for each frame.
    scale_list :
        Scaling factors in case of IFS data (ADI+mSDI cube). Usually, the
        scaling factors are the central channel wavelength divided by the
        shortest wavelength in the cube (more thorough approaches can be used
        to get the scaling factors). This scaling factors are used to re-scale
        the spectral channels and align the speckles.
    radius_int : int, optional
        The radius of the innermost annulus. By default is 0, if >0 then the
        central circular area is discarded.
    fwhm : float, optional
        Known size of the FHWM in pixels to be used. Default is 4.
    asize : float, optional
        The size of the annuli, in pixels.
    n_segments : int or list of ints or 'auto', optional
        The number of segments for each annulus. When a single integer is given
        it is used for all annuli. When set to 'auto', the number of segments is
        automatically determined for every annulus, based on the annulus width.
    delta_rot : float, optional
        Factor for increasing the parallactic angle threshold, expressed in
        FWHM. Default is 1 (excludes 1 FHWM on each side of the considered
        frame). According to Absil+13, a slightly better contrast can be reached
        for the innermost annuli if we consider a ``delta_rot`` condition as
        small as 0.1 lambda/D. This is because at very small separation, the
        effect of speckle correlation is more significant than self-subtraction.
    delta_sep : float or tuple of floats, optional
        The threshold separation in terms of the mean FWHM (for ADI+mSDI data).
        If a tuple of two values is provided, they are used as the lower and
        upper intervals for the threshold (grows as a function of the
        separation).
    ncomp : int or list or 1d numpy array, optional
        How many PCs are used as a lower-dimensional subspace to project the
        target (sectors of) frames. If ``auto`` it will be automatically
        determined. If ``cube`` is a 3d array (ADI), ``ncomp`` can be a list,
        in which case a different number of PCs will be used for each annulus
        (starting with the innermost one). If ``cube`` is a 4d array, then
        ``ncomp`` is the number of PCs obtained from each multi-spectral frame
        (for each sector).
    ncomp2 : int, optional
        Only used for ADI+mSDI (4d) cubes. ``ncomp2`` sets the number of PCs
        used in the second PCA stage (ADI fashion, using the residuals of the
        first stage). If None then the second PCA stage is skipped and the
        residuals are de-rotated and combined.
    mode : {'lapack', 'arpack', 'eigen', 'randsvd', 'cupy', 'eigencupy',
            'randcupy', 'pytorch', 'eigenpytorch', 'randpytorch'}, str optional
        Switch for the SVD method/library to be used. ``lapack`` uses the LAPACK
        linear algebra library through Numpy and it is the most conventional way
        of computing the SVD (deterministic result computed on CPU). ``arpack``
        uses the ARPACK Fortran libraries accessible through Scipy (computation
        on CPU). ``eigen`` computes the singular vectors through the
        eigendecomposition of the covariance M.M' (computation on CPU).
        ``randsvd`` uses the randomized_svd algorithm implemented in Sklearn
        (computation on CPU). ``cupy`` uses the Cupy library for GPU computation
        of the SVD as in the LAPACK version. ``eigencupy`` offers the same
        method as with the ``eigen`` option but on GPU (through Cupy).
        ``randcupy`` is an adaptation of the randomized_svd algorithm, where all
        the computations are done on a GPU (through Cupy). ``pytorch`` uses the
        Pytorch library for GPU computation of the SVD. ``eigenpytorch`` offers
        the same method as with the ``eigen`` option but on GPU (through
        Pytorch). ``randpytorch`` is an adaptation of the randomized_svd
        algorithm, where all the linear algebra computations are done on a GPU
        (through Pytorch).
    nproc : None or int, optional
        Number of processes for parallel computing. If None the number of
        processes will be set to (cpu_count()/2).
    min_frames_lib : int, optional
        Minimum number of frames in the PCA reference library.
    max_frames_lib : int, optional
        Maximum number of frames in the PCA reference library for annuli beyond
        10*FWHM. The more distant/decorrelated frames are removed from the
        library.
    tol : float, optional
        Stopping criterion for choosing the number of PCs when ``ncomp``
        is None. Lower values will lead to smaller residuals and more PCs.
    scaling : {None, 'temp-mean', 'spat-mean', 'temp-standard', 'spat-standard'}
        With None, no scaling is performed on the input data before SVD. With
        "temp-mean" then temporal px-wise mean subtraction is done, with
        "spat-mean" then the spatial mean is subtracted, with "temp-standard"
        temporal mean centering plus scaling to unit variance is done and with
        "spat-standard" spatial mean centering plus scaling to unit variance is
        performed.
    imlib : str, optional
        See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
    interpolation : str, optional
        See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
    collapse : {'median', 'mean', 'sum', 'trimmean'}, str optional
        Sets the way of collapsing the frames for producing a final image.
    full_output: boolean, optional
        Whether to return the final median combined image only or with other
        intermediate arrays.
    verbose : bool, optional
        If True prints to stdout intermediate info.

    Returns
    -------
    frame : array_like, 2d
        Median combination of the de-rotated cube.
    If full_output is True:
    array_out : array_like, 3d
        Cube of residuals.
    array_der : array_like, 3d
        Cube residuals after de-rotation.

    """
    if verbose:
        global start_time
        start_time = time_ini()

    # ADI datacube
    if cube.ndim == 3:
        res = _pca_adi_ann(cube, angle_list, radius_int, fwhm, asize,
                           n_segments, delta_rot, ncomp, svd_mode, nproc,
                           min_frames_lib, max_frames_lib, tol, scaling, imlib,
                           interpolation, collapse, full_output, verbose)

        if verbose:
            print('Done derotating and combining.')
            timing(start_time)
        if full_output:
            cube_out, cube_der, frame = res
            return cube_out, cube_der, frame
        else:
            return res

    # ADI+mSDI (IFS) datacubes
    elif cube.ndim == 4:
        global ARRAY
        ARRAY = cube

        z, n, y_in, x_in = cube.shape
        fwhm = int(np.round(np.mean(fwhm)))
        n_annuli = int((y_in / 2 - radius_int) / asize)

        if scale_list is None:
            raise ValueError('Scaling factors vector must be provided')
        else:
            if np.array(scale_list).ndim > 1:
                raise ValueError('Scaling factors vector is not 1d')
            if not scale_list.shape[0] == z:
                raise ValueError('Scaling factors vector has wrong length')

        if verbose:
            print('First PCA subtraction exploiting the spectral variability')
            print('{} spectral channels per IFS frame'.format(z))
            print('N annuli = {}, mean FWHM = {:.3f}'.format(n_annuli, fwhm))

        res = pool_map(nproc, _pca_sdi_fr, fixed(range(n)), scale_list,
                       radius_int, fwhm, asize, n_segments, delta_sep, ncomp,
                       svd_mode, tol, scaling, imlib, interpolation, collapse,
                       verbose=verbose)
        residuals_cube_channels = np.array(res)

        # Exploiting rotational variability
        if verbose:
            timing(start_time)
            print('{} ADI frames'.format(n))

        if ncomp2 is None:
            if verbose:
                msg = 'Skipping the second PCA subtraction'
                print(msg)

            cube_out = residuals_cube_channels
            cube_der = cube_derotate(cube_out, angle_list, imlib=imlib,
                                     interpolation=interpolation)
            frame = cube_collapse(cube_der, mode=collapse)

        else:
            if verbose:
                msg = 'Second PCA subtraction exploiting the angular '
                msg += 'variability'
                print(msg)

            res = _pca_adi_ann(residuals_cube_channels, angle_list, radius_int,
                               fwhm, asize, n_segments, delta_rot, ncomp2,
                               svd_mode, nproc, min_frames_lib, max_frames_lib,
                               tol, scaling, imlib, interpolation, collapse,
                               full_output, verbose)
            if full_output:
                cube_out, cube_der, frame = res
            else:
                frame = res

        if verbose:
            print('Done derotating and combining.')
            timing(start_time)
        if full_output:
            return cube_out, cube_der, frame
        else:
            return frame

    else:
        raise TypeError('Input array is not a cube or 3d array')

def pca_rdi_annular(cube, angle_list, cube_ref, radius_int=0, asize=1, ncomp=1,
                    svd_mode='lapack', min_corr=0.9, fwhm=4,
                    scaling='temp-standard', imlib='opencv',
                    interpolation='lanczos4', collapse='median',
                    full_output=False, verbose=True):
    """ Annular PCA with Reference Library + Correlation + standardization

    In the case of having a large number of reference images, e.g. for a survey
    on a single instrument, we can afford a better selection of the library by
    constraining the correlation with the median of the science dataset and by
    working on an annulus-wise way. As with other local PCA algorithms in VIP
    the number of principal components can be automatically adjusted by the
    algorithm by minmizing the residuals in the given patch (a la LOCI).

    Parameters
    ----------
    cube : array_like, 3d
        Input science cube.
    angle_list : array_like, 1d
        Corresponding parallactic angle for each frame.
    cube_ref : array_like, 3d
        Reference library cube. For Reference Star Differential Imaging.
    radius_int : int, optional
        The radius of the innermost annulus. By default is 0, if >0 then the
        central circular area is discarded.
    asize : float, optional
        The size of the annuli, in FWHM. Default is 3.
    ncomp : int, optional
        How many PCs are kept. If none it will be automatically determined.
    svd_mode : {'lapack', 'arpack', 'eigen', 'randsvd', 'cupy', 'eigencupy', 'randcupy'}, str
        Switch for the SVD method/library to be used. ``lapack`` uses the LAPACK
        linear algebra library through Numpy and it is the most conventional way
        of computing the SVD (deterministic result computed on CPU). ``arpack``
        uses the ARPACK Fortran libraries accessible through Scipy (computation
        on CPU). ``eigen`` computes the singular vectors through the
        eigendecomposition of the covariance M.M' (computation on CPU).
        ``randsvd`` uses the randomized_svd algorithm implemented in Sklearn
        (computation on CPU). ``cupy`` uses the Cupy library for GPU computation
        of the SVD as in the LAPACK version. ``eigencupy`` offers the same
        method as with the ``eigen`` option but on GPU (through Cupy).
        ``randcupy`` is an adaptation of the randomized_svd algorith, where all
        the computations are done on a GPU.
    min_corr : int, optional
        Level of linear correlation between the library patches and the median
        of the science. Deafult is 0.9.
    fwhm : float, optional
        Known size of the FHWM in pixels to be used. Deafult is 4.
    scaling : {None, 'temp-mean', 'spat-mean', 'temp-standard', 'spat-standard'}
        With None, no scaling is performed on the input data before SVD. With
        "temp-mean" then temporal px-wise mean subtraction is done, with
        "spat-mean" then the spatial mean is subtracted, with "temp-standard"
        temporal mean centering plus scaling to unit variance is done and with
        "spat-standard" spatial mean centering plus scaling to unit variance is
        performed.
    imlib : str, optional
        See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
    interpolation : str, optional
        See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
    collapse : {'median', 'mean', 'sum', 'trimmean'}, str optional
        Sets the way of collapsing the frames for producing a final image.
    full_output: boolean, optional
        Whether to return the final median combined image only or with other
        intermediate arrays.
    verbose : {True, False}, bool optional
        If True prints to stdout intermediate info.

    Returns
    -------
    frame : array_like, 2d
        Median combination of the de-rotated cube.
    If full_output is True:
    array_out : array_like, 3d
        Cube of residuals.
    array_der : array_like, 3d
        Cube residuals after de-rotation.

    """
    def define_annuli(angle_list, ann, n_annuli, fwhm, radius_int,
                      annulus_width,
                      verbose):
        """ Defining the annuli """
        if ann == n_annuli - 1:
            inner_radius = radius_int + (ann * annulus_width - 1)
        else:
            inner_radius = radius_int + ann * annulus_width
        ann_center = (inner_radius + (annulus_width / 2.0))

        if verbose:
            msg2 = 'Annulus {}, Inn radius = {:.2f}, Ann center = {:.2f} '
            print(msg2.format(int(ann + 1), inner_radius, ann_center))
        return inner_radius, ann_center

    def fr_ref_correlation(vector, matrix):
        """ Getting the correlations """
        lista = []
        for i in range(matrix.shape[0]):
            pears, _ = stats.pearsonr(vector, matrix[i])
            lista.append(pears)

        return lista

    def do_pca_annulus(ncomp, matrix, svd_mode, noise_error, data_ref):
        """ PCA for given annulus """
        V = get_eigenvectors(ncomp, matrix, svd_mode,
                             noise_error=noise_error,
                             data_ref=data_ref, debug=False)
        # new variables as linear combinations of the original variables in
        # matrix.T with coefficientes from EV
        transformed = np.dot(V, matrix.T)
        reconstructed = np.dot(V.T, transformed)
        residuals = matrix - reconstructed.T
        return residuals, V.shape[0]

    #---------------------------------------------------------------------------
    array = cube
    array_ref = cube_ref
    if array.ndim != 3:
        raise TypeError('Input array is not a cube or 3d array.')
    if array.shape[0] != angle_list.shape[0]:
        raise TypeError(
            'Input vector or parallactic angles has wrong length.')

    n, y, _ = array.shape
    if verbose:  start_time = time_ini()

    angle_list = check_pa_vector(angle_list)

    annulus_width = asize * fwhm  # equal size for all annuli
    n_annuli = int(np.floor((y / 2 - radius_int) / annulus_width))
    if verbose:
        msg = '# annuli = {}, Ann width = {}, FWHM = {:.3f}\n'
        print(msg.format(n_annuli, annulus_width, fwhm))
        print('PCA will be done locally per annulus and per quadrant.\n')

    cube_out = np.zeros_like(array)
    for ann in range(n_annuli):
        inner_radius, _ = define_annuli(angle_list, ann, n_annuli, fwhm,
                                        radius_int, annulus_width, verbose)
        indices = get_annulus(array[0], inner_radius, annulus_width,
                              output_indices=True)
        yy = indices[0]
        xx = indices[1]

        matrix = array[:, yy, xx]  # shape [nframes x npx_ann]
        matrix_ref = array_ref[:, yy, xx]

        corr = fr_ref_correlation(np.median(matrix, axis=0), matrix_ref)
        indcorr = np.where(np.abs(corr) >= min_corr)
        data_ref = matrix_ref[indcorr]
        nfrslib = data_ref.shape[0]

        if nfrslib < 5:
            msg = 'Too few frames left (<5) fulfill the given correlation'
            msg += ' level. Try decreasing it'
            raise RuntimeError(msg)

        matrix = matrix_scaling(matrix, scaling)
        data_ref = matrix_scaling(data_ref, scaling)

        residuals, ncomps = do_pca_annulus(ncomp, matrix, svd_mode, 10e-3,
                                           data_ref)
        cube_out[:, yy, xx] = residuals

        if verbose in [1, 2]:
            print('# frames in LIB = {}'.format(nfrslib))
            print('# PCs = {}'.format(ncomps))
            print('Done PCA with {} for current annulus'.format(svd_mode))
            timing(start_time)

    cube_der = cube_derotate(cube_out, angle_list, imlib=imlib,
                             interpolation=interpolation)
    frame = cube_collapse(cube_der, mode=collapse)
    if verbose:
        print('Done derotating and combining.')
        timing(start_time)
    if full_output:
        return cube_out, cube_der, frame
    else:
        return frame

################################################################################
# Help functions (encapsulating portions of the main algorithm)
################################################################################

def _pca_sdi_fr(fr, wl, radius_int, fwhm, asize, n_segments, delta_sep,
                ncomp, svd_mode, tol, scaling, imlib, interpolation, collapse):
    """ Optimized PCA subtraction on a multi-spectral frame (IFS data).
    """
    z, n, y_in, x_in = ARRAY.shape

    scale_list = check_scal_vector(wl)
    # rescaled cube, aligning speckles
    multispec_fr = scwave(ARRAY[:, fr, :, :], scale_list,
                          imlib=imlib, interpolation=interpolation)[0]

    # Exploiting spectral variability (radial movement)
    fwhm = int(np.round(np.mean(fwhm)))
    n_annuli = int((y_in / 2 - radius_int) / asize)

    if isinstance(n_segments, int):
        n_segments = [n_segments for _ in range(n_annuli)]
    elif n_segments == 'auto':
        n_segments = list()
        n_segments.append(2)  # for first annulus
        n_segments.append(3)  # for second annulus
        ld = 2 * np.tan(360 / 4 / 2) * asize
        for i in range(2, n_annuli):  # rest of annuli
            radius = i * asize
            ang = np.rad2deg(2 * np.arctan(ld / (2 * radius)))
            n_segments.append(int(np.ceil(360 / ang)))

    cube_res = np.zeros_like(multispec_fr)    # shape (z, resc_y, resc_x)

    if isinstance(delta_sep, tuple):
        delta_sep_vec = np.linspace(delta_sep[0], delta_sep[1], n_annuli)
    else:
        delta_sep_vec = [delta_sep] * n_annuli

    for ann in range(n_annuli):
        if ann == n_annuli - 1:
            inner_radius = radius_int + (ann * asize - 1)
        else:
            inner_radius = radius_int + ann * asize
        ann_center = inner_radius + (asize / 2)

        indices = get_annulus_segments(multispec_fr[0], inner_radius, asize,
                                       n_segments[ann])
        # Library matrix is created for each segment and scaled if needed
        for seg in range(n_segments[ann]):
            yy = indices[seg][0]
            xx = indices[seg][1]
            matrix = multispec_fr[:, yy, xx]  # shape (z, npx_annsegm)
            matrix = matrix_scaling(matrix, scaling)

            for j in range(z):
                indices_left = _find_indices_sdi(wl, ann_center, j,
                                                 fwhm, delta_sep_vec[ann])
                matrix_ref = matrix[indices_left]
                curr_frame = matrix[j]  # current frame
                V = get_eigenvectors(ncomp, matrix_ref, svd_mode,
                                     noise_error=tol, debug=False)
                transformed = np.dot(curr_frame, V.T)
                reconstructed = np.dot(transformed.T, V)
                residuals = curr_frame - reconstructed
                # return residuals, V.shape[0], matrix_ref.shape[0]
                cube_res[j, yy, xx] = residuals

    frame_desc = scwave(cube_res, scale_list, full_output=False, inverse=True,
                        y_in=y_in, x_in=x_in, imlib=imlib,
                        interpolation=interpolation, collapse=collapse)
    return frame_desc

def _pca_adi_ann(cube, angle_list, radius_int=0, fwhm=4, asize=2, n_segments=1,
                 delta_rot=1, ncomp=1, svd_mode='lapack', nproc=None,
                 min_frames_lib=2, max_frames_lib=200, tol=1e-1, scaling=None,
                 imlib='opencv', interpolation='lanczos4', collapse='median',
                 full_output=False, verbose=1):
    """ PCA exploiting angular variability (ADI fashion).
    """
    array = cube
    if array.ndim != 3:
        raise TypeError('Input array is not a cube or 3d array')
    if array.shape[0] != angle_list.shape[0]:
        raise TypeError('Input vector or parallactic angles has wrong length')

    n, y, _ = array.shape

    angle_list = check_pa_vector(angle_list)
    n_annuli = int((y / 2 - radius_int) / asize)

    if isinstance(n_segments, int):
        n_segments = [n_segments for _ in range(n_annuli)]
    elif n_segments == 'auto':
        n_segments = list()
        n_segments.append(2)  # for first annulus
        n_segments.append(3)  # for second annulus
        ld = 2 * np.tan(360 / 4 / 2) * asize
        for i in range(2, n_annuli):  # rest of annuli
            radius = i * asize
            ang = np.rad2deg(2 * np.arctan(ld / (2 * radius)))
            n_segments.append(int(np.ceil(360 / ang)))

    if verbose:
        msg = '# annuli = {}, Ann width = {}, FWHM = {:.3f}'
        print(msg.format(n_annuli, asize, fwhm))
        print('PCA per annulus (or annular sectors):')

    if nproc is None:   # Hyper-threading "duplicates" the cores -> cpu_count/2
        nproc = cpu_count() // 2

    # The annuli are built, and the corresponding PA thresholds for frame
    # rejection are calculated (at the center of the annulus)
    cube_out = np.zeros_like(array)
    for ann in range(n_annuli):
        if isinstance(ncomp, list) or isinstance(ncomp, np.ndarray):
            if len(ncomp) == n_annuli:
                ncompann = ncomp[ann]
            else:
                msge = 'If ncomp is a list, it must match the number of annuli'
                raise TypeError(msge)
        else:
            ncompann = ncomp

        n_segments_ann = n_segments[ann]
        res_ann_par = _define_annuli(angle_list, ann, n_annuli, fwhm,
                                     radius_int, asize, delta_rot,
                                     n_segments_ann, verbose)
        pa_thr, inner_radius, ann_center = res_ann_par
        indices = get_annulus_segments(array[0], inner_radius, asize,
                                       n_segments_ann)
        # Library matrix is created for each segment and scaled if needed
        for j in range(n_segments_ann):
            yy = indices[j][0]
            xx = indices[j][1]
            matrix_segm = array[:, yy, xx]  # shape [nframes x npx_segment]
            matrix_segm = matrix_scaling(matrix_segm, scaling)

            res = pool_map(nproc, do_pca_patch, matrix_segm, fixed(range(n)),
                           angle_list, fwhm, pa_thr, ann_center, svd_mode,
                           ncompann, min_frames_lib, max_frames_lib, tol,
                           verbose=False)

            res = np.array(res)
            residuals = np.array(res[:, 0])
            ncomps = res[:, 1]
            nfrslib = res[:, 2]
            for fr in range(n):
                cube_out[fr][yy, xx] = residuals[fr]

            # number of frames in library printed for each annular quadrant
            # number of PCs printed for each annular quadrant
            if verbose == 2:
                descriptive_stats(nfrslib, verbose=verbose, label='\tLIBsize: ')
                descriptive_stats(ncomps, verbose=verbose, label='\tNum PCs: ')

        if verbose == 2:
            print('Done PCA with {} for current annulus'.format(svd_mode))
        if verbose:
            timing(start_time)

    # Cube is derotated according to the parallactic angle and collapsed
    cube_der = cube_derotate(cube_out, angle_list, imlib=imlib,
                             interpolation=interpolation)
    frame = cube_collapse(cube_der, mode=collapse)
    if verbose:
        print('Done derotating and combining.')
        timing(start_time)
    if full_output:
        return cube_out, cube_der, frame
    else:
        return frame

def do_pca_patch(matrix, frame, angle_list, fwhm, pa_threshold, ann_center,
                 svd_mode, ncomp, min_frames_lib, max_frames_lib, tol):
    """ Does the SVD/PCA for each frame patch (small matrix). For each frame we
    find the frames to be rejected depending on the amount of rotation. The
    library is also truncated on the other end (frames too far or which have
    rotated more) which are more decorrelated to keep the computational cost
    lower. This truncation is done on the annuli after 10*FWHM and the goal is
    to keep min(num_frames/2, 200) in the library.
    """
    if pa_threshold != 0:
        if ann_center > fwhm*10:    # TODO: 10*FWHM optimal? new parameter?
            indices_left = _find_indices_adi(angle_list, frame, pa_threshold,
                                             truncate=True,
                                             max_frames=max_frames_lib)
        else:
            indices_left = _find_indices_adi(angle_list, frame, pa_threshold,
                                             truncate=False)

        data_ref = matrix[indices_left]

        if data_ref.shape[0] <= min_frames_lib:
            msg = 'Too few frames left in the PCA library. '
            msg += 'Try decreasing either delta_rot or min_frames_lib.'
            raise RuntimeError(msg)
    else:
        data_ref = matrix

    data = data_ref
    curr_frame = matrix[frame]                     # current frame

    V = get_eigenvectors(ncomp, data, svd_mode, noise_error=tol, debug=False)

    transformed = np.dot(curr_frame, V.T)
    reconstructed = np.dot(transformed.T, V)
    residuals = curr_frame - reconstructed
    return residuals, V.shape[0], data_ref.shape[0]


#!/usr/bin/env python
# encoding: utf-8

################################################################################
#
#   RMG - Reaction Mechanism Generator
#
#   Copyright (c) 2009-2011 by the RMG Team (rmg_dev@mit.edu)
#
#   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.
#
################################################################################

"""
Contains the :class:`Geometry` class for working with the three-dimensional
geometry of molecules and evaluating properties based on the geometry
information, e.g. moments of inertia.
"""

import numpy
import cython

from rmgpy.quantity import Quantity, constants

################################################################################

class GeometryError(Exception):
    """
    An exception class for errors that occur while working with molecular
    geometries. Pass a string describing the circumstances that caused the
    exceptional behavior.
    """
    pass

################################################################################

class Geometry:
    """
    The three-dimensional geometry of a molecular configuration. The attributes
    are:

    =============== ======================= ====================================
    Attribute       Type                    Description
    =============== ======================= ====================================
    `coordinates`   :class:`numpy.ndarray`  An N x 3 array containing the 3D coordinates of each atom
    `number`        :class:`numpy.ndarray`  An array containing the integer atomic number of each atom
    `mass`          :class:`numpy.ndarray`  An array containing the atomic mass in kg/mol of each atom
    =============== ======================= ====================================

    The integer index of each atom is consistent across all three attributes.
    """

    def __init__(self, coordinates, number, mass):
        self.coordinates = Quantity(coordinates).values
        self.number = Quantity(number).values
        self.mass = Quantity(mass).values

    def __repr__(self):
        """
        Return a string representation that can be used to reconstruct the
        object.
        """
        coordinates = '(['
        for i in range(self.coordinates.shape[0]):
            if i > 0: coordinates += ', '
            coordinates += '[{0}]'.format(','.join(['{0:g}'.format(self.coordinates[i,j]) for j in range(self.coordinates.shape[1])]))
        coordinates += '],"m")'
        number = '[{0}]'.format(','.join(['{0:d}'.format(n) for n in self.number]))
        mass = '([{0}],"g/mol")'.format(','.join(['{0:g}'.format(m * 1000.) for m in self.mass]))
        return 'Geometry(coordinates={0}, number={1}, mass={2})'.format(coordinates, number, mass)

    def __reduce__(self):
        """
        A helper function used when pickling an object.
        """
        return (Geometry, (self.coordinates, self.number, self.mass))

    def getTotalMass(self, atoms=None):
        """
        Calculate and return the total mass of the atoms in the geometry in
        kg/mol. If a list `atoms` of atoms is specified, only those atoms will
        be used to calculate the center of mass. Otherwise, all atoms will be
        used.
        """
        if atoms is None: atoms = range(len(self.mass))
        return sum([self.mass[atom] for atom in atoms])

    def getCenterOfMass(self, atoms=None):
        """
        Calculate and return the [three-dimensional] position of the center of
        mass of the current geometry. If a list `atoms` of atoms is specified,
        only those atoms will be used to calculate the center of mass.
        Otherwise, all atoms will be used.
        """

        cython.declare(center=numpy.ndarray, mass=cython.double, atom=cython.int)

        if atoms is None: atoms = range(len(self.mass))
        center = numpy.zeros(3, numpy.float64); mass = 0.0
        for atom in atoms:
            center += self.mass[atom] * self.coordinates[atom]
            mass += self.mass[atom]
        center /= mass
        return center

    def getMomentOfInertiaTensor(self):
        """
        Calculate and return the moment of inertia tensor for the current
        geometry in kg*m^2. If the coordinates are not at the center of mass,
        they are temporarily shifted there for the purposes of this calculation.
        """

        cython.declare(I=numpy.ndarray, mass=cython.double, atom=cython.int)
        cython.declare(coord0=numpy.ndarray, coord=numpy.ndarray, centerOfMass=numpy.ndarray)

        I = numpy.zeros((3,3), numpy.float64)
        centerOfMass = self.getCenterOfMass()
        for atom, coord0 in enumerate(self.coordinates):
            mass = self.mass[atom] / constants.Na
            coord = coord0 - centerOfMass
            I[0,0] += mass * (coord[1] * coord[1] + coord[2] * coord[2])
            I[1,1] += mass * (coord[0] * coord[0] + coord[2] * coord[2])
            I[2,2] += mass * (coord[0] * coord[0] + coord[1] * coord[1])
            I[0,1] -= mass * coord[0] * coord[1]
            I[0,2] -= mass * coord[0] * coord[2]
            I[1,2] -= mass * coord[1] * coord[2]
        I[1,0] = I[0,1]
        I[2,0] = I[0,2]
        I[2,1] = I[1,2]

        return I

    def getPrincipalMomentsOfInertia(self):
        """
        Calculate and return the principal moments of inertia and corresponding
        principal axes for the current geometry. The moments of inertia are in
        kg*m^2, while the principal axes have unit length.
        """
        I0 = self.getMomentOfInertiaTensor()
        # Since I0 is real and symmetric, diagonalization is always possible
        I, V = numpy.linalg.eig(I0)
        return I, V

    def getInternalReducedMomentOfInertia(self, pivots, top1):
        """
        Calculate and return the reduced moment of inertia for an internal
        torsional rotation around the axis defined by the two atoms in
        `pivots`. The list `top1` contains the atoms that should be considered
        as part of the rotating top; this list should contain the pivot atom
        connecting the top to the rest of the molecule.	The procedure used is
        that of Pitzer [1]_, which is described as :math:`I^{(2,3)}` by East
        and Radom [2]_. In this procedure, the molecule is divided into two
        tops: those at either end of the hindered rotor bond. The moment of
        inertia of each top is evaluated using an axis passing through the
        center of mass of both tops. Finally, the reduced moment of inertia is
        evaluated from the moment of inertia of each top via the formula

        .. math:: \\frac{1}{I^{(2,3)}} = \\frac{1}{I_1} + \\frac{1}{I_2}

        .. [1] Pitzer, K. S. *J. Chem. Phys.* **14**, p. 239-243 (1946).

        .. [2] East, A. L. L. and Radom, L. *J. Chem. Phys.* **106**, p. 6655-6674 (1997).

        """

        cython.declare(Natoms=cython.int, top2=list, top1CenterOfMass=numpy.ndarray, top2CenterOfMass=numpy.ndarray)
        cython.declare(axis=numpy.ndarray, I1=cython.double, I2=cython.double, atom=cython.int, i=cython.int)

        # The total number of atoms in the geometry
        Natoms = len(self.mass)

        # Check that exactly one pivot atom is in the specified top
        if pivots[0] not in top1 and pivots[1] not in top1:
            raise GeometryError('No pivot atom included in top; you must specify which pivot atom belongs with the specified top.')
        elif pivots[0] in top1 and pivots[1] in top1:
            raise GeometryError('Both pivot atoms included in top; you must specify only one pivot atom that belongs with the specified top.')

        # Determine atoms in other top
        top2 = []
        for i in range(Natoms):
            if i not in top1: top2.append(i)

        # Determine centers of mass of each top
        top1CenterOfMass = self.getCenterOfMass(top1)
        top2CenterOfMass = self.getCenterOfMass(top2)

        # Determine axis of rotation
        axis = (top1CenterOfMass - top2CenterOfMass)
        axis /= numpy.linalg.norm(axis)

        # Determine moments of inertia of each top
        I1 = 0.0
        for atom in top1:
            r1 = self.coordinates[atom,:] - top1CenterOfMass
            r1 -= numpy.dot(r1, axis) * axis
            I1 += self.mass[atom] / constants.Na * numpy.linalg.norm(r1)**2
        I2 = 0.0
        for atom in top2:
            r2 = self.coordinates[atom,:] - top2CenterOfMass
            r2 -= numpy.dot(r2, axis) * axis
            I2 += self.mass[atom] / constants.Na * numpy.linalg.norm(r2)**2

        return 1.0 / (1.0 / I1 + 1.0 / I2)


'''
https://github.com/zhangqianhui/Conditional-GAN

Modified by Wei Chen(wchen@cqu.edu.cn), Qiuli Wang(wangqiuli@cqu.edu.cn)
7/12/2020
'''

from utils import save_images, vis_square,sample_label,sample_masks, sample_masks_test
from tensorflow.contrib.layers.python.layers import xavier_initializer
import cv2
from ops import *
import tensorflow as tf
import numpy as np

class CMGAN(object):
    def __init__(self, data_ob, train_dir, eval_dir, test_dir, output_size, learn_rate, batch_size, z_dim, y_dim, log_dir
         , model_path, load = False, gf_dim=64, df_dim = 64, output_c_dim=1, L1_lambda=100):

        self.data_ob = data_ob
        self.train_dir = train_dir
        self.eval_dir = eval_dir
        self.test_dir = test_dir
        self.output_size = output_size
        self.learn_rate = learn_rate
        self.batch_size = batch_size
        self.z_dim = z_dim
        self.y_dim = y_dim
        self.load = load
        self.log_dir = log_dir
        self.model_path = model_path
        self.channel = self.data_ob.shape[2]
        self.images = tf.placeholder(tf.float32, [batch_size, self.output_size, self.output_size, self.channel])
        self.masks = tf.placeholder(tf.float32, [batch_size, self.output_size, self.output_size, self.channel])
        self.lungwindow = tf.placeholder(tf.float32, [batch_size, self.output_size, self.output_size, self.channel])
        self.mediastinumwindow = tf.placeholder(tf.float32, [batch_size, self.output_size, self.output_size, self.channel])

        self.L1_lambda = L1_lambda
        self.z = tf.placeholder(tf.float32, [self.batch_size, self.z_dim])
        self.y = tf.placeholder(tf.float32, [self.batch_size, self.y_dim])
        self.training_step = 50000

        print('image shape: ', self.images.get_shape().as_list())

        self.gf_dim = gf_dim
        self.df_dim = df_dim
        self.output_c_dim = output_c_dim
        self.d_bn1 = batch_norm(name='d_bn1')
        self.d_bn2 = batch_norm(name='d_bn2')
        self.d_bn3 = batch_norm(name='d_bn3')

        self.g_bn_e2 = batch_norm(name='g_bn_e2')
        self.g_bn_e3 = batch_norm(name='g_bn_e3')
        self.g_bn_e4 = batch_norm(name='g_bn_e4')
        self.g_bn_e5 = batch_norm(name='g_bn_e5')
        self.g_bn_e6 = batch_norm(name='g_bn_e6')
        self.g_bn_e7 = batch_norm(name='g_bn_e7')
        self.g_bn_e8 = batch_norm(name='g_bn_e8')

        self.g_bn_d1 = batch_norm(name='g_bn_d1')
        self.g_bn_d2 = batch_norm(name='g_bn_d2')
        self.g_bn_d3 = batch_norm(name='g_bn_d3')
        self.g_bn_d4 = batch_norm(name='g_bn_d4')
        self.g_bn_d5 = batch_norm(name='g_bn_d5')
        self.g_bn_d6 = batch_norm(name='g_bn_d6')
        self.g_bn_d7 = batch_norm(name='g_bn_d7')

	# build model
    def build_model(self):
        print('Building the model:')
        self.real_A = self.masks
        print('shape of real_A: ', self.real_A.get_shape().as_list())

        self.fake_B = self.generator(self.real_A, self.y)
        print('shape of self.y: ', self.y.get_shape().as_list())
        print('shape of fake_B: ', self.fake_B.get_shape().as_list())

        self.lung_logits, self.fake_lungwindow = self.decorator_lung_window(self.fake_B, self.lungwindow, self.y)
        self.mediastinum_logits, self.fake_mediastinumwindow = self.decorator_mediastinum_window(self.fake_B, self.mediastinumwindow, self.y)
        self.real_all_masks_images = tf.concat([self.real_A, self.images], 3)
        self.fake_all_masks_images = tf.concat([self.real_A, self.fake_B], 3)

        self.D_all, self.D_all_logits = self.discriminator_all(self.real_all_masks_images, self.y, reuse = False)
        self.D_all_, self.D_all_logits_ = self.discriminator_all(self.fake_all_masks_images, self.y, reuse = True)

        self.d_all_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits = self.D_all_logits, labels = tf.ones_like(self.D_all)))
        self.d_all_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits = self.D_all_logits_, labels = tf.zeros_like(self.D_all_)))

        self.y_lobu = tf.slice(self.y, [0,0], [64,5])
        self.y_spicu = tf.slice(self.y, [0,5], [64,5])
        self.y_mali = tf.slice(self.y, [0,10], [64,5])
        print('shape of lobu label ', self.y_lobu.get_shape().as_list())
        print('shape of spicu label ', self.y_spicu.get_shape().as_list())
        print('shape of mali label ', self.y_mali.get_shape().as_list())

        # the loss of classify
        self.pre_lung, self.pre_lung_logits = self.classify_lung(self.fake_lungwindow)
        self.d_loss_classify_lung = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.pre_lung, labels=self.y_spicu))

        self.pre_mediastinum, self.pre_mediastinum_logits = self.classify_mediastinum(self.fake_mediastinumwindow)
        self.d_loss_classify_mediastinum = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.pre_mediastinum, labels=self.y_lobu))

        # the loss of malignancy prediction
        self.pre_lung_mali, self.pre_lung_logits_mali = self.classify_lung_mali(self.fake_lungwindow)
        self.d_loss_classify_lung_mali = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.pre_lung_mali, labels=self.y_mali))

        self.pre_mediastinum_mali, self.pre_mediastinum_logits_mali = self.classify_mediastinum_mali(self.fake_mediastinumwindow)
        self.d_loss_classify_mediastinum_mali = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=self.pre_mediastinum_mali, labels=self.y_mali))

        # the loss of generator network
        self.g_lung_loss = self.L1_lambda * tf.reduce_mean(tf.square(self.lungwindow - self.fake_lungwindow))
        self.g_mediastinum_loss = self.L1_lambda * tf.reduce_mean(tf.square(self.mediastinumwindow - self.fake_mediastinumwindow))
        self.g_all_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.D_all_logits_, labels=tf.ones_like(self.D_all_))) \
                        + self.L1_lambda * tf.reduce_mean(tf.square(self.images - self.fake_B))

        # the loss of decorator network
        self.d_all_loss = self.d_all_loss_fake + self.d_all_loss_real + self.d_loss_classify_lung + self.d_loss_classify_mediastinum + self.d_loss_classify_lung_mali + self.d_loss_classify_mediastinum_mali
        self.d_c_loss = self.d_loss_classify_lung + self.d_loss_classify_mediastinum + self.d_loss_classify_lung_mali + self.d_loss_classify_mediastinum_mali
        self.d_loss = self.d_all_loss
        self.g_loss = self.g_lung_loss + self.g_mediastinum_loss + self.g_all_loss

        t_vars = tf.trainable_variables()
        self.d_vars = [var for var in t_vars if 'd_' in var.name]
        self.c_vars = [var for var in t_vars if 'c_' in var.name]
        self.g_vars = [var for var in t_vars if 'g_' in var.name]

        self.saver = tf.train.Saver(keep_checkpoint_every_n_hours=0.5)

    def train(self,args):

        opti_D = tf.train.AdamOptimizer(args.lr, beta1=args.beta1) \
                          .minimize(self.d_loss, var_list=self.d_vars)
        opti_G = tf.train.AdamOptimizer(args.lr, beta1=args.beta1) \
                          .minimize(self.g_loss, var_list=self.g_vars)
        opti_C = tf.train.AdamOptimizer(args.lr, beta1=args.beta1) \
                          .minimize(self.d_c_loss, var_list=self.c_vars)

        init = tf.global_variables_initializer()
        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True

        with tf.Session(config=config) as sess:
            sess.run(init)
            if self.load:
            	# load pretrained model
                print('loading:')
                self.saver = tf.train.import_meta_graph('./model/model.ckpt-20201.meta')  # default to save all variable
                self.saver.restore(sess, tf.train.latest_checkpoint('./model/'))

            self.writer = tf.summary.FileWriter("./logs", sess.graph)
            summary_writer = tf.summary.FileWriter(self.log_dir, graph=sess.graph)

            step = 0
            while step <= self.training_step:
                realbatch_array, real_lungs, real_mediastinums, realmasks, real_labels = self.data_ob.getNext_batch(step,batch_size=self.batch_size)
                batch_z = np.random.uniform(-1, 1, size=[self.batch_size, self.z_dim])
                sess.run([opti_D],feed_dict={self.images: realbatch_array, self.lungwindow: real_lungs, self.mediastinumwindow: real_mediastinums, self.masks: realmasks, self.z: batch_z, self.y: real_labels})
                sess.run([opti_G],feed_dict={self.images: realbatch_array, self.lungwindow: real_lungs, self.mediastinumwindow: real_mediastinums, self.masks: realmasks, self.z: batch_z, self.y: real_labels})
                sess.run([opti_C],feed_dict={self.images: realbatch_array, self.lungwindow: real_lungs, self.mediastinumwindow: real_mediastinums, self.masks: realmasks, self.z: batch_z, self.y: real_labels})

                if np.mod(step, 50) == 1 and step != 0:
                    print('Saving...')
                    sample_images, lungwindow, mediastinumwindow = sess.run([self.fake_B, self.fake_lungwindow, self.fake_mediastinumwindow], feed_dict={self.images: realbatch_array, self.lungwindow: real_lungs, self.mediastinumwindow: real_mediastinums, self.masks: realmasks, self.z: batch_z, self.y: real_labels})
                    save_images(sample_images, [8, 8],
                                './{}/{:04d}_sample.png'.format(self.train_dir, step))
                    save_images(lungwindow, [8, 8],
                                './{}/{:04d}_lung.png'.format(self.train_dir, step))
                    save_images(mediastinumwindow, [8, 8],
                                './{}/{:04d}_mediastinum.png'.format(self.train_dir, step))
                    save_images(realmasks, [8, 8],
                                './{}/{:04d}_mask.png'.format(self.train_dir, step))

                    print('save eval image')

                    real_labels = sample_label()
                    realmasks = sample_masks()
                    sample_images, lungwindow, mediastinumwindow = sess.run([self.fake_B, self.fake_lungwindow, self.fake_mediastinumwindow], feed_dict={self.masks: realmasks, self.y: real_labels})
                    save_images(sample_images, [8, 8],
                                './{}/{:04d}_sample.png'.format(self.eval_dir, step))
                    save_images(lungwindow, [8, 8],
                                './{}/{:04d}_lung.png'.format(self.eval_dir, step))
                    save_images(mediastinumwindow, [8, 8],
                                './{}/{:04d}_mediastinum.png'.format(self.eval_dir, step))
                    save_images(realmasks, [8, 8],
                                './{}/{:04d}_mask.png'.format(self.eval_dir, step))

                    print('save test image')
                    real_labels = sample_label()
                    realmasks = sample_masks_test()
                    sample_images, lungwindow, mediastinumwindow = sess.run([self.fake_B, self.fake_lungwindow, self.fake_mediastinumwindow], feed_dict={self.masks: realmasks, self.y: real_labels})
                    save_images(sample_images, [8, 8],
                                './{}/{:04d}_sample.png'.format(self.test_dir, step))
                    save_images(lungwindow, [8, 8],
                                './{}/{:04d}_lung.png'.format(self.test_dir, step))
                    save_images(mediastinumwindow, [8, 8],
                                './{}/{:04d}_mediastinum.png'.format(self.test_dir, step))
                    save_images(realmasks, [8, 8],
                                './{}/{:04d}_mask.png'.format(self.test_dir, step))
                    # save model each 50 epochs
                    self.saver.save(sess, self.model_path,global_step=step)

                step = step + 1

            save_path = self.saver.save(sess, self.model_path)
            print ("Model saved in file: %s" % save_path)

    def test(self):
    	'''
    	load pretrained model for model testing
    	'''
        init = tf.initialize_all_variables()
        with tf.Session() as sess:
            sess.run(init)
            self.saver.restore(sess, self.model_path)
            sample_z = np.random.uniform(1, -1, size=[self.batch_size, self.z_dim])
            output = sess.run(self.fake_images, feed_dict={self.z: sample_z, self.y: sample_label()})
            save_images(output, [8, 8], './{}/test{:02d}_{:04d}.png'.format(self.train_dir, 0, 0))
            image = cv2.imread('./{}/test{:02d}_{:04d}.png'.format(self.train_dir, 0, 0), 0)
            cv2.imshow("test", image)
            cv2.waitKey(-1)
            print("Test finish!")

    def generator(self, image, y=None):
        with tf.variable_scope("generator") as scope:
            print('generator U-Net:')
            print('shape of y: ', y.get_shape().as_list())
            print('shape of image: ', image.get_shape().as_list())
            y = tf.reshape(y, shape=[self.batch_size, 1, 1, self.y_dim])

            s = self.output_size
            s2, s4, s8, s16, s32, s64, s128 = int(s/2), int(s/4), int(s/8), int(s/16), int(s/32), int(s/64), int(s/128)

            print('shape of image: ', image.get_shape().as_list())

            e1 = conv2d_UNet(image, self.gf_dim, name='g_e1_conv')
            e1 = conv_cond_concat(e1, y)

            e2 = self.g_bn_e2(conv2d_UNet(lrelu(e1), self.gf_dim*2, name='g_e2_conv'))
            e2 = conv_cond_concat(e2, y)

            e3 = self.g_bn_e3(conv2d_UNet(lrelu(e2), self.gf_dim*4, name='g_e3_conv'))
            e3 = conv_cond_concat(e3, y)

            e4 = self.g_bn_e4(conv2d_UNet(lrelu(e3), self.gf_dim*8, name='g_e4_conv'))
            e4 = conv_cond_concat(e4, y)

            e5 = self.g_bn_e5(conv2d_UNet(lrelu(e4), self.gf_dim*8, name='g_e5_conv'))
            e5 = conv_cond_concat(e5, y)

            e6 = self.g_bn_e6(conv2d_UNet(lrelu(e5), self.gf_dim*8, name='g_e6_conv'))
            e6 = conv_cond_concat(e6, y)

            e7 = self.g_bn_e7(conv2d_UNet(lrelu(e6), self.gf_dim*8, name='g_e7_conv'))
            e7 = conv_cond_concat(e7, y)

            e8 = self.g_bn_e8(conv2d_UNet(lrelu(e7), self.gf_dim*8, name='g_e8_conv'))
            e8 = conv_cond_concat(e8, y)

            self.d1, self.d1_w, self.d1_b = deconv2d(tf.nn.relu(e8),
                [self.batch_size, s128, s128, self.gf_dim*8], name='g_d1', with_w=True)
            d1 = tf.nn.dropout(self.g_bn_d1(self.d1), 0.5)
            d1 = tf.concat([d1, e7], 3)
            d1 = conv_cond_concat(d1, y)

            self.d2, self.d2_w, self.d2_b = deconv2d(tf.nn.relu(d1),
                [self.batch_size, s64, s64, self.gf_dim*8], name='g_d2', with_w=True)
            d2 = tf.nn.dropout(self.g_bn_d2(self.d2), 0.5)
            d2 = tf.concat([d2, e6], 3)
            d2 = conv_cond_concat(d2, y)

            self.d3, self.d3_w, self.d3_b = deconv2d(tf.nn.relu(d2),
                [self.batch_size, s32, s32, self.gf_dim*8], name='g_d3', with_w=True)
            d3 = tf.nn.dropout(self.g_bn_d3(self.d3), 0.5)
            d3 = tf.concat([d3, e5], 3)
            d3 = conv_cond_concat(d3, y)

            self.d4, self.d4_w, self.d4_b = deconv2d(tf.nn.relu(d3),
                [self.batch_size, s16, s16, self.gf_dim*8], name='g_d4', with_w=True)
            d4 = self.g_bn_d4(self.d4)
            d4 = tf.concat([d4, e4], 3)
            d4 = conv_cond_concat(d4, y)

            self.d5, self.d5_w, self.d5_b = deconv2d(tf.nn.relu(d4),
                [self.batch_size, s8, s8, self.gf_dim*4], name='g_d5', with_w=True)
            d5 = self.g_bn_d5(self.d5)
            d5 = tf.concat([d5, e3], 3)
            d5 = conv_cond_concat(d5, y)

            self.d6, self.d6_w, self.d6_b = deconv2d(tf.nn.relu(d5),
                [self.batch_size, s4, s4, self.gf_dim*2], name='g_d6', with_w=True)
            d6 = self.g_bn_d6(self.d6)
            d6 = tf.concat([d6, e2], 3)
            d6 = conv_cond_concat(d6, y)

            self.d7, self.d7_w, self.d7_b = deconv2d(tf.nn.relu(d6),
                [self.batch_size, s2, s2, self.gf_dim], name='g_d7', with_w=True)
            d7 = self.g_bn_d7(self.d7)
            d7 = tf.concat([d7, e1], 3)
            d7 = conv_cond_concat(d7, y)

            self.d8, self.d8_w, self.d8_b = deconv2d(tf.nn.relu(d7),
                [self.batch_size, s, s, self.output_c_dim], name='g_d8', with_w=True)
            print('shape of d8: ', self.d8.get_shape().as_list())
            return tf.nn.tanh(self.d8)

    def decorator_lung_window(self, images, masks, y = None):
        with tf.variable_scope('decorator_lungwindow') as scope:
            print('Decorator for lung window')
            y = tf.reshape(y, shape=[self.batch_size, 1, 1, self.y_dim])

            images = conv_cond_concat(images, y)
            h0 = lrelu(conv2d_decorator(images, 32, name='g_deco_lung_h0_conv'))
            h0 = conv_cond_concat(h0, y)
            h1 = lrelu(conv2d_decorator(h0, 64, name='g_deco_lung_h1_conv'))
            h1 = conv_cond_concat(h1, y)
            h2 = lrelu(conv2d_decorator(h1, 128, name = 'g_deco_lung_h2_conv'))
            h2 = conv_cond_concat(h2, y)
            h3 = lrelu(conv2d_decorator(h2, 64, name = 'g_deco_lung_h3_conv'))
            h3 = conv_cond_concat(h3, y)
            h4 = lrelu(conv2d_decorator(h3, 32, name='g_deco_lung_h4_conv'))
            h4 = conv_cond_concat(h4, y)
            h5 = lrelu(conv2d_decorator(h4, self.output_c_dim, name='g_deco_lung_h5_conv'))
            print('shape of h5: ', h5.get_shape().as_list())
            return tf.nn.relu(h5), h5

    def decorator_mediastinum_window(self, images, masks, y = None):
        with tf.variable_scope('decorator_mediastinum') as scope:
            print('Decorator for mediastinum window')
            y = tf.reshape(y, shape=[self.batch_size, 1, 1, self.y_dim])

            images = conv_cond_concat(images, y)
            h0 = lrelu(conv2d_decorator(images, 32, name='g_deco_mediastinum_h0_conv'))
            h0 = conv_cond_concat(h0, y)
            h1 = lrelu(conv2d_decorator(h0, 64, name='g_deco_mediastinum_h1_conv'))
            h1 = conv_cond_concat(h1, y)
            h2 = lrelu(conv2d_decorator(h1, 128, name = 'g_deco_mediastinum_h2_conv'))
            h2 = conv_cond_concat(h2, y)
            h3 = lrelu(conv2d_decorator(h2, 64, name = 'g_deco_mediastinum_h3_conv'))
            h3 = conv_cond_concat(h3, y)
            h4 = lrelu(conv2d_decorator(h3, 32, name='g_deco_mediastinum_h4_conv'))
            h4 = conv_cond_concat(h4, y)
            h5 = lrelu(conv2d_decorator(h4, self.output_c_dim, name='g_deco_mediastinum_h5_conv'))
            print('shape of h5: ', h5.get_shape().as_list())
            return tf.nn.relu(h5), h5

    def decorator(self, images, masks, y = None):
        with tf.variable_scope('decorator_allattributes') as scope:
            print('Decorator for all attributes')
            y = tf.reshape(y, shape=[self.batch_size, 1, 1, self.y_dim])

            images = conv_cond_concat(images, y)

            h0 = lrelu(conv2d_decorator(images, 32, name='g_deco_h0_conv'))
            h0 = conv_cond_concat(h0, y)
            h1 = lrelu(conv2d_decorator(h0, 64, name='g_deco_h1_conv'))
            h1 = conv_cond_concat(h1, y)
            h2 = lrelu(conv2d_decorator(h1, 128, name = 'g_deco_h2_conv'))
            h2 = conv_cond_concat(h2, y)
            h3 = lrelu(conv2d_decorator(h2, 64, name = 'g_deco_h3_conv'))
            h3 = conv_cond_concat(h3, y)
            h4 = lrelu(conv2d_decorator(h3, 32, name='g_deco_h4_conv'))
            h4 = conv_cond_concat(h4, y)
            h5 = lrelu(conv2d_decorator(h4, self.output_c_dim, name='g_deco_h5_conv'))
            print('shape of h5: ', h5.get_shape().as_list())
            return tf.nn.tanh(h5), h5

    def discriminator_all(self, image, y=None, reuse=False):
        with tf.variable_scope("discriminator") as scope:
            print('discriminator:')
            if reuse:
                tf.get_variable_scope().reuse_variables()
            else:
                assert tf.get_variable_scope().reuse == False
            y = tf.reshape(y, shape=[self.batch_size, 1, 1, self.y_dim])

            image = conv_cond_concat(image, y)
            h0 = lrelu(conv2d_UNet(image, self.df_dim, name='d_h0_conv'))
            h0 = conv_cond_concat(h0, y)
            h1 = lrelu(self.d_bn1(conv2d_UNet(h0, self.df_dim*2, name='d_h1_conv')))
            h1 = conv_cond_concat(h1, y)
            h2 = lrelu(self.d_bn2(conv2d_UNet(h1, self.df_dim*4, name='d_h2_conv')))
            h2 = conv_cond_concat(h2, y)
            h3 = lrelu(self.d_bn3(conv2d_UNet(h2, self.df_dim*8, d_h=1, d_w=1, name='d_h3_conv')))
            h3 = conv_cond_concat(h3, y)
            h4 = linear(tf.reshape(h3, [self.batch_size, -1]), 1, 'd_h3_lin')
            print('shape of h4: ', h4.get_shape().as_list())
            return tf.nn.sigmoid(h4), h4

    def classify_lung(self, image, y=None, reuse=False):
        with tf.variable_scope("classify_lung") as scope:
            print('classify_lung:')
            if reuse:
                tf.get_variable_scope().reuse_variables()
            else:
                assert tf.get_variable_scope().reuse == False

            layer = tf.layers.conv2d(image,64,[3,3],padding="same",activation=tf.nn.relu,name='c_classify_1')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_classify_2')
            layer = tf.layers.conv2d(layer,128,[3,3],padding="same",activation=tf.nn.relu,name='c_classify_3')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_classify_4')
            layer = tf.layers.conv2d(layer,256,[3,3],padding="same",activation=tf.nn.relu,name='c_classify_5')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_classify_6')
            layer = tf.reshape(layer, [-1, 16 * 16 * 256])
            layer = tf.layers.dense(layer,100,activation=tf.nn.relu,name='c_classify_7')
            layer = tf.layers.dropout(layer,0.5,name='c_classify_8')
            layer = tf.layers.dense(layer,5,activation=tf.nn.relu,name='c_classify_9')
            logits = tf.nn.softmax(layer)
            return layer, logits

    def classify_mediastinum(self, image, y=None, reuse=False):
        with tf.variable_scope("classify_mediastinum") as scope:
            print('classify_mediastinum:')
            if reuse:
                tf.get_variable_scope().reuse_variables()
            else:
                assert tf.get_variable_scope().reuse == False

            layer = tf.layers.conv2d(image,64,[3,3],padding="same",activation=tf.nn.relu,name='c_mediastinum_1')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_mediastinum_2')
            layer = tf.layers.conv2d(layer,128,[3,3],padding="same",activation=tf.nn.relu,name='c_mediastinum_3')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_mediastinum_4')

            layer = tf.layers.conv2d(layer,256,[3,3],padding="same",activation=tf.nn.relu,name='c_mediastinum_5')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_mediastinum_6')
            layer = tf.reshape(layer, [-1, 16 * 16 * 256])

            layer = tf.layers.dense(layer,100,activation=tf.nn.relu,name='c_mediastinum_7')
            layer = tf.layers.dropout(layer,0.5,name='c_mediastinum_8')
            layer = tf.layers.dense(layer,5,activation=tf.nn.relu,name='c_mediastinum_9')

            logits = tf.nn.softmax(layer)
            return layer, logits

    def classify_lung_mali(self, image, y=None, reuse=False):
        with tf.variable_scope("classify_lung_mali") as scope:
            print('classify_lung_mali:')
            if reuse:
                tf.get_variable_scope().reuse_variables()
            else:
                assert tf.get_variable_scope().reuse == False

            layer = tf.layers.conv2d(image,64,[3,3],padding="same",activation=tf.nn.relu,name='c_classify_mali_1')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_classify_mali_2')
            layer = tf.layers.conv2d(layer,128,[3,3],padding="same",activation=tf.nn.relu,name='c_classify_mali_3')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_classify_mali_4')

            layer = tf.layers.conv2d(layer,256,[3,3],padding="same",activation=tf.nn.relu,name='c_classify_mali_5')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_classify_mali_6')
            layer = tf.reshape(layer, [-1, 16 * 16 * 256])

            layer = tf.layers.dense(layer,100,activation=tf.nn.relu,name='c_classify_mali_7')
            layer = tf.layers.dropout(layer,0.5,name='c_classify_mali_8')
            layer = tf.layers.dense(layer,5,activation=tf.nn.relu,name='c_classify_mali_9')
            logits = tf.nn.softmax(layer)
            return layer, logits

    def classify_mediastinum_mali(self, image, y=None, reuse=False):
        with tf.variable_scope("classify_mediastinum_mali") as scope:
            print('classify_mediastinum_mali:')
            if reuse:
                tf.get_variable_scope().reuse_variables()
            else:
                assert tf.get_variable_scope().reuse == False

            layer = tf.layers.conv2d(image,64,[3,3],padding="same",activation=tf.nn.relu,name='c_mediastinum_mali_1')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_mediastinum_mali_2')
            layer = tf.layers.conv2d(layer,128,[3,3],padding="same",activation=tf.nn.relu,name='c_mediastinum_mali_3')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_mediastinum_mali_4')

            layer = tf.layers.conv2d(layer,256,[3,3],padding="same",activation=tf.nn.relu,name='c_mediastinum_mali_5')
            layer = tf.layers.max_pooling2d(layer,pool_size=[2,2],strides=2,name='c_mediastinum_mali_6')
            layer = tf.reshape(layer, [-1, 16 * 16 * 256])

            layer = tf.layers.dense(layer,100,activation=tf.nn.relu,name='c_mediastinum_mali_7')
            layer = tf.layers.dropout(layer,0.5,name='c_mediastinum_mali_8')
            layer = tf.layers.dense(layer,5,activation=tf.nn.relu,name='c_mediastinum_mali_9')

            logits = tf.nn.softmax(layer)
            return layer, logits


"""
VABY_MODELS_CVR: VABY forward models for CVR

Forward model for CVR measurement using BOLD MRI and PETCo2

Based on Matlab code written by Joana Pinto, December 2020, Oxford
and adapted from Daniel Bulte 2018 script

Python conversion by Martin Craig 2021, Nottingham

(c) 2021 University of Nottingham
"""
import tensorflow as tf

import numpy as np
#import tensorflow_probability as tfp

from vaby.model import Model, ModelOption
from vaby.utils import ValueList
from vaby.parameter import get_parameter

from ._version import __version__

class CvrPetCo2Model(Model):
    """
    Inference forward model for CVR measurement using PETCo2
    """

    def options(self):
        return [
            # Regressor, e.g. physiological data file containing PETCO2 measurements
            ModelOption("regressors", "Regression data (e.g. PETCO2 or O2 time series)", type=str, default=None),
            ModelOption("regressor_types", "Regressor types - comma separated one for each regressor. Supported types: co2, petco2, custom", type=str, default="co2"),
            ModelOption("regressor_trs", "Regressor time resolutions", unit="s", type=ValueList, default=[0.01,]),

            # Protocol parameters
            ModelOption("baseline", "Length of initial baseline block", unit="s", type=int, default=60),
            ModelOption("data_start_time", "Start of MR data relative to start of regressor data - if not provided will be estimated", unit="s", type=float, default=None),
            ModelOption("tr", "Time between MR volumes", unit="s", type=float, default=None),
            ModelOption("air_pressure", "Barometric pressure", unit="mbar", type=int, default=1020),

            # Model options
            ModelOption("infer_sig0", "Infer signal offset", type=bool, default=False),
            ModelOption("infer_delay", "Infer delay shift on regressors", type=bool, default=False),
            #ModelOption("infer_drift", "Infer a linear drift on signal", type=bool, default=False),
            #ModelOption("sigmoid_response", "Use sigmoid relationship between PETCO2 and CVR", type=bool, default=False)
        ]

    def __str__(self):
        return "CVR-PETCO2 model: %s" % __version__

    def __init__(self, data_model, **options):
        Model.__init__(self, data_model, **options)

        if self.regressors is None:
            raise ValueError("A regressor must be provided")

        if isinstance(self.regressors, str):
            fnames = self.regressors.split(",") # Hopefully no commas in file names
            self.regressors = np.array([np.squeeze(np.loadtxt(fname)) for fname in fnames])
        else:
            self.regressors = np.array(self.regressors)

        if self.regressors.ndim < 2:
            self.regressors = self.regressors[np.newaxis, ...]
        if self.regressors.ndim != 2:
            raise ValueError("Regressor must be 1D or 2D")

        self.n_regressors = self.regressors.shape[0]
        self.regressor_types = [s.strip() for s in self.regressor_types.split(",")]
        if len(self.regressor_types) != self.n_regressors:
            raise ValueError("Number of regressors provided (%i) does not match number of regressor types (%i)" % (self.n_regressors, len(self.regressor_types)))
        if len(self.regressor_trs) == 1:
            # Use same value for all regressors
            self.regressor_trs = self.regressor_trs * self.n_regressors
        elif len(self.regressor_trs) != self.n_regressors:
            raise ValueError("Number of regressors provided (%i) does not match number of time resolutions (%i)" % (self.n_regressors, len(self.regressor_trs)))

        # Process regressors and generate regression parameters for each
        self.params = []
        regressors = []
        regressor_tpts = []
        for idx, regressor_type in enumerate(self.regressor_types):
            if regressor_type == "co2":
                # Unprocessed CO2
                regressors.append(self._preproc_co2(self.regressors[idx], self.regressor_trs[idx]))
                regressor_tpts.append(self.tpts())
                self.params.append(get_parameter("cvr%i" % (idx+1), mean=1.0, dist="FoldedNormal", prior_var=2000, post_var=10, **options))
                self.data_start_time = 0
            elif regressor_type == "petco2":
                # Preprocessed end-tidal CO2 - not necessarily aligned with data or at same temporal resolution
                regressors.append(self.regressors[idx].astype(np.float32))
                regressor_tpts.append(np.array(range(len(self.regressors[idx]))) * self.regressor_trs[idx])
                self.params.append(get_parameter("cvr%i" % (idx+1), mean=1.0, dist="FoldedNormal", prior_var=2000, post_var=10, **options))
            elif regressor_type == "custom":
                # Generic regressor, not necessarily aligned to data or at same temporal resolution
                regressors.append(self.regressors[idx].astype(np.float32))
                regressor_tpts.append(np.array(range(len(self.regressors[idx]))) * self.regressor_trs[idx])
                self.params.append(get_parameter("beta%i" % (idx+1), mean=0.0, dist="Normal", prior_var=1e6, post_var=1e3, **options))
            else:
                raise ValueError("Unrecognized regressor type: %s" % regressor_type)
        self.regressors = np.array(regressors)
        self.regressor_tpts = np.array(regressor_tpts)

        # Differences between timepoints for quick interpolation. Given a delay
        # time > 0 we can compute value = regressors[int(delay)] + frac(delay) * regressor_diff[int(delay)]
        self.regressor_diffs = np.zeros(self.regressors.shape, dtype=np.float32)
        self.regressor_diffs[:, :-1] = self.regressors[:, 1:] - self.regressors[:, :-1]

        # Min/max values
        self.regressor_mins = np.min(self.regressors, axis=1)
        self.regressor_maxs = np.max(self.regressors, axis=1)
        self.log.info("Regressor minimum values: %s", self.regressor_mins)
        self.log.info("Regressor maximum values: %s", self.regressor_maxs)

        # Estimate data start time
        if self.data_start_time is None:
            self.estimate_data_start_time(self.regressors[0], self.regressor_tpts[0])

        if self.infer_sig0:
            self.params.append(get_parameter("sig0", mean=1, prior_var=1e9, post_mean=1, post_var=10, post_init=self._init_sig0, **options))
        if self.infer_delay:
            self.params.append(get_parameter("delay", mean=0, prior_var=100, post_var=10, **options))

    def _init_sig0(self, _param, _t, data):
        return np.mean(data, axis=-1), None

    def fit_glm(self, delay_min=-1, delay_max=1, delay_step=1, progress_cb=None):
        self.log.info("GLM: Doing fitting on %i voxels", self.data_model.data_space.size)
        bold_data = self.data_model.data_space.srcdata.flat
        t = self.tpts() # in seconds

        delays = np.arange(delay_min, delay_max+delay_step, delay_step, dtype=np.float32)
        best_resid = np.ones(bold_data.shape[0], dtype=np.float32) * 1e99
        best_delay = np.zeros(bold_data.shape[0], dtype=np.float32)
        best_cvr = np.zeros((bold_data.shape[0], self.n_regressors), dtype=np.float32)
        best_sig0 = np.zeros(bold_data.shape[0], dtype=np.float32)
        best_modelfit = np.zeros(bold_data.shape, dtype=np.float32)
        for idx, delay in enumerate(delays):
            self.log.info("GLM: fitting with delay=%f", delay)
            delayed_tpts = t - delay + self.data_start_time
            x = []
            for idx, regressor in enumerate(self.regressors):
                # FIXME sampling rates, regressor time span...
                delayed = np.interp(delayed_tpts, self.regressor_tpts[idx], regressor)
                regressor_type = self.regressor_types[idx]
                if regressor_type in ("co2", "petco2"):
                    x.append((delayed - self.regressor_mins[idx]) / (self.regressor_maxs[idx] - self.regressor_mins[idx]))
                else:
                    x.append(delayed)

            x.append(np.ones(self.data_model.data_space.srcdata.n_tpts))
            x = np.array(x).T
            for vox in range(bold_data.shape[0]):
                y = bold_data[vox, :]
                beta, resid, _rank, _s = np.linalg.lstsq(x, y)
                model = np.dot(x, beta)
                vox_resid = resid[0]
                if vox_resid < best_resid[vox]:
                    best_delay[vox] = delay
                    best_sig0[vox] = beta[-1]
                    for ridx, regressor_type in enumerate(self.regressor_types):
                        if regressor_type in ("co2", "petco2"):
                            best_cvr[vox, ridx] = beta[ridx]*100/(self.regressor_maxs[ridx] - self.regressor_mins[ridx])/best_sig0[vox]
                        else:
                            best_cvr[vox, ridx] = beta[ridx]/best_sig0[vox]
                    best_modelfit[vox] = model
                    best_resid[vox] = vox_resid
            if progress_cb is not None:
                progress_cb(float(idx)/float(len(delays)))

        self.log.info("GLM: DONE")
        ret = []
        for idx in range(self.n_regressors):
            ret.append(best_cvr[..., idx])
        ret.append(best_delay)
        ret.append(best_sig0)
        ret.append(best_modelfit)
        return tuple(ret)

    def evaluate(self, params, tpts):
        """
        FIXME won't work in SVB batch training because of timepoints

        :param t: Time values tensor of shape [W, 1, N] or [1, 1, N]
        :param params Sequence of parameter values arrays, one for each parameter.
                      Each array is [W, S, 1] tensor where W is the number of nodes and
                      S the number of samples. This
                      may be supplied as a [P, W, S, 1] tensor where P is the number of
                      parameters.

        :return: [W, S, N] tensor containing model output at the specified time values
                 and for each time value using the specified parameter values
        """
        regressor_params = params[:self.n_regressors]

        extra_param = self.n_regressors
        if self.infer_sig0:
            sig0 = params[extra_param]
            extra_param += 1
        else:
            sig0 = 0

        if self.infer_delay:
            delay = params[extra_param] - self.data_start_time
            extra_param += 1

            # Apply time delay [W, (S), N] FIXME what is length of regressor
            t_delayed = (tpts - delay) / self.tr
            t_delayed = tf.clip_by_value(t_delayed, 0, len(self.regressors[0])-1)
            t_base = tf.floor(t_delayed)

            # Integer index into the CO2 and diff arrays
            t_base_idx = tf.cast(t_base, tf.int32)

            # Fractional distance to next array index, or 0 if base index was < 0
            t_frac = tf.clip_by_value(t_delayed - t_base, 0, 1)
        else:
            t_base_idx = tf.cast(tf.floor(tpts / self.tr), tf.int32)
            t_frac = None

        fit = 1
        for idx, regressor in enumerate(self.regressors):
            # Tile regressor arrays over all nodes so we can use tf.gather
            regressor = tf.tile(regressor[np.newaxis, ...], (tf.shape(t_base_idx)[0], 1))

            # Get value of regressor at integer part of time points
            delayed = tf.gather(regressor, t_base_idx, axis=1, batch_dims=1)

            if t_frac is not None:
                # If we have a delay, use the differenced regressor to do linear interpolation on the
                # fractional part of the time points
                regressor_diff = tf.tile(self.regressor_diffs[idx][np.newaxis, ...], (tf.shape(t_base_idx)[0], 1))
                delayed_diff = tf.gather(regressor_diff, t_base_idx, axis=1, batch_dims=1)
                delayed += t_frac * delayed_diff

            # Sigmoid response
            #return sig0 + (b/(1+c.(e^(-(delayed_co2-c)/d))))/100

            if self.regressor_types[idx] in ("petco2", "co2"):
                # Regressor parameter is CVR
                fit += regressor_params[idx] * (delayed - self.regressor_mins[idx]) / 100
            elif self.regressor_types[idx] == "custom":
                # Regressor parameter is generic coefficient
                fit += regressor_params[idx] * delayed

        fit = sig0 * fit
        return fit

    def tpts(self):
        """
        Get the full set of timeseries time values

        :return: Either a Numpy array of shape [N] or a Numpy array of shape
                 [W, N] for nodewise timepoints.
        """
        return np.linspace(0, self.data_model.data_space.srcdata.n_tpts, num=self.data_model.data_space.srcdata.n_tpts, endpoint=False, dtype=np.float32) * self.tr

    def estimate_data_start_time(self, regressor, regressor_tpts):
        # Mean time series
        bold_data_average = np.mean(self.data_model.data_space.srcdata.flat, axis=0)

        # Interpolate BOLD timeseries onto regressor
        mr_timings = self.tpts()
        bold_data_interp = np.interp(regressor_tpts, mr_timings, bold_data_average)

        _cc, delay_vols = self._cross_corr(bold_data_interp, regressor)
        self.data_start_time = -delay_vols * regressor_tpts[1] # to seconds assuming uniform spacing
        self.log.info("Cross correlation estimated data start time: %f", self.data_start_time)

        # Calculate the latest possible start time of the MR data
        # in case the cross correlation method returns something silly
        regressor_duration = len(regressor) * regressor_tpts[1]
        mr_duration = mr_timings[-1]
        max_time_begin = regressor_duration - mr_duration
        self.log.debug("Regressor duration: %f", regressor_duration)
        self.log.debug("MR duration: %f", mr_duration)
        self.log.info("Absolute latest start time: %f", max_time_begin)
        self.data_start_time = min(self.data_start_time, max_time_begin)
        self.log.info("Final estimated data start time: %f", self.data_start_time)

    def _cross_corr(self, y1, y2):
        """
        Calculates the cross correlation and lags

        :param y1: 1D Numpy array
        :param y2: 1D Numpy array, same length as y1

        :return: Tuple of Maximum correlation, lag in terms of the index
        """
        if len(y1) != len(y2):
            raise ValueError('The lengths of the inputs should be the same.')

        corr = np.correlate(y1 - np.mean(y1),
                            y2 - np.mean(y2),
                            mode='full')
        lag = corr.argmax() - (len(y1) - 1)
        return np.max(corr), lag

    def _preproc_co2(self, co2, regressor_tr):
        """
        Preprocess CO2 measurements from physiological data file
        """
        co2 = np.squeeze(co2)
        samp_rate = 1/regressor_tr

        # FIXME we need the data start time to have been already estimated,
        # although we probably shouldn't since there is no need to trim the CO2 trace
        if self.data_start_time is None:
            self.estimate_data_start_time(co2, np.array(range(len(co2))) * regressor_tr)
            data_start = self.data_start_time
            self.data_start_time = None

        # FIXME should we really trim the CO2 trace?
        co2_trim = co2[int(data_start * samp_rate):]

        # Determined respiratory frequency during baseline and use info to
        # determine size of end-tidal search window
        baseline_vols = int(self.baseline * samp_rate)
        baseline_fft = np.fft.fft(co2_trim[:baseline_vols])
        p2 = np.abs(baseline_fft/baseline_vols)
        p1 = np.array(p2[:int(baseline_vols/2)+1])
        p1[1:-2] = 2*p1[1:-2]
        f = np.linspace(0, samp_rate/2, int(baseline_vols/2)+1)

        loc = np.argmax(p1[1:])

        pkloc = loc+1
        harm = f[pkloc]
        resp_period = round(1/harm) # e.g. 8s

        # Search window = 1 second more than the respiratory period
        nsearch_vols = int((resp_period+1)*samp_rate)
        windows = int(np.floor(co2_trim.shape[0]/nsearch_vols))

        # Find peak PETCO2 in each window - it's value and index position
        posmax = np.zeros(windows, dtype=np.int)
        winmax = np.zeros(windows)
        for i in range(windows):
            for j in range(nsearch_vols):
                if j == 0 or co2_trim[i*nsearch_vols+j] > winmax[i]:
                    winmax[i] = co2_trim[i*nsearch_vols+j]
                    posmax[i] = i*nsearch_vols+j

        # Make new full sample ET time course where the PETCO2 changes linearly
        # between window maxima
        co2_resamp = np.zeros((co2_trim.shape[0], 1))
        for x in range(windows-1):
            dist_c = posmax[x+1] - posmax[x]
            step_c = winmax[x+1] - winmax[x]
            ramp_c = step_c / dist_c
            for g in range(dist_c+1):
                co2_resamp[posmax[x]+g] = winmax[x] + (ramp_c * g)

        # Pad the start and end with repeats of first and last value to maintain
        # length and phase
        co2_resamp[:posmax[0]] = co2_resamp[posmax[0]]
        co2_resamp[posmax[-1]:] = co2_resamp[posmax[-1]]

        # Create a timecourse of the end tidal CO2 values at the TR's for use with CVR sigmoids
        # Make new time course at the TR resolution and normalise timecourse betwwen 0 and 1 to create EV
        block = int(round(self.tr*samp_rate))
        ev_co2 = np.zeros((self.data_model.data_space.srcdata.n_tpts,), dtype=np.float32)
        for i in range(self.data_model.data_space.srcdata.n_tpts):
            ev_co2[i] = co2_resamp[block * i + block-1]

        # Convert to mmHg
        air_pressure_mmhg = self.air_pressure/1.33322387415 # pressure mbar
        co2_mmHg = (ev_co2 * air_pressure_mmhg) / 100 # div by 100 as values are in percent

        # Calculation of normo/hypercapnea from on/off volumes
        # Convert time periods to number of volumes
        #baseline_vols = self.baseline/self.tr
        #blocksize_on_vols = self.blocksize_on/self.tr
        #blocksize_off_vols = self.blocksize_off/self.tr

        # Average all of first baseline block
        #self.normocap = np.mean(co2_mmHg[:int(baseline_vols+self.delay)])

        #s1 = (baseline_vols+self.delay+blocksize_on_vols/2)
        #s2 = (baseline_vols+self.delay+blocksize_on_vols)
        #s3 = (baseline_vols+self.delay+blocksize_on_vols+blocksize_off_vols+blocksize_on_vols/2)
        #s4 = (baseline_vols+self.delay+blocksize_on_vols+blocksize_off_vols+blocksize_on_vols)
        #s1, s2, s3, s4 = int(s1), int(s2), int(s3), int(s4)
        # Select 2nd half of each hypercapnic block to average
        #hyperblock = np.concatenate([co2_mmHg[s1-1:s2], co2_mmHg[s3-1:s4]])
        #self.hypercap = np.mean(hyperblock)

        return co2_mmHg


from model.abstract_VAE import VAE
import numpy as np
import torch
import torch.nn as nn
from torch.autograd import Variable
from scipy.stats import norm

class StableBCELoss(nn.modules.Module):
    def __init__(self):
        super(StableBCELoss, self).__init__()
    def forward(self, input, target):
        neg_abs = - input.abs()
        loss = input.clamp(min=0) - input * target + (1 + neg_abs.exp()).log()
        return loss.sum()

class CNNEncodeLayer(nn.Module):
    def __init__(self, input, output, zdim, batchnorm, activacation, out_img_dims):
        super(CNNEncodeLayer, self).__init__()
        if activacation == "lrelu":
            self.act = nn.LeakyReLU()
        else:
            self.act = nn.ReLU()
        if batchnorm:
            main = nn.Sequential(
                nn.Conv2d(input, output,  4, stride=2, padding=1),
                nn.BatchNorm2d(output),
                self.act,
            )
            main2 = nn.Sequential(
                nn.Conv2d(input, output,  4, stride=2, padding=1),
                nn.BatchNorm2d(output),
                self.act,
            )
        else:
            main = nn.Sequential(
                nn.Conv2d(input, output,  4, stride=2, padding=1),
                self.act,
            )
            main2 = nn.Sequential(
                nn.Conv2d(input, output,  4, stride=2, padding=1),
                self.act,
            )
        self.main = main
        self.main2 = main2
        self.fc1 = nn.Linear(output * out_img_dims[0] * out_img_dims[1], zdim)
        self.fc2 = nn.Linear(output * out_img_dims[0] * out_img_dims[1], zdim)
        self.out_img_dims = out_img_dims
        # print ("Not implemented now...")
        return
    def forward(self, x):

        h = self.main(x)
        h2 = self.main2(x)
        return h, self.fc1(h2.view(h2.size(0), -1)), self.fc2(h2.view(h2.size(0), -1))

class CNNDecodeLayer(nn.Module):
    def __init__(self, input, output, zdim, batchnorm, activacation, input_img_dims):
        super(CNNDecodeLayer, self).__init__()

        if activacation == "lrelu":
            self.act = nn.LeakyReLU()
        else:
            self.act = nn.ReLU()
        if input == 0:
            input = output
            self.fc = nn.Linear(zdim, input * input_img_dims[0] * input_img_dims[0])
        else:
            self.fc = nn.Linear(zdim, input * input_img_dims[0] * input_img_dims[0])
            input *= 2
        if batchnorm:
            main = nn.Sequential(
                nn.ConvTranspose2d(input, output,  4, stride=2, padding=1),
                nn.BatchNorm2d(output),
                self.act,
            )
        else:
            main = nn.Sequential(
                nn.ConvTranspose2d(input, output,  4, stride=2, padding=1),
                self.act,
            )
        self.main = main
        self.input_img_dims = input_img_dims
    def forward(self, input, z):
        if input is None:
            input = self.act(self.fc(z).view(z.size(0), -1, self.input_img_dims[0], self.input_img_dims[1]))
        else:
            input = torch.cat([input, self.fc(z).view(z.size(0), -1, self.input_img_dims[0], self.input_img_dims[1])], 1)
        return self.main(input)

class EncodeLayer(nn.Module):
    def __init__(self, input, output, zdim, batchnorm, activacation):
        super(EncodeLayer, self).__init__()
        if activacation == "lrelu":
            self.act = nn.LeakyReLU()
        else:
            self.act = nn.ReLU()
        if batchnorm:
            main = nn.Sequential(
                nn.Linear(input, output),
                nn.BatchNorm1d(output),
                self.act,
            )
        else:
            main = nn.Sequential(
                nn.Linear(input, output),
                self.act,
            )
        self.main = main
        self.fc1 = nn.Linear(output, zdim)
        self.fc2 = nn.Linear(output, zdim)
    def forward(self, x):
        h = self.main(x)
        return self.main(x),self.fc1(h), self.fc2(h)

class DecodeLayer(nn.Module):
    def __init__(self, input, output, zdim, batchnorm, activacation):
        super(DecodeLayer, self).__init__()

        if activacation == "lrelu":
            self.act = nn.LeakyReLU()
        else:
            self.act = nn.ReLU()
        if input == 0:
            input = output
            self.fc = nn.Linear(zdim, input)
        else:
            self.fc = nn.Linear(zdim, input)
            input *= 2
        if batchnorm:
            main = nn.Sequential(
                nn.Linear(input, output),
                nn.BatchNorm1d(output),
                self.act,
            )
        else:
            main = nn.Sequential(
                nn.Linear(input, output),
                self.act,
            )
        self.main = main

    def forward(self, input, z):
        if input is None:
            input = self.act(self.fc(z))
        else:
            input = torch.cat([input, self.act(self.fc(z))], 1)
        return self.main(input)

class VLAE(VAE):

    def __init__(self, input_dims, code_dims, beta=1.0,
                 hidden=400, activacation="lrelu",
                 decoder="Bernoulli", batchnorm=False):

        super(VLAE, self).__init__(input_dims, code_dims)
        self.name = "VLAE"
        self.nx = int(np.prod(input_dims))
        self.nz = int(np.prod(code_dims))
        self.beta = beta
        if activacation == "lrelu":
            self.act = nn.LeakyReLU()
        else:
            self.act = nn.ReLU()
        if decoder == "Bernoulli":
            self.reconstruct_loss = StableBCELoss()
        else:
            self.reconstruct_loss = nn.MSELoss()

        self.encode_layers = nn.ModuleList([EncodeLayer(self.nx, hidden, code_dims[1], batchnorm, activacation)])
        self.decode_layers = nn.ModuleList([])
        for i in range(code_dims[0]-1):
            el = EncodeLayer(hidden, hidden, code_dims[1], batchnorm, activacation)
            dl = DecodeLayer(hidden, hidden, code_dims[1], batchnorm, activacation)
            self.encode_layers.append(el)
            self.decode_layers.append(dl)

        self.fc1 = nn.Linear(code_dims[1], hidden)
        self.fc2 = nn.Linear(hidden, self.nx)

    def encode(self, x):
        h = x.view(x.size(0), -1)
        mu_list = []
        logvar_list = []
        for fc in self.encode_layers:
            h, mu, logvar = fc(h)
            mu_list.append(mu)
            logvar_list.append(logvar)
        return torch.cat(mu_list, dim=1), torch.cat(logvar_list, dim=1)

    def reparametrize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        if isinstance(mu, torch.cuda.FloatTensor):
            eps = torch.cuda.FloatTensor(std.size()).normal_()
        else:
            eps = torch.FloatTensor(std.size()).normal_()
#        eps[:,-2:-1] = (eps[:,-2:-1] - mu.data[:,-2:-1]) / std.data[:,-2:-1]
        eps = Variable(eps)
        return eps.mul(std).add_(mu)

    def decode(self, z):
        zcode = list(torch.chunk(z, self.code_dims[0], dim=1))[::-1]
        h = self.act(self.fc1(zcode[0]))
        for z, fc in zip(zcode[1:], self.decode_layers):
            h = fc(h, z)
        return self.fc2(h)

    def forward(self, x):
        mu, logvar = self.encode(x.view(x.size(0), -1))
        z = self.reparametrize(mu, logvar)
        return self.decode(z), mu, logvar, z

    def loss(self, recon_x, x, mu, logvar, z):
        x = x.view(x.size(0), -1)
        BCE = self.reconstruct_loss(recon_x, x) / x.size(0)
        # see Appendix B from VAE paper:
        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
        # https://arxiv.org/abs/1312.6114
        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
        KLD = torch.sum(KLD_element).mul_(-0.5) / x.size(0)
        return BCE + self.beta * KLD, BCE, KLD

    def mutual_info_q(self, x):
        mu, logvar = self.encode(x.view(x.size(0), -1))
        z = self.reparametrize(mu, logvar)
        l = z.size(0)
        z = z.repeat(l, 1, 1)
        mu = mu.unsqueeze(2).repeat(1,1,l).transpose(1,2)
        logvar = logvar.unsqueeze(2).repeat(1,1,l).transpose(1,2)
        p_matrix =  ( - torch.sum((z - mu) ** 2  / logvar.exp(), dim=2) / 2.0 - 0.5 * torch.sum(logvar, dim=2)).exp_()
        p_split_matrix = (- (z - mu) ** 2  / logvar.exp() / 2.0 - 0.5 * logvar ).exp_()
        p_split_vector = torch.sum(p_split_matrix, dim=1)
        p_vector =  torch.sum(p_matrix, dim=1)
        I = torch.FloatTensor([np.log(l)])
        I_split = torch.FloatTensor([np.log(l)] * int(z.size(2)))
        for i in range(l):
            I += (p_matrix[i][i].log() - p_vector[i].log()).data / l
            I_split += (p_split_matrix[i][i].log() - p_split_vector[i].log()).data / l
        # q(z_i) is not independent..
        # assert np.allclose(I.numpy(), np.sum(I_split.numpy()))
        return I, I_split

class MMDVLAE(VLAE):
    def compute_kernel(self, x, y):
        x_size = x.size(0)
        y_size = y.size(0)
        dim = x.size(1)
        tiled_x = x.unsqueeze(1).repeat(1, y_size, 1)
        tiled_y = y.unsqueeze(0).repeat(x_size, 1, 1)
        return ((-(tiled_x - tiled_y) ** 2).mean(dim=2) / float(dim)).exp_()

    def compute_mmd(self, x, y, sigma_sqr=1.0):
        x_kernel = self.compute_kernel(x, x)
        y_kernel = self.compute_kernel(y, y)
        xy_kernel = self.compute_kernel(x, y)
        return torch.mean(x_kernel) + torch.mean(y_kernel) - 2 * torch.mean(xy_kernel)
    def loss(self, recon_x, x, mu, logvar, z):
        x = x.view(x.size(0), -1)
        BCE = self.reconstruct_loss(recon_x, x) / (x.size(0) * x.size(1))

        true_samples = Variable(torch.FloatTensor(x.size(0), self.nz).normal_())
        MMD = self.compute_mmd(true_samples, z)
        return BCE + self.beta *  MMD , BCE, MMD

class CNNVLAE(VAE):
    def __init__(self, input_dims, code_dims, beta=1.0,
                 hidden=400, activacation="lrelu",
                 decoder="Bernoulli", batchnorm=True):

        super(CNNVLAE, self).__init__(input_dims, code_dims)
        self.name = "CNNVLAE"
        self.nx = input_dims[0]
        self.nz = int(np.prod(code_dims))
        self.beta = beta
        if activacation == "lrelu":
            self.act = nn.LeakyReLU()
        else:
            self.act = nn.ReLU()
        if decoder == "Bernoulli":
            self.reconstruct_loss = StableBCELoss()
        else:
            self.reconstruct_loss = nn.MSELoss()

        assert(input_dims[1] == input_dims[2])
        l = input_dims[1]
        l = int(l/2)
        self.encode_layers = [CNNEncodeLayer(self.nx, hidden, code_dims[1], batchnorm, activacation, (l, l))]
        self.decode_layers = [CNNDecodeLayer(hidden*2, hidden, code_dims[1], batchnorm, activacation, (l, l))]
        self.conv2 = nn.ConvTranspose2d(hidden, 1, 1, 1)
        for i in range(code_dims[0]-2):
            l = int(l/2)
            el = CNNEncodeLayer(hidden, hidden * 2, code_dims[1], batchnorm, activacation, (l, l))
            hidden *= 2
            dl = CNNDecodeLayer(hidden * 2, hidden, code_dims[1], batchnorm, activacation, (l, l))
            self.encode_layers.append(el)
            self.decode_layers.insert(0, dl)
        self.encode_layers = nn.ModuleList(self.encode_layers)
        self.decode_layers = nn.ModuleList(self.decode_layers)
        l = int(l/2)
        self.encode_layers.append(CNNEncodeLayer(hidden, hidden*2, code_dims[1], batchnorm, activacation, (l, l)))
        hidden *= 2
        self.conv1 = CNNDecodeLayer(0, hidden, code_dims[1], batchnorm, activacation, (l, l))

    def encode(self, x):
        h = x
        mu_list = []
        logvar_list = []
        for conv in self.encode_layers:
            h, mu, logvar = conv(h)
            mu_list.append(mu)
            logvar_list.append(logvar)
        return torch.cat(mu_list, dim=1), torch.cat(logvar_list, dim=1)

    def reparametrize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        if isinstance(mu, torch.cuda.FloatTensor):
            eps = torch.cuda.FloatTensor(std.size()).normal_()
        else:
            eps = torch.FloatTensor(std.size()).normal_()
#        eps[:,-2:-1] = (eps[:,-2:-1] - mu.data[:,-2:-1]) / std.data[:,-2:-1]
        eps = Variable(eps)
        return eps.mul(std).add_(mu)

    def decode(self, z):

        zcode = list(torch.chunk(z, self.code_dims[0], dim=1))[::-1]
        h = self.act(self.conv1(None, zcode[0]))
        for z, conv in zip(zcode[1:], self.decode_layers):
            h = conv(h, z)
        return self.conv2(h)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        return self.decode(z), mu, logvar, z

    def loss(self, recon_x, x, mu, logvar, z):
        x = x.view(x.size(0), -1)
        recon_x = recon_x.view(recon_x.size(0), -1)
        BCE = self.reconstruct_loss(recon_x, x) / x.size(0)
        # see Appendix B from VAE paper:
        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
        # https://arxiv.org/abs/1312.6114
        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
        KLD = torch.sum(KLD_element).mul_(-0.5) / x.size(0)
        return BCE + self.beta * KLD, BCE, KLD

    def mutual_info_q(self, x):
        mu, logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        z = z.view(z.size(0), -1)
        mu = mu.view(mu.size(0), -1)
        logvar = logvar.view(logvar.size(0), -1)
        l = z.size(0)
        z = z.repeat(l, 1, 1)
        mu = mu.unsqueeze(2).repeat(1,1,l).transpose(1,2)
        logvar = logvar.unsqueeze(2).repeat(1,1,l).transpose(1,2)
        p_matrix =  ( - torch.sum((z - mu) ** 2  / logvar.exp(), dim=2) / 2.0 - 0.5 * torch.sum(logvar, dim=2)).exp_()
        p_split_matrix = (- (z - mu) ** 2  / logvar.exp() / 2.0 - 0.5 * logvar ).exp_()
        p_split_vector = torch.sum(p_split_matrix, dim=1)
        p_vector =  torch.sum(p_matrix, dim=1)
        I = torch.FloatTensor([np.log(l)])
        I_split = torch.FloatTensor([np.log(l)] * int(z.size(2)))
        for i in range(l):
            I += (p_matrix[i][i].log() - p_vector[i].log()).data / l
            I_split += (p_split_matrix[i][i].log() - p_split_vector[i].log()).data / l
        # q(z_i) is not independent..
        # assert np.allclose(I.numpy(), np.sum(I_split.numpy()))
        return I, I_split