# -*- coding: utf-8 -*-
"""
Contains helper code for running spatial nulls models
"""

from pathlib import Path
import warnings

import numpy as np
from scipy import optimize, spatial
try:  # scipy >= 1.8.0
    from scipy.ndimage._measurements import _stats, labeled_comprehension
except ImportError:  # scipy < 1.8.0
    from scipy.ndimage.measurements import _stats, labeled_comprehension
from sklearn.utils.validation import check_random_state

from neuromaps.images import load_gifti, PARCIGNORE
from neuromaps.points import _geodesic_parcel_centroid

def load_spins(fn, n_perm=None):
    """
    Loads spins from `fn`

    Parameters
    ----------
    fn : os.PathLike
        Filepath to file containing spins to load
    n_perm : int, optional
        Number of spins to retain (i.e., subset data)

    Returns
    -------
    spins : (N, P) array_like
        Loaded spins
    """

    try:
        npy = Path(fn).with_suffix('.npy')
        if npy.exists():
            spins = np.load(npy, allow_pickle=False, mmap_mode='c')
        else:
            spins = np.loadtxt(fn, delimiter=',', dtype='int32')
    except TypeError:
        spins = np.asarray(fn, dtype='int32')

    if n_perm is not None:
        spins = spins[..., :n_perm]

    return spins

def get_parcel_centroids(surfaces, parcellation=None, method='surface',
                         drop=None):
    """
    Returns vertex coordinates corresponding to parcel centroids

    If `parcellation` is not specified then returned `centroids` are vertex
    coordinates of `surfaces`

    Parameters
    ----------
    surfaces : (2,) list-of-str
        Surfaces on which to compute parcel centroids; generally spherical
        surfaces are recommended. Surfaces should be (left, right) hemisphere.
        If no parcellations are provided then returned `centroids` represent
        all vertices in `surfaces`
    parcellation : (2,) list-of-str, optional
        Path to GIFTI label files containing labels of parcels on the
        (left, right) hemisphere. If not specified then vertex coordinates from
        `surfaces` are returned instead. Default: None
    method : {'average', 'surface', 'geodesic'}, optional
        Method for calculation of parcel centroid. See Notes for more
        information. Default: 'surface'
    drop : list, optional
        Specifies regions in `parcellation` for which the parcel centroid
        should not be calculated. If not specified, centroids for parcels
        defined in `PARCIGNORE` are not calculated. Default: None

    Returns
    -------
    centroids : (N, 3) numpy.ndarray
        Coordinates of parcel centroids. If `parcellation` is not specified
        these are simply the vertex coordinates
    hemiid : (N,) numpy.ndarray
        Array denoting hemisphere designation of coordinates in `centroids`,
        where `hemiid=0` denotes the left and `hemiid=1` the right hemisphere

    Notes
    -----
    The following methods can be used for finding parcel centroids:

    1. ``method='average'``

       Uses the arithmetic mean of the coordinates for the vertices in each
       parcel. Note that in this case the calculated centroids will not act
       actually fall on the surface of `surf`.

    2. ``method='surface'``

       Calculates the 'average' coordinates and then finds the closest vertex
       on `surf`, where closest is defined as the vertex with the minimum
       Euclidean distance.

    3. ``method='geodesic'``

       Uses the coordinates of the vertex with the minimum average geodesic
       distance to all other vertices in the parcel. Note that this is slightly
       more time-consuming than the other two methods, especially for
       high-resolution meshes.
    """

    methods = ['average', 'surface', 'geodesic']
    if method not in methods:
        raise ValueError('Provided method for centroid calculation {} is '
                         'invalid. Must be one of {}'.format(methods, methods))

    if drop is None:
        drop = PARCIGNORE
    if parcellation is None:
        parcellation = (None, None)

    centroids, hemiid = [], []
    for n, (parc, surf) in enumerate(zip(parcellation, surfaces)):
        vertices, faces = load_gifti(surf).agg_data()
        if parc is not None:
            labels = load_gifti(parc).agg_data()
            labeltable = parc.labeltable.get_labels_as_dict()

            for lab in np.unique(labels):
                if labeltable.get(lab) in drop:
                    continue

                mask = labels == lab
                if method in ('average', 'surface'):
                    roi = np.atleast_2d(vertices[mask].mean(axis=0))
                    if method == 'surface':  # find closest vertex on surf
                        idx = np.argmin(spatial.distance_matrix(vertices, roi),
                                        axis=0)[0]
                        roi = vertices[idx]
                elif method == 'geodesic':
                    inds, = np.where(mask)
                    roi = _geodesic_parcel_centroid(vertices, faces, inds)

                centroids.append(roi)
                hemiid.append(n)
        else:
            centroids.append(vertices)
            hemiid.extend([n] * len(vertices))

    return np.row_stack(centroids), np.asarray(hemiid)

def _gen_rotation(seed=None):
    """
    Generates random matrix for rotating spherical coordinates

    Parameters
    ----------
    seed : {int, np.random.RandomState instance, None}, optional
        Seed for random number generation

    Returns
    -------
    rotate_{l,r} : (3, 3) numpy.ndarray
        Rotations for left and right hemisphere coordinates, respectively
    """

    rs = check_random_state(seed)

    # for reflecting across Y-Z plane
    reflect = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])

    # generate rotation for left
    rotate_l, temp = np.linalg.qr(rs.normal(size=(3, 3)))
    rotate_l = rotate_l @ np.diag(np.sign(np.diag(temp)))
    if np.linalg.det(rotate_l) < 0:
        rotate_l[:, 0] = -rotate_l[:, 0]

    # reflect the left rotation across Y-Z plane
    rotate_r = reflect @ rotate_l @ reflect

    return rotate_l, rotate_r

def gen_spinsamples(coords, hemiid, n_rotate=1000, check_duplicates=True,
                    method='original', seed=None, verbose=False,
                    return_cost=False):
    """
    Returns a resampling array for `coords` obtained from rotations / spins

    Using the method initially proposed in [ST1]_ (and later modified + updated
    based on findings in [ST2]_ and [ST3]_), this function applies random
    rotations to the user-supplied `coords` in order to generate a resampling
    array that preserves its spatial embedding. Rotations are generated for one
    hemisphere and mirrored for the other (see `hemiid` for more information).

    Due to irregular sampling of `coords` and the randomness of the rotations
    it is possible that some "rotations" may resample with replacement (i.e.,
    will not be a true permutation). The likelihood of this can be reduced by
    either increasing the sampling density of `coords` or changing the
    ``method`` parameter (see Notes for more information on the latter).

    Parameters
    ----------
    coords : (N, 3) array_like
        X, Y, Z coordinates of `N` nodes/parcels/regions/vertices defined on a
        sphere
    hemiid : (N,) array_like
        Array denoting hemisphere designation of coordinates in `coords`, where
        values should be {0, 1} denoting the different hemispheres. Rotations
        are generated for one hemisphere and mirrored across the y-axis for the
        other hemisphere.
    n_rotate : int, optional
        Number of rotations to generate. Default: 1000
    check_duplicates : bool, optional
        Whether to check for and attempt to avoid duplicate resamplings. A
        warnings will be raised if duplicates cannot be avoided. Setting to
        True may increase the runtime of this function! Default: True
    method : {'original', 'vasa', 'hungarian'}, optional
        Method by which to match non- and rotated coordinates. Specifying
        'original' will use the method described in [ST1]_. Specfying 'vasa'
        will use the method described in [ST4]_. Specfying 'hungarian' will use
        the Hungarian algorithm to minimize the global cost of reassignment
        (will dramatically increase runtime). Default: 'original'
    seed : {int, np.random.RandomState instance, None}, optional
        Seed for random number generation. Default: None
    verbose : bool, optional
        Whether to print occasional status messages. Default: False
    return_cost : bool, optional
        Whether to return cost array (specified as Euclidean distance) for each
        coordinate for each rotation Default: True

    Returns
    -------
    spinsamples : (N, `n_rotate`) numpy.ndarray
        Resampling matrix to use in permuting data based on supplied `coords`.
    cost : (N, `n_rotate`,) numpy.ndarray
        Cost (specified as Euclidean distance) of re-assigning each coordinate
        for every rotation in `spinsamples`. Only provided if `return_cost` is
        True.

    Notes
    -----
    By default, this function uses the minimum Euclidean distance between the
    original coordinates and the new, rotated coordinates to generate a
    resampling array after each spin. Unfortunately, this can (with some
    frequency) lead to multiple coordinates being re-assigned the same value:

        >>> from neuromaps.nulls.spins import gen_spinsamples
        >>> coords = [[0, 0, 1], [1, 0, 0], [0, 0, 1], [1, 0, 0]]
        >>> hemi = [0, 0, 1, 1]
        >>> gen_spinsamples(coords, hemi, n_rotate=1, seed=1,
        ...                 check_duplicates=False)
        array([[0],
               [0],
               [2],
               [3]])

    While this is reasonable in most circumstances, if you feel incredibly
    strongly about having a perfect "permutation" (i.e., all indices appear
    once and exactly once in the resampling), you can set the ``method``
    parameter to either 'vasa' or 'hungarian':

        >>> gen_spinsamples(coords, hemi, n_rotate=1, seed=1,
        ...                 method='vasa', check_duplicates=False)
        array([[1],
               [0],
               [2],
               [3]])
        >>> gen_spinsamples(coords, hemi, n_rotate=1, seed=1,
        ...                 method='hungarian', check_duplicates=False)
        array([[0],
               [1],
               [2],
               [3]])

    Note that setting this parameter may increase the runtime of the function
    (especially for `method='hungarian'`). Refer to [ST1]_ for information on
    why the default suffices in most cases.

    For the original MATLAB implementation of this function refer to [ST5]_.

    References
    ----------
    .. [ST1] Alexander-Bloch, A., Shou, H., Liu, S., Satterthwaite, T. D.,
       Glahn, D. C., Shinohara, R. T., Vandekar, S. N., & Raznahan, A. (2018).
       On testing for spatial correspondence between maps of human brain
       structure and function. NeuroImage, 178, 540-51.

    .. [ST2] Blaser, R., & Fryzlewicz, P. (2016). Random Rotation Ensembles.
       Journal of Machine Learning Research, 17(4), 1–26.

    .. [ST3] Lefèvre, J., Pepe, A., Muscato, J., De Guio, F., Girard, N.,
       Auzias, G., & Germanaud, D. (2018). SPANOL (SPectral ANalysis of Lobes):
       A Spectral Clustering Framework for Individual and Group Parcellation of
       Cortical Surfaces in Lobes. Frontiers in Neuroscience, 12, 354.

    .. [ST4] Váša, F., Seidlitz, J., Romero-Garcia, R., Whitaker, K. J.,
       Rosenthal, G., Vértes, P. E., ... & Jones, P. B. (2018). Adolescent
       tuning of association cortex in human structural brain networks.
       Cerebral Cortex, 28(1), 281-294.

    .. [ST5] https://github.com/spin-test/spin-test
    """

    methods = ['original', 'vasa', 'hungarian']
    if method not in methods:
        raise ValueError('Provided method "{}" invalid. Must be one of {}.'
                         .format(method, methods))

    seed = check_random_state(seed)

    coords = np.asanyarray(coords)
    hemiid = np.squeeze(np.asanyarray(hemiid, dtype='int8'))

    # check supplied coordinate shape
    if coords.shape[-1] != 3 or coords.squeeze().ndim != 2:
        raise ValueError('Provided `coords` must be of shape (N, 3), not {}'
                         .format(coords.shape))

    # ensure hemisphere designation array is correct
    if hemiid.ndim != 1:
        raise ValueError('Provided `hemiid` array must be one-dimensional.')
    if len(coords) != len(hemiid):
        raise ValueError('Provided `coords` and `hemiid` must have the same '
                         'length. Provided lengths: coords = {}, hemiid = {}'
                         .format(len(coords), len(hemiid)))
    if np.max(hemiid) > 1 or np.min(hemiid) < 0:
        raise ValueError('Hemiid must have values in {0, 1} denoting left and '
                         'right hemisphere coordinates, respectively. '
                         + 'Provided array contains values: {}'
                         .format(np.unique(hemiid)))

    # empty array to store resampling indices
    spinsamples = np.zeros((len(coords), n_rotate), dtype=int)
    cost = np.zeros((len(coords), n_rotate))
    inds = np.arange(len(coords), dtype=int)

    # generate rotations and resampling array!
    msg, warned = '', False
    for n in range(n_rotate):
        count, duplicated = 0, True

        if verbose:
            msg = 'Generating spin {:>5} of {:>5}'.format(n, n_rotate)
            print(msg, end='\r', flush=True)

        while duplicated and count < 500:
            count, duplicated = count + 1, False
            resampled = np.zeros(len(coords), dtype='int32')

            # rotate each hemisphere separately
            for h, rot in enumerate(_gen_rotation(seed=seed)):
                hinds = (hemiid == h)
                coor = coords[hinds]
                if len(coor) == 0:
                    continue

                # if we need an "exact" mapping (i.e., each node needs to be
                # assigned EXACTLY once) then we have to calculate the full
                # distance matrix which is a nightmare with respect to memory
                # for anything that isn't parcellated data.
                # that is, don't do this with vertex coordinates!
                if method == 'vasa':
                    dist = spatial.distance_matrix(coor, coor @ rot)
                    # min of max a la Vasa et al., 2018
                    col = np.zeros(len(coor), dtype='int32')
                    for r in range(len(dist)):
                        # find parcel whose closest neighbor is farthest away
                        # overall; assign to that
                        row = dist.min(axis=1).argmax()
                        col[row] = dist[row].argmin()
                        cost[inds[hinds][row], n] = dist[row, col[row]]
                        # set to -inf and inf so they can't be assigned again
                        dist[row] = -np.inf
                        dist[:, col[row]] = np.inf
                # optimization of total cost using Hungarian algorithm. this
                # may result in certain parcels having higher cost than with
                # `method='vasa'` but should always result in the total cost
                # being lower #tradeoffs
                elif method == 'hungarian':
                    dist = spatial.distance_matrix(coor, coor @ rot)
                    row, col = optimize.linear_sum_assignment(dist)
                    cost[hinds, n] = dist[row, col]
                # if nodes can be assigned multiple targets, we can simply use
                # the absolute minimum of the distances (no optimization
                # required) which is _much_ lighter on memory
                # huge thanks to https://stackoverflow.com/a/47779290 for this
                # memory-efficient method
                elif method == 'original':
                    dist, col = spatial.cKDTree(coor @ rot).query(coor, 1)
                    cost[hinds, n] = dist

                resampled[hinds] = inds[hinds][col]

            # if we want to check for duplicates ensure that we don't have any
            if check_duplicates:
                if np.any(np.all(resampled[:, None] == spinsamples[:, :n], 0)):
                    duplicated = True
                # if our "spin" is identical to the input then that's no good
                elif np.all(resampled == inds):
                    duplicated = True

        # if we broke out because we tried 500 rotations and couldn't generate
        # a new one, warn that we're using duplicate rotations and give up.
        # this should only be triggered if check_duplicates is set to True
        if count == 500 and not warned:
            warnings.warn('Duplicate rotations used. Check resampling array '
                          'to determine real number of unique permutations.')
            warned = True

        spinsamples[:, n] = resampled

    if verbose:
        print(' ' * len(msg) + '\b' * len(msg), end='', flush=True)

    if return_cost:
        return spinsamples, cost

    return spinsamples

