# pulsar.py

# Class containing pulsar data from timing package [tempo2/PINT].

from __future__ import (absolute_import, division,
                        print_function, unicode_literals)

import enterprise
import numpy as np
from ephem import Ecliptic, Equatorial
import os
import json
from enterprise.signals import utils
import logging

try:
    import cPickle as pickle
except:
    import pickle

try:
    import libstempo as t2
except ImportError:
    print('Ooh, no libstempo?')
    t2 = None

try:
    import pint
    import pint.toa as toa
    import pint.models.model_builder as mb
    from pint.models import TimingModel
    from pint.residuals import Residuals as resids
except ImportError:
    print('Cannot import PINT? Meh...')
    pint = None

import astropy.units as u

if pint is None and t2 is None:
    err_msg = 'Must have either PINT or libstempo timing package installed'
    raise ImportError(err_msg)

logging.basicConfig(format='%(levelname)s: %(name)s: %(message)s',
                    level=logging.INFO)
logger = logging.getLogger(__name__)

def get_maxobs(timfile):
    """Utility function to return number of lines in tim file.

    :param timfile:
        Full path to tim-file. For tim-files that use INCLUDEs this
        should be the base tim file.

    :returns: Number of lines in tim-file
    """

    maxobs = 0
    with open(timfile) as tfile:
        flines = tfile.readlines()
        lines = [ln for ln in flines if not ln.startswith('C')]
        if any(map(lambda x: 'INCLUDE' in x, lines)):
            for line in filter(lambda x: 'INCLUDE' in x, lines):
                maxobs += get_maxobs(line.split()[-1])
        else:
            maxobs = sum(1 for line in lines if line.rstrip('\n'))
    return maxobs

class BasePulsar(object):
    """Abstract Base Class for Pulsar objects."""

    def _get_pdist(self):
        dfile = enterprise.__path__[0] + '/datafiles/pulsar_distances.json'
        with open(dfile, 'r') as fl:
            pdict = json.load(fl)

        if self.name[0] not in ['J', 'B']:
            if 'J' + self.name in pdict:
                pdist = tuple(pdict.get('J'+self.name))
            elif 'B' + self.name in pdict:
                pdist = tuple(pdict.get('B'+self.name))
        else:
            pdist = tuple(pdict.get(self.name, (1.0, 0.2)))

        if pdist == (1.0, 0.2):
            msg = 'WARNING: Could not find pulsar distance for '
            msg += 'PSR {0}.'.format(self.name)
            msg += ' Setting value to 1 with 20% uncertainty.'
            logger.warning(msg)
        return pdist

    def _get_radec_from_ecliptic(self, elong, elat):
        # convert via pyephem
        try:
            ec = Ecliptic(elong, elat)

            # check for B name
            if 'B' in self.name:
                epoch = '1950'
            else:
                epoch = '2000'
            eq = Equatorial(ec, epoch=str(epoch))
            raj = np.double(eq.ra)
            decj = np.double(eq.dec)

        except TypeError:
            msg = 'WARNING: Cannot fine sky location coordinates '
            msg += 'for PSR {0}. '.format(self.name)
            msg += 'Setting values to 0.0'
            logger.warning(msg)
            raj = 0.0
            decj = 0.0

        return raj, decj

    def _get_pos(self):
        return np.array([np.cos(self._raj) * np.cos(self._decj),
                         np.sin(self._raj) * np.cos(self._decj),
                         np.sin(self._decj)])

    def sort_data(self):
        """Sort data by time."""
        if self._sort:
            self._isort = np.argsort(self._toas, kind='mergesort')
            self._iisort = np.zeros(len(self._isort), dtype=np.int)
            for ii, p in enumerate(self._isort):
                self._iisort[p] = ii
        else:
            self._isort = slice(None, None, None)
            self._iisort = slice(None, None, None)

    def filter_data(self, start_time=None, end_time=None):
        """Filter data to create a time-slice of overall dataset."""
        if start_time is None and end_time is None:
            mask = np.ones(self._toas.shape, dtype=bool)
        else:
            mask = np.logical_and(self._toas >= start_time * 86400,
                                  self._toas <= end_time * 86400)

        self._toas = self._toas[mask]
        self._toaerrs = self._toaerrs[mask]
        self._residuals = self._residuals[mask]
        self._ssbfreqs = self._ssbfreqs[mask]

        self._designmatrix = self._designmatrix[mask, :]
        dmx_mask = np.sum(self._designmatrix, axis=0) != 0.0
        self._designmatrix = self._designmatrix[:, dmx_mask]

        for key in self._flags:
            self._flags[key] = self._flags[key][mask]

        if self._planetssb is not None:
            self._planetssb = self.planetssb[mask, :, :]

        self.sort_data()

    def to_pickle(self, outdir=None):
        """Save object to pickle file."""

        # drop t2pulsar object
        if hasattr(self, 't2pulsar'):
            del self.t2pulsar

        if outdir is None:
            outdir = os.getcwd()

        if not os.path.exists(outdir):
            os.makedirs(outdir)

        with open(outdir + '/{0}.pkl'.format(self.name), 'wb') as f:
            pickle.dump(self, f)

    @property
    def isort(self):
        """Return sorting indices."""
        return self._isort

    @property
    def iisort(self):
        """Return inverse of sorting indices."""
        return self._iisort

    @property
    def toas(self):
        """Return array of TOAs in seconds."""
        return self._toas[self._isort]

    @property
    def stoas(self):
        """Return array of observatory TOAs in seconds."""
        return self._stoas[self._isort]

    @property
    def residuals(self):
        """Return array of residuals in seconds."""
        return self._residuals[self._isort]

    @property
    def toaerrs(self):
        """Return array of TOA errors in seconds."""
        return self._toaerrs[self._isort]

    @property
    def freqs(self):
        """Return array of radio frequencies in MHz."""
        return self._ssbfreqs[self._isort]

    @property
    def Mmat(self):
        """Return ntoa x npar design matrix."""
        return self._designmatrix[self._isort, :]

    @property
    def pdist(self):
        """Return tuple of pulsar distance and uncertainty in kpc."""
        return self._pdist

    @property
    def dm(self):
        """Return DM parameter from parfile."""
        return self._dm

    @property
    def dmx(self):
        """Return a dictionary of DMX-parameter values and stoa ranges
        from parfile."""
        return self._dmx

    @property
    def flags(self):
        """Return a dictionary of tim-file flags."""

        return dict((k, v[self._isort]) for k, v in self._flags.items())

    @property
    def backend_flags(self):
        """Return array of backend flags.

        Not all TOAs have the same flags for all data sets. In order to
        facilitate this we have a ranked ordering system that will look
        for flags. The order is `group`, `g`, `sys`, `i`, `f`, `fe`+`be`.

        """

        nobs = len(self._toas)
        bflags = ['flag'] * nobs
        check = lambda i, fl: fl in self._flags and self._flags[fl][i] != ''
        flags = [['group'], ['g'], ['sys'], ['i'], ['f'], ['fe', 'be']]
        for ii in range(nobs):
            # TODO: make this cleaner
            for f in flags:
                if np.all(list(map(lambda xx: check(ii, xx), f))):
                    bflags[ii] = '_'.join(self._flags[x][ii] for x in f)
                    break
        return np.array(bflags)[self._isort]

    @property
    def theta(self):
        """Return polar angle of pulsar in radians."""
        return np.pi / 2 - self._decj

    @property
    def phi(self):
        """Return azimuthal angle of pulsar in radians."""
        return self._raj

    @property
    def pos(self):
        """Return unit vector to pulsar."""
        return self._pos

    @property
    def pos_t(self):
        """Return unit vector to pulsar as function of time."""
        return self._pos_t[self._isort, :]

    @property
    def planetssb(self):
        """Return planetary position vectors at all timestamps"""
        return self._planetssb[self._isort, :, :]

class PintPulsar(BasePulsar):

    def __init__(self, toas, model, sort=True, planets=True):

        self._sort = sort
        self.planets = planets
        self.name = model.PSR.value

        self._toas = np.array(toas.table['tdbld'], dtype='float64') * 86400
        self._residuals = np.array(resids(toas, model).time_resids.to(u.s),
                                   dtype='float64')
        self._toaerrs = np.array(toas.get_errors().to(u.s), dtype='float64')
        self._designmatrix = model.designmatrix(toas)[0]
        self._ssbfreqs = np.array(model.barycentric_radio_freq(toas),
                                  dtype='float64')

        # fitted parameters
        self.fitpars = ['Offset'] + [par for par in model.params
                                     if not getattr(model, par).frozen]

        # set parameters
        spars = [par for par in model.params]
        self.setpars = [sp for sp in spars if sp not in self.fitpars]

        self._flags = {}
        for ii, obsflags in enumerate(toas.get_flags()):
            for jj, flag in enumerate(obsflags):

                if flag not in list(self._flags.keys()):
                    self._flags[flag] = [''] * toas.ntoas

                self._flags[flag][ii] = obsflags[flag]

        # convert flags to arrays
        # TODO probably better way to do this
        for key, val in self._flags.items():
            if isinstance(val[0], u.quantity.Quantity):
                self._flags[key] = np.array([v.value for v in val])
            else:
                self._flags[key] = np.array(val)

        self._pdist = self._get_pdist()
        self._raj, self._decj = self._get_radec(model)
        self._pos = self._get_pos()
        self._planetssb = self._get_planetssb()

        # TODO: pos_t not currently implemented
        self._pos_t = np.zeros((len(self._toas), 3))

        self.sort_data()

    def _get_radec(self, model):
        if hasattr(model, 'RAJ') and hasattr(model, 'DECJ'):
            return (model.RAJ.value, model.DECJ.value)
        else:
            # TODO: better way of dealing with units
            d2r = np.pi / 180
            elong, elat = model.ELONG.value, model.ELAT.value
            return self._get_radec_from_ecliptic(elong*d2r, elat*d2r)

    def _get_planetssb(self):
        return np.zeros((len(self._toas), 9, 6))