def spin_parcels(surfaces, parcellation, method='surface', n_rotate=1000,
                 spins=None, verbose=False, **kwargs):
    """
    Rotates parcels in `parcellation` and re-assigns based on maximum overlap

    Vertex labels are rotated and a new label is assigned to each *parcel*
    based on the region maximally overlapping with its boundaries.

    Parameters
    ----------
    surfaces : (2,) list-of-str
        Surfaces to use for rotating parcels; generally spherical surfaces
        are recommended. Surfaces should be (left, right) hemisphere
    parcellation : (2,) list-of-str, optional
        Path to GIFTI label files containing parcel labels on the (left, right)
        hemisphere of `surfaces`
    n_rotate : int, optional
        Number of rotations to generate. Default: 1000
    spins : array_like, optional
        Pre-computed spins to use instead of generating them on the fly. If not
        provided will use other provided parameters to create them. Default:
        None
    seed : {int, np.random.RandomState instance, None}, optional
        Seed for random number generation. Default: None
    verbose : bool, optional
        Whether to print occasional status messages. Default: False
    return_cost : bool, optional
        Whether to return cost array (specified as Euclidean distance) for each
        coordinate for each rotation. Default: True
    kwargs : key-value pairs
        Keyword arguments passed to :func:`~.gen_spinsamples`

    Returns
    -------
    spinsamples : (N, `n_rotate`) numpy.ndarray
        Resampling matrix to use in permuting data parcellated with labels from
        `parcellation`, where `N` is the number of parcels. Indices of -1
        indicate that the parcel was completely encompassed by regions in
        `drop` and should be ignored.
    """

    def overlap(vals):
        """ Returns most common positive value in `vals`; -1 if all negative
        """
        vals = np.asarray(vals)
        vals, counts = np.unique(vals[vals > 0], return_counts=True)
        try:
            return vals[counts.argmax()] - 1
        except ValueError:
            return -1

    # get vertex-level labels (set drop labels to - values)
    vertices = np.hstack([
        load_gifti(parc).agg_data() for parc in parcellation
    ])
    labels = np.unique(vertices)
    mask = labels != 0

    # get spins + cost (if requested)
    if spins is None:
        coords, hemiid = get_parcel_centroids(surfaces, method=method)
        spins = gen_spinsamples(coords, hemiid, n_rotate=n_rotate,
                                verbose=verbose, **kwargs)
        if kwargs.get('return_cost'):
            spins, cost = spins
    spins = load_spins(spins)

    if len(vertices) != len(spins):
        raise ValueError('Provided annotation files have a different '
                         'number of vertices than the specified fsaverage '
                         'surface.\n    ANNOTATION: {} vertices\n     '
                         'FSAVERAGE:  {} vertices'
                         .format(len(vertices), len(spins)))

    # spin and assign regions based on max overlap
    regions = np.zeros((len(labels[mask]), n_rotate), dtype='int32')
    for n in range(n_rotate):
        if verbose:
            msg = f'Calculating parcel overlap: {n:>5}/{n_rotate}'
            print(msg, end='\b' * len(msg), flush=True)
        regions[:, n] = labeled_comprehension(vertices[spins[:, n]], vertices,
                                              labels, overlap, int, -1)[mask]

    if kwargs.get('return_cost'):
        return regions, cost

    return regions

def parcels_to_vertices(data, parcellation):
    """
    Projects parcellated `data` to vertices as defined by `parcellation`

    Parameters
    ----------
    data : (N,) numpy.ndarray
        Parcellated data to be projected to vertices
    parcellation : tuple-of-str or os.PathLike
        Filepaths to parcellation images to project `data` to vertices

    Reurns
    ------
    projected : numpy.ndarray
        Vertex-level data
    """

    data = np.vstack(data).astype(float)
    vertices = np.hstack([
        load_gifti(parc).agg_data() for parc in parcellation
    ])
    expected = np.unique(vertices)[1:].size
    n_vert = vertices.shape[0]
    if expected != len(data):
        raise ValueError('Number of parcels in provided annotation files '
                         'differs from size of parcellated data array.\n'
                         '    EXPECTED: {} parcels\n'
                         '    RECEIVED: {} parcels'
                         .format(expected, len(data)))

    projected = np.zeros((n_vert, data.shape[-1]), dtype=data.dtype)
    n_vert = 0
    for parc in parcellation:
        labels = load_gifti(parc).agg_data().astype('int')
        currdata = np.append([[np.nan]], data, axis=0)
        projected[n_vert:n_vert + len(labels), :] = currdata[labels, :]
        n_vert += len(labels)

    return np.squeeze(projected)

def vertices_to_parcels(data, parcellation):
    """
    Reduces vertex-level `data` to parcels defined by `parcellation`

    Takes average of vertices within each parcel (excluding NaN values).
    Assigns NaN to parcels for which *all* vertices are NaN.

    Parameters
    ----------
    data : (N,) numpy.ndarray
        Vertex-level data to be reduced to parcels
    parcellation : tuple-of-str or os.PathLike
        Filepaths to parcellation images to parcellate `data`

    Reurns
    ------
    reduced : numpy.ndarray
        Parcellated `data`
    """

    data = np.vstack(data)
    vertices = np.hstack([
        load_gifti(parc).agg_data() for parc in parcellation
    ])
    n_parc = np.unique(vertices).size
    expected = vertices.shape[0]
    if expected != len(data):
        raise ValueError('Number of vertices in provided annotation files '
                         'differs from size of vertex-level data array.\n'
                         '    EXPECTED: {} vertices\n'
                         '    RECEIVED: {} vertices'
                         .format(expected, len(data)))

    numerator = np.zeros((n_parc, data.shape[-1]), dtype=data.dtype)
    denominator = np.zeros((n_parc, data.shape[-1]), dtype=data.dtype)
    start = end = 0
    for parc in parcellation:
        labels = load_gifti(parc).agg_data().astype('int')
        indices = np.unique(labels)
        end += len(labels)

        for idx in range(data.shape[-1]):
            currdata = np.squeeze(data[start:end, idx])
            counts, sums = _stats(np.nan_to_num(currdata), labels, indices)
            _, nacounts = _stats(np.isnan(currdata), labels, indices)
            counts = (np.asanyarray(counts, dtype=float)
                      - np.asanyarray(nacounts, dtype=float))

            numerator[indices, idx] += sums
            denominator[indices, idx] += counts

        start = end

    with np.errstate(divide='ignore', invalid='ignore'):
        reduced = np.squeeze(numerator / denominator)[1:]

    return reduced

def spin_data(data, surfaces, parcellation, method='surface', n_rotate=1000,
              spins=None, verbose=False, **kwargs):
    """
    Projects parcellated `data` to `surfaces`, rotates, and re-parcellates

    Projection of `data` to `surfaces` uses provided `parcellation` files.
    Re-parcellated data will not be exactly identical to original values due to
    re-averaging process. Parcels subsumed by regions in `drop` will be listed
    as NaN.

    Parameters
    ----------
    data : (N,) numpy.ndarray
        Parcellated data to be rotated. Parcels should be ordered by [left,
        right] hemisphere; ordering within hemisphere should correspond to the
        provided `parcellation` files.
    surfaces : (2,) list-of-str
        Surfaces to use for rotating parcels; generally spherical surfaces
        are recommended. Surfaces should be (left, right) hemisphere
    parcellation : (2,) list-of-str, optional
        Path to GIFTI label files containing parcel labels on the (left, right)
        hemisphere of `surfaces` mapping `data` to vertices in `surfaces`
    n_rotate : int, optional
        Number of rotations to generate. Default: 1000
    spins : array_like, optional
        Pre-computed spins to use instead of generating them on the fly. If not
        provided will use other provided parameters to create them. Default:
        None

    verbose : bool, optional
        Whether to print occasional status messages. Default: False
    kwargs : key-value pairs
        Keyword arguments passed to function used to generate rotations

    Returns
    -------
    rotated : (N, `n_rotate`) numpy.ndarray
        Rotated `data
    """

    # get coordinates and hemisphere designation for spin generation
    vertices = parcels_to_vertices(data, parcellation)

    if spins is None:
        coords, hemiid = get_parcel_centroids(surfaces, method=method)
        spins = gen_spinsamples(coords, hemiid, n_rotate=n_rotate,
                                verbose=verbose, **kwargs)
        if kwargs.get('return_cost'):
            spins, cost = spins
    spins = load_spins(spins)

    if len(vertices) != len(spins):
        raise ValueError('Provided parcellation files have a different '
                         'number of vertices than the specified surfaces.\n'
                         '    ANNOTATION: {} vertices\n'
                         '     FSAVERAGE: {} vertices'
                         .format(len(vertices), len(spins)))

    spun = np.zeros(data.shape + (n_rotate,))
    for n in range(n_rotate):
        if verbose:
            msg = f'Reducing vertices to parcels: {n:>5}/{n_rotate}'
            print(msg, end='\b' * len(msg), flush=True)
        spun[..., n] = vertices_to_parcels(vertices[spins[:, n]], parcellation)

    if verbose:
        print(' ' * len(msg) + '\b' * len(msg), end='', flush=True)

    if kwargs.get('return_cost'):
        return spun, cost

    return spun


from pyscf import scf, dft, gto, ao2mo, df, lib, cc
from pyscf.dft.numint import eval_ao, eval_rho
from pyscf.dft.gen_grid import Grids
from pyscf.pbc.tools.pyscf_ase import atoms_from_ase
import numpy as np
import logging

CALC_TYPES = {
    'RHF'   : scf.hf.RHF,
    'UHF'   : scf.uhf.UHF,
    'RKS'   : dft.rks.RKS,
    'UKS'   : dft.uks.UKS,
    'CCSD'  : cc.ccsd.CCSD,
    'UCCSD' : cc.uccsd.UCCSD
}

SCF_TYPES = {
    'RHF'  : scf.hf.RHF,
    'ROHF' : scf.rohf.ROHF,
    'UHF'  : scf.uhf.UHF,
    'RKS'  : dft.RKS,
    'UKS'  : dft.UKS
}

########################################################
# BASIC HELPER ROUTINES FOR RUNNING PYSCF CALCULATIONS #
########################################################

def mol_from_ase(atoms, basis, spin=0, charge=0):
    """
    Get a pyscf gto.Mole object from an ase Atoms object (atoms).
    Assign it the atomic basis set (basis).
    Return the Mole object.
    """
    mol = gto.Mole()
    mol.atom = atoms_from_ase(atoms)
    mol.basis = basis
    mol.spin = spin
    mol.charge = charge
    mol.build()
    return mol

def setup_rks_calc(mol, xc, grid_level=3, vv10=False, **kwargs):
    """
    Set up a PySCF RKS calculation with sensible defaults.
    """
    rks = dft.RKS(mol)
    rks.xc = xc
    rks.grids.level = grid_level
    rks.grids.build()
    logging.warning('xc: {}, grid level: {}'.format(xc, grid_level))
    if vv10:
        logging.warning('Using VV10 in UKS setup')
        rks.nlc = 'VV10'
        if np.array([gto.charge(mol.atom_symbol(i)) <= 18 for i in range(mol.natm)]).all():
            rks.nlcgrids.prune = dft.gen_grid.sg1_prune
        else:
            rks.nlcgrids.prune = None
        rks.nlcgrids.level = 1
    return rks

def setup_uks_calc(mol, xc, grid_level=3, vv10=False, **kwargs):
    """
    Set up a PySCF UKS calculation with sensible defaults.
    """
    uks = dft.UKS(mol)
    uks.xc = xc
    uks.grids.level = grid_level
    uks.grids.build()
    logging.warning('xc: {}, grid level: {}'.format(xc, grid_level))
    if vv10:
        logging.warning('Using VV10 in UKS setup')
        uks.nlc = 'VV10'
        if np.array([gto.charge(mol.atom_symbol(i)) <= 18 for i in range(mol.natm)]).all():
            uks.nlcgrids.prune = dft.gen_grid.sg1_prune
        else:
            uks.nlcgrids.prune = None
        uks.nlcgrids.level = 1
    return uks

def run_scf(mol, calc_type, functional=None, remove_ld=False, dm0=None):
    """
    Run an SCF calculation on a gto.Mole object (Mole)
    of a given calc_type in SCF_TYPES. Return the calc object.
    Note, if RKS or UKS is the calc_type, default functional is used.
    """
    if not calc_type in SCF_TYPES:
        raise ValueError('Calculation type must be in {}'.format(list(SCF_TYPES.keys())))

    calc = SCF_TYPES[calc_type](mol)
    if remove_ld:
        logging.info("Removing linear dependence from overlap matrix")
        calc = scf.addons.remove_linear_dep_(calc)
    if 'KS' in calc_type and functional is not None:
        calc.xc = functional
        if 'MN' in functional:
            logging.info('MN grid level 4')
            calc.grids.level = 4
        if functional == 'wB97M_V':
            logging.info('Using Specialized wB97M-V params')
            calc.nlc = 'VV10'
            calc.grids.prune = None
            calc.grids.level = 4
            if np.array([gto.charge(mol.atom_symbol(i)) <= 18 for i in range(mol.natm)]).all():
                calc.nlcgrids.prune = dft.gen_grid.sg1_prune
            else:
                calc.nlcgrids.prune = None
            calc.nlcgrids.level = 1

    calc.kernel(dm0 = dm0)
    return calc

def run_cc(hf):
    """
    Run and return a restricted CCSD calculation on mol,
    with HF molecular orbital coefficients in the RHF object hf.
    """
    if type(hf) == SCF_TYPES['RHF']:
        calc_cls = cc.CCSD
    elif type(hf) == SCF_TYPES['UHF']:
        calc_cls = cc.UCCSD
    else:
        raise NotImplementedError('HF type {} not supported'.format(type(hf)) +\
            '\nSupported Types: {}'.format(SCF_TYPES['RHF'], SCF_TYPES['UHF']))
    calc = calc_cls(hf)
    calc.kernel()
    return calc

#############################################
# HELPER FUNCTIONS FOR THE analyzers MODULE #
#############################################

def get_grid(mol, level=3):
    """
    Get the real-space grid of a molecule for numerical integration.
    """
    grid = Grids(mol)
    grid.level = level
    grid.kernel()
    return grid

def get_ha_total(rdm1, eeint):
    return np.sum(np.sum(eeint * rdm1, axis=(2,3)) * rdm1)

def get_hf_coul_ex_total(mol, hf):
    rdm1 = hf.make_rdm1()
    jmat, kmat = hf.get_jk(mol, rdm1)
    return np.sum(jmat * rdm1) / 2, -np.sum(kmat * rdm1) / 4

def get_hf_coul_ex_total2(rdm1, jmat, kmat):
    if len(rdm1.shape) == 2:
        return np.sum(jmat * rdm1) / 2, -np.sum(kmat * rdm1) / 4
    else:
        return np.sum(jmat * np.sum(rdm1, axis=0)) / 2, -np.sum(kmat * rdm1) / 2

def get_hf_coul_ex_total_unrestricted(mol, hf):
    rdm1 = hf.make_rdm1()
    jmat, kmat = hf.get_jk(mol, rdm1)
    return np.sum(jmat * np.sum(rdm1, axis=0)) / 2, -np.sum(kmat * rdm1) / 2

def transform_basis_1e(mat, coeff):
    """
    Transforms the 1-electron matrix mat into the basis
    described by coeff (with the basis vectors being the columns).
    To transform AO operator to MO operator, pass mo_coeff.
    To transform MO operator to AO operator, pass inv(mo_coeff).
    To transform AO density matrix to MO density matrix, pass inv(transpose(mo_coeff)).
    To transform MO density matrix to AO density matrix, pass transpose(mo_coeff).
    """
    if len(coeff.shape) == 2:
        return np.matmul(coeff.transpose(), np.matmul(mat, coeff))
    else:
        if len(coeff) != 2 or len(mat) != 2:
            raise ValueError('Need two sets of orbitals, two mats for unrestricted case.')
        part0 = np.matmul(coeff[0].transpose(), np.matmul(mat[0], coeff[0]))
        part1 = np.matmul(coeff[1].transpose(), np.matmul(mat[1], coeff[1]))
        return np.array([part0, part1])

def make_rdm2_from_rdm1(rdm1):
    """
    For an RHF calculation, return the 2-RDM from
    a given 1-RDM. Given D2(ijkl)=<psi| i+ k+ l j |psi>,
    and D(ij)=<psi| i+ j |psi>, then
    D2(ijkl) = D(ij) * D(kl) - 0.5 * D(lj) * D(ki)
    """
    rdm1copy = rdm1.copy()
    part1 = np.einsum('ij,kl->ijkl', rdm1, rdm1copy)
    part2 = np.einsum('lj,ki->ijkl', rdm1, rdm1copy)
    return part1 - 0.5 * part2

def make_rdm2_from_rdm1_unrestricted(rdm1):
    """
    For a UHF calculation, return the 2-RDM from
    a given 1-RDM. Given D2(ijkl)=<psi| i+ k+ l j |psi>,
    and D(ij)=<psi| i+ j |psi>, then:
    For like spin, D2(ijkl) = D(ij) * D(kl) - D(lj) * D(ki).
    For opposite spin, D2(ijkl) = D(ij) * D(kl)
    Return D(uu,ijkl), D(ud,ijkl), D(dd,ijkl)
    """
    spinparts = []
    rdm1copy = rdm1.copy()
    for s in [0,1]:
        part1 = np.einsum('ij,kl->ijkl', rdm1[s], rdm1copy[s])
        part2 = np.einsum('lj,ki->ijkl', rdm1[s], rdm1copy[s])
        spinparts.append(part1 - part2)
    mixspinpart = np.einsum('ij,kl->ijkl', rdm1[0], rdm1copy[1])
    return np.array([spinparts[0], mixspinpart, spinparts[1]])

def get_ao_vals(mol, points):
    return eval_ao(mol, points)

def get_mgga_data(mol, grid, rdm1):
    """
    Get atomic orbital and density data.
    See eval_ao and eval_rho docs for details.
    Briefly, returns 0-3 derivatives of the atomic orbitals
    in ao_data;
    and the density, first derivatives of density,
    Laplacian of density, and kinetic energy density
    in rho_data.
    """
    ao_data = eval_ao(mol, grid.coords, deriv=3)
    if len(rdm1.shape) == 2:
        rho_data = eval_rho(mol, ao_data, rdm1, xctype='mGGA')
    else:
        part0 = eval_rho(mol, ao_data, rdm1[0], xctype='mGGA')
        part1 = eval_rho(mol, ao_data, rdm1[1], xctype='mGGA')
        rho_data = np.array([part0, part1])
    return ao_data, rho_data