class Tempo2Pulsar(BasePulsar):

    def __init__(self, t2pulsar, sort=True,
                 drop_t2pulsar=True, planets=True):

        self._sort = sort
        self.t2pulsar = t2pulsar
        self.planets = planets
        self.name = str(t2pulsar.name)

        self._toas = np.double(t2pulsar.toas()) * 86400
        # saving also stoas (e.g., for DMX comparisons)
        self._stoas = np.double(t2pulsar.stoas) * 86400
        self._residuals = np.double(t2pulsar.residuals())
        self._toaerrs = np.double(t2pulsar.toaerrs) * 1e-6
        self._designmatrix = np.double(t2pulsar.designmatrix())
        self._ssbfreqs = np.double(t2pulsar.ssbfreqs()) / 1e6

        # fitted parameters
        self.fitpars = ['Offset'] + list(map(str, t2pulsar.pars()))

        # set parameters
        spars = list(map(str, t2pulsar.pars(which='set')))
        self.setpars = [sp for sp in spars if sp not in self.fitpars]

        self._flags = {}
        for key in t2pulsar.flags():
            self._flags[key] = t2pulsar.flagvals(key)

        self._pdist = self._get_pdist()
        self._raj, self._decj = self._get_radec(t2pulsar)
        self._pos = self._get_pos()
        self._planetssb = self._get_planetssb(t2pulsar)

        # gather DM/DMX information if available
        self._set_dm(t2pulsar)

        self._pos_t = t2pulsar.psrPos.copy()
        if 'ELONG' and 'ELAT' in np.concatenate((t2pulsar.pars(which='fit'),
                                                 t2pulsar.pars(which='set'))):
            self._pos_t = utils.ecl2eq_vec(self._pos_t)

        self.sort_data()

        if drop_t2pulsar:
            del self.t2pulsar

    # gather DM/DMX information if available
    def _set_dm(self, t2pulsar):
        pars = t2pulsar.pars(which='set')

        if 'DM' in pars:
            self._dm = t2pulsar['DM'].val

        dmx = {par: {'DMX': t2pulsar[par].val,
                     'DMXerr': t2pulsar[par].err,
                     'DMXR1': t2pulsar[par[:3] + 'R1' + par[3:]].val,
                     'DMXR2': t2pulsar[par[:3] + 'R2' + par[3:]].val,
                     'fit': par in pars}
               for par in pars if 'DMX_' in par}

        if dmx:
            self._dmx = dmx

    def _get_radec(self, t2pulsar):
        if 'RAJ' in np.concatenate((t2pulsar.pars(which='fit'),
                                    t2pulsar.pars(which='set'))):
            return (np.double(t2pulsar['RAJ'].val),
                    np.double(t2pulsar['DECJ'].val))

        else:
            # use ecliptic coordinates
            elong = t2pulsar['ELONG'].val
            elat = t2pulsar['ELAT'].val
            return self._get_radec_from_ecliptic(elong, elat)

    def _get_planetssb(self, t2pulsar):
        planetssb = None
        if self.planets:
            for ii in range(1, 10):
                tag = 'DMASSPLANET' + str(ii)
                self.t2pulsar[tag].val = 0.0
            self.t2pulsar.formbats()
            planetssb = np.zeros((len(self._toas), 9, 6))
            planetssb[:, 0, :] = self.t2pulsar.mercury_ssb
            planetssb[:, 1, :] = self.t2pulsar.venus_ssb
            planetssb[:, 2, :] = self.t2pulsar.earth_ssb
            planetssb[:, 3, :] = self.t2pulsar.mars_ssb
            planetssb[:, 4, :] = self.t2pulsar.jupiter_ssb
            planetssb[:, 5, :] = self.t2pulsar.saturn_ssb
            planetssb[:, 6, :] = self.t2pulsar.uranus_ssb
            planetssb[:, 7, :] = self.t2pulsar.neptune_ssb
            planetssb[:, 8, :] = self.t2pulsar.pluto_ssb

            if 'ELONG' and 'ELAT' in np.concatenate((t2pulsar.pars(),
                                                     t2pulsar.pars(
                                                         which='set'))):
                for ii in range(9):
                    planetssb[:,ii,:3] = utils.ecl2eq_vec(planetssb[:,ii,:3])
                    planetssb[:,ii,3:] = utils.ecl2eq_vec(planetssb[:,ii,3:])
        return planetssb

def Pulsar(*args, **kwargs):

    ephem = kwargs.get('ephem', None)
    clk = kwargs.get('clk', None)
    planets = kwargs.get('planets', True)
    sort = kwargs.get('sort', True)
    drop_t2pulsar = kwargs.get('drop_t2pulsar', True)
    timing_package = kwargs.get('timing_package', 'tempo2')

    if pint is not None:
        toas = list(filter(lambda x: isinstance(x, toa.TOAs), args))
        model = list(filter(lambda x: isinstance(x, TimingModel), args))

    if t2 is not None:
        t2pulsar = list(filter(lambda x: isinstance(x, t2.tempopulsar), args))

    parfile = list(filter(lambda x: isinstance(x, str) and
                          x.split('.')[-1] == 'par', args))
    timfile = list(filter(lambda x: isinstance(x, str) and
                          x.split('.')[-1] in ['tim', 'toa'], args))

    if pint and toas and model:
        return PintPulsar(toas[0], model[0], sort=sort, planets=planets)
    elif t2 and t2pulsar:
        return Tempo2Pulsar(t2pulsar[0], sort=sort,
                            drop_t2pulsar=drop_t2pulsar,
                            planets=planets)
    elif parfile and timfile:
        # Check whether the two files exist
        if not os.path.isfile(parfile[0]) or not os.path.isfile(timfile[0]):
            msg = 'Cannot find parfile {0} or timfile {1}!'.format(
                parfile[0], timfile[0])
            raise IOError(msg)

        # Obtain the directory name of the timfile, and change to it
        timfiletup = os.path.split(timfile[0])
        dirname = timfiletup[0] or './'
        reltimfile = timfiletup[-1]
        relparfile = os.path.relpath(parfile[0], dirname)

        # get current directory
        cwd = os.getcwd()

        # Change directory to the base directory of the tim-file to deal with
        # INCLUDE statements in the tim-file
        os.chdir(dirname)

        if timing_package.lower() == 'pint':
            if ephem is None:
                ephem = 'DE421'
            if clk is None:
                bipm_version = 'BIPM2015'
            else:
                bipm_version = clk.split('(')[1][:-1]
            toas = toa.get_TOAs(reltimfile, ephem=ephem, planets=planets,
                                bipm_version=bipm_version)
            model = mb.get_model(relparfile)
            os.chdir(cwd)
            return PintPulsar(toas, model, sort=sort, planets=planets)

        elif timing_package.lower() == 'tempo2':

            # hack to set maxobs
            maxobs = get_maxobs(reltimfile) + 100
            t2pulsar = t2.tempopulsar(relparfile, reltimfile,
                                      maxobs=maxobs, ephem=ephem, clk=clk)
            os.chdir(cwd)
            return Tempo2Pulsar(t2pulsar, sort=sort,
                                drop_t2pulsar=drop_t2pulsar,
                                planets=planets)
    else:
        print('Unknown arguments {}'.format(args))


#################################################################################
# Copyright (c) 2011-2013, Pacific Biosciences of California, Inc.
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in the
#   documentation and/or other materials provided with the distribution.
# * Neither the name of Pacific Biosciences nor the names of its
#   contributors may be used to endorse or promote products derived from
#   this software without specific prior written permission.
#
# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY
# THIS LICENSE.  THIS SOFTWARE IS PROVIDED BY PACIFIC BIOSCIENCES AND ITS
# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL PACIFIC BIOSCIENCES OR
# ITS CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
# BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
# IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#################################################################################

from math import sqrt
import math
import scipy.stats as s
import array as a
import sys

from numpy import log, pi, log10, e, log1p, exp
import numpy as np
import re

log10e = log10(e)

canonicalBaseMap = {'A': 'A', 'C': 'C', 'G': 'G', 'T': 'T', 'H': 'A', 'I': 'C', 'J': 'C', 'K': 'C'}
modNames = {'H': 'm6A', 'I': 'm5C', 'J': 'm4C', 'K': 'm5C'}

m5CCode = 'I'

iupacMap = {
    'A': 'A',
    'C': 'C',
    'G': 'G',
    'T': 'T',
    'K': 'GT',
    'M': 'AC',
    'R': 'AG',
    'Y': 'CT',
    'S': 'CG',
    'W': 'AT',
    'B': 'CGT',
    'D': 'AGT',
    'H': 'ACT',
    'V': 'ACG',
    'N': 'ACGT'
}

def findMotifPositions(seq, motifs):
    regexs = []

    # Generate a regex for each motif, honouring degenerate bases
    for m in motifs:
        regex = ''

        for c in m:
            regex = regex + "[" + iupacMap[c] + "]"

        regexs.append(regex)

    allMatches = []

    # Return a list of matching positions in the sequence
    for r in regexs:
        rr = re.compile(r)
        matches = [x.start() for x in rr.finditer(seq)]
        allMatches.extend(matches)

    allMatches.sort()

    return allMatches