def get_tau_and_grad_helper(mol, grid, rdm1, ao_data):
    """
    Passes the derivatives of the atomic orbitals
    to eval_rho to get the kinetic energy density and its
    derivatives. Not sure if this works.
    """
    # 0 1 2 3 4  5  6  7  8  9
    # 0 x y z xx xy xz yy yz zz
    aox = ao_data[[1, 4, 5, 6]]
    aoy = ao_data[[2, 5, 7, 8]]
    aoz = ao_data[[3, 6, 8, 9]]
    tau  = eval_rho(mol, aox, rdm1, xctype='GGA')
    tau += eval_rho(mol, aoy, rdm1, xctype='GGA')
    tau += eval_rho(mol, aoz, rdm1, xctype='GGA')
    return 0.5 * tau

def get_tau_and_grad(mol, grid, rdm1, ao_data):
    if len(rdm1.shape) == 2:
        return get_tau_and_grad_helper(mol, grid, rdm1, ao_data)
    else:
        return np.array([get_tau_and_grad_helper(mol, grid, rdm1[0], ao_data),\
                        get_tau_and_grad_helper(mol, grid, rdm1[1], ao_data)])

def get_rho_second_deriv_helper(mol, grid, dm, ao):
    from pyscf.dft.numint import _contract_rho, _dot_ao_dm
    from pyscf.dft.gen_grid import make_mask, BLKSIZE

    nao = mol.nao_nr()
    N = grid.weights.shape[0]
    non0tab = np.ones(((N+BLKSIZE-1)//BLKSIZE, mol.nbas),
                         dtype=np.uint8)
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    c0 = _dot_ao_dm(mol, ao[0], dm, non0tab, shls_slice, ao_loc)
    c1 = np.zeros((3, N, nao))
    # 0 1 2 3 4  5  6  7  8  9
    # 0 x y z xx xy xz yy yz zz
    # - - - - 0  1  2  3  4  5
    # - - - - 11 12 13 22 23 33
    ddrho = np.zeros((6, N))
    alphas = [0, 0, 0, 1, 1, 2]
    betas =  [0, 1, 2, 1, 2, 2]
    for i in range(3):
        c1[i] = _dot_ao_dm(mol, ao[i+1], dm.T, non0tab, shls_slice, ao_loc)
    for i in range(6):
        term1 = _contract_rho(c0, ao[i + 4])
        term2 = _contract_rho(c1[alphas[i]], ao[betas[i]+1])
        total = term1 + term2
        ddrho[i] = total + total.conj()
    return ddrho

def get_rho_second_deriv(mol, grid, rdm1, ao_data):
    if len(rdm1.shape) == 2:
        return get_rho_second_deriv_helper(mol, grid, rdm1, ao_data)
    else:
        return np.array([get_rho_second_deriv_helper(mol, grid, rdm1[0], ao_data),\
                        get_rho_second_deriv_helper(mol, grid, rdm1[1], ao_data)])

def get_vele_mat(mol, points):
    """
    Return shape (N, nao, nao)
    """
    auxmol = gto.fakemol_for_charges(points)
    vele_mat = df.incore.aux_e2(mol, auxmol)
    return np.ascontiguousarray(np.transpose(vele_mat, axes=(2,0,1)))

def get_mo_vals(ao_vals, mo_coeff):
    """
    Args:
        ao_vals shape (N,nao)
        mo_coeff shape (nao,nao)
    Returns
        shape (N,nao)
    """
    return np.matmul(ao_vals, mo_coeff)

def get_mo_vele_mat(vele_mat, mo_coeff):
    """
    Convert the return value of get_vele_mat to the MO basis.
    """
    if len(mo_coeff.shape) == 2:
        return np.matmul(mo_coeff.transpose(),
            np.matmul(vele_mat, mo_coeff))
    else:
        tmp = np.einsum('puv,svj->spuj', vele_mat, mo_coeff)
        return np.einsum('sui,spuj->spij', mo_coeff, tmp)

def get_vele_mat_chunks(mol, points, num_chunks, orb_vals, mo_coeff=None):
    """
    Generate chunks of vele_mat on the fly to reduce memory load.
    """
    num_pts = points.shape[0]
    for i in range(num_chunks):
        start = (i * num_pts) // num_chunks
        end = ((i+1) * num_pts) // num_chunks
        auxmol = gto.fakemol_for_charges(points[start:end])
        orb_vals_chunk = orb_vals[start:end]
        vele_mat_chunk = df.incore.aux_e2(mol, auxmol)
        vele_mat_chunk = np.ascontiguousarray(np.transpose(
                                vele_mat_chunk, axes=(2,0,1)))
        if mo_coeff is not None:
            vele_mat_chunk = get_mo_vele_mat(vele_mat_chunk, mo_coeff)
        yield vele_mat_chunk, orb_vals_chunk

def get_vele_mat_generator(mol, points, num_chunks, mo_coeff=None):
    get_generator = lambda orb_vals: get_vele_mat_chunks(mol, points,
                                num_chunks, orb_vals, mo_coeff)
    return get_generator

def get_ha_energy_density(mol, rdm1, vele_mat, ao_vals):
    """
    Get the classical Hartree energy density on a real-space grid,
    for a given molecular structure with basis set (mol),
    for a given 1-electron reduced density matrix (rdm1).
    Returns the Hartree energy density.
    """
    if len(rdm1.shape) == 2:
        Vele = np.einsum('pij,ij->p', vele_mat, rdm1)
    else:
        rdm1 = np.array(rdm1)
        Vele = np.einsum('pij,sij->p', vele_mat, rdm1)
    rho = eval_rho(mol, ao_vals, rdm1)
    return 0.5 * Vele * rho

def get_fx_energy_density(mol, mo_occ, mo_vele_mat, mo_vals):
    """
    Get the Hartree Fock exchange energy density on a real-space grid,
    for a given molecular structure with basis set (mol),
    for a given atomic orbital (AO) 1-electron reduced density matrix (rdm1).
    Returns the exchange energy density, which is negative.
    """
    A = mo_occ * mo_vals
    tmp = np.einsum('pi,pij->pj', A, mo_vele_mat)
    return -0.25 * np.sum(A * tmp, axis=1)

# The following functions are helpers that check whether vele_mat
# is a numpy array or a generator before passing to the methods
# above. This allows one to integrate the memory-saving (but slower)
# chunk-generating approach smoothly.

def get_ha_energy_density2(mol, rdm1, vele_mat, ao_vals):
    if isinstance(vele_mat, np.ndarray):
        return get_ha_energy_density(mol, rdm1, vele_mat, ao_vals)
    else:
        ha_energy_density = np.array([])
        for vele_mat_chunk, orb_vals_chunk in vele_mat(ao_vals):
            ha_energy_density = np.append(ha_energy_density,
                                    get_ha_energy_density(mol, rdm1,
                                        vele_mat_chunk, orb_vals_chunk))
        return ha_energy_density

def get_fx_energy_density2(mol, mo_occ, mo_vele_mat, mo_vals):
    # make sure to test that the grids end up the same
    if isinstance(mo_vele_mat, np.ndarray):
        return get_fx_energy_density(mol, mo_occ, mo_vele_mat, mo_vals)
    else:
        fx_energy_density = np.array([])
        for vele_mat_chunk, orb_vals_chunk in mo_vele_mat(mo_vals):
            fx_energy_density = np.append(fx_energy_density,
                                    get_fx_energy_density(mol, mo_occ,
                                        vele_mat_chunk, orb_vals_chunk))
        return fx_energy_density

def mol_from_dict(mol_dict):
    for item in ['charge', 'spin', 'symmetry', 'verbose']:
        if type(mol_dict[item]).__module__ == np.__name__:
            mol_dict[item] = mol_dict[item].item()
    mol = gto.mole.unpack(mol_dict)
    mol.build()
    return mol

def get_scf(calc_type, mol, calc_data = None):
    calc = CALC_TYPES[calc_type](mol)
    calc.__dict__.update(calc_data)
    return calc

def get_ccsd(calc_type, mol, calc_data = None):
    if calc_type == 'CCSD':
        hf = scf.hf.RHF(mol)
    else:
        hf = scf.uhf.UHF(mol)
    hf.e_tot = calc_data.pop('e_tot') - calc_data['e_corr']
    calc = CALC_TYPES[calc_type](hf)
    calc.__dict__.update(calc_data)
    return calc

def load_calc(fname):
    analyzer_dict = lib.chkfile.load(fname, 'analyzer')
    mol = mol_from_dict(analyzer_dict['mol'])
    calc_type = analyzer_dict['calc_type']
    if 'CCSD' in calc_type:
        return get_ccsd(calc_type, mol, analyzer_dict['calc']), calc_type
    else:
        return get_scf(calc_type, mol, analyzer_dict['calc']), calc_type

def load_analyzer_data(fname):
    data_file = os.path.join(dirname, fname)
    return lib.chkfile.load(data_file, 'analyzer/data')

##################################################
# HELPER FUNCTIONS FOR COMPUTING DFT INGREDIENTS #
##################################################

def get_ws_radii(rho):
    return (3.0 / (4 * np.pi * rho + 1e-16))**(1.0/3)

def get_gradient_magnitude(rho_data):
    return np.linalg.norm(rho_data[1:4,:], axis=0)

def get_normalized_grad(rho, mag_grad):
    sprefac = 2 * (3 * np.pi * np.pi)**(1.0/3)
    n43 = rho**(4.0/3)
    s = mag_grad / (sprefac * n43 + 1e-16)
    return s

def get_single_orbital_tau(rho, mag_grad):
    return mag_grad**2 / (8 * rho + 1e-16)

def get_uniform_tau(rho):
    return (3.0/10) * (3*np.pi**2)**(2.0/3) * rho**(5.0/3)

def get_regularized_tau(tau, tau_w, tau_unif):
    alpha = (tau - tau_w) / (tau_unif + 1e-4)
    return alpha**3 / (alpha**2 + 1e-3)

def get_normalized_tau(tau, tau_w, tau_unif):
    return (tau - tau_w) / (tau_unif + 1e-16)

def get_dft_input(rho_data):
    rho = rho_data[0,:]
    r_s = get_ws_radii(rho)
    mag_grad = get_gradient_magnitude(rho_data)
    s = get_normalized_grad(rho, mag_grad)
    tau_w = get_single_orbital_tau(rho, mag_grad)
    tau_unif = get_uniform_tau(rho)
    alpha = get_regularized_tau(rho_data[5], tau_w, tau_unif)
    return rho, s, alpha, tau_w, tau_unif

def get_dft_input2(rho_data):
    rho = rho_data[0,:]
    r_s = get_ws_radii(rho)
    mag_grad = get_gradient_magnitude(rho_data)
    s = get_normalized_grad(rho, mag_grad)
    tau_w = get_single_orbital_tau(rho, mag_grad)
    tau_unif = get_uniform_tau(rho)
    alpha = get_normalized_tau(rho_data[5], tau_w, tau_unif)
    return rho, s, alpha, tau_w, tau_unif

def squish_density(rho_data, coords, weights, alpha):
    new_coords = coords / alpha
    new_weights = weights / alpha**3
    rho_data = rho_data.copy()
    rho_data[0,:] *= alpha**3
    rho_data[1:4,:] *= alpha**4
    rho_data[4:6,:] *= alpha**5
    return new_coords, new_weights, rho_data

def squish_tau(tau_data, alpha):
    tau_data = tau_data.copy()
    tau_data[0,:] *= alpha**5
    tau_data[1:4] *= alpha**6
    return tau_data


#########################################################################
# Dicomifier - Copyright (C) Universite de Strasbourg
# Distributed under the terms of the CeCILL-B license, as published by
# the CEA-CNRS-INRIA. Refer to the LICENSE file or to
# http://www.cecill.info/licences/Licence_CeCILL-B_V1-en.html
# for details.
#########################################################################

"""
Extract diffusion-related information from JSON meta-data and convert it to
other formats. Unless otherwise specified, all b-values extracted from meta-data
are expressed in :math:`s/m^2` (i.e. SI units).
"""

import base64
import binascii
import re

import numpy

from .. import dicom_to_nifti
from .. import logger

def from_standard(data):
    """ Extract diffusion gradient direction and b-value from standard DICOM
        elements (MR Diffusion Sequence).
    """

    diffusion_data = data["MRDiffusionSequence"]

    scheme = []
    for entry in diffusion_data:
        entry = entry[0]
        if "DiffusionBValue" not in entry:
            raise Exception("Missing b-value")
        b_value = entry["DiffusionBValue"][0]
        # Convert from s/mm^2 to s/m^2
        b_value *= 1e6

        if "DiffusionGradientDirectionSequence" not in entry:
            if b_value != 0:
                raise Exception("Missing direction")
            else:
                entry["DiffusionGradientDirectionSequence"] = [{
                    "DiffusionGradientOrientation": [0,0,0]}]
        direction = entry["DiffusionGradientDirectionSequence"][0]
        if "DiffusionGradientOrientation" not in direction:
            raise Exception("Missing direction")
        direction = numpy.array(direction["DiffusionGradientOrientation"])
        norm = numpy.linalg.norm(direction)
        if norm > 0:
            direction /= norm

        scheme.append((b_value, direction))

    return scheme

def from_siemens_csa(data):
    """ Extract diffusion gradient direction and b-value from Siemens-specific
        elements (CSA Image Header Info (0029,xx10)).
    """

    logger.warning(
        "The coordinate system of the gradient direction is unspecified. "
        "Results may be wrong on non-axial images.")

    scheme = []

    # Look for "SIEMENS CSA HEADER" private creator and get the concrete tag
    # of CSA Image Header Info (0029,xx10)
    element = None
    for tag, item in data.items():
        match = re.match(r"([\da-f]{4})00([\da-f]{2})", tag)
        if match:
            try:
                item = [base64.b64decode(item[0]).decode()]
            except binascii.Error:
                pass
            except UnicodeDecodeError:
                pass
            if item[0] == "SIEMENS CSA HEADER":
                element = match.group(1)+match.group(2)+"10"
                break

    item = data[element]
    for entry in numpy.ravel(item):
        siemens_data = dicom_to_nifti.siemens.parse_csa(base64.b64decode(entry))

        b_value = siemens_data["B_value"][0]
        # Convert from s/mm^2 to s/m^2
        b_value *= 1e6

        direction = siemens_data["DiffusionGradientDirection"]
        if len(direction) == 0:
            direction = [0,0,0]
        norm = numpy.linalg.norm(direction)
        if norm > 0:
            direction /= norm

        scheme.append((b_value, direction))
    return scheme

def from_ge_private(data):
    """ Extract diffusion gradient direction and b-value from GE-specific
        elements (0019,xxbb, 0019,xxbc, 0019,xxbd, and 0043,xx39).
    """

    # Look for "GEMS_ACQU_01" and "GEMS_PARM_01" private creators and build base
    # tags.
    gems_acq = None
    gems_parm = None
    for tag, item in sorted(data.items()):
        if tag[:4] == "0019" and tag[-4:-2] == "00":
            if item and item[0] == "GEMS_ACQU_01":
                gems_acq = "0019"+tag[-2:]
        if tag[:4] == "0043" and tag[-4:-2] == "00":
            if item and item[0] == "GEMS_PARM_01":
                gems_parm = "0043"+tag[-2:]
        if tag>"004300ff" or None not in [gems_acq, gems_parm]:
            break

    directions = None
    if gems_acq is not None:
        directions = numpy.squeeze(
            numpy.transpose([data[gems_acq+x] for x in ["bb", "bc", "bd"]]))
        norm = numpy.maximum(1e-30, numpy.linalg.norm(directions, axis=1))
        directions /= norm[:,None]

    b_values = None
    if gems_parm is not None:
        b_values = data.get(gems_parm+"39")
        if b_values is not None:
            b_values = [x[0] for x in b_values]
    if b_values is None:
        b_values = [x[0] if x else 0 for x in data.get("DiffusionBValue")]
    # Convert from s/mm^2 to s/m^2
    b_values = [x*1e6 for x in b_values]

    return list(zip(b_values, directions))

def to_mrtrix(scheme, fd):
    """ Save a diffusion scheme in MRtrix format to a file-like object.
    """

    # https://mrtrix.readthedocs.io/en/latest/concepts/dw_scheme.html
    # > the direction vectors are assumed to be provided with respect to real
    # > or scanner coordinates. This is the same convention as is used in the
    # > DICOM format.
    # NOTE: DICOM uses *patient* coordinates, not *scanner* coordinates
    for b_value, direction in scheme:
        # Convert from s/m^2 to s/mm^2
        b_value /= 1e6
        print(*direction, b_value, file=fd)

def to_fsl(scheme, transform, bvecs_fd, bvals_fd):
    """ Save a diffusion scheme in FSL bvecs+bvals format. A reference
        transform is required as the bvecs are store in image coordinates, not
        in patient coordinates. This transform must correspond to an
        image-to-patient transform, e.g. what is stored in the *affine* member
        of nibabel images.
    """

    # https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FDT/UserGuide#Diffusion_data_in_FSL
    # https://fsl.fmrib.ox.ac.uk/fsl/fslwiki/FDT/FAQ#What_conventions_do_the_bvecs_use.3F

    directions = numpy.array([direction for b_value, direction in scheme])
    # Convert from patient coordinates to image coordinates
    # WARNING: for highly anisotropic images and non axis-aligned transforms,
    # this seem to introduce a slight bias in the directions.
    bvecs = numpy.array([numpy.linalg.inv(transform) @ d for d in directions])

    if numpy.linalg.det(transform)>0:
        bvecs[:,0] *= -1

    bvecs[bvecs == numpy.NZERO] = 0

    # Re-normalize (not required by FSL)
    norm = numpy.maximum(1e-20, numpy.linalg.norm(bvecs, axis=1))
    bvecs /= norm[:,None]
    for row in bvecs.T:
        print(*row, file=bvecs_fd)

    # Convert from s/m^2 to s/mm^2
    b_values = numpy.array([b_value for b_value, direction in scheme]) / 1e6
    print(*b_values, file=bvals_fd)


import sys
# from bfuncs import *
from math import *
import copy
import re
import numpy

from Scientific.Geometry import Vector, isVector, Tensor, isTensor
from Scientific.indexing import index_expression
from Scientific import N

from MMTK import *
from MMTK.MoleculeFactory import *
from MMTK.ForceFields import Amber99ForceField
from MMTK import Universe
from MMTK.InternalCoordinates import *
from MMTK.ChemicalObjects import *

########################################
########## AUX FUNCTIONS ###############
########################################

# START FUNC bsubstr @@@@@
def bsubstr(astring, pos, nchar):
	"""Awk substr equivalent"""
	result = 0;
	l = len(astring);
	if((pos > 0) & (nchar <= l)):
		f = pos+nchar-1;
		if(f <= l):
			return astring[(pos-1):f];
	else:
		return result;

# END FUNC   bsubstr @@@@@

# START FUNC bsubstr @@@@@
def bcontains(basket, apple):
	"""
	Checks if apple is in the basket
	@basket: container
	@basket type: list
	@apple: subset
	@apple type type: list
	@returns 1 if apple is contained in basket
	"""
	found = 0
	lapple = len(apple)
	lbasket = len(basket)

	if lapple >= lbasket : return 0

	for i in range(lapple) :
		found = 0
		for j in range(lbasket) :
			if apple[i] == basket[j]:
				found = 1
				break
		if found == 0 :
			return 0

	return 1

# START FUNC bdiff @@@@@
def bPDBprint_atom(record, name, position):
	serial = 1
	altLoc = ' '
	resName = 'LIG'
	chainID = 'X'
	resSeq = 1
	occupancy = 1.0
	tempFactor = 0.0
	print '{0:<6}{1:>5}{2:>4}{3} {4} {5}{6:>4}    {7:8.3f}{8:8.3f}{9:8.3f}  {10:4.2f}{11:6.2f}'.format\
		(record, serial, name, altLoc, resName, chainID, resSeq, \
		position.x(), position.y(), position.z(), occupancy, tempFactor)
# END FUNC   bdiff  @@@@@

# START FUNC bisSaturated @@@@@
def bisSaturated(element, bonds):
	"""
	Checks if an atom is involved in multiple bonds
	@element: the atom
	@element type: string
	@bonds: how many bonds is the atom involved in
	@bonds type: integer
	@rvalue: difference between the valence and bonds (unsaturated value)
	@rtype: integer
	"""
	if element == 'C' :
		return 4 - bonds
	elif element == 'Si' :
		return 4 - bonds
	elif element == 'N' :
		return 3 - bonds
	elif element == 'O' :
		return 2 - bonds
	elif element == 'S' :
		return 2 - bonds
	elif element == 'P' :
		return 5 - bonds
	elif element == 'Se' :
		return 2 - bonds
	else :
		return 0

####################################
########## FUNCTIONS ###############
####################################

def anchor(universe):
	"""Anchor a universe"""
	if(universe.numberOfAtoms() >=3):
		#set anchor atoms indexes
		a1i	= 0
		for i in range(len(universe.atomList()[a1i].bondedTo())):
			for j in range(len(universe.atomList()[a1i].bondedTo()[i].bondedTo())):
				if universe.atomList()[a1i] is not \
					universe.atomList()[a1i].bondedTo()[i].bondedTo()[j] :
					a2i = i;
					a3i = j;
					break
		print 'anchor using: ', a1i, a2i, a3i

		v1 = universe.atomList()[a1i].position();
		v2 = universe.atomList()[a1i].bondedTo()[a2i].position();
		v3 = universe.atomList()[a1i].bondedTo()[a2i].bondedTo()[a3i].position();

		vT01 = -v1;
		tra01 = Translation(vT01);
		universe.applyTransformation(tra01);

		x = universe.atomList()[a1i].bondedTo()[a2i].position().x();
		y = 0;
		z = universe.atomList()[a1i].bondedTo()[a2i].position().z();
		vR01 = Vector(x,y,z);
		vR02 = Vector(1,0,0);

		sign = 1.0
		if ((z > 0) and (x > 0)) or ((z < 0) and (x < 0)):
			sign = -1.0

		rot01 = Rotation(Vector(0,1,0), sign*vR01.angle(vR02));
		universe.applyTransformation(rot01);

		v4 = universe.atomList()[a1i].bondedTo()[a2i].position();

		sign = -1.0
		if ((x > 0) and (y > 0)) or ((x < 0) and (y < 0)):
			sign = 1.0
		rot02 = Rotation(Vector(0,0,1), sign*v4.angle(Vector(1,0,0)));
		universe.applyTransformation(rot02);

		x = 0;
		y = universe.atomList()[a1i].bondedTo()[a2i].bondedTo()[a3i].position().y();
		z = universe.atomList()[a1i].bondedTo()[a2i].bondedTo()[a3i].position().z();

		sign = 1.0
		if ((z > 0) and (y > 0)) or ((z < 0) and (y < 0)):
			sign = -1.0

		vR03 = Vector(x,y,z);
		vR04 = Vector(0,1,0);
		rot03 = Rotation(Vector(1,0,0), sign*vR03.angle(vR04));
		universe.applyTransformation(rot03);

	else:
		print "Molecule too little to be anchored";
#--------------------------------------------------

####################################
############ CLASSES ###############
####################################

class bConvertor():
	"""
	Converts from cartesian coordinates
	found in universe to bonds, angles, dihedrals
	lists found within self
	"""
	def __init__(self, universe):
		"""
		:param universe: an MMTK universe
		:type universe: MMTK universe
		"""
		#Variables
		self.bonds = []		#list of bonds
		self.angles = []	#list of angles
		self.dihedrals = [] #list of dihedrals

	def cart2all_internals(self):
		k = -1
		# Bonds
		for i in universe[0].bonds :
			k += 1;
			x = universe.distance(i.a1, i.a2);
			self.bonds.append([i.a1, i.a2, x]);

		# Angles
		k = -1;
		for i in self.bonds :
			for j in self.bonds:
				if j != i :
					if (i[0] == j[0]):
						k += 1;
						x = universe.angle(i[1], i[0], j[1]);
						self.angles.append([[i[1], i[0], i[2]],\
											[i[0], j[1], j[2]],\
										 	x]);
					elif (i[0] == j[1]):
						k += 1;
						x = universe.angle(i[1], i[0], j[0]);
						self.angles.append([[i[1], i[0], i[2]],\
											[i[0], j[0], j[2]],\
										 	x]);
					elif (i[1] == j[0]):
						k += 1;
						x = universe.angle(i[0], i[1], j[1]);
						self.angles.append([[i[0], i[1], i[2]],\
											[i[1], j[1], j[2]],\
										 	x]);
					elif (i[1] == j[1]):
						k += 1;
						x = universe.angle(i[0], i[1], j[0]);
						self.angles.append([[i[0], i[1], i[2]],\
											[i[1], j[0], j[2]],\
										 	x]);

		# Dihedrals
		k = -1;
		for i in self.angles :
			for j in self.angles:
				if j < i :
					if (i[0] == j[1]) :
						k += 1;
						x = universe.dihedral(j[0][0], j[0][1], j[1][1], i[1][1]);
						self.dihedrals.append([j,i,x])
					if (i[1] == j[0]) :
						k += 1;
						x = universe.dihedral(i[0][0], i[0][1], i[1][1], j[1][1]);
						self.dihedrals.append([i,j,x])
	#---------------------------------------------

	def print_internals(self):
		"""Writes Bonds, Angles & Dihedrals to stdout"""
		print "Bonds";
		for i in self.bonds :
			print i;
		print "Angles";
		for i in self.angles :
			print i
		print "Dihedrals"
		for i in self.dihedrals :
			print i;
	#---------------------------------------------

	def swap_bond(self, bond):
		"""Reverses the atoms order"""
		"""in a [a1,a2,val] type bond"""
		rbond = copy.deepcopy(bond)
		x = rbond[0]
		rbond[0] = rbond[1]
		rbond[1] = x
		return rbond
	#---------------------------------------------

	def swap_angle(self, angle):
		"""Reverses the atoms order in a"""
		"""[[a1,a2,bval], [a2,a3,bval], aval] type angle"""
		rangle = copy.deepcopy(angle)
		a1 =  angle[0][0]
		a2l = angle[0][1]
		bvall = angle[0][2]
		a2r = angle[1][0]
		a3  = angle[1][1]
		bvalr = angle[1][2]

		rangle[0][0] = a3
		rangle[0][1] = a2r
		rangle[0][2] = bvalr
		rangle[1][0] = a2l
		rangle[1][1] = a1
		rangle[1][2] = bvall
		return rangle
	#---------------------------------------------

	def swap_dihe(self, dihe):
		"""Reverses the atoms order in a"""
		"""[ [[a1,a2,bv1],[a2l,a3l,bv2l], av1],"""
		"""  [[a2r,a3r,bv2r],[a3,a4,bv3], av2],"""
		"""  dv ] type dihedral angle"""
		rdihe = copy.deepcopy(dihe)

		a1   = dihe[0][0][0]
		a2   = dihe[0][0][1]
		bv1  = dihe[0][0][2]
		a2l  = dihe[0][1][0]
		a3l  = dihe[0][1][1]
		bv2l = dihe[0][1][2]
		av1  = dihe[0][2]
		a2r  = dihe[1][0][0]
		a3r  = dihe[1][0][1]
		bv2r = dihe[1][0][2]
		a3   = dihe[1][1][0]
		a4   = dihe[1][1][1]
		bv3  = dihe[1][1][2]
		av2  = dihe[1][2]

		rdihe[0][0][0] = a4
		rdihe[0][0][1] = a3
		rdihe[0][0][2] = bv3
		rdihe[0][1][0] = a3r
		rdihe[0][1][1] = a2r
		rdihe[0][1][2] = bv2r
		rdihe[0][2] = av2
		rdihe[1][0][0] = a3l
		rdihe[1][0][1] = a2l
		rdihe[1][0][2] = bv2l
		rdihe[1][1][0] = a2
		rdihe[1][1][1] = a1
		rdihe[1][1][2] = bv1
		rdihe[1][2] = av1
		return rdihe
	#---------------------------------------------

	def sprint_zmat_dihe(self, dihe):
		"""Write dihedral to stdout"""
		""" in Z-matrix format"""
		a1   = dihe[0][0][0]
		a2   = dihe[0][0][1]
		bv1  = dihe[0][0][2]
		a2l  = dihe[0][1][0]
		a3l  = dihe[0][1][1]
		bv2l = dihe[0][1][2]
		av1  = dihe[0][2]
		a2r  = dihe[1][0][0]
		a3r  = dihe[1][0][1]
		bv2r = dihe[1][0][2]
		a3   = dihe[1][1][0]
		a4   = dihe[1][1][1]
		bv3  = dihe[1][1][2]
		av2  = dihe[1][2]
		return '{0:>4} {1:>4} {2:>4} {3:>4} {4:>10.6f} {5:>4} {6:>10.6f} {7:>4} {8:>10.6f}'.format(a1.index+1, a1.name, a1.name[0], a2.index+1, bv1*10.0, a3.index+1, av1, a4.index+1, dihe[2])
	#---------------------------------------------

	def print_zmat_dihe(self, dihe):
		"""Write dihedral to stdout"""
		""" in Z-matrix format"""
		a1   = dihe[0][0][0]
		a2   = dihe[0][0][1]
		bv1  = dihe[0][0][2]
		a2l  = dihe[0][1][0]
		a3l  = dihe[0][1][1]
		bv2l = dihe[0][1][2]
		av1  = dihe[0][2]
		a2r  = dihe[1][0][0]
		a3r  = dihe[1][0][1]
		bv2r = dihe[1][0][2]
		a3   = dihe[1][1][0]
		a4   = dihe[1][1][1]
		bv3  = dihe[1][1][2]
		av2  = dihe[1][2]
		print '{0:>4} {1:>4} {2:>4} {3:>4} {4:>10.6f} {5:>4} {6:>10.6f} {7:>4} {8:>10.6f}'.format(a1.index+1, a1.name, a1.name[0], a2.index+1, bv1*10.0, a3.index+1, av1, a4.index+1, dihe[2])
	#---------------------------------------------

	def NeRFinte2cart(self, dihe):	# Needs debug
		"""
		Reconstructs atomlist positions
		2005 Parsons et al - NeRF algorithm
		:param universe: MMTK universe
		:type universe: MMTK universe
		:param dihe: list of dihedrals
		:type dihe: [[[],[],aval] , [[],[],aval] , dval] where [] is a bond
		"""
		known_atoms = []
		unknown_atoms = copy.deepcopy(universe.atomList())
		C = Vector(0,0,0)
		D = Vector(0,0,0)
		D2 = Vector(0,0,0)
		l = len(dihe)
		for di in range(l):	# Create the START stack
			dihe.append(self.swap_dihe(dihe[di]))

		#for i in dihe:
		#	print_zmat_dihe(i)

		# Assign initial values
		i = dihe[0]
		a1 = i[0][0][0]
		a2 = i[0][0][1]
		a3 = i[1][1][0]
		a4 = i[1][1][1]
		R1 = i[0][0][2]
		R2 = i[1][1][2]
		theta1 = i[0][2]
		phi = i[2]
		universe.atomList()[a1.index].setPosition(Vector(0,0,0))
		universe.atomList()[a2.index].setPosition(Vector(R1,0,0))
		universe.atomList()[a3.index].setPosition(Vector(R1-R2*cos(theta1), R2*sin(theta1), 0))
		known_atoms.append(a1)
		known_atoms.append(a2)
		known_atoms.append(a3)
		#print 'initialize at: ', a1, a2, a3
		#bPDBprint_atom('HETATM', a1.name, 10*a1.position())
		#bPDBprint_atom('HETATM', a2.name, 10*a2.position())
		#bPDBprint_atom('HETATM', a3.name, 10*a3.position())
		# Main loop
		cnt=0;
		while len(known_atoms) < len(universe.atomList()):
			cnt += 1
			if cnt > 9999:
				print 'NeRFinte2cart() cnt EXCEEDED: 9999'
				break
			ai = -1
			for a in unknown_atoms:
				ai += 1
				for i in dihe:				# Search if atom a is a4 in any dihe
					a1 = i[0][0][0]
					a2 = i[0][0][1]
					a3 = i[1][0][1]
					a4 = i[1][1][1]
					a1known = 0
					a2known = 0
					a3known = 0
					a4known = 0
					if a4.index == a.index:
						for k in known_atoms:	# Check if atom is in KNOWN stack
							if a4.index == k.index:
								a4known = 1
								break
							if a1.index == k.index :
								a1known = 1
							if a2.index == k.index :
								a2known = 1
							if a3.index == k.index :
								a3known = 1
						if ((a4known == 0) and \
							(a1known == 1) and (a2known == 1) and (a3known == 1)):
								# Move everything so that a3 is in the center
								M = Translation(-1*a3.position())
								universe.applyTransformation(M)
								for k in known_atoms:
									M(k.position())
								# Solve a4
								#print 'SOLVE: ', a4
								print_zmat_dihe(i)
								A = a1.position()
								B = a2.position()
								C = a3.position()
								#print 'A = ', A, '(', a1.name , ')'
								#print 'B = ', B, '(', a2.name , ')'
								#print 'C = ', C, '(', a3.name , ')'
								R = i[1][1][2]
								#print 'R = ', R
								theta = i[1][2]
								if theta > pi/2:
									theta = pi - theta
								#print 'theta = ', theta, numpy.rad2deg(theta)
								phi = i[2]
								if phi < 0:
									phi = pi - phi
								#print 'phi = ', phi, numpy.rad2deg(phi)
								AB = B - A	# B - A
								#print 'AB = ', AB
								BC = C - B	# C - B
								#print 'BC = ', BC
								bc = BC.normal()
								#print 'bc = ', bc
								n = AB.cross(bc)
								n = n.normal()
								#print 'n = ', n
								Mx = bc
								#print 'Mx = ', Mx
								My = n.cross(bc)
								My = My.normal()
								#print 'My = ', My
								Mz = n
								#print 'Mz = ', Mz
								M = Tensor([ [Mx.x(), Mx.y(), Mx.z()],\
											 [My.x(), My.y(), My.z()],\
											 [Mz.x(), Mz.y(), Mz.z()] ])
								#print 'M: '
								#print M
								D2 = Vector(R*cos(theta), R*cos(phi)*sin(theta), R*sin(phi)*sin(theta))
								#print 'D2: '
								#print D2
								D = M*D2
								D = D + C
								#print 'D: '
								#print D
								a.setPosition(D)
								universe.atomList()[a4.index].setPosition(D)
								# Add a4 to KNOWN
								known_atoms.append(a4)
								unknown_atoms.pop(ai)
								#print '========================'
								break

		for k in known_atoms:
			bPDBprint_atom('HETATM', k.name, 10*k.position())

	# -----------------------------------

	def inte2cart(self, dihe):
		"""
		Reconstructs atomlist positions
		Classic
		:param universe: MMTK universe
		:type universe: MMTK universe
		:param dihe: list of dihedrals
		:type dihe: [[[],[],aval] , [[],[],aval] , dval] where [] is a bond
		"""
		known_atoms = []
		unknown_atoms = copy.deepcopy(universe.atomList())
		C = Vector(0,0,0)
		D0 = Vector(0,0,0)
		D1 = Vector(0,0,0)
		D = Vector(0,0,0)
		D2 = Vector(0,0,0)
		l = len(dihe)
		for di in range(l):	# Create the START stack
			dihe.append(self.swap_dihe(dihe[di]))

		# Assign initial values
		i = dihe[0]
		a1 = i[0][0][0]
		a2 = i[0][0][1]
		a3 = i[1][1][0]
		a4 = i[1][1][1]
		R1 = i[0][0][2]
		R2 = i[1][1][2]
		theta1 = i[0][2]
		phi = i[2]
		universe.atomList()[a1.index].setPosition(Vector(0,0,0))
		universe.atomList()[a2.index].setPosition(Vector(R1,0,0))
		universe.atomList()[a3.index].setPosition(Vector(R1-R2*cos(theta1), R2*sin(theta1), 0))
		known_atoms.append(a1)
		known_atoms.append(a2)
		known_atoms.append(a3)

		# Main loop
		cnt=0;
		while len(known_atoms) < len(universe.atomList()):
			cnt += 1
			if cnt > 9999:
				print 'inte2cart() cnt EXCEEDED: 9999'
				break
			ai = -1
			for a in unknown_atoms:
				ai += 1
				for i in dihe:				# Search if atom a is a4 in any dihe
					a1 = i[0][0][0]
					a2 = i[0][0][1]
					a3 = i[1][0][1]
					a4 = i[1][1][1]
					a1known = 0
					a2known = 0
					a3known = 0
					a4known = 0
					if a4.index == a.index:
						for k in known_atoms:	# Check if atom is in KNOWN stack
							if a4.index == k.index:
								a4known = 1
								break
							if a1.index == k.index :
								a1known = 1
							if a2.index == k.index :
								a2known = 1
							if a3.index == k.index :
								a3known = 1
						if ((a4known == 0) and \
							(a1known == 1) and (a2known == 1) and (a3known == 1)):
								# Move everything so that a3 is in the center
								M = Translation(-1*a3.position())
								universe.applyTransformation(M)
								for k in known_atoms:
									M(k.position())
								# Solve a4
								A = a1.position()
								B = a2.position()
								C = a3.position()
								#print 'SOLVE: ', a4
								R = i[1][1][2]
								theta = pi-i[1][2]
								phi = i[2]
								AB = B - A	# B - A
								BC = C - B	# C - B
								bc = BC.normal()
								n = AB.cross(bc)
								n = n.normal()
								D0 = C + (bc*R)
								M = Rotation(n, theta)
								D1 = M(D0)
								M = Rotation(bc, phi)
								D = M(D1)
								a.setPosition(D)
								universe.atomList()[a4.index].setPosition(D)

								# Add a4 to KNOWN
								known_atoms.append(a4)
								unknown_atoms.pop(ai)
								break

		for k in known_atoms:
			bPDBprint_atom('HETATM', k.name, 10*k.position())

	# -----------------------------------

class bRBSelector():
	"""
	Finds and writes rigid bodies as lists
	"""
	def __init__(self, universe, convertor):
		"""
		:param universe: an MMTK universe
		:type universe: MMTK universe
		:param bonds: lists of precomputed bonds
		:type bonds: a list of 3 element lists: [a1,a2,bval]
		"""
		#Variables
		self.bonds = convertor.bonds	#
		self.dihedrals = convertor.dihedrals
		self.convertor = convertor #
		self.cnt = 0		#checks if infinite loop
		self.bibonds = []	#an element is a bond with an index attached
		self.bibonds_dict = {}
		self.wostack = []	#stack which uses bibonds element types (work)
		self.rings = []		#stack which uses bibonds element types	(rings deposit)
		self.trash = []		#stack which uses bibonds element types (recovery)
		self.index = 0		#used to identify bonds !not the same as Atom.index
		self.inbuff = []	#indexes buffer used in inrings
		self.inrings = []	#list of sorted lists of indexes
		self.buff = []
		self.conv = 0
		self.atring = []	#list of lists of atom indexes in rings
		self.atbond = []	#list of 2-tuples atom indexes in rigid bonds
		self.atatom = []	#list of 2 tuples [atom, valence]
		self.atconj = []	#list of conjugate multiple bonds
		self.atrb = []
		self.stickbonds = []

	def __getstate__(self):
		return self.atring

	def findPaths(self):
		# Fill the START stack (bibonds)
		self.index = 0
		for i in self.bonds :
			self.index += 1
			self.bibonds.append([i, self.index])
			self.bibonds.append([convertor.swap_bond(i), self.index])
			self.bibonds_dict[self.index] = i
		# Kickstart - put the first element in WORK stack (wostack)
		self.wostack.append(self.bibonds[-1])
		self.bibonds.pop()
		# Main
		while True:
			if not self.bibonds:
				break
			elif not self.wostack:
				self.wostack.append(self.bibonds[-1])
				self.bibonds.pop()
			else:
				self.cnt += 1
				if self.cnt > (universe.numberOfAtoms()*universe.numberOfAtoms()):
					print 'findPaths(): cnt EXCEEDED!'
					break
				if not self.wostack:
					break
				else:
					found_link = 0
					for i in range(len(self.bibonds)) :
						if ( (self.wostack[-1][0][1].index == self.bibonds[i][0][0].index) and \
							 (self.wostack[-1][1] != self.bibonds[i][1])):
							found_link = 1
							self.wostack.append(self.bibonds[i])
							self.bibonds.pop(i)
							# Check for rings
							wolen = len(self.wostack)
							if(wolen >=3):
								for j in range(-3, -1 * wolen, -1) :
									#print 'wostack[',j,']:',wostack[j]
									#Check for circular path
									if ((self.wostack[-1][0][1].index == self.wostack[j][0][0].index) and\
										(self.wostack[-1][1] != self.wostack[j][1])):
										self.inbuff = []
										for k in range(wolen+j, wolen):
											self.inbuff.append(self.wostack[k][1])
										#Check if path has two way bonds
										bonds_doubled = 0
										if len(self.inbuff) > len(set(self.inbuff)) :
											bonds_doubled = 1
										#Check if ring is already in the rings stack
										self.inbuff.sort()
										ring_already_in = 0
										for t in self.inrings :
											if self.inbuff == t:
												ring_already_in = 1
												break
										#Copy a new ring to rings
										if (ring_already_in == 0) and (bonds_doubled == 0):
											for k in range(wolen+j, wolen):
												self.rings.append(self.wostack[k])
											self.inrings.append(self.inbuff)
										break

							break
					if found_link == 0:
						self.trash.append(self.wostack[-1])
						self.wostack.pop()

		# Prepare a list of names
		i = 0
		for m in self.inrings :
			self.buff = []
			for n in m :
				self.buff.append(self.bibonds_dict[n][0].index)
				self.buff.append(self.bibonds_dict[n][1].index)
			self.buff = set(self.buff)
			self.atring.append(list(self.buff))

		# ---------------------------------------------------------------

	def check_for_cond_rings(self, bonds, conv) :
		for mi in range(len(self.atring)) :
			for ni in range(len(self.atring)) :
				self.buff = []
				if bcontains(self.atring[mi], self.atring[ni]) == 1 :
					# Extract the hidden ring
					self.buff = list(set(self.atring[mi])-set(self.atring[ni]))
					buffbuff = []
					for o in self.atring[ni] :	# closed ring
						for p in self.buff :		# diff
							found_bond = 0
							if (p != o) :
								for r in bonds :
									if ( ((r[0].index == o) and (r[1].index == p)) or\
										 ((r[1].index == o) and (r[0].index == p)) ):
										buffbuff.append(o)
										found_bond = 1
										break
							if found_bond == 1 :
								break
					for s in list(set(buffbuff)) :
						self.buff.append(s)

					self.atring.append(self.buff)
					self.atring.pop(mi)

					c = 0
					for zi in (range(len(self.atring))) :
						c += len(self.atring[zi])
					if c == conv : # check for convergence
						return

					self.check_for_cond_rings( bonds, c)	#recursivity
					#break

	def untangleCondensed(self):
		self.conv = 0	#convergence
		self.check_for_cond_rings(self.bonds, self.conv)
		# Sort atring
		for i in self.atring:
			i.sort()
		self.atring.sort()

	def find_rings(self):
		self.findPaths()
		self.untangleCondensed()

	def print_atring(self):
		"""Prints atring list to stdout"""
		if(len(self.atring) > 1) :
			if(self.atring[0] != self.atring[1]) :
				print self.atring[0]
			for i in range(1,len(self.atring)):
				if self.atring[i] != self.atring[i-1] :
					print self.atring[i]
		elif (len(self.atring) == 1):
			print self.atring[0]
		else:
			print 'No rings found'
	#---------------------------------------------

	def print_rings(self):
		"""Prints ring atom names to stdout"""
		fin_string = ""
		fin_string += 'rings = ['
		t = 0
		if(len(self.atring) > 1) :
			if(self.atring[0] != self.atring[1]) :
				#print self.atring[0]
				if t != 0: fin_string += ', '
				t += 1
				fin_string += '['
				k = 0
				for z in self.atring[0]:
					if k != 0: fin_string += ', '
					fin_string += universe.atomList()[z].name
					k += 1
				fin_string += ']'
			for i in range(1,len(self.atring)):
				if self.atring[i] != self.atring[i-1] :
					#print self.atring[i]
					if t != 0: fin_string += ', '
					t += 1
					fin_string += '['
					k = 0
					for z in self.atring[i]:
						if k != 0: fin_string += ', '
						fin_string += universe.atomList()[z].name
						k += 1
					fin_string += ']'
		elif (len(self.atring) == 1):
			#print self.atring[0]
			if t != 0: fin_string += ', '
			t += 1
			fin_string += '['
			k = 0
			for z in self.atring[0]:
				if k != 0: fin_string += ', '
				fin_string += universe.atomList()[z].name
				k += 1
			fin_string += ']'
		else:
			pass
			#print 'No rings found'
		fin_string += ']'
		print fin_string
	#---------------------------------------------

	def growRB(self):
		"""
		Grows rigid bodies starting from
		rigid bonds (C-H, X=X) and rings
		"""
#		print 'BEGIN growRB ============='
		# Build atatom list ([atom, unsaturated rank])
		for ai in universe.atomList():
			unsat = bisSaturated(ai.type.symbol, len(ai.bondedTo()))
			if unsat > 0:
				self.atatom.append([ai.index, unsat])

		# Build atbond list (bonds with sp2 or sp atoms)
		self.atbond = []
		for ata in self.atatom:
			ring_found1 = 0
			for ri in self.atring:		# Check if ata is in any ring
				for rj in ri:
					if ata[0] == rj:
						ring_found1 = 1
						#break			# Atom 1 touches a ring
				if ring_found1 == 1:
					pass
					#break				# Checked
			for atb in self.atatom:	# Check if atb is in any ring
				ring_found2 = 0
				for ri in self.atring:
					for rj in ri:
						if atb[0] == rj:
							ring_found2 = 1
							#break	# Atom 2 touches a ring
				if (ring_found2 == 0) or (ring_found1 == 0):
					if ata.index < atb.index:
						for bo in universe.atomList()[ata[0]].bondedTo():
							if bo.index == atb[0]:
								self.atbond.append([ata[0],atb[0]])
								break

		fin_str = ""
		fin_str += 'non_ring_pi_bonds = [' # len(self.atbond)
		k = 0
		for bo in self.atbond:
			#print universe.atomList()[bo[0]].name, universe.atomList()[bo[1]].name
			if k != 0: fin_str += ', '
			fin_str += 'Bond('
			fin_str += universe.atomList()[bo[0]].name
			fin_str += ', '
			fin_str += universe.atomList()[bo[1]].name
			fin_str += ')'
			k += 1
		fin_str += ']'
		print fin_str
#			print bo
#		print '------------------'

# THE FOLLOWING CODE IS NOT NECESSARY
#		for ri in self.atring:	# Eliminates bonds sticked into rings BUG !!
#			for rj in ri:
#				le = len(self.atbond)
#				for bi in range(le):
#					if (self.atbond[bi][0] == rj) or (self.atbond[bi][1] == rj):
#						self.atbond.pop(bi)
#						break
#
#		print len(self.atbond), ': multiple bonds without sticked: '
#		for bo in self.atbond:
#			print universe.atomList()[bo[0]].name, universe.atomList()[bo[1]].name
#			#print bo
#		print '------------------'

		# Build atconj - merge multiple bonds in conjugated systems
		self.atconj = copy.deepcopy(self.atbond)
		buff = []
		cnt = 0
		merge_found = 1
		while (merge_found == 1) and (cnt < 100):
			cnt += 1
			merge_found = 0
			for coi in range(len(self.atconj)):
				for coj in range(len(self.atconj)):
					if coj != coi:
						# Check if they have a common atom (to merge them)
						for i in self.atconj[coi]:
							for j in self.atconj[coj]:
								if i == j:
									merge_found = 1
									break
							if merge_found == 1:
								break
						if merge_found == 1:	# Add coj to coi (merge)
							buff = []
							for j in self.atconj[coj]:
								common = 0
								for i in self.atconj[coi]:
									if i == j:
										common = 1
										break
								if common == 0:
									buff.append(j)
							for t in buff:
								self.atconj[coi].append(t)
							# Print the new conj
							# Remove bond coj
							self.atconj.pop(coj) # Eliminate coj from atconj
							break
					if merge_found == 1:
						break
				if merge_found == 1:
					break

#		print 'Conjugated systems: '
#		for coi in self.atconj:
#				for k in coi:
#					print universe.atomList()[k].name, ' ',
#				print

#		print 'END growRB ================'

	def diffPseudoDihe(self):
		"""
		Eliminates correlated dihedrals by
		replacing them with differences
		"""
		print 'BEGIN diffPseudoDihe() =========='
		#for ai in universe.atomList():
		#	print ai.position()

		#for bi in self.atbond:
		#	print 'atbond', universe.atomList()[bi[0]].name, universe.atomList()[bi[1]].name

		#for di in self.dihedrals:
		#	convertor.print_zmat_dihe(di)

		#Stickbonds
		self.stickbonds = []
		for bi in self.bonds:
			at1found =  -1
			at2found =  -1
			for ri in self.atring: # Search the rings
				for rj in ri:
					if (bi[0].index == rj) and \
						(universe.atomList()[bi[1].index].type.symbol != 'H'):
						at1found = rj
					elif (bi[1].index == rj)  and \
						(universe.atomList()[bi[0].index].type.symbol != 'H'):
						at2found = rj
			if ((at1found >= 0) and (at2found == -1) or\
				(at1found == -1) and (at2found >= 0)):
				self.stickbonds.append([bi[0].index, bi[1].index])

#		print 'Sticked bonds:'
#		for bi in self.stickbonds:
#			print universe.atomList()[bi[0]].name, universe.atomList()[bi[1]].name

		for ri in self.atring:
			print ri
			for rj in ri:
				stick = -1
				for bi in self.stickbonds:
					if bi[0] == rj:
						stick = 1
					elif bi[1] == rj:
						stick = 0
					if stick >= 0:
						print '[',universe.atomList()[bi[0]].name,\
								universe.atomList()[bi[1]].name, '] :'
						# Search dihedrals
						for di in self.dihedrals:
							a2 = di[0][0][1]
							a3 = di[1][0][1]
							if ((a2.index == bi[0] and a3.index == bi[1]) or\
								(a2.index == bi[1] and a3.index == bi[0])):
									convertor.print_zmat_dihe(di)
						print '+++++++++++++++'
						break

		print 'END   diffPseudoDihe() =========='

#-------------------------------------------------------------------------

########
# MAIN #
########

# BUILD MOLECULE #

# "Declare" some variables
natm = 0;
atmi = 0;
x = .0;
y = .0;
z = .0;
elem = '--';
name = [];
##################

fpo1 = open(sys.argv[1], 'r');

factory = MoleculeFactory();
factory.createGroup('main');

linecnt = 0;
while True:
	line = fpo1.readline();
	if not line: break;
	linecnt += 1;
	if linecnt == 4:
		natm = int(bsubstr(line, 1, 3));
	if (linecnt >= 5) and (linecnt < (5 + natm)):   # Read cart coords
		atmi += 1;
		x = float(bsubstr(line,  2, 9));
		y = float(bsubstr(line, 12, 9));
		z = float(bsubstr(line, 22, 9));
		elem = bsubstr(line, 32, 2);
		if elem[1] == ' ':
			elem = elem[0];
		name.append(elem + str(atmi));
		factory.addAtom('main',name[atmi-1],elem);
		factory.setPosition('main',name[atmi-1],Vector\
			(x*Units.Ang, y*Units.Ang, z*Units.Ang));
	if linecnt >= (5 + natm):
		if line[0] == 'M': break
		factory.addBond('main',\
			name[int(bsubstr(line,1,3))-1], name[int(bsubstr(line,4,3))-1]);

fpo1.close();

newmol = factory.retrieveMolecule('main');
universe = InfiniteUniverse(Amber99ForceField(mod_files=['frcmod.ff99SB']));
universe.addObject(newmol);
universe.configuration();

# -------------------------------------------

#anchor(universe)	# BUG pseudo01.sdf when first atom has only one bond

#universe.writeXML(file('out.xml', 'w'));
#universe.writeToFile('out.pdb');

convertor = bConvertor(universe)
convertor.cart2all_internals()
#convertor.print_internals()

#print '\n++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++\n'
#for i in convertor.dihedrals:
#	convertor.print_zmat_dihe(i)
#convertor.inte2cart(convertor.dihedrals)

# Print rings
rb_selector = bRBSelector(universe, convertor)
rb_selector.find_rings()
#rb_selector.print_atring()
rb_selector.print_rings()
rb_selector.growRB()
#rb_selector.diffPseudoDihe()


# %% [markdown]
"""
calculate ODNP using DNPLab
===========================

This example demonstrates how to use the dnplab.dnpHydration module

"""
# %%

# %% [markdown]
# First import dnplab and numpy,
import dnplab
import numpy as np

# %%

# %% [markdown]
# To use the dnpHydration module first create a dictionary with the necessary inputs. Start by creating a workspace and assinging an inputs dictionary to the key **'hydration_inputs'**. For example,

Enhancements = [] # list of signal enhancements
Enhancement_powers = [] # list of powers in Watts corresponding to Enhancements
T1s = [] # list of T1 values in seconds
T1_powers = [] # list of powers in Watts corresponding to T1s

inputs = {
          'E_array' : np.array(Enhancements),
          'E_powers' : np.array(Enhancement_powers),
          'T1_array' : np.array(T1s),
          'T1_powers' : np.array(T1_powers),
          'T10': 2.0, # T1 measured with power=0
          'T100': 2.5, # T1 measured with SL=0 and power=0
          'spin_C': 100, # spin concentration in micromolar
          'field': 350, # magnetic field in mT
          'smax_model': 'tethered', # choice of smax model
          'interpolate_method': 'second_order' # choice of interpolation method
          }
# %%

# %% [markdown]
# Now you can either create a workspace and add the dictionary under the key **'hydration_inputs'**,
workspace = dnplab.create_workspace('hydration_inputs', inputs)
# %%

# %% [markdown]
# Or add to an existing workspace,
workspace.add('hydration_inputs', inputs)
# %%

# %% [markdown]
# In rare cases the bulk water or second order T1 interpolation constants may need to be altered. This is not necessary for the odnp module to operate, but if needed this can be done by adding the dictionary **'hydration_constants'** to the workspace. For example,
constants = {
             'ksigma_bulk': 95.4, # bulk ksigma value
             'krho_bulk': 353.4, # bulk krho value
             'klow_bulk': 366, # bulk klow value
             'tcorr_bulk': 54, # bulk tcorr value
             'D_H2O': 2.3e-9, # bulk water diffusivity
             'D_SL': 4.1e-10, # diffusivity of spin probe in bulk water
             'delta_T1_water': 1, # change in water proton T1 due to microwaves
             'T1_water': 2.5, # T1 of bulk water protons
             'macro_C': 100, # concentration of macromolecule in uM
             }

workspace.add('hydration_constants', constants)
# %%

# %% [markdown]
# Next, pass the workspace to dnplab.dnpHydration.hydration to perform calculations using,
hydration_results = dnplab.dnpHydration.hydration(workspace)
# %%

# %% [markdown]
# or operate in-place with:
dnplab.dnpHydration.hydration(workspace)
# %%

# %% [markdown]
# For use without creating a DNPLab workspace simply skip the above steps and pass the dictionaries to dnpHydration directly,
hydration_results = dnplab.dnpHydration.odnp(inputs=inputs, constants=constants)
# %%