class MultiSiteDetection(object):

    def __init__(self, gbmModel, sequence, rawKinetics, callBounds, methylMinCov, motifs=['CG']):
        """

        """

        self.methylMinCov = methylMinCov
        self.motifs = motifs

        self.gbmModel = gbmModel
        self.sequence = sequence

        self.callStart = callBounds[0]
        self.callEnd = callBounds[1]

        # Extents that we will attempt to call a modification
        self.callRange = xrange(self.callStart, self.callEnd)

        # These switch because we changing viewpoints
        self.pre = gbmModel.post
        self.post = gbmModel.pre

        self.lStart = self.pre
        self.lEnd = len(self.sequence) - self.post

        # Extents that we will use for likelihoods
        self.likelihoodRange = xrange(self.lStart, self.lEnd)

        self.alternateBases = dict((x, list(sequence[x])) for x in xrange(len(sequence)))

        self.rawKinetics = rawKinetics

    def getConfigs(self, centerIdx):
        ''' Enumerate all the contexts centered at centerIdx with one
            modification added '''
        start = centerIdx - self.pre
        end = centerIdx + self.post
        return self._possibleConfigs(start, end)

    def _possibleConfigs(self, start, end):
        ''' Enumerate all the contexts coming from the substring self.sequence[start,end] with one
            modification added '''

        if start == end:
            return self.alternateBases[start]
        else:
            r = []
            allSuffixes = self._possibleConfigs(start + 1, end)

            # The first suffix is alway the one with no modifications
            # Only add the alternate to that one -- that way we only
            # get configurations with a single modification, not all combos

            noModsSuffix = allSuffixes[0]
            if len(allSuffixes) > 1:
                    restSuffixes = allSuffixes[1:]
            else:
                    restSuffixes = []

            # The noMods suffix get the alternates
            for c in self.alternateBases[start]:
                    r.append(c + noModsSuffix)

            # the other suffixes already have mods -- they just get the unmodified base
            for suffix in restSuffixes:
                    r.append(self.alternateBases[start][0] + suffix)

            return r

        # Compute something for all the windows in [start, end]
    def getContexts(self, start, end, sequence):
        contexts = []

        for pos in xrange(start, end + 1):
            ctx = sequence[(pos - self.pre):(pos + self.post + 1)].tostring()
            contexts.append(ctx)

        return contexts

    def computeContextMeans(self):
        """Generate a hash of the mean ipd for all candidate contexts"""

        allContexts = []

        for pos in self.motifPositions:
            for offsetPos in xrange(pos - self.post, pos + self.pre + 1):
                cfgs = self.getConfigs(offsetPos)
                allContexts.extend(cfgs)

        predictions = self.gbmModel.getPredictions(allContexts)
        self.contextMeanTable = dict(zip(allContexts, predictions))

    def decode(self):
        """Use this method to do the full modification finding protocol"""

        # Find sites matching the desired motif
        self.findMotifs()

        # Compute all the required mean ipds under all possible composite hypotheses
        self.computeContextMeans()

        # Compute a confidence for each mod and return results
        return self.scorePositions()

    def findMotifs(self):
        """ Mark all the positions matching the requested motif """

        # Generate list of matching positions
        allMotifPositions = findMotifPositions(self.sequence, self.motifs)
        self.motifPositions = []

        for pos in allMotifPositions:
            # Only use bases that are inside the callBounds
            if self.callStart <= pos < self.callEnd:
                self.alternateBases[pos].append('I')
                self.motifPositions.append(pos)

    def multiSiteDetection(self, positions, nullPred, modPred, centerPosition):
        ''' kinetics, nullPred, and modifiedPred are parallel arrays
            containing the observations and predictions surrounding a
            single candidate motif site.  Estimate the p-value of
            modification and the modified fraction here'''

        # Apply the error model to the predictions
        nullErr = 0.01 + 0.03 * nullPred + 0.06 * nullPred ** (1.7)
        modErr = 0.01 + 0.03 * modPred + 0.06 * modPred ** (1.7)

        obsMean = np.zeros(nullPred.shape)
        obsErr = np.zeros(nullPred.shape)

        # Get the observations into the same array format
        for i in xrange(len(positions)):
            position = positions[i]

            if position in self.rawKinetics:
                siteObs = self.rawKinetics[position]
                obsMean[i] = siteObs['tMean']
                obsErr[i] = siteObs['tErr']
            else:
                # Crank up the variance -- we don't have an observation at this
                # position, so we should ignore it.
                obsMean[i] = 0.0
                obsErr[i] = 999999999

        # Subtract off the background model from the observations and the modified prediction
        dObs = obsMean - nullPred
        # Error of observation and prediction are uncorrelated
        obsSigma = obsErr ** 2 + nullErr ** 2
        invObsSigma = 1.0 / obsSigma

        # Error of null prediction and mod prediction are probably correlated -- need a better estimate of the error of the difference!!
        dPred = modPred - nullPred
        dPredSigma = (obsErr ** 2 + nullErr ** 2) / 2  # Just stubbing in a factor of 2 here...

        weightsNumerator = invObsSigma * dPred
        weights = weightsNumerator / (dPred * weightsNumerator).sum()

        signalEstimate = (weights * dObs).sum()
        varianceEstimate = (np.abs(weights) * obsSigma).sum()

        maxSignal = (weights * dPred).sum()
        maxSignalVariance = (np.abs(weights) * dPredSigma).sum()

        # Now just run the standard erf on this Gaussian to quantify the probability that there is some signal
        # What we want now:
        #
        # 1. p-value that dObs * dPred (dot product) is greater than 0.
        # 2. Distribution of \alpha, where dObs = \alpha dPred, where \alpha \in [0,1], with appropriate error propagation
        # 2a. Is it possible to summarize 2 with a Beta distribution?

        pvalue = s.norm._cdf(-signalEstimate / varianceEstimate)
        pvalue = max(sys.float_info.min, pvalue)
        score = -10.0 * log10(pvalue)

        centerPosition['MSscore'] = score
        centerPosition['MSpvalue'] = pvalue

        centerPosition['signal'] = signalEstimate
        centerPosition['variance'] = varianceEstimate

        centerPosition['modelSignal'] = maxSignal
        centerPosition['modelVariance'] = maxSignalVariance

        centerPosition['Mask'] = []

        return centerPosition

    def scorePositions(self):
        """
        Score each motif site in the sequence.
        """

        qvModCalls = dict()

        dnaSeq = a.array('c')
        dnaSeq.fromstring(self.sequence)

        for pos in self.motifPositions:
            if pos in self.rawKinetics:

                # Fetch unmodified positions
                nullPred = self.getRegionPredictions(pos - self.post, pos + self.pre, dnaSeq)

                # Fetch modified positions and reset sequence
                originalBase = dnaSeq[pos]
                dnaSeq[pos] = m5CCode
                modifiedPred = self.getRegionPredictions(pos - self.post, pos + self.pre, dnaSeq)
                dnaSeq[pos] = originalBase

                # Position that contribute to this call
                positions = xrange(pos - self.post, pos + self.pre + 1)

                # Run the multi-site detection and save the results
                centerStats = self.rawKinetics[pos]
                centerStats = self.multiSiteDetection(positions, nullPred, modifiedPred, centerStats)

                qvModCalls[pos] = centerStats

        return qvModCalls

    def getRegionPredictions(self, start, end, sequence):
        predictions = np.zeros(end - start + 1)

        for pos in xrange(start, end + 1):
            ctx = sequence[(pos - self.pre):(pos + self.post + 1)].tostring()
            predictions[pos - start] = self.contextMeanTable[ctx]

        return predictions


# Copyright 2020 Arthur Coqué, Guillaume Morin, Pôle OFB-INRAE ECLA, UR RECOVER
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""This module gathers wc algorithms used for estimating SPM concentrations.

Each class of this module correspond to one algorithm. An algorithm can have
several calibrations (a calibration is a set of parameters), either
packaged within SISPPEO (these default calibrations are located in
'resources/wc_algo_calibration') or provided by the user.
Before its utilisation, an algorithm has to be instantiate with specific
settings like the product_type of further input products, the calibration
used, the band used (if needed), etc.

Example:

    algo1 = SPMNechad('S2_GRS', 'B4', 'Nechad_2016')
    out_array2 = algo1(input_array, 'rho')

    algo2 = SPMGet('L8_GRS', 'GET_2018')
    out_array2 = algo2(red_array, nir_array, 'rrs')
"""

from pathlib import Path
from typing import Optional, Union

import numpy as np
import xarray as xr

from sisppeo.utils.algos import load_calib, producttype_to_sat
from sisppeo.utils.config import wc_algo_config as algo_config, wc_calib
from sisppeo.utils.exceptions import InputError

# pylint: disable=invalid-name
# Ok for a custom type.
P = Union[str, Path]
N = Union[int, float]

def _nechad(rho, a, c):
    ssc = a * rho / (1 - (rho / c))
    return ssc

class SPMNechad:
    """Semi-analytical algorithm to retrieve SPM concentration (in mg/l) from reflectance.

    Semi-analytical algorithm to retrieve SPM concentrations (in mg/L) from
    surface reflectances (rho, unitless) or remote sensing reflectances (Rrs,
    in sr-1).
    This algorithm was presented in Nechad et al., 2010 and 2016.

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_band = 'B4'
    _default_calibration_file = wc_calib / 'spm_nechad.yaml'
    _default_calibration_name = 'Nechad_2016'
    name = 'spm-nechad'

    def __init__(self,
                 product_type: str,
                 requested_band: str = _default_band,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'SPMNechad' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT).
            requested_band: Optional; The band used by the algorithm (
                default=_default_band).
            calibration: The calibration (set of parameters) used by the
                algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        self.requested_bands = [requested_band]
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._valid_limit = calibration_dict['validity_limit']
        try:
            params = calibration_dict[producttype_to_sat(product_type)][
                requested_band]
        except KeyError as invalid_input:
            msg = (f'{product_type} or {requested_band} is not allowed with '
                   f'{self.name}/this calibration')
            raise InputError(msg) from invalid_input
        self.__dict__.update(params)
        self.meta = {'band': requested_band,
                     'calibration': calibration_name,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 rho: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('rho').

        Args:
            rho: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'rho' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of SPM concentration (in mg/L).
        """

        if data_type == 'rrs':
            rho = np.pi * rho

        np.warnings.filterwarnings('ignore')
        # pylint: disable=no-member
        # Loaded in __init__ with "__dict__.update".
        spm = _nechad(rho, self.a, self.c)
        spm = spm.where((rho >= 0) & (spm >= 0) & (spm < self._valid_limit))
        return spm

class SPMHan:
    """Switching Semi-analytical algorithm to retrieve SPM (in mg/l) from reflectance.

    Switching Semi-analytical algorithm to retrieve suspended particulate
    matter (in mg/l) from surface reflectances (rho, unitless) or remote
    sensing reflectances (Rrs, in sr-1).
    This algorithm was published in Han et al., 2016

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_calibration_file = wc_calib / 'spm_han.yaml'
    _default_calibration_name = 'Han_2016'
    name = 'spm-han'

    def __init__(self,
                 product_type: str,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'SPMHan' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT)
            calibration: The calibration (set of parameters) used by the
                algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from invalid_product
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._valid_limit = calibration_dict['validity_limit']
        try:
            params = calibration_dict[producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with this calibration'
            raise InputError(msg) from invalid_product
        self.__dict__.update(params)
        self.meta = {'calibration': calibration_name,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 refl_red: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('rho').

        Args:
            refl_red: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'rho' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of SPM concentration (in mg/L).
        """
        if data_type == 'rho':
            refl_red = refl_red / np.pi
            print(data_type)

        np.warnings.filterwarnings('ignore')
        rrs_red = refl_red.where(refl_red >= 0)
        # pylint: disable=no-member
        # Loaded in __init__ whit "__dict__.update".
        spm_low = _nechad(rrs_red, self.a_low, self.c_low)
        spm_high = _nechad(rrs_red, self.a_high, self.c_high)

        w_low = np.log10(self.switch_sup) - np.log10(rrs_red)
        w_high = np.log10(rrs_red) - np.log10(self.switch_inf)
        spm_mixing = (w_low * spm_low + w_high * spm_high) / (w_low + w_high)

        spm = rrs_red.where(rrs_red > self.switch_inf, spm_low)
        spm = spm.where(rrs_red < self.switch_sup, spm_high)
        spm = spm.where((rrs_red <= self.switch_inf)
                        | (rrs_red >= self.switch_sup), spm_mixing)
        spm = spm.where((spm >= 0) & (spm <= self._valid_limit))
        return spm

class SPMGet:
    """Switching Semi-analytical algorithm to retrieve SPM (in mg/l) from reflectance.

    Switching Semi-analytical algorithm to retrieve suspended particulate
    matter (in mg/l) from surface reflectances (rho, unitless) or remote
    sensing reflectances (Rrs, in sr-1).
    This algorithm was calibrated on GET radiometric database in 2018.

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_calibration_file = wc_calib / 'spm_get.yaml'
    _default_calibration_name = 'GET_2018'
    name = 'spm-get'

    def __init__(self,
                 product_type: str,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'SPMGet' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT)
            calibration: The calibration (set of parameters) used by the
                algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from invalid_product
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._switch_inf = calibration_dict['switch_inf']
        self._switch_sup = calibration_dict['switch_sup']
        self._valid_limit = calibration_dict['validity_limit']
        try:
            params = calibration_dict[producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with this calibration'
            raise InputError(msg) from invalid_product
        self.__dict__.update(params)
        self.meta = {'calibration': calibration_name,
                     'switch inf': self._switch_inf,
                     'switch sup': self._switch_sup,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 refl_red: xr.DataArray,
                 refl_nir: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('rho').

        Args:
            refl_red: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'rho' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of SPM concentration (in mg/L).
        """
        if data_type == 'rrs':
            refl_red = np.pi * refl_red
            refl_nir = np.pi * refl_nir

        np.warnings.filterwarnings('ignore')
        rho_red = refl_red.where(refl_red >= 0)
        rho_nir = refl_red.where(refl_nir >= 0)

        # pylint: disable=no-member
        # Loaded in __init__ whit "__dict__.update".
        spm_low = _nechad(rho_red, self.a_nechad, self.c_nechad)
        spm_high = self.coef_br * np.power((rho_nir / rho_red), self.exp_br)

        w = ((rho_red - self._switch_inf)
             / (self._switch_sup - self._switch_inf))
        spm_mixing = (1 - w) * spm_low + w * spm_high

        spm = rho_red.where(rho_red > self._switch_inf, spm_low)
        spm = spm.where(rho_red < self._switch_sup, spm_high)
        spm = spm.where((rho_red <= self._switch_inf)
                        | (rho_red >= self._switch_sup), spm_mixing)
        spm = spm.where((spm >= 0) & (spm <= self._valid_limit))
        return spm

class TURBIDogliotti:
    """Switching Semi-analytical algorithm to retrieve Turbidity (in FNU) from reflectance.

    Switching Semi-analytical algorithm to retrieve Turbidity (in FNU) from
    surface reflectances (rho, unitless) or remote sensing reflectances (Rrs,
    in sr-1).
    This algorithm was published in Dogliotti et al., 2015

    Attributes:
        name: The name of the algorithm used. This is the key used by
            L3AlgoBuilder and that you must provide in config or when using
            the CLI.
        requested_bands: A list of bands further used by the algorithm.
        meta: A dict of metadata (calibration name, model coefficients, etc).
    """
    _default_calibration_file = wc_calib / 'turbi_dogliotti.yaml'
    _default_calibration_name = 'Dogliotti_2015'
    name = 'turbi-dogliotti'

    def __init__(self,
                 product_type: str,
                 calibration: Optional[P] = None,
                 **_ignored) -> None:
        """Inits an 'TURBIDogliotti' instance with specific settings.

        Args:
            product_type: The type of the input satellite product (e.g.
                S2_ESA_L2A or L8_USGS_L1GT)
            calibration: The calibration (set of parameters) used by the
                algorithm (default=_default_calibration_name).
            **_ignored: Unused kwargs sent to trash.
        """
        try:
            self.requested_bands = algo_config[self.name][
                producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with {self.name}'
            raise InputError(msg) from invalid_product
        calibration_dict, calibration_name = load_calib(
            calibration,
            self._default_calibration_file,
            self._default_calibration_name
        )
        self._switch_inf = calibration_dict['switch_inf']
        self._switch_sup = calibration_dict['switch_sup']
        self._valid_limit = calibration_dict['validity_limit']
        try:
            params = calibration_dict[producttype_to_sat(product_type)]
        except KeyError as invalid_product:
            msg = f'{product_type} is not allowed with this calibration'
            raise InputError(msg) from invalid_product
        self.__dict__.update(params)
        self.meta = {'calibration': calibration_name,
                     'switch inf': self._switch_inf,
                     'switch sup': self._switch_sup,
                     'validity_limit': self._valid_limit,
                     **params}

    def __call__(self,
                 rho_red: xr.DataArray,
                 rho_nir: xr.DataArray,
                 data_type: str,
                 **_ignored) -> xr.DataArray:
        """Runs the algorithm on the input array ('rho').

        Args:
            rho_red: An array (dimension 1 * N * M) of 'data_type'.
            rho_nir: An array (dimension 1 * N * M) of 'data_type'.
            data_type: Either 'rho' or 'rrs' (respectively surface reflectance
                and remote sensing reflectance).
            **_ignored: Unused kwargs sent to trash.

        Returns:
            An array (dimension 1 * N * M) of Turbidity (in FNU).
        """

        if data_type == 'rrs':
            rho_red = np.pi * rho_red
            rho_nir = np.pi * rho_nir

        np.warnings.filterwarnings('ignore')
        rho_red = rho_red.where(rho_red >= 0)
        rho_nir = rho_red.where(rho_nir >= 0)
        # pylint: disable=no-member
        # Loaded in __init__ whit "__dict__.update".
        t_low = _nechad(rho_red, self.a_low, self.c_low)
        t_high = _nechad(rho_nir, self.a_high, self.c_high)
        w = ((rho_red - self._switch_inf)
             / (self._switch_sup - self._switch_inf))
        t_mixing = (1 - w) * t_low + w * t_high

        turb = rho_red.where(rho_red > self._switch_inf, t_low)
        turb = turb.where(rho_red < self._switch_sup, t_high)
        turb = turb.where((rho_red <= self._switch_inf)
                          | (rho_red >= self._switch_sup), t_mixing)
        turb = turb.where((turb >= 0) & (turb <= self._valid_limit))
        return turb


# Author: Roman Goj <roman.goj@gmail.com>
#
# License: BSD (3-clause)

import copy as cp

import numpy as np
from scipy.fftpack import fftfreq

from ..io.pick import pick_types
from ..utils import logger, verbose, warn
from ..time_frequency.multitaper import (dpss_windows, _mt_spectra,
                                         _csd_from_mt, _psd_from_mt_adaptive)

class CrossSpectralDensity(object):
    """Cross-spectral density

    Parameters
    ----------
    data : array of shape (n_channels, n_channels)
        The cross-spectral density matrix.
    ch_names : list of string
        List of channels' names.
    projs :
        List of projectors used in CSD calculation.
    bads :
        List of bad channels.
    frequencies : float | list of float
        Frequency or frequencies for which the CSD matrix was calculated. If a
        list is passed, data is a sum across CSD matrices for all frequencies.
    n_fft : int
        Length of the FFT used when calculating the CSD matrix.
    """
    def __init__(self, data, ch_names, projs, bads, frequencies, n_fft):
        self.data = data
        self.dim = len(data)
        self.ch_names = cp.deepcopy(ch_names)
        self.projs = cp.deepcopy(projs)
        self.bads = cp.deepcopy(bads)
        self.frequencies = np.atleast_1d(np.copy(frequencies))
        self.n_fft = n_fft

    def __repr__(self):
        s = 'frequencies : %s' % self.frequencies
        s += ', size : %s x %s' % self.data.shape
        s += ', data : %s' % self.data
        return '<CrossSpectralDensity  |  %s>' % s

@verbose
def compute_epochs_csd(epochs, mode='multitaper', fmin=0, fmax=np.inf,
                       fsum=True, tmin=None, tmax=None, n_fft=None,
                       mt_bandwidth=None, mt_adaptive=False, mt_low_bias=True,
                       projs=None, verbose=None):
    """Estimate cross-spectral density from epochs

    Note: Baseline correction should be used when creating the Epochs.
          Otherwise the computed cross-spectral density will be inaccurate.

    Note: Results are scaled by sampling frequency for compatibility with
          Matlab.

    Parameters
    ----------
    epochs : instance of Epochs
        The epochs.
    mode : str
        Spectrum estimation mode can be either: 'multitaper' or 'fourier'.
    fmin : float
        Minimum frequency of interest.
    fmax : float | np.inf
        Maximum frequency of interest.
    fsum : bool
        Sum CSD values for the frequencies of interest. Summing is performed
        instead of averaging so that accumulated power is comparable to power
        in the time domain. If True, a single CSD matrix will be returned. If
        False, the output will be a list of CSD matrices.
    tmin : float | None
        Minimum time instant to consider. If None start at first sample.
    tmax : float | None
        Maximum time instant to consider. If None end at last sample.
    n_fft : int | None
        Length of the FFT. If None the exact number of samples between tmin and
        tmax will be used.
    mt_bandwidth : float | None
        The bandwidth of the multitaper windowing function in Hz.
        Only used in 'multitaper' mode.
    mt_adaptive : bool
        Use adaptive weights to combine the tapered spectra into PSD.
        Only used in 'multitaper' mode.
    mt_low_bias : bool
        Only use tapers with more than 90% spectral concentration within
        bandwidth. Only used in 'multitaper' mode.
    projs : list of Projection | None
        List of projectors to use in CSD calculation, or None to indicate that
        the projectors from the epochs should be inherited.
    verbose : bool, str, int, or None
        If not None, override default verbose level (see mne.verbose).

    Returns
    -------
    csd : instance of CrossSpectralDensity
        The computed cross-spectral density.
    """
    # Portions of this code adapted from mne/connectivity/spectral.py

    # Check correctness of input data and parameters
    if fmax < fmin:
        raise ValueError('fmax must be larger than fmin')
    tstep = epochs.times[1] - epochs.times[0]
    if tmin is not None and tmin < epochs.times[0] - tstep:
        raise ValueError('tmin should be larger than the smallest data time '
                         'point')
    if tmax is not None and tmax > epochs.times[-1] + tstep:
        raise ValueError('tmax should be smaller than the largest data time '
                         'point')
    if tmax is not None and tmin is not None:
        if tmax < tmin:
            raise ValueError('tmax must be larger than tmin')
    if epochs.baseline is None and epochs.info['highpass'] < 0.1:
        warn('Epochs are not baseline corrected or enough highpass filtered. '
             'Cross-spectral density may be inaccurate.')

    if projs is None:
        projs = cp.deepcopy(epochs.info['projs'])
    else:
        projs = cp.deepcopy(projs)

    picks_meeg = pick_types(epochs[0].info, meg=True, eeg=True, eog=False,
                            ref_meg=False, exclude='bads')
    ch_names = [epochs.ch_names[k] for k in picks_meeg]

    # Preparing time window slice
    tstart, tend = None, None
    if tmin is not None:
        tstart = np.where(epochs.times >= tmin)[0][0]
    if tmax is not None:
        tend = np.where(epochs.times <= tmax)[0][-1] + 1
    tslice = slice(tstart, tend, None)
    n_times = len(epochs.times[tslice])
    n_fft = n_times if n_fft is None else n_fft

    # Preparing frequencies of interest
    sfreq = epochs.info['sfreq']
    orig_frequencies = fftfreq(n_fft, 1. / sfreq)
    freq_mask = (orig_frequencies > fmin) & (orig_frequencies < fmax)
    frequencies = orig_frequencies[freq_mask]
    n_freqs = len(frequencies)

    if n_freqs == 0:
        raise ValueError('No discrete fourier transform results within '
                         'the given frequency window. Please widen either '
                         'the frequency window or the time window')

    # Preparing for computing CSD
    logger.info('Computing cross-spectral density from epochs...')
    if mode == 'multitaper':
        # Compute standardized half-bandwidth
        if mt_bandwidth is not None:
            half_nbw = float(mt_bandwidth) * n_times / (2 * sfreq)
        else:
            half_nbw = 2

        # Compute DPSS windows
        n_tapers_max = int(2 * half_nbw)
        window_fun, eigvals = dpss_windows(n_times, half_nbw, n_tapers_max,
                                           low_bias=mt_low_bias)
        n_tapers = len(eigvals)
        logger.info('    using multitaper spectrum estimation with %d DPSS '
                    'windows' % n_tapers)

        if mt_adaptive and len(eigvals) < 3:
            warn('Not adaptively combining the spectral estimators due to a '
                 'low number of tapers.')
            mt_adaptive = False
    elif mode == 'fourier':
        logger.info('    using FFT with a Hanning window to estimate spectra')
        window_fun = np.hanning(n_times)
        mt_adaptive = False
        eigvals = 1.
        n_tapers = None
    else:
        raise ValueError('Mode has an invalid value.')

    csds_mean = np.zeros((len(ch_names), len(ch_names), n_freqs),
                         dtype=complex)

    # Picking frequencies of interest
    freq_mask_mt = freq_mask[orig_frequencies >= 0]

    # Compute CSD for each epoch
    n_epochs = 0
    for epoch in epochs:
        epoch = epoch[picks_meeg][:, tslice]

        # Calculating Fourier transform using multitaper module
        x_mt, _ = _mt_spectra(epoch, window_fun, sfreq, n_fft)

        if mt_adaptive:
            # Compute adaptive weights
            _, weights = _psd_from_mt_adaptive(x_mt, eigvals, freq_mask,
                                               return_weights=True)
            # Tiling weights so that we can easily use _csd_from_mt()
            weights = weights[:, np.newaxis, :, :]
            weights = np.tile(weights, [1, x_mt.shape[0], 1, 1])
        else:
            # Do not use adaptive weights
            if mode == 'multitaper':
                weights = np.sqrt(eigvals)[np.newaxis, np.newaxis, :,
                                           np.newaxis]
            else:
                # Hack so we can sum over axis=-2
                weights = np.array([1.])[:, None, None, None]

        x_mt = x_mt[:, :, freq_mask_mt]

        # Calculating CSD
        # Tiling x_mt so that we can easily use _csd_from_mt()
        x_mt = x_mt[:, np.newaxis, :, :]
        x_mt = np.tile(x_mt, [1, x_mt.shape[0], 1, 1])
        y_mt = np.transpose(x_mt, axes=[1, 0, 2, 3])
        weights_y = np.transpose(weights, axes=[1, 0, 2, 3])
        csds_epoch = _csd_from_mt(x_mt, y_mt, weights, weights_y)

        # Scaling by number of samples and compensating for loss of power due
        # to windowing (see section 11.5.2 in Bendat & Piersol).
        if mode == 'fourier':
            csds_epoch /= n_times
            csds_epoch *= 8 / 3.

        # Scaling by sampling frequency for compatibility with Matlab
        csds_epoch /= sfreq

        csds_mean += csds_epoch
        n_epochs += 1

    csds_mean /= n_epochs

    logger.info('[done]')

    # Summing over frequencies of interest or returning a list of separate CSD
    # matrices for each frequency
    if fsum is True:
        csd_mean_fsum = np.sum(csds_mean, 2)
        csd = CrossSpectralDensity(csd_mean_fsum, ch_names, projs,
                                   epochs.info['bads'],
                                   frequencies=frequencies, n_fft=n_fft)
        return csd
    else:
        csds = []
        for i in range(n_freqs):
            csds.append(CrossSpectralDensity(csds_mean[:, :, i], ch_names,
                                             projs, epochs.info['bads'],
                                             frequencies=frequencies[i],
                                             n_fft=n_fft))
        return csds


#!/usr/bin/python

###Sterimol (and Tolman CA) Calculator###

###############################################################
#                       sterimoltools.py                      #
#                                                             #
###############################################################

#Python Libraries
import subprocess, sys, os
from numpy import *
from scipy import *
from math import *
import numpy as np
#from vpython import *

#Chemistry Libaries
#from radialdata import *
#from pars import *

#Avoid number error warnings
import warnings
warnings.filterwarnings("ignore")

#Chemistry Arrays
periodictable = ["Bq","H","He","Li","Be","B","C","N","O","F","Ne","Na","Mg","Al","Si","P","S","Cl","Ar","K","Ca","Sc","Ti","V","Cr","Mn","Fe","Co","Ni","Cu","Zn","Ga","Ge","As","Se","Br","Kr","Rb","Sr","Y","Zr",
             "Nb","Mo","Tc","Ru","Rh","Pd","Ag","Cd","In","Sn","Sb","Te","I","Xe","Cs","Ba","La","Ce","Pr","Nd","Pm","Sm","Eu","Gd","Tb","Dy","Ho","Er","Tm","Yb","Lu","Hf","Ta","W","Re","Os","Ir","Pt","Au","Hg","Tl",
             "Pb","Bi","Po","At","Rn","Fr","Ra","Ac","Th","Pa","U","Np","Pu","Am","Cm","Bk","Cf","Es","Fm","Md","No","Lr","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Uub","Uut","Uuq","Uup","Uuh","Uus","Uuo"]

metals = ["Li","Be","Na","Mg","Al","K","Ca","Sc","Ti","V","Cr","Mn","Fe","Co","Ni","Cu","Zn","Ga","Rb","Sr","Y","Zr","Nb","Mo","Tc","Ru","Rh","Pd","Ag","Cd","In","Sn","Cs","Ba","La","Ce","Pr","Nd","Pm","Sm","Eu","Gd","Tb","Dy","Ho","Er","Tm","Yb","Lu","Hf","Ta","W","Re","Os","Ir","Pt","Au","Hg","Tl","Pb","Bi","Po","Fr","Ra","Ac","Th","Pa","U","Np","Pu","Am","Cm","Bk","Cf","Es","Fm","Md","No","Lr","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Cn","Uut","Fl","Uup","Lv"]

# Verloop's original Sterimol parameters use CPK atomic VdW radii based on atom-type definitions
sterimol_atomtypes = ["C", "C2", "C3", "C4", "C5/N5", "C6/N6", "C7", "C8", "H", "N", "C66", "N4", "O", "O2", "P", "S", "S1", "F", "C1", "S4", "B1", "I"]

# CPK VdW radii in pm
cpk_radii = [150,160,160,150,170,170,170,150,100,150,170,145,135,135,140,170,100,135,180,140,195,215]

def getfragment(atom,molcart):
   bondlist=[atom]
   for a in range(len(molcart)):
      if calcdist(atom,a,molcart)<1.92 and a not in bondlist:bondlist.append(a)

      for b in range(len(bondlist)):
         for c in range(len(molcart)):

            if calcdist(bondlist[b],c,molcart)<1.92 and c not in bondlist:bondlist.append(c)
   return bondlist

def connectivity(atom,molcart,aty):
   con=[]
   for a in range(len(molcart)):
      if aty[a]in metals and molcart[a] != molcart[atom] and 0.1<calcdist(a,atom,molcart)<2:con.append(a)
      if molcart[a] != molcart[atom] and 0.1<calcdist(a,atom,molcart)<1.7:con.append(a)
   return len(con)
def genradii(atom,molcart,aty):
   #molcart=fileData.CARTESIANS
   con=connectivity(atom,molcart,aty)
   if con==0:con=1
   type=aty[atom]
   arow=periodictable.index(type)
   radius=molmod[arow][con]
   if radius==0:radius=1;print "Warning: No atomic radii found", arow, con
   return radius

def rotrel(vect1,vect2,vect3):
   ax=np.cross(vect1,vect2)
   ang=math.acos((np.dot(vect1,vect2))/(np.linalg.norm(vect1)*np.linalg.norm(vect2)))
   norm=1/(np.linalg.norm(ax))
   axnorm=np.dot(ax,norm)
   ux=axnorm[0]
   uy=axnorm[1]
   uz=axnorm[2]
   a=math.cos(ang)+((ux*ux)*(1-math.cos(ang)))
   b=(ux*uy*(1-math.cos(ang)))-(uz*math.sin(ang))
   c=(ux*uz*(1-math.cos(ang)))+(uy*math.sin(ang))
   d=(uy*ux*(1-math.cos(ang)))+(uz*math.sin(ang))
   e=(math.cos(ang))+(uy*uy*(1-math.cos(ang)))
   f=(uy*uz*(1-math.cos(ang)))-(ux*math.sin(ang))
   g=(uz*ux*(1-math.cos(ang)))-(uy*math.sin(ang))
   h=(uz*uy*(1-math.cos(ang)))+(ux*math.sin(ang))
   i=math.cos(ang)+(uz*uz*(1-math.cos(ang)))
   bigmat=([[a,b,c],[d,e,f,],[g,h,i]])
   vect=np.dot(bigmat,vect3)
   return vect

def calcdist(a,b,carts):
   return np.linalg.norm(np.subtract(carts[a],carts[b]))

def elementID(massno):
   if massno < len(periodictable): return periodictable[massno]
   else: return "XX"

def bondiRadius(massno):
   #Bondi van der Waals radii for all atoms from: Bondi, A. J. Phys. Chem. 1964, 68, 441-452, except hydrogen, which is taken from Rowland, R. S.; Taylor, R. J. Phys. Chem. 1996, 100, 7384-7391
   #Radii that are not available in either of these publications have RvdW = 2.00 Angstrom

   bondi = [0.0,1.09, 1.40, 1.82,2.00,2.00,1.70,1.55,1.52,1.47,1.54,2.27,1.73,2.00,2.10,1.80,1.80,1.75,1.88,2.75,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,1.63,1.40,1.39,1.87,2.00,1.85,1.90,
            1.85,2.02,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,1.63,1.72,1.58,1.93,2.17,2.00,2.06,1.98,2.16,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,1.72,1.66,1.55,1.96,2.02,2.00,2.00,2.00,
            2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,2.00,1.86]
   if massno<len(bondi): radius = bondi[massno]
   else: radius = 2.0
   return radius

def calcopposite(atom1,atom2,angle,molcart):
   h=calcdist(atom1,atom2,molcart)
   d=h*math.sin(angle)
   return d
def calcadj(atom1,atom2,angle,molcart):
   h=calcdist(atom1,atom2,molcart)
   d=h*math.cos(angle)
   return d

def getcoords(atom,molcart):
   coords=[]
   for i in range(3):
      coords.append(molcart[atom][i])
   return coords

def avpoints(atomnos,molcart):
   xcoords=[]
   ycoords=[]
   zcoords=[]
   for a in atomnos:
      xcoords.append(molcart[a][0])
      ycoords.append(molcart[a][1])
      zcoords.append(molcart[a][2])
   syslength=len(xcoords)
   x=0;y=0;z=0
   for i in range(syslength):
      x=x+xcoords[i]
      y=y+ycoords[i]
      z=z+zcoords[i]
   x=x/syslength; y=y/syslength; z=z/syslength
   return round(x,8),round(y,8),round(z,8)

def distcalc(atom1,atom2):
   x=atom1[0]-atom2[0]
   y=atom1[1]-atom2[1]
   z=atom1[2]-atom2[2]
   dist = (x**2+y**2+z**2)**0.5
   return dist

def dprod(v1, v2): return sum((a*b) for a, b in zip(v1, v2))

def length(v): return math.sqrt(dprod(v, v))

def angle(v1, v2):
   val = dprod(v1, v2) / length(v1) / length(v2)
   if val > 0.999999: val = 1.0
   if val < -0.999999: val = -1.0
   return math.acos(val)

def dihedral(atoma,atomb,atomc,atomd):
   x1=atoma[0]
   y1=atoma[1]
   z1=atoma[2]
   x2=atomb[0]
   y2=atomb[1]
   z2=atomb[2]
   x3=atomc[0]
   y3=atomc[1]
   z3=atomc[2]
   x4=atomd[0]
   y4=atomd[1]
   z4=atomd[2]
   ax= (y2-y1)*(z2-z3)-(z2-z1)*(y2-y3)
   ay= (z2-z1)*(x2-x3)-(x2-x1)*(z2-z3)
   az= (x2-x1)*(y2-y3)-(y2-y1)*(x2-x3)
   bx= (y3-y2)*(z3-z4)-(z3-z2)*(y3-y4)
   by= (z3-z2)*(x3-x4)-(x3-x2)*(z3-z4)
   bz= (x3-x2)*(y3-y4)-(y3-y2)*(x3-x4)
   nbx= (y2-y3)*(z4-z3)-(z2-z3)*(y4-y3)
   nby= (z2-z3)*(x4-x3)-(x2-x3)*(z4-z3)
   nbz= (x2-x3)*(y4-y3)-(y2-y3)*(x4-x3)
   torsion=180.0/math.pi*math.acos((ax*bx+ay*by+az*bz)/(math.sqrt(ax*ax+ay*ay+az*az)*math.sqrt(bx*bx+by*by+bz*bz)))
   sign=180.0/math.pi*math.acos((nbx*(x2-x1)+nby*(y2-y1)+nbz*(z2-z1))/(math.sqrt(nbx*nbx+nby*nby+nbz*nbz)*math.sqrt((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)+(z2-z1)*(z2-z1))))
   if sign<90.0:
      torsion=torsion*-1.0
   return torsion

#Get PDB data?

class getinData:
   def __init__(self, file):
      if not os.path.exists(file+".com"):
         print ("\nFATAL ERROR: Input file [ %s ] does not exist"%file)
      def getATOMTYPES(self, inlines):
         self.ATOMTYPES = []
         self.LEVELTYPES = []
         for i in range(0,len(inlines)):
            if inlines[i].find("#") > -1:
               if len(inlines[i+1].split()) == 0: start = i+5
               if len(inlines[i+2].split()) == 0: start = i+6
               break
         for i in range(start,len(inlines)):
            if len(inlines[i].split()) ==0:
               break
            else:
               self.ATOMTYPES.append(inlines[i].split()[0].split("-")[0])
               for oniomlevel in ["H", "M", "L"]:
                  if inlines[i].rfind(oniomlevel)>1:
                     self.LEVELTYPES.append(inlines[i][inlines[i].rfind("H"):])
                     break
      def getCARTESIANS(self, inlines, natoms):
         self.CARTESIANS = []
         for i in range(0,len(inlines)):
            if inlines[i].find("#") > -1:
               start = i+5
               break

         for i in range(start,len(inlines)):
            if len(inlines[i].split()) == 0:
               break
            elif len(inlines[i].split()) == 4:
               self.CARTESIANS.append([float(inlines[i].split()[1]), float(inlines[i].split()[2]), float(inlines[i].split()[3])])

      infile = open(file+".com","r")
      inlines = infile.readlines()
      getATOMTYPES(self, inlines)
      self.NATOMS=len(self.ATOMTYPES)

      getCARTESIANS(self, inlines, self.NATOMS)

class getinData2:
   def __init__(self, file):
      start=2
      if not os.path.exists(file+".xyz"):
         print ("\nFATAL ERROR: Input file [ %s ] does not exist"%file)
      def getATOMTYPES(self, inlines):
         self.ATOMTYPES = []
         self.LEVELTYPES = []
         for i in range(2,len(inlines)):
            if len(inlines[i].split()) ==0:
               break
            else:
               self.ATOMTYPES.append(inlines[i].split()[0].split("-")[0])
               for oniomlevel in ["H", "M", "L"]:
                  if inlines[i].rfind(oniomlevel)>1:
                     self.LEVELTYPES.append(inlines[i][inlines[i].rfind("H"):])
                     break
      def getCARTESIANS(self, inlines, natoms):
         self.CARTESIANS = []
#         for i in range(0,len(inlines)):
#            if inlines[i].find("#") > -1:
#               start = i+5
#               break
         for i in range(start,len(inlines)):
            if len(inlines[i].split()) == 0:
               break
            elif len(inlines[i].split()) == 4:
               self.CARTESIANS.append([float(inlines[i].split()[1]), float(inlines[i].split()[2]), float(inlines[i].split()[3])])

#         print self.CARTESIANS
#         for i in range(0,len(inlines)):
      infile = open(file+".xyz","r")
      inlines = infile.readlines()
      getATOMTYPES(self, inlines)
      self.NATOMS=len(self.ATOMTYPES)

      getCARTESIANS(self, inlines, self.NATOMS)

class getoutData:
   def __init__(self, file):
      if not os.path.exists(file+".out"):
         if not os.path.exists(file+".log"):
            print ("\nFATAL ERROR: Output file [ %s ] does not exist"%file)
      def getFORMAT(self, outlines):
         for i in range(0,len(outlines)):
            if outlines[i].find("Gaussian") > -1: self.FORMAT = "Gaussian"; break
      def getATOMTYPES(self, outlines, format):
         self.ATOMTYPES = []
         self.CARTESIANS = []
         if format == "Gaussian":
            for i in range(0,len(outlines)):

               if outlines[i].find("Standard orientation") > -1:
                  standor = i
                  arb=0
               if outlines[i].find("Input orientation") > -1:
                  standor2 = i
                  arb=1
               if outlines[i].find("Rotational constants") > -1 and outlines[i-1].find("-------") > -1 and arb==0:
                  self.NATOMS = i-standor-6
               if outlines[i].find("Rotational constants") > -1 and outlines[i-1].find("-------") > -1 and arb==1:
                  self.NATOMS = i-standor2-6
                  arb=2
               if outlines[i].find("Distance matrix") > -1 and outlines[i-1].find("-------") > -1:
                  self.NATOMS = i-standor2-6
            try: standor, standor2
            except NameError: pass
            else:
               if standor2>standor:del standor
               if standor>standor2:del standor2
            try: standor
            except NameError: pass
            else:
               for i in range (standor+5,standor+5+self.NATOMS):
                  self.ATOMTYPES.append(elementID(int(outlines[i].split()[1])))
                  if len(outlines[i].split())==6:self.CARTESIANS.append([float(outlines[i].split()[3]),float(outlines[i].split()[4]),float(outlines[i].split()[5])])
                  else: self.CARTESIANS.append([float(outlines[i].split()[2]),float(outlines[i].split()[3]),float(outlines[i].split()[4])])
            try: standor2
            except NameError: pass
            else:
               for i in range (standor2+5,standor2+5+self.NATOMS):
                     self.ATOMTYPES.append(elementID(int(outlines[i].split()[1])))
                     self.CARTESIANS.append([float(outlines[i].split()[3]),float(outlines[i].split()[4]),float(outlines[i].split()[5])])

      if os.path.exists(file+".out"):outfile = open(file+".out","r")
      else: outfile = open(file+".log","r")
      outlines = outfile.readlines()
      getFORMAT(self, outlines)
      getATOMTYPES(self, outlines, self.FORMAT)
      self.NATOMS=len(self.ATOMTYPES)

def concheck(conpar,val):
   cons=[]
   for a in range(len(conpar)):
      for b in range(len(conpar[a])):
         if val ==conpar[a][0]:
            for c in range(len(conpar[a])-1):
               cons.append(conpar[a][c+1])
            return cons

def twod_dist(a,b,c):
   vect1=np.subtract(a,b)
   vect2=np.subtract(b,c)
   ang=angle(vect1,vect2)
   return math.sin(ang)*np.linalg.norm(vect1)

def twod_vect(a,b,c):
   vect1=np.subtract(a,b)
   vect2=np.subtract(b,c)
   ang=angle(vect1,vect2)
   nvect2=vect2/np.linalg.norm(vect2)
   return ((math.cos(ang)*np.linalg.norm(vect1))*nvect2)+b

def twod_rot(vect,theta):
   a=math.cos(theta)
   b=math.sin(theta)
   mat=[[a,-b],[b,a]]
   vect=np.dot(mat,vect)
   return vect

# Generate Sterimol atom type from connectivity data
def generate_atom_types(atomtype, cn):
   st_types = []
   for i in range(0,len(atomtype)):
      atom = atomtype[i]
      if atom == "H": st_types.append("H")
      elif atom == "P": st_types.append("P")
      elif atom == "F": st_types.append("F")
      elif atom == "Cl": st_types.append("C1")
      elif atom == "Br": st_types.append("B1")
      elif atom == "I": st_types.append("I")
      elif atom == "O": #Sterimol distinguishes between "normal", and double-bonded O atoms
         if cn[i] < 1.5: st_types.append("O2")
         if cn[i] > 1.5: st_types.append("O")
      elif atom == "S": #Sterimol distinguishes between "normal", tetrahedral, and octohedral S atoms
         if cn[i] < 2.5: st_types.append("S")
         if 5.5 > cn[i] > 2.5: st_types.append("S4")
         if cn[i] > 5.5: st_types.append("S1")
      elif atom == "N": #Sterimol distinguishes between tetrahedral and planar (amide) N atoms
         if cn[i] > 2.5: st_types.append("N")
         if cn[i] < 2.5: st_types.append("C6/N6")
      elif atom == "C": #Sterimol distinguishes between myriad types of C atoms ...
         if cn[i] < 2.5: st_types.append("C3")
         if 3.5 > cn[i] > 2.5: # need to differentiate between sp2 carbon and aromatic carbon ...
            st_types.append("C6/N6") # assumes aromatic rather than sp2
         if cn[i] > 3.5: st_types.append("C")
   return st_types

# Calculation of atomic coordination numbers (taken from Grimme's DFTD3 definitions)
def ncoord(natom, rcov, atomtype, coords):
   max_elem = 94
   k1 = 16.0
   k2 = 4.0/3.0
   cn =[]
   for i in range(0,natom):
      xn = 0.0
      for iat in range(0,natom):
         if iat != i:
            dx = coords[iat][0] - coords[i][0]
            dy = coords[iat][1] - coords[i][1]
            dz = coords[iat][2] - coords[i][2]
            r2 = dx*dx+dy*dy+dz*dz
            r = math.pow(r2,0.5)
            r = r
            for k in range(0,max_elem):
               if atomtype[i].find(elements[k])>-1:Zi=k
               if atomtype[iat].find(elements[k])>-1:Ziat=k

            rco = rcov[Zi]+rcov[Ziat]
            rco = rco*k2
            rr=rco/r
            damp=1.0/(1.0+math.exp(-k1*(rr-1.0)))
            xn=xn+damp
      cn.append(xn)
   return cn

def linearcheck(carts):
   ans=0;xgrad=[];ygrad=[]
   for row in carts:xgrad.append(round(np.gradient(row)[0],4));ygrad.append(round(np.gradient(row)[1],4))
   if min(xgrad)==max(xgrad) and min(ygrad)==max(ygrad):ans=1
   return ans

class calcSterimol:
   def __init__(self, file, radii, atomA, atomB,verbose):
      if len(file.split(".com"))>1 or len(file.split(".gjf"))>1:
            fileData = getinData(file.split(".")[0])
      if len(file.split(".xyz"))>1:
            fileData = getinData2(file.split(".")[0])
      if len(file.split(".out"))>1 or len(file.split(".log"))>1: fileData = getoutData(file.split(".")[0])

      # initialize the array of atomic vdw radii
      molcart = fileData.CARTESIANS; atomtype = fileData.ATOMTYPES; natoms = len(molcart); vdw_radii = []
#      print fileData.ATOMTYPES

      if radii == "cpk":
         atomic_co_no = ncoord(natoms, rcov, atomtype, molcart)
         sterimol_types = generate_atom_types(atomtype, atomic_co_no)
         print(sterimol_types)
         #print sterimol_types
         for i in range(0,natoms):
            for j in range(0,len(sterimol_atomtypes)):
               if sterimol_types[i] == sterimol_atomtypes[j]: vdw_radii.append(cpk_radii[j]/100.00)

      if radii == "bondi":
         for i in range(0,natoms):
               vdw_radii.append(bondiRadius(periodictable.index(fileData.ATOMTYPES[i])))

# Define vector along the L-axis connecting base atom and the next attached atom
# subtract one since the array starts from zero not one
      atomA = atomA - 1; atomB = atomB - 1
      next_atom = molcart[atomB]
      vect1=np.subtract(getcoords(atomA,molcart),next_atom)
      if verbose == True:
            pass
          #print "   Atoms", atomA, "and", atomB, "define the L-axis and direction", vect1

#          print "\n", "   Atom ".ljust(9), "  Xco/A".rjust(9), "  Yco/A".rjust(9), "  Zco/A".rjust(9), " VdW/pm".rjust(9)
          #print "   ##############################################"
      # Remove the base atom from the list of atoms to be considered for sterics (after printing all)
      atomlist = list(xrange(0,natoms))
      if verbose == True:
          for atom in atomlist:
             pass
#             if radii == "cpk": print "  ", sterimol_types[atom].ljust(6),
#             if radii == "bondi": print "  ", atomtype[atom].ljust(6),
             for coord in molcart[atom]:
                pass
#                if coord < 0.0: print "   %.3f".rjust(6) % coord,
#                else: print "    %.3f".rjust(6) % coord,
#             print "    %.1f" % round(vdw_radii[atom]*100)
      atomlist.remove(atomA)

      adjlist=[]; opplist=[]; theta=[]
      for i in atomlist:
         vect2=np.subtract(getcoords(atomA,molcart),getcoords(i,molcart))
         oppdist=calcopposite(atomA,i,angle(vect1,vect2),molcart)
         opplist.append(oppdist+vdw_radii[i])
         adjdist=calcadj(atomA,i,angle(vect1,vect2),molcart)
         #minadjlist.append(adjdist-vdw_radii[i])
         adjlist.append(adjdist+vdw_radii[i])

      B5=max(opplist)
   #self.lval=max(adjlist)-minval
   # A bit weird, but seems like original sterimol adds on the difference between the bond length and vdw radius of atom B. For a C-H bond this is 1.50 - 1.10 = 0.40 Angstrom)
      self.lval=max(adjlist)+0.40

      ###Useful - do not delete!
      #print "   B5 atom", atomlist[opplist.index(max(opplist))]+1, "distance", max(opplist)
      #print "   Highest atom", atomlist[adjlist.index(max(adjlist))]+1,"distance", max(adjlist),"\n   Lowest atom", atomlist[minadjlist.index(min(minadjlist))]+1,"distance", min(minadjlist)

      zcarts=[]#zeroed carts
      for i in atomlist: zcarts.append(np.subtract(molcart[i],molcart[atomA]))
      zvect=[0,0,1]
      zcent=np.subtract(next_atom,molcart[atomA])
      for cart in range(len(zcarts)):
         zcoord= rotrel(zcent,zvect,zcarts[cart])
         zcarts[cart]=zcoord
      twodcarts=[]
      for row in zcarts: twodcarts.append([row[0],row[1]])
      fragrad=[]#radii of fragment atoms
      for t in atomlist: fragrad.append(vdw_radii[t])
      singledist=[]
      for t in range(len(fragrad)):
         d=np.linalg.norm(twodcarts[t])#;print d
         d=d+fragrad[t]
         singledist.append(d)
      self.newB5=max(singledist) #This is the same as the 3D calculated value from above

      center=[0,0]
      vlist=[]#list of distances from the origin to the tangential vectors
      alist=[]#list of atoms between which the tangential vectors pass through no other atoms
      iav=[]#interatomic vectors
      sym=symcheck(twodcarts)
      for x in range(len(twodcarts)):
         if sym==1:
            twodcarts[x][0]=twodcarts[x][0]+0.000001
            twodcarts[x][1]=twodcarts[x][1]+0.000001
         for y in range(len(twodcarts)):
            if x!=y:
               try:nvect= (twod_vect(center,twodcarts[x],twodcarts[y]))#origin normal vector to connecting atomic centers vector
               except ValueError:nvect=[0,0]
               iav=np.subtract(twodcarts[x],twodcarts[y])#interatomic vector
               iad=np.linalg.norm(iav)#interatomic distance
               try:theta=math.asin((fragrad[y]-fragrad[x])/iad)#calculates angle by which to rotate vdw radii before adding
               except ValueError: theta=np.pi/2
               try:unvect=nvect/np.linalg.norm(nvect)
               except RuntimeWarning:pass#unvect=[0,0]
               xradv=twod_rot(unvect*fragrad[x],theta)
               yradv=twod_rot(unvect*fragrad[y],theta)
               mvect= (twod_vect(center,twodcarts[x]-xradv,twodcarts[y]-yradv))
               nvect= (twod_vect(center,twodcarts[x]+xradv,twodcarts[y]+yradv))#origin normal vector to connecting atomic surfaces tangential vector
               newx=twodcarts[x]+xradv
               newy=twodcarts[y]+yradv
               mewx=twodcarts[x]-xradv
               mewy=twodcarts[y]-yradv
               if np.cross(nvect,xradv)<0.000000001 and theta!=np.pi/2:
                  satpoint=[]#Satisfied points not within range of tangential vector
                  for z in range(len(twodcarts)):
                     pvdist=twod_dist(twodcarts[z],newx,newy)
                     if z!=x and z!=y and pvdist>(fragrad[z]-0.0001):satpoint.append(pvdist)
                  if len(satpoint)==len(atomlist)-2:vlist.append(np.linalg.norm(nvect));alist.append([x,y]);#print x,y
                  satpoint=[]
                  for z in range(len(twodcarts)):
                     pvdist=twod_dist(twodcarts[z],mewx,mewy)
                     if z!=x and z!=y and pvdist>(fragrad[z]-0.0001):satpoint.append(pvdist)
                  if len(satpoint)==len(atomlist)-2:vlist.append(np.linalg.norm(mvect));alist.append([x,y])
      if linearcheck(twodcarts)==1:self.B1 = max(fragrad)
      elif len(vlist) > 0: self.B1=min(vlist)
      else: self.B1 = max(fragrad)

def symcheck(carts):#Add symmetry criteria
   center=[0,0]
   distlist=[]
   distlist.append(10)
   for a in range(len(carts)):
      for b in range(len(carts)):
         if a!=b:
            dist=np.linalg.norm(twod_vect(center,carts[a],carts[b]))
            distlist.append(dist)
   if min(distlist)<0.0000000001:ans=1
   else:ans=0
   return ans

def calcSandwich(file):
   metalatoms=[]
   if file.split(".")[1]=="log" or file.split(".")[1]=="out":fileData=getoutData(file.split(".")[0])
#   if file.split(".")[1]=="xyz":
#    fileData=getinData2(file.split(".")[0])
   if file.split(".")[1]=="com" or file.split(".")[1]=="gjf":fileData=getinData(file.split(".")[0])
   for i in range(len(fileData.ATOMTYPES)):
      if fileData.ATOMTYPES[i] in metals:metalatoms.append(i)

   ivals=[]
   jvals=[]
   for i in range(len(fileData.ATOMTYPES)):
      for j in range(len(fileData.ATOMTYPES)):
         dist = ((fileData.CARTESIANS[i][0]-fileData.CARTESIANS[j][0])**2 +(fileData.CARTESIANS[i][1]-fileData.CARTESIANS[j][1])**2+(fileData.CARTESIANS[i][2]-fileData.CARTESIANS[j][2])**2)**0.5
         if 0.01<dist<1.511 and fileData.ATOMTYPES[j] == "C" and fileData.ATOMTYPES[i] == "C":
            ivals.append(i)
            jvals.append(j)
   conpar=[]
   for a in range(len(ivals)):
      rar=[]
      rar.append(ivals[a])

      for b in range(len(ivals)):
         if ivals[a]==ivals[b]:rar.append(jvals[b])
      if rar not in conpar:conpar.append(rar)

   allrings=[]
   for a in range(len(conpar)):
      z=conpar[a][0]
      for b in concheck(conpar,z):
         y=b
         for c in concheck(conpar,y):
            x=c
            for d in concheck(conpar,x):
               w=d
               for e in concheck(conpar,w):
                  v=e
                  rar=[]
                  rar.extend([z,y,x,w,v])
                  if z in concheck(conpar,v) and sorted(rar) not in allrings and len(set(rar))==5:allrings.append(sorted(rar))
                  for f in concheck(conpar,v):
                     u=f
                     tar=[]
                     tar.extend([z,y,x,w,v,u])
                     if z in concheck(conpar,u) and sorted(tar) not in allrings and len(set(tar))==6:allrings.append(sorted(tar))

   if not allrings:
      for ma in metalatoms:
         for s in range(len(fileData.CARTESIANS)):
            if 0.1<np.linalg.norm(np.subtract(fileData.CARTESIANS[ma],fileData.CARTESIANS[s]))<2.1:allrings.append([s,s,s,s,s])

   mcdists=[]
   mcdist=9999
   for ring in allrings:

      if len(ring)==5:
         tolman=[]
         cent=avpoints(ring,fileData.CARTESIANS)
         m=fileData.CARTESIANS[metalatoms[0]]
         tempmcdist=mcdist
         mcdist=distcalc(m,cent)
         for b in metalatoms:#find closest metal to ring
            m=fileData.CARTESIANS[b]
            if mcdist>=distcalc(m,cent):mcdist=distcalc(m,cent);metal=b
         mcdists.append([mcdist,metal])
         frag=getfragment(ring[0],fileData.CARTESIANS)
         vect1=np.subtract(getcoords(metal,fileData.CARTESIANS),cent)
         if tempmcdist==mcdist:break#Stops if dealing with identical ring system as before (intended for symmetric dimers)
         adjlist=[]
         minadjlist=[]
         opplist=[]
         alpha=[]
         beta=[]
         theta=[]#Candidate Tolman angle substituent
         omega=[]#standardised atom "dihedral" orientation
         ringang=[]
         for i in frag:
            vect2=np.subtract(getcoords(metal,fileData.CARTESIANS),getcoords(i,fileData.CARTESIANS))
            oppdist=calcopposite(metal,i,angle(vect1,vect2),fileData.CARTESIANS)
            opplist.append(oppdist+genradii(i,fileData.CARTESIANS,fileData.ATOMTYPES))
            adjdist=calcadj(metal,i,angle(vect1,vect2),fileData.CARTESIANS)
            minadjlist.append(adjdist-genradii(i,fileData.CARTESIANS,fileData.ATOMTYPES))
            adjlist.append(adjdist+genradii(i,fileData.CARTESIANS,fileData.ATOMTYPES))
            alpha.append(angle(vect1,vect2))
            hyp=distcalc(getcoords(i,fileData.CARTESIANS),getcoords(metal,fileData.CARTESIANS))
            beta.append(math.asin(genradii(i,fileData.CARTESIANS,fileData.ATOMTYPES)/hyp))
            theta.append(alpha[-1]+beta[-1])
            if ring[0]!=ring[1]:omega.append(dihedral([10,10,10],getcoords(metal,fileData.CARTESIANS),cent,getcoords(i,fileData.CARTESIANS)))
            if ring[0]!=ring[1] and i in ring:ringang.append(dihedral([10,10,10],getcoords(metal,fileData.CARTESIANS),cent,getcoords(i,fileData.CARTESIANS)))
         B5=max(opplist)#Bondi
         lval=max(adjlist)-min(minadjlist)
         interval=180/len(ring)
         if ring[0]!=ring[1]:
            for k in ringang:
               tlist=[];tang=[];tfrag=[]

               for h in range(len(frag)):
                  if k-interval<omega[h]<k+interval:tlist.append(frag[h])
                  if k>(180-interval) and k-interval<omega[h]+360<k+interval:tlist.append(frag[h])
                  if k<-(180-interval) and k-interval<omega[h]-360<k+interval:tlist.append(frag[h])
               for t in range(len(frag)):
                  if frag[t] in tlist: tang.append(theta[t]);tfrag.append(frag[t])
               tolman.append(math.degrees(max(tang)))
            x=0
            for c in tolman:
               x=x+c
            tolmanCA=round(2*(x/len(tolman)),3)
         else:tolmanCA=0
         smcdist=round(mcdist,3);lval=round(lval,3);lval=round(lval,3);B5=round(B5,3)
         molcart=fileData.CARTESIANS
         zcarts=[]
         for i in frag:
            zcarts.append(np.subtract(molcart[i],molcart[metal]))
         zvect=[0,0,1]
         zcent=np.subtract(cent,molcart[metal])
         for cart in range(len(zcarts)):
            zcoord= rotrel(zcent,zvect,zcarts[cart])
            zcarts[cart]=zcoord
         twodcarts=[]
         for row in zcarts:
            twodcarts.append([row[0],row[1]])
         fragrad=[]#radii of fragment atoms
         for t in frag:
            fragrad.append(genradii(t,fileData.CARTESIANS,fileData.ATOMTYPES))
         singledist=[]
         for t in range(len(fragrad)):
            d=np.linalg.norm(twodcarts[t])#;print d
            d=d+fragrad[t]
            singledist.append(d)
         newB5=round(max(singledist),3)#This is the same as the 3D calculated value from above

         center=[0,0]
         vlist=[]#list of distances from the origin to the tangential vectors
         alist=[]#list of atoms between which the tangential vectors pass through no other atoms
         iav=[]#interatomic vectors

         for x in range(len(twodcarts)):
            for y in range(len(twodcarts)):
               if x!=y:
                  try:nvect= (twod_vect(center,twodcarts[x],twodcarts[y]))#origin normal vector to connecting atomic centers vector
                  except ValueError:nvect=[0,0]
                  iav=np.subtract(twodcarts[x],twodcarts[y])#interatomic vector
                  iad=np.linalg.norm(iav)#interatomic distance
                  try:theta=math.asin((fragrad[y]-fragrad[x])/iad)#calculates angle by which to rotate vdw radii before adding
                  except ValueError: theta=np.pi/2
                  try:unvect=nvect/np.linalg.norm(nvect)
                  except RuntimeWarning:pass#unvect=[0,0]
                  xradv=twod_rot(unvect*fragrad[x],theta)
                  yradv=twod_rot(unvect*fragrad[y],theta)
                  nvect= (twod_vect(center,twodcarts[x]+xradv,twodcarts[y]+yradv))#origin normal vector to connecting atomic surfaces tangential vector
                  newx=twodcarts[x]+xradv
                  newy=twodcarts[y]+yradv
                  if np.cross(nvect,xradv)<0.000000001 and theta!=np.pi/2:
                     satpoint=[]#Satisfied points not within range of tangential vector
                     for z in range(len(twodcarts)):
                        pvdist=twod_dist(twodcarts[z],newx,newy)
                        if z!=x and z!=y and pvdist>fragrad[z]:satpoint.append(pvdist)
                     if len(satpoint)==len(frag)-2:vlist.append(np.linalg.norm(nvect));alist.append([x,y])#;print x,y
         B1=round(min(vlist),3)
#         print "   "+file.ljust(25),str(tolmanCA).rjust(9), str(smcdist).rjust(9), str(lval).rjust(9),str(B1).rjust(9), str(newB5).rjust(9)

molmod=[['Bq', 0, 0, 0, 0],
        ['H', 1, 1, 1, 1],
        ['He', 0, 0, 0, 0],
        ['Li', 0, 0, 0, 0],
        ['Be', 0, 0, 0, 0],
        ['B', 0, 0, 0, 0],
        ['C', 0, 1.6, 1.6, 1.5],
        ['N', 1.45, 1.45, 1.5, 1.25],
        ['O', 1.35, 1.35, 1.35, 0],
        ['F', 1.35, 1.35, 0, 0],
        ['Ne', 0, 0, 0, 0],
        ['Na', 0, 0, 0, 0],
        ['Mg', 0, 0, 0, 0],
        ['Al', 0, 0, 0, 0],
        ['Si', 2.1, 2.1, 2.1, 2.1],
        ['P', 0, 0, 0, 0],
        ['S', 0, 0, 0, 0],
        ['Cl', 1.8, 0, 0, 0],
        ['Ar', 0, 0, 0, 0],
        ['K', 0, 0, 0, 0],
        ['Ca', 0, 0, 0, 0],
        ['Sc', 0, 0, 0, 0],
        ['Ti', 0, 0, 0, 0],
        ['V', 0, 0, 0, 0],
        ['Cr', 0, 0, 0, 0],
        ['Mn', 0, 0, 0, 0],
        ['Fe', 0, 0, 0, 0],
        ['Co', 0, 0, 0, 0],
        ['Ni', 0, 0, 0, 0],
        ['Cu', 0, 0, 0, 0],
        ['Zn', 0, 0, 0, 0],
        ['Ga', 0, 0, 0, 0],
        ['Ge', 0, 0, 0, 0],
        ['As', 0, 0, 0, 0],
        ['Se', 0, 0, 0, 0],
        ['Br', 1.95, 0, 0, 0],
        ['Kr', 0, 0, 0, 0],
        ['Rb', 0, 0, 0, 0],
        ['Sr', 0, 0, 0, 0],
        ['Y', 0, 0, 0, 0],
        ['Zr', 0, 0, 0, 0],
        ['Nb', 0, 0, 0, 0],
        ['Mo', 0, 0, 0, 0],
        ['Tc', 0, 0, 0, 0],
        ['Ru', 0, 0, 0, 0],
        ['Rh', 0, 0, 0, 0],
        ['Pd', 0, 0, 0, 0],
        ['Ag', 0, 0, 0, 0],
        ['Cd', 0, 0, 0, 0],
        ['In', 0, 0, 0, 0],
        ['Sn', 0, 0, 0, 0],
        ['Sb', 0, 0, 0, 0],
        ['Te', 0, 0, 0, 0],
        ['I', 2.15, 0, 0, 0],
        ['Xe', 0, 0, 0, 0],
        ['Cs', 0, 0, 0, 0],
        ['Ba', 0, 0, 0, 0],
        ['La', 0, 0, 0, 0],
        ['Ce', 0, 0, 0, 0],
        ['Pr', 0, 0, 0, 0],
        ['Nd', 0, 0, 0, 0],
        ['Pm', 0, 0, 0, 0],
        ['Sm', 0, 0, 0, 0],
        ['Eu', 0, 0, 0, 0],
        ['Gd', 0, 0, 0, 0],
        ['Tb', 0, 0, 0, 0],
        ['Dy', 0, 0, 0, 0],
        ['Ho', 0, 0, 0, 0],
        ['Er', 0, 0, 0, 0],
        ['Tm', 0, 0, 0, 0],
        ['Yb', 0, 0, 0, 0],
        ['Lu', 0, 0, 0, 0],
        ['Hf', 0, 0, 0, 0],
        ['Ta', 0, 0, 0, 0],
        ['W', 0, 0, 0, 0],
        ['Re', 0, 0, 0, 0],
        ['Os', 0, 0, 0, 0],
        ['Ir', 0, 0, 0, 0],
        ['Pt', 0, 0, 0, 0],
        ['Au', 0, 0, 0, 0],
        ['Hg', 0, 0, 0, 0],
        ['Tl', 0, 0, 0, 0],
        ['Pb', 0, 0, 0, 0],
        ['Bi', 0, 0, 0, 0],
        ['Po', 0, 0, 0, 0],
        ['At', 0, 0, 0, 0],
        ['Rn', 0, 0, 0, 0],
        ['Fr', 0, 0, 0, 0],
        ['Ra', 0, 0, 0, 0],
        ['Ac', 0, 0, 0, 0],
        ['Th', 0, 0, 0, 0],
        ['Pa', 0, 0, 0, 0],
        ['U', 0, 0, 0, 0],
        ['Np', 0, 0, 0, 0],
        ['Pu', 0, 0, 0, 0],
        ['Am', 0, 0, 0, 0],
        ['Cm', 0, 0, 0, 0],
        ['Bk', 0, 0, 0, 0],
        ['Cf', 0, 0, 0, 0],
        ['Es', 0, 0, 0, 0],
        ['Fm', 0, 0, 0, 0],
        ['Md', 0, 0, 0, 0],
        ['No', 0, 0, 0, 0],
        ['Lr', 0, 0, 0, 0],
        ['Rf', 0, 0, 0, 0],
        ['Db', 0, 0, 0, 0],
        ['Sg', 0, 0, 0, 0],
        ['Bh', 0, 0, 0, 0],
        ['Hs', 0, 0, 0, 0],
        ['Mt', 0, 0, 0, 0],
        ['Ds', 0, 0, 0, 0],
        ['Rg', 0, 0, 0, 0],
        ['Uub', 0, 0, 0, 0],
        ['Uut', 0, 0, 0, 0],
        ['Uuq', 0, 0, 0, 0],
        ['Uup', 0, 0, 0, 0],
        ['Uuh', 0, 0, 0, 0],
        ['Uus', 0, 0, 0, 0],
        ['Uuo', 0, 0, 0, 0],]

elements = ["H","He","Li","Be","B","C","N","O","F","Ne","Na","Mg","Al","Si",
			"P","S","Cl","Ar","K","Ca","Sc","Ti","V","Cr","Mn","Fe","Co","Ni",
			"Cu","Zn","Ga","Ge","As","Se","Br","Kr","Rb","Sr","Y","Zr","Nb","Mo",
			"Tc","Ru","Rh","Pd","Ag","Cd","In","Sn","Sb","Te","I","Xe","Cs","Ba",
			"La","Ce","Pr","Nd","Pm","Sm","Eu","Gd","Tb","Dy","Ho","Er","Tm","Yb",
			"Lu","Hf","Ta","W","Re","Os","Ir","Pt","Au","Hg","Tl","Pb","Bi","Po","At",
			"Rn","Fr","Ra","Ac","Th","Pa","U","Np","Pu","Am","Cm","Bk","Cf","Es","Fm",
			"Md","No","Lr","Rf","Db","Sg","Bh","Hs","Mt","Ds","Rg","Uub","Uut","Uuq",
			"Uup","Uuh","Uus","Uuo"]

## covalent radii (taken from Pyykko and Atsumi, Chem. Eur. J. 15, 2009, 188-197 ##
## values for metals decreased by 10 % ##
rcov = [0.32, 0.46, 1.20, 0.94, 0.77, 0.75, 0.71, 0.63, 0.64, 0.67,
		1.40, 1.25, 1.13, 1.04, 1.10, 1.02, 0.99, 0.96, 1.76, 1.54,
		1.33, 1.22, 1.21, 1.10, 1.07, 1.04, 1.00, 0.99, 1.01, 1.09,
		1.12, 1.09, 1.15, 1.10, 1.14, 1.17, 1.89, 1.67, 1.47, 1.39,
		1.32, 1.24, 1.15, 1.13, 1.13, 1.08, 1.15, 1.23, 1.28, 1.26,
		1.26, 1.23, 1.32, 1.31, 2.09, 1.76, 1.62, 1.47, 1.58, 1.57,
		1.56, 1.55, 1.51, 1.52, 1.51, 1.50, 1.49, 1.49, 1.48, 1.53,
		1.46, 1.37, 1.31, 1.23, 1.18, 1.16, 1.11, 1.12, 1.13, 1.32,
		1.30, 1.30, 1.36, 1.31, 1.38, 1.42, 2.01, 1.81, 1.67, 1.58,
		1.52, 1.53, 1.54, 1.55]