"""
Item-based k-NN collaborative filtering.
"""

from sys import intern
import logging
import warnings

import pandas as pd
import numpy as np
import scipy.sparse as sps
import scipy.sparse.linalg as spla
from numba import jit, njit, prange
from numba.typed import List

from lenskit import util, matrix, DataWarning
from lenskit.sharing import in_share_context
from lenskit.util.parallel import is_mp_worker
from lenskit.util.accum import kvp_minheap_insert, kvp_minheap_sort
from . import Predictor

_logger = logging.getLogger(__name__)
_mkl_ops = matrix.mkl_ops()
_empty_csr = matrix._impl_mod()._empty_csr
_subset_rows = matrix._impl_mod()._subset_rows

if _mkl_ops is not None:
    # we have to import LK CFFI utils into this module
    for lkv in dir(_mkl_ops):
        if lkv.startswith('_lk_mkl'):
            globals()[lkv] = getattr(_mkl_ops, lkv)

def _make_blocks(n, size):
    "Create blocks for the range 0..n."
    return [(s, min(s + size, n)) for s in range(0, n, size)]

@njit(nogil=True)
def _count_nbrs(mat: matrix._CSR, thresh: float):
    "Count the number of neighbors passing the threshold for each row."
    counts = np.zeros(mat.nrows, dtype=np.int32)
    cs = mat.colinds
    vs = mat.values
    for i in range(mat.nrows):
        sp, ep = mat.row_extent(i)
        for j in range(sp, ep):
            c = cs[j]
            v = vs[j]
            if c != i and v >= thresh:
                counts[i] += 1

    return counts

@njit
def _insert(dst, used, limits, i, c, v):
    "Insert one item into a heap"
    sp = dst.rowptrs[i]
    ep = sp + used[i]
    ep = kvp_minheap_insert(sp, ep, limits[i], c, v, dst.colinds, dst.values)
    used[i] = ep - sp

@njit(nogil=True)
def _copy_nbrs(src: matrix._CSR, dst: matrix._CSR, limits, thresh: float):
    "Copy neighbors into the output matrix."
    used = np.zeros(dst.nrows, dtype=np.int32)

    for i in range(src.nrows):
        sp, ep = src.row_extent(i)

        for j in range(sp, ep):
            c = src.colinds[j]
            v = src.values[j]
            if c != i and v >= thresh:
                _insert(dst, used, limits, i, c, v)

    return used

@njit(nogil=True, parallel=not is_mp_worker())
def _sort_nbrs(smat):
    for i in prange(smat.nrows):
        sp, ep = smat.row_extent(i)
        kvp_minheap_sort(sp, ep, smat.colinds, smat.values)

@njit(nogil=True)
def _make_sim_block(nitems, bsp, bitems, r_sp, r_ep, r_cs, r_vs, min_sim, max_nbrs):
    # pass 1: compute the size of each row
    sizes = np.zeros(bitems, np.int32)
    for i in range(nitems):
        for j in range(r_sp[i], r_ep[i]):
            # we accept the neighbor if it passes threshold and isn't a self-similarity
            r = r_cs[j]
            if i != bsp + r and r_vs[j] >= min_sim:
                sizes[r] += 1

    if max_nbrs > 0:
        for i in range(bitems):
            if sizes[i] > max_nbrs:
                sizes[i] = max_nbrs

    # if bnc == 0:
    #     # empty resulting matrix, oops
    #     return _empty_csr(bitems, nitems, np.zeros(bitems, np.int32))

    # allocate a matrix
    block_csr = _empty_csr(bitems, nitems, sizes)

    # pass 2: truncate each row into the matrix
    eps = block_csr.rowptrs[:-1].copy()
    for c in range(nitems):
        for j in range(r_sp[c], r_ep[c]):
            v = r_vs[j]
            r = r_cs[j]
            sp, lep = block_csr.row_extent(r)
            lim = lep - sp
            if c != bsp + r and v >= min_sim:
                eps[r] = kvp_minheap_insert(sp, eps[r], lim, c, v,
                                            block_csr.colinds, block_csr.values)
        # we're done!
    return block_csr

@njit(nogil=True)
def _mkl_sim_block(block, bsp, bep, rmh, min_sim, max_nbrs, nitems):
    "Compute a single block of the similarity matrix"
    # assert block.nrows == bep - bsp

    bitems = block.nrows
    if block.nnz == 0:
        return _empty_csr(bitems, nitems, np.zeros(bitems, np.int32))

    # create a matrix handle for the subset matrix
    amh = _mkl_ops._from_csr(block)
    _lk_mkl_spopt(amh)

    smh = _lk_mkl_spmabt(rmh, amh)

    _lk_mkl_spfree(amh)

    _lk_mkl_sporder(smh)  # for reproducibility

    block = _lk_mkl_spexport_p(smh)
    bnr = _lk_mkl_spe_nrows(block)
    bnc = _lk_mkl_spe_ncols(block)
    # bnr and bnc should be right
    # assert bnc == bep - bsp
    # assert bnr == nitems

    r_sp = _lk_mkl_spe_row_sp(block)
    r_ep = _lk_mkl_spe_row_ep(block)
    r_cs = _lk_mkl_spe_colinds(block)
    r_vs = _lk_mkl_spe_values(block)

    block_csr = _make_sim_block(nitems, bsp, bitems, r_sp, r_ep, r_cs, r_vs, min_sim, max_nbrs)

    _lk_mkl_spe_free(block)
    _lk_mkl_spfree(smh)
    return block_csr

@njit(nogil=True, parallel=not is_mp_worker())
def _mkl_sim_blocks(trmat, blocks, ptrs, min_sim, max_nbrs):
    "Compute the similarity matrix with blocked MKL calls"
    nitems = trmat.nrows
    nblocks = len(blocks)

    null = _empty_csr(1, 1, np.zeros(1, dtype=np.int32))
    res = [null for i in range(nblocks)]

    rmat_h = _mkl_ops._from_csr(trmat)
    _lk_mkl_sporder(rmat_h)
    _lk_mkl_spopt(rmat_h)

    for bi in prange(nblocks):
        b = blocks[bi]
        p = ptrs[bi]
        bs, be = p
        bres = _mkl_sim_block(b, bs, be, rmat_h, min_sim, max_nbrs, nitems)
        res[bi] = bres

    _lk_mkl_spfree(rmat_h)

    return res

def _scipy_sim_block(block, bsp, bep, rmat, min_sim, max_nbrs, nitems):
    "Compute a single block of the similarity matrix"
    assert block.nrows == bep - bsp

    _logger.debug('processing block %d:%d (%d nnz)', bsp, bep, block.nnz)

    if rmat.nnz == 0:
        return _empty_csr(block.nrows, nitems, np.zeros(block.nrows, np.int32))

    sims = rmat @ block.to_scipy().transpose()
    sims = matrix.CSR.from_scipy(sims)

    r_sp = sims.rowptrs[:-1]
    r_ep = sims.rowptrs[1:]
    r_cs = sims.colinds
    r_vs = sims.values

    block_csr = _make_sim_block(nitems, bsp, block.nrows, r_sp, r_ep, r_cs, r_vs, min_sim, max_nbrs)
    _logger.debug('umm %d %d', block_csr.nrows, block.nrows)
    assert block_csr.nrows == block.nrows
    assert block_csr.ncols == nitems
    _logger.debug('block %d:%d has %d similarities', bsp, bep, block_csr.nnz)
    _logger.debug('block: %s', matrix.CSR(N=block_csr))

    return block_csr

# we compile this in object mode so we can use numba's thread pool with scipy
@jit(parallel=not is_mp_worker(), forceobj=True)
def _scipy_sim_blocks(trmat, blocks, ptrs, min_sim, max_nbrs):
    "Compute the similarity matrix with blocked SciPy calls"
    nitems, nusers = trmat.shape
    nblocks = len(blocks)

    null = _empty_csr(1, 1, np.zeros(1, dtype=np.int32))
    res = [null for i in range(nblocks)]

    for bi in prange(nblocks):
        b = blocks[bi]
        bs, be = ptrs[bi]
        bres = _scipy_sim_block(b, bs, be, trmat, min_sim, max_nbrs, nitems)
        res[bi] = bres
        assert bres.nrows == be - bs

    return res

@njit(nogil=True)
def _predict_weighted_average(model, nitems, nrange, ratings, rated, targets):
    "Weighted average prediction function"
    min_nbrs, max_nbrs = nrange
    scores = np.full(nitems, np.nan, dtype=np.float_)

    for i in prange(targets.shape[0]):
        iidx = targets[i]
        rptr = model.rowptrs[iidx]
        rend = model.rowptrs[iidx + 1]

        num = 0
        denom = 0
        nnbrs = 0

        for j in range(rptr, rend):
            nidx = model.colinds[j]
            if not rated[nidx]:
                continue

            nnbrs = nnbrs + 1
            num = num + ratings[nidx] * model.values[j]
            denom = denom + np.abs(model.values[j])

            if max_nbrs > 0 and nnbrs >= max_nbrs:
                break

        if nnbrs < min_nbrs:
            continue

        scores[iidx] = num / denom

    return scores

@njit(nogil=True)
def _predict_sum(model, nitems, nrange, ratings, rated, targets):
    "Sum-of-similarities prediction function"
    min_nbrs, max_nbrs = nrange
    scores = np.full(nitems, np.nan, dtype=np.float_)

    for i in prange(targets.shape[0]):
        iidx = targets[i]
        rptr = model.rowptrs[iidx]
        rend = model.rowptrs[iidx + 1]

        score = 0
        nnbrs = 0

        for j in range(rptr, rend):
            nidx = model.colinds[j]
            if not rated[nidx]:
                continue

            nnbrs = nnbrs + 1
            score = score + model.values[j]

            if max_nbrs > 0 and nnbrs >= max_nbrs:
                break

        if nnbrs < min_nbrs:
            continue

        scores[iidx] = score

    return scores

_predictors = {
    'weighted-average': _predict_weighted_average,
    'sum': _predict_sum
}

class ItemItem(Predictor):
    """
    Item-item nearest-neighbor collaborative filtering with ratings. This item-item implementation
    is not terribly configurable; it hard-codes design decisions found to work well in the previous
    Java-based LensKit code.

    Args:
        nnbrs(int):
            the maximum number of neighbors for scoring each item (``None`` for unlimited)
        min_nbrs(int): the minimum number of neighbors for scoring each item
        min_sim(double): minimum similarity threshold for considering a neighbor
        save_nbrs(double):
            the number of neighbors to save per item in the trained model
            (``None`` for unlimited)
        center(bool):
            whether to normalize (mean-center) rating vectors.  Turn this off when working
            with unary data and other data types that don't respond well to centering.
        aggregate:
            the type of aggregation to do. Can be ``weighted-average`` or ``sum``.

    Attributes:
        item_index_(pandas.Index): the index of item IDs.
        item_means_(numpy.ndarray): the mean rating for each known item.
        item_counts_(numpy.ndarray): the number of saved neighbors for each item.
        sim_matrix_(matrix.CSR): the similarity matrix.
        user_index_(pandas.Index): the index of known user IDs for the rating matrix.
        rating_matrix_(matrix.CSR): the user-item rating matrix for looking up users' ratings.
    """
    AGG_SUM = intern('sum')
    AGG_WA = intern('weighted-average')
    _use_mkl = True

    def __init__(self, nnbrs, min_nbrs=1, min_sim=1.0e-6, save_nbrs=None,
                 center=True, aggregate='weighted-average'):
        self.nnbrs = nnbrs
        if self.nnbrs is not None and self.nnbrs < 1:
            self.nnbrs = -1
        self.min_nbrs = min_nbrs
        if self.min_nbrs is not None and self.min_nbrs < 1:
            self.min_nbrs = 1
        self.min_sim = min_sim
        self.save_nbrs = save_nbrs
        self.center = center
        self.aggregate = aggregate

    def fit(self, ratings, **kwargs):
        """
        Train a model.

        The model-training process depends on ``save_nbrs`` and ``min_sim``, but *not* on other
        algorithm parameters.

        Args:
            ratings(pandas.DataFrame):
                (user,item,rating) data for computing item similarities.
        """
        # Training proceeds in 2 steps:
        # 1. Normalize item vectors to be mean-centered and unit-normalized
        # 2. Compute similarities with pairwise dot products
        self._timer = util.Stopwatch()

        _logger.debug('[%s] beginning fit, memory use %s', self._timer, util.max_memory())

        init_rmat, users, items = matrix.sparse_ratings(ratings)
        n_items = len(items)
        _logger.info('[%s] made sparse matrix for %d items (%d ratings from %d users)',
                     self._timer, len(items), init_rmat.nnz, len(users))
        _logger.debug('[%s] made matrix, memory use %s', self._timer, util.max_memory())

        rmat, item_means = self._mean_center(ratings, init_rmat, items)
        _logger.debug('[%s] centered, memory use %s', self._timer, util.max_memory())

        rmat = self._normalize(rmat)
        _logger.debug('[%s] normalized, memory use %s', self._timer, util.max_memory())

        _logger.info('[%s] computing similarity matrix', self._timer)
        smat = self._compute_similarities(rmat)
        _logger.debug('[%s] computed, memory use %s', self._timer, util.max_memory())

        _logger.info('[%s] got neighborhoods for %d of %d items',
                     self._timer, np.sum(np.diff(smat.rowptrs) > 0), n_items)

        _logger.info('[%s] computed %d neighbor pairs', self._timer, smat.nnz)

        self.item_index_ = items
        self.item_means_ = item_means
        self.item_counts_ = np.diff(smat.rowptrs)
        self.sim_matrix_ = smat
        self.user_index_ = users
        self.rating_matrix_ = init_rmat
        # create an inverted similarity matrix for efficient scanning
        self._sim_inv_ = smat.transpose()
        _logger.info('[%s] transposed matrix for optimization', self._timer)
        _logger.debug('[%s] done, memory use %s', self._timer, util.max_memory())

        return self

    def _mean_center(self, ratings, rmat, items):
        if not self.center:
            return rmat, None

        item_means = ratings.groupby('item').rating.mean()
        item_means = item_means.reindex(items).values
        mcvals = rmat.values - item_means[rmat.colinds]
        nmat = matrix.CSR(rmat.nrows, rmat.ncols, rmat.nnz,
                          rmat.rowptrs.copy(), rmat.colinds.copy(), mcvals)
        _logger.info('[%s] computed means for %d items', self._timer, len(item_means))
        return nmat, item_means

    def _normalize(self, rmat):
        rmat = rmat.to_scipy()
        # compute column norms
        norms = spla.norm(rmat, 2, axis=0)
        # and multiply by a diagonal to normalize columns
        recip_norms = norms.copy()
        is_nz = recip_norms > 0
        recip_norms[is_nz] = np.reciprocal(recip_norms[is_nz])
        norm_mat = rmat @ sps.diags(recip_norms)
        assert norm_mat.shape[1] == rmat.shape[1]
        # and reset NaN
        norm_mat.data[np.isnan(norm_mat.data)] = 0
        _logger.info('[%s] normalized rating matrix columns', self._timer)
        return matrix.CSR.from_scipy(norm_mat, False)

    def _compute_similarities(self, rmat):
        trmat = rmat.transpose()
        nitems = trmat.nrows
        m_nbrs = self.save_nbrs
        if m_nbrs is None or m_nbrs < 0:
            m_nbrs = 0

        bounds = _make_blocks(nitems, 1000)
        _logger.info('[%s] splitting %d items (%d ratings) into %d blocks',
                     self._timer, nitems, trmat.nnz, len(bounds))
        blocks = [trmat.subset_rows(sp, ep) for (sp, ep) in bounds]

        if self._use_mkl and _mkl_ops is not None:
            _logger.info('[%s] computing similarities with MKL', self._timer)
            ptrs = List(bounds)
            nbs = List(b.N for b in blocks)
            if not nbs:
                # oops, this is the bad place
                # in non-JIT node, List doesn't actually make the list
                nbs = [b.N for b in blocks]
                ptrs = bounds
            s_blocks = _mkl_sim_blocks(trmat.N, nbs, ptrs, self.min_sim, m_nbrs)
        else:
            s_blocks = _scipy_sim_blocks(trmat.to_scipy(), blocks, bounds, self.min_sim, m_nbrs)

        s_blocks = [matrix.CSR(N=b) for b in s_blocks]
        nnz = sum(b.nnz for b in s_blocks)
        tot_rows = sum(b.nrows for b in s_blocks)
        _logger.info('[%s] computed %d similarities for %d items in %d blocks',
                     self._timer, nnz, tot_rows, len(s_blocks))
        row_nnzs = np.concatenate([b.row_nnzs() for b in s_blocks])
        assert len(row_nnzs) == nitems, \
            'only have {} rows for {} items'.format(len(row_nnzs), nitems)

        smat = matrix.CSR.empty((nitems, nitems), row_nnzs, rpdtype=np.int64)
        start = 0
        for bi, b in enumerate(s_blocks):
            bnr = b.nrows
            end = start + bnr
            v_sp = smat.rowptrs[start]
            v_ep = smat.rowptrs[end]
            _logger.debug('block %d (%d:%d) has %d entries, storing in %d:%d',
                          bi, start, end, b.nnz, v_sp, v_ep)
            smat.colinds[v_sp:v_ep] = b.colinds
            smat.values[v_sp:v_ep] = b.values
            start = end

        _logger.info('[%s] sorting similarity matrix with %d entries', self._timer, smat.nnz)
        _sort_nbrs(smat.N)

        return smat

    def predict_for_user(self, user, items, ratings=None):
        _logger.debug('predicting %d items for user %s', len(items), user)
        if ratings is None:
            if user not in self.user_index_:
                _logger.debug('user %s missing, returning empty predictions', user)
                return pd.Series(np.nan, index=items)
            upos = self.user_index_.get_loc(user)
            ratings = pd.Series(self.rating_matrix_.row_vs(upos),
                                index=pd.Index(self.item_index_[self.rating_matrix_.row_cs(upos)]))

        if not ratings.index.is_unique:
            wmsg = 'user {} has duplicate ratings, this is likely to cause problems'.format(user)
            warnings.warn(wmsg, DataWarning)

        # set up rating array
        # get rated item positions & limit to in-model items
        n_items = len(self.item_index_)
        ri_pos = self.item_index_.get_indexer(ratings.index)
        m_rates = ratings[ri_pos >= 0]
        ri_pos = ri_pos[ri_pos >= 0]
        rate_v = np.full(n_items, np.nan, dtype=np.float_)
        rated = np.zeros(n_items, dtype='bool')
        # mean-center the rating array
        if self.center:
            rate_v[ri_pos] = m_rates.values - self.item_means_[ri_pos]
        else:
            rate_v[ri_pos] = m_rates.values
        rated[ri_pos] = True

        _logger.debug('user %s: %d of %d rated items in model', user, len(ri_pos), len(ratings))
        assert np.sum(np.logical_not(np.isnan(rate_v))) == len(ri_pos)
        assert np.all(np.isnan(rate_v) == np.logical_not(rated))

        # set up item result vector
        # ipos will be an array of item indices
        i_pos = self.item_index_.get_indexer(items)
        i_pos = i_pos[i_pos >= 0]
        _logger.debug('user %s: %d of %d requested items in model', user, len(i_pos), len(items))

        # now we take a first pass through the data to count _viable_ targets
        # This computes the number of neighbors (and their weight sum) for
        # each target item based on the user's ratings, allowing us to fast-path
        # other computations and avoid as many neighbor truncations as possible
        i_cts, i_sums, i_nbrs = self._count_viable_targets(i_pos, ri_pos)
        viable = i_cts >= self.min_nbrs
        i_pos = i_pos[viable]
        i_cts = i_cts[viable]
        i_sums = i_sums[viable]
        i_nbrs = i_nbrs[viable]
        _logger.debug('user %s: %d of %d requested items possibly reachable',
                      user, len(i_pos), len(items))

        # look for some fast paths
        if self.aggregate == self.AGG_SUM and self.min_sim >= 0:
            # similarity sums are all we need
            if self.nnbrs >= 0:
                fast_mask = i_cts <= self.nnbrs
                fast_items = i_pos[fast_mask]
                fast_scores = i_sums[fast_mask]
                slow_items = i_pos[~fast_mask]
            else:
                fast_items = i_pos
                fast_scores = i_sums
                slow_items = np.array([], dtype='i4')

            _logger.debug('user %s: using fast-path similarity sum for %d items',
                          user, len(fast_items))

            if len(slow_items):
                iscores = _predict_sum(self.sim_matrix_.N, len(self.item_index_),
                                       (self.min_nbrs, self.nnbrs),
                                       rate_v, rated, slow_items)
            else:
                iscores = np.full(len(self.item_index_), np.nan)
            iscores[fast_items] = fast_scores

        elif self.aggregate == self.AGG_WA and self.min_nbrs == 1:
            # fast-path single-neighbor targets - common in sparse data
            fast_mask = i_cts == 1
            fast_items = i_pos[fast_mask]
            fast_scores = rate_v[i_nbrs[fast_mask]]
            if self.min_sim < 0:
                fast_scores *= np.sign(i_sums[fast_mask])
            _logger.debug('user %s: fast-pathed %d scores', user, len(fast_scores))

            slow_items = i_pos[i_cts > 1]
            iscores = _predict_weighted_average(self.sim_matrix_.N, len(self.item_index_),
                                                (self.min_nbrs, self.nnbrs),
                                                rate_v, rated, slow_items)
            iscores[fast_items] = fast_scores
        else:
            # now compute the predictions
            _logger.debug('user %s: taking the slow path', user)
            agg = _predictors[self.aggregate]
            iscores = agg(self.sim_matrix_.N, len(self.item_index_), (self.min_nbrs, self.nnbrs),
                          rate_v, rated, i_pos)

        if self.center:
            iscores += self.item_means_

        results = pd.Series(iscores, index=self.item_index_)
        results = results.reindex(items, fill_value=np.nan)

        _logger.debug('user %s: predicted for %d of %d items',
                      user, results.notna().sum(), len(items))

        return results

    def _count_viable_targets(self, targets, rated):
        "Count upper-bound on possible neighbors for target items and rated items."
        # initialize counts to zero
        counts = np.zeros(len(self.item_index_), dtype=np.int32)
        sums = np.zeros(len(self.item_index_))
        last_nbrs = np.full(len(self.item_index_), -1, 'i4')
        # count the number of times each item is reachable from the neighborhood
        for ri in rated:
            nbrs = self._sim_inv_.row_cs(ri)
            counts[nbrs] += 1
            sums[nbrs] += self._sim_inv_.row_vs(ri)
            last_nbrs[nbrs] = ri

        # we want the reachability counts for the target items
        return counts[targets], sums[targets], last_nbrs[targets]

    def __getstate__(self):
        state = dict(self.__dict__)
        if '_sim_inv_' in state and not in_share_context():
            del state['_sim_inv_']
        return state

    def __setstate__(self, state):
        self.__dict__.update(state)
        if hasattr(self, 'sim_matrix_') and not hasattr(self, '_sim_inv_'):
            self._sim_inv_ = self.sim_matrix_.transpose()

    def __str__(self):
        return 'ItemItem(nnbrs={}, msize={})'.format(self.nnbrs, self.save_nbrs)


#!/usr/bin/env python
#    Under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see http://www.gnu.org/licenses/.
#
# Renamed to svt to keep separate from the wav2png.py script
#
# Based on wav2png.py by Bram de Jong <bram.dejong at domain.com where domain in gmail>
# From http://freesound.iua.upf.edu/blog/?p=10
#
"""
svt.py [options] filename.wav
The options the script takes are:

-a = output waveform image (default input filename + _w.png)

-s = output spectrogram image (default input filename + _s.png)

-w = image width in pixels (default 500)

-h = image height in pixels (default 170)

-f = fft size, power of 2 (default 2048)

-m = maximum frequency to draw, in Hz (default 22050)

-o = 1 if you want to draw the waveform too

"""

import optparse, math, sys
import scikits.audiolab as audiolab
import ImageFilter, ImageChops, Image, ImageDraw, ImageColor
import numpy

#For music visualization, this class is more or less irrelevant.
class TestAudioFile(object):
    """A class that mimics audiolab.sndfile but generates noise instead of reading
    a wave file. Additionally it can be told to have a "broken" header and thus crashing
    in the middle of the file. Also useful for testing ultra-short files of 20 samples."""

    def __init__(self, num_frames, has_broken_header=False):
        self.seekpoint = 0
        self.num_frames = num_frames
        self.has_broken_header = has_broken_header

    def seek(self, seekpoint):
        self.seekpoint = seekpoint

    def get_nframes(self):
        return self.num_frames

    def get_samplerate(self):
        return 44100

    def get_channels(self):
        return 1

    def read_frames(self, frames_to_read):
        if self.has_broken_header and self.seekpoint + frames_to_read > self.num_frames / 2:
            raise IOError()

        num_frames_left = self.num_frames - self.seekpoint
        will_read = num_frames_left if num_frames_left < frames_to_read else frames_to_read
        self.seekpoint += will_read
        return numpy.random.random(will_read)*2 - 1

class AudioProcessor(object):
    def __init__(self, audio_file, fft_size, channel, window_function=numpy.ones):
        self.fft_size = fft_size
        self.window = window_function(self.fft_size)
        self.audio_file = audio_file
        self.frames = audio_file.get_nframes()
        self.samplerate = audio_file.get_samplerate()
        self.channels = audio_file.get_channels()
        self.spectrum_range = None
        self.lower = 100
        self.higher = 22050
        self.lower_log = math.log10(self.lower)
        self.higher_log = math.log10(self.higher)
        self.clip = lambda val, low, high: min(high, max(low, val))
        self.channel = channel

    #returns a number of samples from the audio file.
    def read(self, start, size, resize_if_less=False):
        """ read size samples starting at start, if resize_if_less is True and less than size
        samples are read, resize the array to size and fill with zeros """

        # number of zeros to add to start and end of the buffer
        add_to_start = 0
        add_to_end = 0

        if start < 0:
            # the first FFT window starts centered around zero
            if size + start <= 0:
                return numpy.zeros(size) if resize_if_less else numpy.array([])
            else:
                self.audio_file.seek(0)

                add_to_start = -start # remember: start is negative!
                to_read = size + start

                if to_read > self.frames:
                    add_to_end = to_read - self.frames
                    to_read = self.frames
        else:
            self.audio_file.seek(start)

            to_read = size
            if start + to_read >= self.frames:
                to_read = self.frames - start
                add_to_end = size - to_read

        try:
            samples = self.audio_file.read_frames(to_read)
        except IOError:
            # this can happen for wave files with broken headers...
            return numpy.zeros(size) if resize_if_less else numpy.zeros(2)

        # select which channel to draw
        if self.channels > 1:
            if self.channel==1:
                samples = samples[:,0]
            if self.channel==2:
                samples = samples[:,1]

        if resize_if_less and (add_to_start > 0 or add_to_end > 0):
            if add_to_start > 0:
                samples = numpy.concatenate((numpy.zeros(add_to_start), samples), axis=1)

            if add_to_end > 0:
                samples = numpy.resize(samples, size)
                samples[size - add_to_end:] = 0

        return samples

    """
    The spectral centroid is a measure used in digital signal processing to characterise
    a spectrum. It indicates where the "center of mass" of the spectrum is. Perceptually,
    it has a robust connection with the impression of "brightness" of a sound. It is calculated
    as the weighted mean of the frequencies present in the signal, determined using a Fourier
    transform, with their magnitudes as the wiehgts.
    The spectral centroid is widely used in digital audio and music processing as an automatic
    measure of musical timbre. -Wikipedia

    Probably extremely useful for visualization.
    """
    def spectral_centroid(self, seek_point, spec_range=120.0):
        """ starting at seek_point read fft_size samples, and calculate the spectral centroid """

        samples = self.read(seek_point - self.fft_size/2, self.fft_size, True)

        samples *= self.window
        fft = numpy.fft.fft(samples)
        spectrum = numpy.abs(fft[:fft.shape[0] / 2 + 1]) / float(self.fft_size) # normalized abs(FFT) between 0 and 1
        length = numpy.float64(spectrum.shape[0])

        # scale the db spectrum from [- spec_range db ... 0 db] > [0..1]
        db_spectrum = ((20*(numpy.log10(spectrum + 1e-30))).clip(-spec_range, 0.0) + spec_range)/spec_range

        energy = spectrum.sum()
        spectral_centroid = 0

        if energy > 1e-20:
            # calculate the spectral centroid

            if self.spectrum_range == None:
                self.spectrum_range = numpy.arange(length)

            spectral_centroid = (spectrum * self.spectrum_range).sum() / (energy * (length - 1)) * self.samplerate * 0.5

            # clip > log10 > scale between 0 and 1
            spectral_centroid = (math.log10(self.clip(spectral_centroid, self.lower, self.higher)) - self.lower_log) / (self.higher_log - self.lower_log)

        return (spectral_centroid, db_spectrum)

    #Goes through the samples and finds the min and max amplitudes; used to draw the waveform.
    #This function is probably what I need to use to output amplitudes to a visualizer..
    def peaks(self, start_seek, end_seek):
        """ read all samples between start_seek and end_seek, then find the minimum and maximum peak
        in that range. Returns that pair in the order they were found. So if min was found first,
        it returns (min, max) else the other way around. """

        # larger blocksizes are faster but take more mem...
        # Aha, Watson, a clue, a tradeof!
        block_size = 4096

        max_index = -1
        max_value = -1
        min_index = -1
        min_value = 1

        if end_seek > self.frames:
            end_seek = self.frames

        if block_size > end_seek - start_seek:
            block_size = end_seek - start_seek

        if block_size <= 1:
            samples = self.read(start_seek, 1)
            return samples[0], samples[0]
        elif block_size == 2:
            samples = self.read(start_seek, True)
            return samples[0], samples[1]

        for i in range(start_seek, end_seek, block_size):
            samples = self.read(i, block_size)

            local_max_index = numpy.argmax(samples)
            local_max_value = samples[local_max_index]

            if local_max_value > max_value:
                max_value = local_max_value
                max_index = local_max_index

            local_min_index = numpy.argmin(samples)
            local_min_value = samples[local_min_index]

            if local_min_value < min_value:
                min_value = local_min_value
                min_index = local_min_index

        return (min_value, max_value) if min_index < max_index else (max_value, min_value)

#not relevant to visualization
def interpolate_colors(colors, flat=False, num_colors=256):
    """ given a list of colors, create a larger list of colors interpolating
    the first one. If flatten is True a list of numers will be returned. If
    False, a list of (r,g,b) tuples. num_colors is the number of colors wanted
    in the final list """

    palette = []

    for i in range(num_colors):
        index = (i * (len(colors) - 1))/(num_colors - 1.0)
        index_int = int(index)
        alpha = index - float(index_int)

        if alpha > 0:
            r = (1.0 - alpha) * colors[index_int][0] + alpha * colors[index_int + 1][0]
            g = (1.0 - alpha) * colors[index_int][1] + alpha * colors[index_int + 1][1]
            b = (1.0 - alpha) * colors[index_int][2] + alpha * colors[index_int + 1][2]
        else:
            r = (1.0 - alpha) * colors[index_int][0]
            g = (1.0 - alpha) * colors[index_int][1]
            b = (1.0 - alpha) * colors[index_int][2]

        if flat:
            palette.extend((int(r), int(g), int(b)))
        else:
            palette.append((int(r), int(g), int(b)))

    return palette

class WaveformImage(object):
    def __init__(self, image_width, image_height, palette):
        self.image = Image.new("RGB", (image_width, image_height))

        self.image_width = image_width
        self.image_height = image_height

        self.draw = ImageDraw.Draw(self.image)
        self.previous_x, self.previous_y = None, None

        if palette==2:
	        colors = [
	                    (255,255,255),
	                    (255,255,255),
	                    (255,255,255),
	                    (225,248,255),
	                    (210,241,255),
	                    (195,232,255),
	                    (180,221,255),
	                    (165,208,255),
	                    (150,193,255),
	                    (135,175,255),
	                    (120,156,255),
	                    (105,134,255),
	                    (90,110,255),
	                    (75,85,255),
	                    (63,60,255),
	                    (64,45,255),
	                    (66,30,255),
	                    (76,0,255),
	                    (0,128,13),
	                    (8,138,0),
	                    (20,143,0),
	                    (33,148,0),
	                    (46,153,0),
	                    (60,158,0),
	                    (91,168,0),
	                    (108,173,0),
	                    (125,179,0),
	                    (143,184,0),
	                    (162,189,0),
	                    (182,194,0),
	                    (199,195,0),
	                    (204,184,0),
	                    (255,230,128),
	                    (255,221,119),
	                    (255,213,111),
	                    (255,203,102),
	                    (255,193,94),
	                    (255,181,85),
	                    (255,169,77),
	                    (255,157,68),
	                    (255,143,60),
	                    (255,129,51),
	                    (255,113,43),
	                    (255,97,34),
	                    (255,81,25),
	                    (255,63,17),
	                    (255,45,8),
	                    (255,26,0)
	                 ]
        elif palette==1:
	        colors = [
	                    (0, 0, 0),
	                    (58/4,68/4,65/4),
	                    (80/2,100/2,153/2),
	                    (90,180,100),
	                    (224,224,44),
	                    (255,60,30),
	                    (255,255,255)
	                 ]

        # this line gets the old "screaming" colors back...
        # colors = [self.color_from_value(value/29.0) for value in range(0,30)]

        self.color_lookup = interpolate_colors(colors)
        self.pix = self.image.load()

    def color_from_value(self, value):
        """ given a value between 0 and 1, return an (r,g,b) tuple """

        return ImageColor.getrgb("hsl(%d,%d%%,%d%%)" % (int( (1.0 - value) * 360 ), 80, 50))

    def draw_peaks(self, x, peaks, spectral_centroid):
        """ draw 2 peaks at x using the spectral_centroid for color """

        y1 = self.image_height * 0.5 - peaks[0] * (self.image_height - 4) * 0.5
        y2 = self.image_height * 0.5 - peaks[1] * (self.image_height - 4) * 0.5

        line_color = self.color_lookup[int(spectral_centroid*255.0)]

        if self.previous_y != None:
            self.draw.line([self.previous_x, self.previous_y, x, y1, x, y2], line_color)
        else:
            self.draw.line([x, y1, x, y2], line_color)

        self.previous_x, self.previous_y = x, y2

        self.draw_anti_aliased_pixels(x, y1, y2, line_color)

    def draw_anti_aliased_pixels(self, x, y1, y2, color):
        """ vertical anti-aliasing at y1 and y2 """

        y_max = max(y1, y2)
        y_max_int = int(y_max)
        alpha = y_max - y_max_int

        if alpha > 0.0 and alpha < 1.0 and y_max_int + 1 < self.image_height:
            current_pix = self.pix[x, y_max_int + 1]

            r = int((1-alpha)*current_pix[0] + alpha*color[0])
            g = int((1-alpha)*current_pix[1] + alpha*color[1])
            b = int((1-alpha)*current_pix[2] + alpha*color[2])

            self.pix[x, y_max_int + 1] = (r,g,b)

        y_min = min(y1, y2)
        y_min_int = int(y_min)
        alpha = 1.0 - (y_min - y_min_int)

        if alpha > 0.0 and alpha < 1.0 and y_min_int - 1 >= 0:
            current_pix = self.pix[x, y_min_int - 1]

            r = int((1-alpha)*current_pix[0] + alpha*color[0])
            g = int((1-alpha)*current_pix[1] + alpha*color[1])
            b = int((1-alpha)*current_pix[2] + alpha*color[2])

            self.pix[x, y_min_int - 1] = (r,g,b)

    def save(self, filename):
        # draw a zero "zero" line
        a = 25
        for x in range(self.image_width):
            self.pix[x, self.image_height/2] = tuple(map(lambda p: p+a, self.pix[x, self.image_height/2]))

        self.image.save(filename)

class SpectrogramImage(object):
    def __init__(self, image_width, image_height, fft_size, f_max, f_min, nyquist_freq, palette):
        self.image = Image.new("P", (image_height, image_width))

        self.image_width = image_width
        self.image_height = image_height
        self.fft_size = fft_size
        self.f_max = f_max
        self.f_min = f_min
        self.nyquist_freq = nyquist_freq
        self.palette = palette

        if nyquist_freq<f_max:
            print "\nWarning: The specified maximum frequency to draw (%d Hz) is higher that what the digital file allows, which is %d Hz. The image file will have black areas on top that correspond to empty data.\n" % (f_max,nyquist_freq)

        if palette==2:
	        colors = [
	                    (255,255,255),
	                    (255,255,255),
	                    (255,255,255),
	                    (225,248,255),
	                    (210,241,255),
	                    (195,232,255),
	                    (180,221,255),
	                    (165,208,255),
	                    (150,193,255),
	                    (135,175,255),
	                    (120,156,255),
	                    (105,134,255),
	                    (90,110,255),
	                    (75,85,255),
	                    (63,60,255),
	                    (64,45,255),
	                    (66,30,255),
	                    (76,0,255),
	                    (0,128,13),
	                    (8,138,0),
	                    (20,143,0),
	                    (33,148,0),
	                    (46,153,0),
	                    (60,158,0),
	                    (91,168,0),
	                    (108,173,0),
	                    (125,179,0),
	                    (143,184,0),
	                    (162,189,0),
	                    (182,194,0),
	                    (199,195,0),
	                    (204,184,0),
	                    (255,230,128),
	                    (255,221,119),
	                    (255,213,111),
	                    (255,203,102),
	                    (255,193,94),
	                    (255,181,85),
	                    (255,169,77),
	                    (255,157,68),
	                    (255,143,60),
	                    (255,129,51),
	                    (255,113,43),
	                    (255,97,34),
	                    (255,81,25),
	                    (255,63,17),
	                    (255,45,8),
	                    (255,26,0)
	                 ]
        elif palette==1:
	        colors = [
	                    (0, 0, 0),
	                    (58/4,68/4,65/4),
	                    (80/2,100/2,153/2),
	                    (90,180,100),
	                    (224,224,44),
	                    (255,60,30),
	                    (255,255,255)
	                 ]

        self.image.putpalette(interpolate_colors(colors, True))

        # generate the lookup which translates y-coordinate to fft-bin
        self.y_to_bin = []
        y_min = math.log10(f_min)
        y_max = math.log10(f_max)
        for y in range(self.image_height):
#            log scale
#            freq = math.pow(10.0, y_min + y / (image_height - 1.0) *(y_max - y_min))
#            arithmetic scale
            freq = f_min + y / (image_height - 1.0) *(f_max - f_min)
#            uses the nyquist frequency to allow files of different sampling rate
            bin = freq / nyquist_freq * (self.fft_size/2 + 1)
#            bin = freq / 22050.0 * (self.fft_size/2 + 1)

            if bin < self.fft_size/2:
                alpha = bin - int(bin)

                self.y_to_bin.append((int(bin), alpha * 255))

        # this is a bit strange, but using image.load()[x,y] = ... is
        # a lot slower than using image.putadata and then rotating the image
        # so we store all the pixels in an array and then create the image when saving
        self.pixels = []

    def draw_spectrum(self, x, spectrum):
        for (index, alpha) in self.y_to_bin:
            self.pixels.append( int( ((255.0-alpha) * spectrum[index] + alpha * spectrum[index + 1] )) )

        for y in range(len(self.y_to_bin), self.image_height):
            self.pixels.append(0)

    def save(self, filename):
        self.image.putdata(self.pixels)
        self.image.transpose(Image.ROTATE_90).save(filename)

def create_png(input_filename, output_filename_w, output_filename_s, image_width, image_height, fft_size, f_max, f_min, wavefile, palette, channel):
    """
    Given command line arguments this basically does everything.

    WHAT I HAVE GATHERED:
    db_spectrum has the frequencies of the sound file.
    spectral_centroid tells us what the color of the sound is.
    peaks tell us what the amplitude of the sound is.

    Should be trivial to adapt this from image output to output
    to our JavaScript visualizer now.
    """

    print "processing file %s:\n\t" % input_file,

    audio_file = audiolab.sndfile(input_filename, 'read')  #opens the wavfile; audio_file is an object now

    samples_per_pixel = audio_file.get_nframes() / float(image_width)
    nyquist_freq = (audio_file.get_samplerate() / 2) + 0.0
    """
    Initializes AudioProcessor class, which does FFT analysis and spits
    out amplitudes and frequencies to the SpectrogramImage and WaveformImage
    classes below later. For a stereo wav file, this selects a single channel
    to analyze. We might want to analyze both channels to give more input to
    the visualizer,though.
    """
    processor = AudioProcessor(audio_file, fft_size, channel, numpy.hanning)

    if wavefile==1:
        waveform = WaveformImage(image_width, image_height, palette)
    spectrogram = SpectrogramImage(image_width, image_height, fft_size, f_max, f_min, nyquist_freq, palette)

    for x in range(image_width):
        #shows progress
        if x % (image_width/10) == 0:
            sys.stdout.write('.')
            sys.stdout.flush()

        seek_point = int(x * samples_per_pixel)
        next_seek_point = int((x + 1) * samples_per_pixel)

        (spectral_centroid, db_spectrum) = processor.spectral_centroid(seek_point)

        #let's have a look at the spectral centroid and the db_spectrum
        #print "Spectral Centroid:" + str(spectral_centroid)
        #print "DB Spectrum:" + str(db_spectrum)

        if wavefile==1:
            #aha! The peaks and spectral centroid make up the waveform.
            #Since the spectral centroid indicates timbre (often referred to as color),
            #it's probably what colors the waveform.
            peaks = processor.peaks(seek_point, next_seek_point)
            #let's have a look at these peaks
            #print "Peaks:" + str(peaks)
            waveform.draw_peaks(x, peaks, spectral_centroid)

        spectrogram.draw_spectrum(x, db_spectrum)

    if wavefile==1:
        waveform.save(output_filename_w)
    spectrogram.save(output_filename_s)

    print " done"

def processWav(filename, channel):
    """
    filename: path to a wav file
    Channel: 1 for left, 2 for right
    Returns centroids, frequencies, volumes
    """
    #open file
    audio_file = audiolab.sndfile(filename, 'read')
    #should be length of audiofile in seconds * 60. will fix this later

    import contextlib
    import wave
    with contextlib.closing(wave.open(filename, 'r')) as f:
        frames = f.getnframes()
        rate = f.getframerate()
        duration = frames / float(rate)
    duration *= 30 #30 data points for every second of audio yay
    duration = int(duration) #can only return an integer number of frames so yeah
    #print duration
    #Not really samples per pixel but I'll let that slide
    samples_per_pixel = audio_file.get_nframes() / float(duration)
    #some rule says this frequency has to be half of the sample rate
    nyquist_freq = (audio_file.get_samplerate() / 2) + 0.0
    #fft_size stays 2048; smaller size == more efficient, fewer frequency samples
    processor = AudioProcessor(audio_file, 2048, channel, numpy.hanning)

    centroids = []
    frequencies = []
    volumes = []

    for x in range(duration):
        seek_point = int(x * samples_per_pixel)
        next_seek_point = int((x + 1) * samples_per_pixel)
        (spectral_centroid, db_spectrum) = processor.spectral_centroid(seek_point)
        peaks = processor.peaks(seek_point, next_seek_point)

        centroids.append(spectral_centroid)
        frequencies.append(db_spectrum)
        volumes.append(peaks)
    #print "Centroids:" + str(centroids)
    #print "Frequencies:" + str(frequencies)
    #print "Volumes:" + str(volumes)

    #convert volumes[] from peaks to actual volumes
    for i in range(len(volumes)):
        volumes[i] = abs(volumes[i][0]) + abs(volumes[i][1])
    #round frequencies to save resources
    for i in range(len(frequencies)):
        for j in range(len(frequencies[i])):
            frequencies[i][j] = round(frequencies[i][j], 4)
    return centroids, frequencies, volumes


"""
Draft of a force-field.
At this point only a Lennard Jones and a solvation term are implemented.

Functions assign_atom_types(), assign_params(), get_neighbors() and,
get_sasa() adapted from
https://github.com/BIOS-IMASL/Azahar/blob/master/Azahar/energy.py
"""

import openbabel as ob
import numpy as np
from scipy.spatial.distance import cdist
from scipy.spatial import cKDTree
from numba import jit
from ..utils.geometry import dist
from ..utils.constants import par_s_ij, par_eps_ij

def assign_atom_types(selection='all'):
    """
    Assign properties at each atom in the selection.

    Using openbabel-2.4.1
    read http://openbabel.org/dev-api/classOpenBabel_1_1OBAtom.shtml#ae09ed28481ac044dab3f31c8605b44a9
    for available functions provided by openbabel to extract atom properties.
    There is a GetType() function but is more limited.

    Parameters
    -----------
    selection: str
        Atoms selection

    Returns
    --------
    atom_types: list of tuples, a tuple for each description of an atom
        [(at_index, elem, heavy_nb, in_ring, is_arom,
          ring_membership, neighbors)]
    """
    atom_types = []
    pdb_string = selection.dump_pdb(filename=False, b_factors=None, to_file=False)
    mol = ob.OBMol()
    obconversion = ob.OBConversion()
    obconversion.SetInAndOutFormats('pdb', 'pdb')
    obconversion.ReadString(mol, pdb_string)
    rings = mol.GetSSSR()
    for at in ob.OBMolAtomIter(mol):
        ring_member = [ring.IsMember(at) for ring in rings]
        neighbors = [neighbor.GetAtomicNum()
                     for neighbor in ob.OBAtomAtomIter(at)]
        atom_types.append(
            (at.GetIndex(),
             at.GetAtomicNum(),
             at.GetHvyValence(),
             any(ring_member),
                at.IsAromatic(),
                at.MemberOfRingCount(),
                (neighbors)))
    return atom_types

def assign_params(atom_types):
    """
    Assign solvation parameters for each atom.
    For now, simplified from: DOI: 10.1002/prot.340140112

    SRA_model = {'atom_types': (Radius, theta)}
    Radius: Hydrated radius in Angstroms
    theta: Atomic solvatation parameter (kcal/mol/A²)

    Parameters
    ----------
    atom_types: list of tuples, a tuple for each description of atoms
        [(at_index, elem, heavy_nb, in_ring, is_arom,
          ring_membership, neighbors)]

    Result
    --------
    params: Dictionary of atoms index with the corresponding Radius and theta
        {'atom_index': (Radius, theta)}
    """
    SRA_model = {'hydroxyl_carboxyl H': (2.85, -0.0487),
                 'amine_amide H': (2.85, -0.0487),
                 'thiol H': (2.85, -0.0487),

                 'aliphatic CH3': (0.946, 0.0676),
                 'aliphatic CH2': (0.946, 0.0676),
                 'aliphatic CH': (0.946, 0.0676),
                 'aliphatic_alicyclic C': (0.946, 0.0676),
                 'alicyclic CH2': (0.946, 0.0676),
                 'alicyclic CH': (0.946, 0.0676),
                 'aromatic CH': (0.946, 0.0676),
                 'aromatic C': (0.946, 0.0676),
                 'aromatic C of fused ring': (0.946, 0.0676),
                 'aromatic CH of fused ring': (0.946, 0.0676),
                 'carbonyl_carboxylic C': (0.946, 0.0676),

                 'N primary amine': (4.10, -0.0225),
                 'N secondary amine': (4.10, -0.0225),
                 'N cyclic amine': (4.10, -0.0225),
                 'aromatic N': (4.10, -0.0225),
                 'N amide': (4.10, -0.0225),

                 'ether_hydroxyl O': (2.83, -0.0282),
                 'carboxylic O': (2.83, -0.0282),
                 'carbonyl O': (2.83, -0.0282),
                 'amide carbonyl O': (2.83, -0.0282),

                 'thiol or sulfide S': (7.37, -0.0020)}
    params = {}
    for atom in atom_types:
        at_index, elem, heavy_nb, in_ring, is_arom, ring_membership, neighbors = atom
        if elem == 1:
            if (neighbors[0] == 6 or neighbors[0] == 8):
                params[at_index] = SRA_model['hydroxyl_carboxyl H']
            elif neighbors[0] == 7:
                params[at_index] = SRA_model['amine_amide H']
            elif neighbors[0] == 16:
                params[at_index] = SRA_model['thiol H']
        elif elem == 6:
            if neighbors.count(1) == 3:
                params[at_index] = SRA_model['aliphatic CH3']
            elif neighbors.count(1) == 2:
                if in_ring:
                    params[at_index] = SRA_model['alicyclic CH2']
                else:
                    params[at_index] = SRA_model['aliphatic CH2']
            elif neighbors.count(1) == 1:
                if in_ring:
                    params[at_index] = SRA_model['alicyclic CH']
                else:
                    params[at_index] = SRA_model['aliphatic CH']
                if is_arom:
                    if ring_membership > 1:
                        params[at_index] = SRA_model['aromatic CH of fused ring']
                    else:
                        params[at_index] = SRA_model['aromatic CH']
            elif neighbors.count(1) == 0:
                if is_arom:
                    if ring_membership > 1:
                        params[at_index] = SRA_model['aromatic C of fused ring']
                    else:
                        params[at_index] = SRA_model['aromatic C']
                else:
                    params[at_index] = SRA_model['aliphatic_alicyclic C']
            elif neighbors.count(8) > 1:
                params[at_index] = SRA_model['carbonyl_carboxylic C']
        elif elem == 7:
            if neighbors.count(1) == 2:
                if is_arom:
                    params[at_index] = SRA_model['N cyclic amine']
                else:
                    params[at_index] = SRA_model['N primary amine']
            elif neighbors.count(1) == 1:
                params[at_index] = SRA_model['N secondary amine']
            elif is_arom:
                params[at_index] = SRA_model['aromatic N']
            # elif # N de laamide
                #params.append(SRA_model['aromatic N'])
        elif elem == 8:
            if len(neighbors) == 2:
                params[at_index] = SRA_model['ether_hydroxyl O']
            else:
                params[at_index] = SRA_model['carboxylic O']
            # C=O
            # C=O of ester
            # O of amide
        if elem == 16:
            params[at_index] = SRA_model['thiol or sulfide S']
    return params

def get_neighbors(coord, probe, k, params):
    """
    Returns list of index of neighbors.

    TODO: Combine with compute_neighbors()

    Parameters
    ----------
    coord : Array (n,3)
        Cartesian coordinates
    probe : float
        Radius of the solvent
    k: int
        Atoms index from which to compute neighbors
    """
    dist = cdist(coord, coord, metric='euclidean')
    neighbor_indices = []
    radius = params[k][0] + probe * 2
    for key, values in params.items():
        if dist[key, k] < radius + values[0]:
            neighbor_indices.append(key)
    return neighbor_indices

def get_sasa(params, points, selection='all', probe=1.4):
    """
    Returns the solvent-accessible-surface area and empirical solvation term.

    Parameters
    ----------
    params: dictionary
        {atom: [atom's radius, atom's energy solvatation]}
    points: interger
        Number of points per sphere
    selection: str
        Atoms selection
    probe : float
        The radius of the solvent

    Return
    ------
    energy: float
        The energy of solvation from the atoms selection
    areas: float
        The total solvent-accessible-surface area from the atoms selection
    """
    # compute the area each point represents
    areas = []
    energies = []
    coord = selection.coords
    const = 4.0 *(np.pi/points)
    for key, values in params.items():
        # scale the points to the correct radius
        radius = values[0] + probe
        points_scaled = (coord[key] + points * radius).reshape(1,-1)
        # get all the indices of neighbors of the i residue
        neighbors = get_neighbors(coord, probe, key, params)
        # compute the distance between points and neighbors
        d_matrix = cdist(points_scaled, coord[neighbors],
                         metric='euclidean')
        # create a matrix and store the vdW radii for each neighbor
        nb_matrix = np.zeros((len(points_scaled), len(neighbors)))
        for nb_i, nb in enumerate(neighbors):
            nb_matrix[:, nb_i] = values[0]
        # compute the number of buried points, we have to be carefull
        # because we have counted how many times a point is buried
        # and we only need how many points are buried
        buried = np.sum(np.sum(d_matrix < nb_matrix + probe, axis=1) > 0)
        exposed = len(points_scaled) - buried
        area_per_atom = const * exposed * radius**2
        energy_per_atom = area_per_atom * values[1]
        areas.append(area_per_atom)
        energies.append(energy_per_atom)
    return sum(energies), sum(areas)

def compute_neighbors(coords, exclusions, cut_off):
    """
    Use a KD-tree (from scipy) to compute the neighbors atoms for each atom.

    Parameters
    ----------
    coords : array (m, 3)
        Cartesian coordinates of a molecule
    exclusions : set of tuples
        Pairs of atoms excluded from the computation of the neighbors
    cut_off : float
        Only pairs of atoms closer than cut_off will be used to compute the
        neighbors.

    Results
    -------
    neighbors: set of tuples
        Pairs of neighbors atoms within a given "cut_off" and excluding
        "exclusions"
    """
    tree_c = cKDTree(coords)
    all_pairs = tree_c.query_pairs(cut_off)
    return all_pairs - exclusions

def LJ(neighbors, xyz, elements):
    """
    Lennard Jones energy term

    .. math::

    LJ_{ij} = \\epsilon \\left [ \\left (\\frac{\\sigma_{ij}}{r_{ij}} \\right)^{12}
     - 2 \\left (\\frac{\\sigma_{ij}}{r_{ij}} \\right)^{6} \\right]

    \\sigma_{ij} is the distance at which the potential reaches its minimum
    \\epsilon_{ij} is the depth of the potential well
    r_{ij} is the distance between the particles

    Parameters
    ----------
    neighbors : set of tuples
        Pairs of neighbors atoms
    xyz : array (m, 3)
        Cartesian coordinates of a molecule
    elements : list of strings
        list of atoms in the molecule, used has key dictionary in
        par_s_ij and par_eps_ij

    Results
    -------
    E_LJ: float
        Lennard Jones energy contribution
    """

    E_vdw = 0.
    for i, j in neighbors:
        key_ij = elements[i] + elements[j]
        # use par_vdw without precomputing values
        #sigma_ij = par_vdw[key_i][0] + par_vdw[key_j][0]
        #epsilon_ij = (par_vdw[key_i][1] * par_vdw[key_j][1])**0.5
        # or use precomputed values
        sigma_ij = par_s_ij[key_ij]
        epsilon_ij = par_eps_ij[key_ij]

        E_vdw += _LJ(xyz, i, j, sigma_ij, epsilon_ij)
    return E_vdw

# convenient function just to speed up computation by 3x
@jit
def _LJ(xyz, i, j, sigma_ij, epsilon_ij):

    r_ij = dist(xyz[i], xyz[j])
    C6 = (sigma_ij / r_ij) ** 6

    return epsilon_ij * (C6 * C6 - 2 * C6)


'''
LEAKY INTEGRATE AND FIRE NETWORK W. STIMULATION
----------------------
Runs a simulation (from 0s to 9s), with a stimulation period from
3s to 6s. During this period, the current slowly ramps up, mimicking
the optogenetic stimulation (or silencing) implemented experimentally
----------------------
'''

from brian import *
import numpy as np
import numpy.matlib as ml
from numpy.random import rand as rand
from numpy.random import randn as randn
import h5py
import os.path
from subprocess import call
from sys import platform
import elephant
from neo.core import SpikeTrain
from elephant.conversion import BinnedSpikeTrain
import quantities as pq
from scipy import signal

root_dir = '/home/tpfeffer/neonates/'

from make_circuit import make_circuit

if __name__ == '__main__':

    #------------------------------------------------------------------------------
    # Simulation parameters
    #------------------------------------------------------------------------------
    # Timing
    runtime  = 9000.0 * ms
    stim_on  = 3000.0 * ms
    stim_off = 6000.0 * ms
    #------------------------------------------------------------------------------
    # VERSION 1: Parameters from run_network.py
    #------------------------------------------------------------------------------
    v = 1
    # Inputs: stimululus, AMPA, NMDA, GABA
    inputs      = np.array([0.9])
    AMPA_mods   = np.array([5.1])
    NMDA_mods   = np.array([1])
    GABA_mods   = np.linspace(0.7,6.2,56)
    ntrls       = 1
    stim_curr   = 0.2 # max stimulation current
    stim_frac   = 0.2 # fraciton of stimulated /inhibited neurons
    #------------------------------------------------------------------------------
    # preallocate
    #resp = np.zeros([len(AMPA_mods), len(NMDA_mods), len(GABA_mods), len(inputs)])
    #mean_corr = np.zeros([len(AMPA_mods), len(NMDA_mods), len(GABA_mods), len(inputs)])
    myclock = EventClock(dt=1.0*ms,t=0.0*ms,order=1,makedefaultclock=False)
    voltage_clock = EventClock(dt=5.0*ms,t=0.0*ms,order=2,makedefaultclock=False)
    #------------------------------------------------------------------------------
    # Run simulation
    #------------------------------------------------------------------------------

    if not(os.path.exists('/home/tpfeffer/neonates/proc/v%d' %v)):
      os.makedirs('/home/tpfeffer/neonates/proc/v%d' %v)

    for itr in range(ntrls):
      # Loop through exp parameters
      for igaba in range(0,GABA_mods.size):
          for iinp in range(0,inputs.size):
              for iampa in range(0,AMPA_mods.size):
                  for inmda in range(0,NMDA_mods.size):
                      for stim_protocol in range(1,3): # 1 = inhibit IN / 2 = excite IN

                        fn = os.path.expanduser(root_dir + 'proc/opto/v%d/neonates_network_opto_iampa%d_inmda%d_gaba%d_inp%d_stim%d_tr%d_v%d_processing.txt') % (v,iampa, inmda, igaba, iinp,stim_protocol,itr,v)
                        if os.path.isfile(fn)==False:
                            call(['touch', fn])
                        else:
                            continue

                        #  initialize
                        defaultclock.reinit(t=0.0*ms)
                        myclock.reinit(t=0.0*ms)

                        clear(True)

                        print("Computing INPUT%d, GABA%d, AMPA%d, trial%d ...") % (iinp, igaba,iampa,itr)

                        inp = inputs[0]
                        AMPA_mod = AMPA_mods[iampa]
                        NMDA_mod = NMDA_mods[inmda]
                        GABA_mod = GABA_mods[igaba]

                        Dgroups, Dconnections, Dnetfunctions, subgroups = make_circuit(inp,GABA_mod,AMPA_mod,NMDA_mod)

                        # get populations from the integrations circuit
                        popE = Dgroups['DE']
                        popI = Dgroups['DI']

                        #------------------------------------------------------------------------------
                        # Set up external stimulation (optogenetics protocol)
                        #------------------------------------------------------------------------------

                        if stim_protocol == 1: # inhibit inhibitory cells
                          stim = np.zeros([len(popI),1000 * int(stim_off-stim_on)])
                          rand_idx = np.random.choice(np.linspace(0,len(popI)-1,len(popI)),np.int(stim_frac*len(popI)),replace=False)
                          f = np.vectorize(np.int); rand_idx = f(rand_idx)
                          stim[rand_idx,:] = np.matlib.repmat(np.linspace(0,-stim_curr,1000 * int(stim_off-stim_on)),np.int(stim_frac*len(popI)),1)
                        elif stim_protocol == 2: # excite inhibitory cells
                          stim = np.zeros([len(popI),1000 * int(stim_off-stim_on)])
                          rand_idx = np.random.choice(np.linspace(0,len(popI)-1,len(popI)),np.int(stim_frac*len(popI)),replace=False)
                          f = np.vectorize(np.int); rand_idx = f(rand_idx)
                          stim[rand_idx,:] = np.matlib.repmat(np.linspace(0,stim_curr,1000 * int(stim_off-stim_on)),np.int(stim_frac*len(popI)),1)

                        @network_operation(myclock,stim_protocol)

                        def update_input():
                          if myclock.t >= stim_on and myclock.t < stim_off:
                            if stim_protocol == 1: # inhibit inhibitory cells
                              popI.I = stim[: ,int( (myclock.t - stim_on) / (1 * ms))] * nA
                            else: # excite inhibitory cells
                              popI.I = stim[: ,int( (myclock.t - stim_on) / (1 * ms))] * nA
                          else:
                            popI.I = 0.0 * nA
                            popE.I = 0.0 * nA

                        #------------------------------------------------------------------------------
                        # ---- set initial conditions (random)
                        popE.gen = popE.gen * (1 + 0.2 * rand(popE.__len__()))
                        popI.gen = popI.gen * (1 + 0.2 * rand(popI.__len__()))
                        popE.V = popE.V + rand(popE.__len__()) * 2 * mV
                        popI.V = popI.V + rand(popI.__len__()) * 2 * mV
                        #popE.I = 0.0 * nA
                        #popI.I = 0.0 * nA

                        # record spikes of excitatory/inhibitory neurons
                        Sp_E = SpikeMonitor(popE, record=True)
                        Sp_I = SpikeMonitor(popI, record=True)

                        # record instantaneous excitatory/inhibitory population activity
                        R_E = PopulationRateMonitor(popE, bin=5*ms)
                        R_I = PopulationRateMonitor(popI, bin=5*ms)

                        # record voltage
                        Vm_E = StateMonitor(popE, 'V', record=True, clock=voltage_clock)
                        Vm_I = StateMonitor(popI, 'V', record=True, clock=voltage_clock)

                        gE = StateMonitor(popE, 'gea', record=True, clock=voltage_clock)
                        gI = StateMonitor(popE, 'gi', record=True, clock=voltage_clock)

                        # record current
                        I_E = StateMonitor(popE, 'I', record=True, clock=myclock)
                        I_I = StateMonitor(popI, 'I', record=True, clock=myclock)

                        #------------------------------------------------------------------------------
                        # Run the simulation
                        #------------------------------------------------------------------------------
                        print("Running simulation...")
                        net = Network(Dgroups.values(),  Dconnections.values(), Dnetfunctions, update_input, Sp_E, Sp_I, R_E, R_I, Vm_E, Vm_I, gE, gI, I_E, I_I)
                        net.prepare()
                        net.run(runtime)

                        print("Computing power spectra ...")
                        LFP = np.abs(np.concatenate([gE.values,gI.values])).sum(axis=0)

                        fxx, pxx = signal.welch(LFP,200,window='hann')

                        # convert to array and save output has .h5
                        spt_E = []; spt_E_idx = []
                        spt_I = []; spt_I_idx = []

                        for ineuron in range(0,len(Sp_E.spiketimes)):
                            spt_E = np.append(spt_E,Sp_E.spiketimes.values()[ineuron], axis=None)
                            spt_E_idx = np.append(spt_E_idx,ml.repmat(ineuron,1,len(Sp_E.spiketimes.values()[ineuron])), axis=None)

                        for ineuron in range(0,len(Sp_I.spiketimes)):
                            spt_I = np.append(spt_I,Sp_I.spiketimes.values()[ineuron], axis=None)
                            spt_I_idx = np.append(spt_I_idx,ml.repmat(ineuron,1,len(Sp_I.spiketimes.values()[ineuron])), axis=None)

                        spt_E = np.vstack((spt_E,spt_E_idx))
                        spt_I = np.vstack((spt_I,spt_I_idx))

                        # -------------------------------------------
                        # COMPUTE SPIKE COUNT CORRELATIONS
                        # -------------------------------------------
                        print("Computing spike count correlations...")
                        spikes = dict(); frE = np.zeros([320,3]); frI = np.zeros([80,3])

                        first_spike = 0 # start analysis at t = 0
                        for ineuron in range(0,320):
                            spikes[ineuron] = SpikeTrain(spt_E[0][(spt_E[1]==ineuron) & ((spt_E[0]>0) & (spt_E[0]<runtime))]*pq.s, t_start = 0, t_stop = runtime)
                            frE[ineuron,0] = np.sum((spikes[ineuron]>0) & (spikes[ineuron]<3))
                            frE[ineuron,1] = np.sum((spikes[ineuron]>3) & (spikes[ineuron]<6))
                            frE[ineuron,2] = np.sum((spikes[ineuron]>6) & (spikes[ineuron]<9))
                        for ineuron in range(0,80):
                            spikes[ineuron+320] = SpikeTrain(spt_I[0][(spt_I[1]==ineuron) & ((spt_I[0]>0) & (spt_I[0]<runtime))]*pq.s, t_start = 0, t_stop = runtime)
                            frI[ineuron,0] = np.sum((spikes[ineuron+320]>0) & (spikes[ineuron+320]<3))
                            frI[ineuron,1] = np.sum((spikes[ineuron+320]>3) & (spikes[ineuron+320]<6))
                            frI[ineuron,2] = np.sum((spikes[ineuron+320]>6) & (spikes[ineuron+320]<9))

                        frE=np.mean(frE,axis=0)/int(runtime)
                        frI=np.mean(frI,axis=0)/int(runtime)

                        st_baseline = []; st_stim = []; st_post = []
                        subsamp = 1
                        matidx = np.triu_indices(len(range(0,len(spikes),subsamp)),1)
                        for isp in range(0,len(spikes),subsamp):
                            st_baseline.append(spikes[isp][(spikes[isp]>0) & (spikes[isp]<=3)])
                            st_stim.append(spikes[isp][(spikes[isp]>3) & (spikes[isp]<=6)])
                            st_post.append(spikes[isp][(spikes[isp]>6) & (spikes[isp]<=9)])

                        stc=np.zeros([3,1])
                        sts_corr_baseline=BinnedSpikeTrain(st_baseline, binsize=100*pq.ms)
                        tmp1=elephant.spike_train_correlation.corrcoef(sts_corr_baseline)
                        sts_corr_stim=BinnedSpikeTrain(st_stim, binsize=100*pq.ms)
                        tmp2=elephant.spike_train_correlation.corrcoef(sts_corr_stim)
                        sts_corr_post=BinnedSpikeTrain(st_post, binsize=100*pq.ms)
                        tmp3=elephant.spike_train_correlation.corrcoef(sts_corr_post)

                        # average correlations across upper triangular part
                        stc[0] = np.triu(tmp1,1).sum()/((tmp1.shape[0]**2-tmp1.shape[0])/2)
                        stc[1] = np.triu(tmp2,1).sum()/((tmp2.shape[0]**2-tmp2.shape[0])/2)
                        stc[2] = np.triu(tmp3,1).sum()/((tmp3.shape[0]**2-tmp3.shape[0])/2)

                        spike_array = np.zeros([400,9000])
                        if itr==0:
                          sts_corr_baseline=BinnedSpikeTrain(st_baseline, binsize=1*pq.ms)
                          spike_array[:,0:3000]=sts_corr_baseline.to_array()[:,0:3000]
                          sts_corr_stim=BinnedSpikeTrain(st_stim, binsize=1*pq.ms)
                          spike_array[:,3000:6000]=sts_corr_stim.to_array()[:,3000:6000]
                          sts_corr_post=BinnedSpikeTrain(st_post, binsize=1*pq.ms)
                          spike_array[:,6000:9000]=sts_corr_post.to_array()[:,6000:9000]

                          hf = h5py.File(os.path.expanduser(root_dir + 'proc/opto/v%d/neonates_opto_spikes_iampa%d_inmda%d_gaba%d_inp%d_prot%d_tr%d_v%d.h5') % (v,iampa, inmda, igaba, iinp,stim_protocol,itr,v), 'w')
                          hf.create_dataset('spike_array', data=spike_array)
                          hf.close()

                        # stE = []; stI = []
                        # subsamp = 1
                        # matidx = np.triu_indices(len(range(0,len(spikesE),subsamp)),1)
                        # for isp in range(0,len(spikesE),subsamp):
                        #     stE.append(spikesE[isp])
                        #     stI.append(spikesI[isp])

                        # stsE=BinnedSpikeTrain(stE, binsize=1*pq.ms)
                        # stsI=BinnedSpikeTrain(stI, binsize=1*pq.ms)

                        # stsE_corr=BinnedSpikeTrain(stE, binsize=10*pq.ms)
                        # corr=elephant.spike_train_correlation.corrcoef(stsE_corr)

                        # hf = h5py.File(os.path.expanduser(root_dir + 'proc/opto/v%d/neonates_network_opto_spiketrain_iampa%d_inmda%d_gaba%d_inp%d_stim%d_tr%d_v%d.h5') % (v,iampa, inmda, igaba, iinp,stim_protocol,itr,v), 'w')
                        # hf.create_dataset('spike_train_E', data= stsE.to_array())
                        # hf.create_dataset('spike_train_I', data= stsI.to_array())
                        # hf.close()

                        # mean_corr = corr[matidx].mean()
                        curr = I_I.values
                        hf = h5py.File(os.path.expanduser(root_dir + 'proc/opto/v%d/neonates_opto_iampa%d_inmda%d_gaba%d_inp%d_prot%d_tr%d_v%d.h5') % (v,iampa, inmda, igaba, iinp,stim_protocol,itr,v), 'w')
                        # hf.create_dataset('sttc_all', data=sttc_all)
                        hf.create_dataset('stc', data=stc)
                        # hf.create_dataset('sttcE', data=sttcE)
                        # hf.create_dataset('sttcI', data=sttcI)
                        hf.create_dataset('frE', data=frE)
                        hf.create_dataset('frI', data=frI)
                        hf.create_dataset('pxx', data=pxx)
                        hf.create_dataset('fxx', data=fxx)
                        hf.create_dataset('LFP', data=LFP)
                        hf.create_dataset('curr',data=curr)
                        hf.close()


"""Classes for odorants, mixtures, chemical orders, etc."""

import base64
import io
import json
import re
import time
import warnings
from collections import OrderedDict
from datetime import datetime
from urllib.parse import quote

import numpy as np
import pandas as pd
import pubchempy as pcp
import requests
from IPython.display import display
from PIL import Image

import quantities as pq
from pyrfume import load_data, logger, tqdm, trange
from pyrfume.physics import mackay
from quantities.constants.statisticalmechanics import R
from typing import Dict
try:
    from rdkit import Chem
    from rdkit.Chem import Draw, AllChem, SaltRemover
    from rdkit import RDLogger
    rdkit_logger = RDLogger.logger()
    RDKIT = True
except ImportError:
    warnings.warn(
        "Parts of mordred and/or rdkit could not be imported; try installing rdkit via conda",
        UserWarning,
    )
    RDKIT = False

ROOM_TEMP = (22 + 273.15) * pq.Kelvin
ROOM_PRESSURE = 1 * pq.atm
GAS_MOLAR_DENSITY = ROOM_PRESSURE / (R * ROOM_TEMP)

ODORANTS_BASIC_INFO_PATH = "molecules/all-cids-properties.csv"
ODORANT_SOURCES_PATH = "molecules/all-cids.csv"

PUBCHEM_KINDS = ['name', 'smiles', 'inchi', 'inchikey', 'formula', 'sdf', None]

class Solution:
    components: Dict["Compound", pq.quantity.Quantity] = None
    date_created: datetime = None

    def __init__(self, components: dict, date_created: datetime=None):
        self.total_volume = 0 * pq.mL
        assert isinstance(components, dict), "Components must be a dict"
        for component, volume in components.items():
            assert isinstance(
                component, (Compound, Solution)
            ), "Each component must be a Compound or a Solution"
            try:
                volume = volume.rescale(pq.mL)
            except ValueError:
                raise ValueError("Components must be provided with volumes")
            self.total_volume += volume  # Assume that volume is conserved
        self.components = components
        if not date_created:
            date_created = str(datetime.now())[:-7]
        self.date_created = date_created if date_created else datetime.now()

    @property
    def compounds(self):
        return self._compounds()

    def _compounds(self, result: dict = None):
        if result is None:
            result = {}
        for component, volume in self.components.items():
            if isinstance(component, Compound):
                if component in result:
                    result[component] += volume
                else:
                    result[component] = volume
            else:  # If it is a Solution
                component._compounds(result=result)
        return result

    @property
    def molecules(self):
        """Returns a dictionary with the moles of each Molecule"""
        compounds = self.compounds
        assert all([c.density for c, v in compounds.items() if v and not c.is_solvent]), (
            "All non-solvent compounds must have a known density " "in order to compute moles"
        )
        assert all([c.molecular_weight for c, v in compounds.items() if v and not c.is_solvent]), (
            "All non-solvent compounds must have a known molecular weight "
            "in order to compute moles"
        )
        return {
            c.molecule: (v * c.molarity).rescale(pq.mol) for c, v in self.compounds.items() if v
        }

    @property
    def molarities(self):
        """Returns a dictionary with the molarity of each Molecule"""
        return {m: mol / self.total_volume for m, mol in self.molecules.items() if mol}

    @property
    def mole_fractions(self):
        """Returns a dictionary with the mole fraction of each Molecule"""
        molecules = self.molecules
        assert [moles for molecule, moles in molecules.items()], (
            "All compounds must have a known number of moles " "in order to compute mole fraction"
        )
        # A Quantities bug prevents me from simply summing molecules.values()
        total_moles = 0 * pq.mol
        for moles in molecules.values():
            total_moles += moles
        return {molecule: moles / total_moles for molecule, moles in molecules.items()}

    def mole_fraction(self, molecule):
        return self.mole_fractions[molecule] if molecule in self.mole_fractions else 0

    @property
    def dilutions(self):
        return {
            c.molecule: self.total_volume / c.volume for c in self.compounds if not c.is_solvent
        }

    @property
    def partial_pressures(self):
        """Computes partial pressures for each odorant
        in the mixture using Raoult's law"""
        return {
            m: self.mole_fraction(m) * m.vapor_pressure for m in self.molecules if m.vapor_pressure
        }

    def partial_pressure(self, molecule):
        return self.partial_pressures[molecule]

    @property
    def total_pressure(self):
        """Computes total pressure of the vapor using Dalton's law"""
        preferred_units = pq.Pa
        partial_pressures = [pressure.rescale(preferred_units) for pressure in self.partial_pressures.values()]
        return preferred_units * np.sum(partial_pressures)

    @property
    def vapor_concentrations(self):
        """Concentrations of each component in the vapor phase at steady state.
        Units are fraction of volume. Air is assumed to make up the balance"""
        pp = self.partial_pressures
        result = {}
        for m, p in pp.items():
            ratio = (p / pq.atm).simplified
            assert ratio.units == pq.dimensionless
            result[m] = float(ratio)
        return result

    def vapor_concentration(self, molecule):
        return self.vapor_concentrations[molecule]

    @property
    def molar_evaporation_rates(self):
        mf = self.mole_fractions
        result = {
            molecule: mole_fraction * molecule.molar_evaporation_rate
            for molecule, mole_fraction in mf.items()
        }
        return result

    def molar_evaporation_rate(self, molecule):
        return self.molar_evaporation_rates[molecule]

class Molecule:
    def __init__(self, cid: int, name: str=None, fill: bool=False):
        self.cid = cid

        if fill:
            self.fill_details()
        if name:
            self.name = name

    # Integer Chemical ID number (CID) from PubChem
    cid: int = 0
    # Chemical Abstract Service (CAS) number
    cas: str = ""
    # Principal name
    name: str = ""
    # Synonyms
    synonyms: str = ""
    # IUPAC name (long, unique name)
    iupac: str = ""
    # Density (pq.g / pq.ml)
    density: float = None
    # Vapor pressure (pq.Pa)
    vapor_pressure: float = None
    # Molecular weight (pq.g / pq.mol)
    molecular_weight: float = None

    @property
    def molarity(self):
        if not (self.molecular_weight and self.density):
            result = None
        else:
            result = self.density / self.molecular_weight
            result = result.rescale(pq.mol / pq.L)
        return result

    @property
    def molar_evaporation_rate(self):
        return mackay(self.vapor_pressure)

    def fill_details(self):
        assert self.cid is not None
        url_template = (
            "https://pubchem.ncbi.nlm.nih.gov/" "rest/pug/compound/cid/%d/property/" "%s/JSON"
        )
        property_list = ["MolecularWeight", "IsomericSMILES"]
        url = url_template % (self.cid, ",".join(property_list))
        json_data = url_to_json(url)
        details = json_data["PropertyTable"]["Properties"][0]

        def convert(name):
            s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
            return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()

        for key, value in details.items():
            if key == "CID":
                assert value == self.cid, "REST API CID does not match provided CID"
            key = convert(key)
            if key == "molecular_weight":
                value = float(value) * pq.g / pq.mol
            setattr(self, key, value)

        if not self.name:
            self.name = self.get_name_from_api()

    def get_name_from_api(self):
        url_template = "https://pubchem.ncbi.nlm.nih.gov/" "rest/pug/compound/cid/%d/synonyms/JSON"
        url = url_template % (self.cid)
        json_data = url_to_json(url)
        name = None
        if json_data:
            information = json_data["InformationList"]["Information"][0]
            synonyms = information["Synonym"]
            name = synonyms[0].lower()
        return name

    def get_cid_from_api(self):
        url_template = "https://pubchem.ncbi.nlm.nih.gov/" "rest/pug/compound/%s/%s/cids/JSON"
        options = [getattr(self, x) for x in ("cas", "name") if len(getattr(self, x))]
        cid = None
        query = self.name
        for option in options:
            url = url_template % (option, query)
            json_data = url_to_json(url)
        cid = json_data["IdentifierList"]["CID"][0]
        return cid

    def __eq__(self, other):
        if self.cid:
            return self.cid == other.cid
        else:
            return self.name == self.name

    def __lt__(self, other):
        if self.cid:
            return self.cid < other.cid
        else:
            return self.name < self.name

    def __hash__(self):
        return id(self)

    def __repr__(self):
        if self.cid and self.name:
            result = "%d (%s)" % (self.cid, self.name)
        elif self.cid:
            result = "%d" % self.cid
        elif self.name:
            result = "%s" % self.name
        else:
            result = "Unknown"
        return result

class Vendor:
    def __init__(self, name: str, url: str):
        self.name = name
        self.url = url

    name: str = ""
    url: str = ""

class ChemicalOrder:
    def __init__(self, molecule: "Molecule", vendor: "Vendor", part_id: str, purity: float=1, known_impurities: list=None):
        self.molecule = molecule
        self.vendor = vendor
        self.part_id = part_id
        self.purity = purity
        self.known_impurities = known_impurities

    # Molecule
    molecule: Molecule = None
    # Vendor, e.g. Sigma-Aldrich
    vendor: Vendor = None
    # ID number of compound at vendor
    part_id: str = ""
    # Reported purity as a fraction
    purity: float = 1
    # List of known impurities (Molecules)
    known_impurities: list = None

class Compound:
    def __init__(
        self, chemical_order: ChemicalOrder, stock: str="", date_arrived: datetime=None,
        date_opened: datetime=None, is_solvent: bool=False
    ):

        self.chemical_order = chemical_order
        self.stock = stock
        self.date_arrived = date_arrived if date_arrived else datetime.now
        self.date_opened = date_opened
        self.is_solvent = is_solvent

    # ChemicalOrder
    chemical_order: ChemicalOrder = None
    # Stock number (supplied by vendor, usually on bottle)
    stock: str = ""
    # Date arrived at the lab/clinic
    date_arrived: datetime = None
    # Date opened
    date_opened: datetime = None
    # Is it a solvent?
    is_solvent: bool = False

    def __getattr__(self, attr):
        """If no attribute is found, try looking up on the
        ChemicalOrder or the Molecule"""
        try:
            return getattr(self.chemical_order, attr)
        except AttributeError:
            return getattr(self.chemical_order.molecule, attr)

def url_to_json(url, verbose=True) -> str:
    json_data = None
    response = requests.get(url)
    if response.status_code == 200:
        string = response.content.decode("utf-8")
        json_data = json.loads(string)
    else:
        msg = "HTTP Status Code %d for %s" % (response.status_code, url)
        print(msg)
        if verbose:
            logger.error(msg)
    return json_data

def is_kind(identifier: str, kind: str) -> bool:
    if kind == 'smiles':
        if RDKIT:
            rdkit_logger.setLevel(RDLogger.CRITICAL)
            result = Chem.MolFromSmiles(identifier) is not None
            rdkit_logger.setLevel(RDLogger.WARNING)
        else:
            result = None
    elif kind == 'inchikey':
        result = len(identifier) == 27 and identifier[14]=='-' and identifier[25]=='-'
    elif kind == 'inchi':
        if RDKIT:
            result = Chem.inchi.MolFromInchi(identifier, logLevel=None) is not None
        else:
            result = None
    elif kind == 'name':
        result = True
    else:
        result = False
    return result

def get_kind(identifier: str):
    kinds = [kind for kind in PUBCHEM_KINDS if is_kind(identifier, kind)]
    return kinds[-1]  # Return most sophisticated kind ('name' will always be in the list)

def deisomerize_smiles(smiles: str) -> str:
    mol = Chem.MolFromSmiles(smiles)
    if mol:  # If a mol object was successfully create (i.e. not `None`)
        smiles = Chem.MolToSmiles(mol, isomericSmiles=False)
    else:
        smiles = smiles.replace('@','').replace('@@','').replace('/','').replace('\\','')
    return smiles

def canonical_smiles(smiles: str, kekulize: bool = False) -> str:
    """Use rdkit to convert the `smiles` string to canonical form"""
    mol = Chem.MolFromSmiles(smiles)
    if mol:  # If a mol object was successfully create (i.e. not `None`)
        if kekulize:
            Chem.Kekulize(mol)
        smiles = Chem.MolToSmiles(mol, isomericSmiles=True)
    else:  # No mol object means the `smiles` string was invalid
        smiles = ""
    return smiles

def get_cids(
    identifiers: list,
    kind: str = None,
    verbose: bool = True,
    wait: float = 0,
    results: dict = None,
) -> dict:
    """Return CIDs for molecule based on any synonym,
    including a chemical name or a CAS"""
    if isinstance(kind, str):
        kind = kind.lower()
    assert kind in PUBCHEM_KINDS
    if results is None:
        results = {}
    p = tqdm(identifiers)
    for identifier in p:
        #if not isinstance(identifier, str):
        #    logger.warning("%s is not a string" % identifier)
        #    continue
        p.set_description(str(identifier))
        cid = get_cid(identifier, kind=kind, verbose=verbose)
        if not cid:
            logger.warning("Could not find %s" % identifier)
        results[identifier] = cid
        if wait:
            time.sleep(wait)
    return results

def get_cid(
    identifier: str, kind: str = None, verbose: bool = True, fix_smiles_on_error: bool = True, attempt=0
) -> int:
    """
    Return data about a molecule from any synonym,
    including a chemical name or a CAS.
    """
    if isinstance(identifier, float) and np.isnan(identifier):
        return 0
    replace = [('α', 'alpha'), ('β', 'beta'), ('γ', 'gamma'), ('δ', 'delta')]
    for a, b in replace:
        identifier = identifier.replace(a, b)
    if kind is None:
        kind = get_kind(identifier)
    else:
        kind = kind.lower()
    try:
        result = pcp.get_cids(identifier, namespace=kind)
    except pcp.BadRequestError:
        logger.warning('Request Error for "%s"' % identifier)
        result = []
    except pcp.PubChemHTTPError as e:
        if attempt == 0:
            import time
            time.sleep(10)
            return get_cid(identifier,kind,verbose, fix_smiles_on_error, 1)
        else:
            raise e
    if not len(result):
        cid = 0
    else:
        if (len(result) > 1) and verbose:
            logger.warning("Multiple CIDs for %s: %s" % (identifier, result))
        cid = result[0]
    if not cid and kind == "smiles" and fix_smiles_on_error:
        # Retry with canonical SMILES
        identifier = canonical_smiles(identifier)
        if identifier:
            cid = get_cid(identifier, kind=kind, verbose=verbose, fix_smiles_on_error=False)
    return cid

def from_cids(cids: list, property_list: bool = None) -> list:
    if property_list is None:
        property_list = ["MolecularWeight", "IsomericSMILES", "IUPACName"]
    result = []
    chunk_size = 100
    for start in trange(0, len(cids), chunk_size):
        stop = min(start + chunk_size, len(cids))
        logger.info("Retrieving %d through %d" % (start, stop - 1))
        cid_subset = [
            str(x) for x in [cid for cid in cids[start:stop] if int(cid) > 0 and cid is not None]
        ]
        cid_subset = ",".join(cid_subset)
        properties_template = (
            "https://pubchem.ncbi.nlm.nih.gov/" "rest/pug/compound/cid/%s/property/" "%s/JSON"
        )
        url = properties_template % (cid_subset, ",".join(property_list))
        json_data = url_to_json(url)
        data = json_data["PropertyTable"]["Properties"]

        synonyms_template = (
            "https://pubchem.ncbi.nlm.nih.gov/" "rest/pug/compound/cid/%s/synonyms/JSON"
        )
        url = synonyms_template % (cid_subset)
        json_data = url_to_json(url)
        information = json_data["InformationList"]["Information"]
        for i, d in enumerate(data):
            try:
                synonyms = information[i]["Synonym"]
            except KeyError:
                try:
                    d["name"] = d["IUPACName"]
                except KeyError:
                    d["name"] = ""
            else:
                d["name"] = synonyms[0].lower()
        result += data
    return result

def cids_to_smiles(cids: list) -> dict:
    """Returns an ordered dictionary of SMILES strings with CIDs as keys"""
    info = from_cids(cids, property_list=["IsomericSMILES"])
    smiles = {item["CID"]: item["IsomericSMILES"] for item in info}
    return smiles

def cids_to_cas(cids: list) -> OrderedDict:
    result = OrderedDict()
    chunk_size = 100
    for start in trange(0, len(cids), chunk_size):
        stop = min(start + chunk_size, len(cids))
        # msg = "Retrieving %d through %d" % (start, stop-1)
        cid_subset = [str(x) for x in cids[start:stop]]
        cid_subset = ",".join(cid_subset)
        synonyms_template = (
            "https://pubchem.ncbi.nlm.nih.gov/" "rest/pug/compound/cid/%s/synonyms/JSON"
        )
        url = synonyms_template % (cid_subset)
        json_data = url_to_json(url)
        information = json_data["InformationList"]["Information"]
        for i, info in enumerate(information):
            cid = info["CID"]
            try:
                synonyms = info["Synonym"]
            except KeyError:
                result[cid] = []
            else:
                result[cid] = cas_from_synonyms(synonyms)
        result
    return result

def cas_from_synonyms(synonyms: list) -> list:
    result = []
    for s in synonyms:
        if re.match(r"^[0-9]+\-[0-9]+\-[0-9]+$", s):
            result.append(s)
    return result

def cactus(identifier: str, output: str = "cas") -> str:
    url_template = "https://cactus.nci.nih.gov/chemical/structure/%s/%s"
    identifier = identifier.replace(' ', '%20')
    url = url_template % (identifier, output)
    response = requests.get(url)
    if response.status_code == 200:
        result = response.content.decode("utf-8")
    else:
        logger.error("HTTP Status Code %d for %s" % (response.status_code, url))
        result = None
    return result

def cactus_image(smiles: str) -> None:
    url_template = "https://cactus.nci.nih.gov/chemical/structure/%s/image"
    smiles = smiles.replace(' ', '%20')
    url = url_template % smiles
    response = requests.get(url)
    if response.status_code == 200:
        image_data = response.content.decode("utf-8")
        image = Image(image_data)
        display(image)
    else:
        logger.error("HTTP Status Code %d for %s" % (response.status_code, url))

def get_compound_summary(cid: int, heading: str):
    """Get summary info about `heading` from PubChem for the compound
    given by `cid`.  Example heading: 'Physical Description'"""
    url_template = (
        "https://pubchem.ncbi.nlm.nih.gov/" "rest/pug_view/data/compound/%d/JSON?heading=%s"
    )
    escaped_heading = quote(heading)  # Escape the string
    url = url_template % (cid, escaped_heading)
    json_data = url_to_json(url, verbose=False)
    return json_data

def get_compound_odor(cid, raw=False):
    info = []
    for heading in ["Odor", "Physical Description"]:
        json_data = get_compound_summary(cid, heading)
        if raw:
            info += [] if json_data is None else [json_data]
        else:
            info += _parse_odor_info(json_data)
    return info

def _parse_odor_info(info, odors=None, any_string=False):
    if odors is None:
        odors = []
    if isinstance(info, dict):
        for key, value in info.items():
            if key == "TOCHeading" and value == "Odor":
                any_string = True
            if key == "String" and (any_string or "odor" in value.lower()):
                odors.append(value)
            else:
                _parse_odor_info(value, odors=odors, any_string=any_string)
    elif isinstance(info, list):
        for value in info:
            _parse_odor_info(value, odors=odors, any_string=any_string)
    return odors

def _parse_other_info(info, records=None):
    if records is None:
        records = []
    if isinstance(info, dict):
        for key, value in info.items():
            if key == "String":
                records.append(value)
            elif key == "Value" and "Number" in value:
                records.append(value)
            else:
                _parse_other_info(value, records=records)
    elif isinstance(info, list):
        for value in info:
            _parse_other_info(value, records=records)
    return records

def display_molecules(molecules: pd.DataFrame, no_of_columns=5, figsize=(15, 15)):
    import matplotlib.pyplot as plt
    from IPython.display import display
    fig = plt.figure(figsize=figsize)
    column = 0
    for i, (cid, info) in enumerate(molecules.iterrows()):
        column += 1
        #  check for end of column and create a new figure
        if column == no_of_columns+1:
            fig = plt.figure(figsize=figsize)
            column = 1
        fig.add_subplot(1, no_of_columns, column)
        image = smiles_to_image(info['IsomericSMILES'], png=False)
        plt.imshow(image)
        plt.axis('off')
        plt.title("%d: %s" % (cid, info['name']))

def embed_molecules(molecules: pd.DataFrame):
    import matplotlib.pyplot as plt
    plt.figure(figsize=(6, 6))
    ax = plt.gca()
    embedding = load_data('embedding/pf_umap.pkl')
    ax = embedding.plot.scatter(x=0, y=1, alpha=0.05, c='k', ax=ax)
    smiles = molecules['IsomericSMILES']
    embedding_ = embedding.loc[smiles]
    embedding_.plot.scatter(x=0, y=1, alpha=1, c='r', s=100, ax=ax)
    ax.set_xlabel('Dimension 1')
    ax.set_ylabel('Dimension 2')

def smiles_to_image(smiles, png=True, b64=False, crop=True, padding=10, size=300):
    """
    png: Whether to convert to .png data (or to leave as a PIL image)
    b64: Whether to base64 encode (only possible for .png data)
    """
    buffer = io.BytesIO()
    mol = Chem.MolFromSmiles(smiles)
    image = Draw.MolToImage(mol, fitImage=True, size=(size, size))
    if crop:
        image = crop_image(image, padding=padding)
    if png:
        image.save(buffer, format="PNG")
        image = buffer.getvalue()
    if b64:
        assert png, "Can only base64 encode PNG data, not a PIL image"
        image = base64.b64encode(image).decode("utf8")
    return image

def smiles_to_mol(smiles: list, max_attempts: int=10, use_random_coords: bool=False, deisomerize=False) -> dict:
    if deisomerize:
        f = deisomerize_smiles
    else:
        f = lambda x: x
    mols_raw = [Chem.MolFromSmiles(f(smi)) for smi in smiles]
    logger.info("Computing 3D coordinates...")
    s = SaltRemover.SaltRemover()
    mols = {}
    n = len(mols_raw)
    pbar = tqdm(total=n)
    for i, mol in enumerate(mols_raw):
        pbar.update()
        logger.debug("Embedding %s" % smiles[i])
        try:
            mol = s.StripMol(mol, dontRemoveEverything=True)
            mol = Chem.AddHs(mol)
            AllChem.Compute2DCoords(mol)
            AllChem.EmbedMolecule(mol, maxAttempts=max_attempts, useRandomCoords=use_random_coords)
            AllChem.UFFOptimizeMolecule(mol)  # Is this deterministic?
        except Exception as e:
            logger.warning("Exception for %s: %s" % (smiles[i], str(e)))
        else:
            mols[smiles[i]] = mol
    logger.info("Finished embedding all molecules")
    return mols

def crop_image(img, padding=0):
    """Crop white out of a PIL image."""
    as_array = np.array(img)  # N x N x (r,g,b,a)
    if as_array.shape[2] == 4:
        as_array[as_array[:, :, 3] == 0] = [255, 255, 255, 255]
    has_content = np.sum(as_array, axis=2, dtype=np.uint32) != 255 * 4
    xs, ys = np.nonzero(has_content)
    x_range = max([min(xs) - padding, 0]), min([max(xs) + padding, as_array.shape[0]])
    y_range = max([min(ys) - padding, 0]), min([max(ys) + padding, as_array.shape[1]])
    as_array_cropped = as_array[x_range[0] : x_range[1], y_range[0] : y_range[1], 0:3]
    img = Image.fromarray(as_array_cropped, mode="RGB")
    return img

def all_odorants():
    """All CIDs, SMILES, Names, and Molecular Weights found in the
    file at ODORANTS_BASIC_INFO_PATH"""
    df = load_data(ODORANTS_BASIC_INFO_PATH)
    df = df.sort_index()
    return df

def all_sources():
    """Whether or not each odorant (by CID) is in each of the data sources"""
    df = load_data(ODORANT_SOURCES_PATH)
    df = df.sort_index()
    return df

def all_cids():
    """All CIDs found in the file at ODORANTS_BASIC_INFO_PATH"""
    df = all_odorants()
    return list(df.index)

def all_smiles():
    """All SMILES found in the file at ODORANTS_BASIC_INFO_PATH.
    May contain duplicates (if two CIDs give the same SMILES)"""
    df = all_odorants()
    return list(df["SMILES"])

if __name__ == "__main__":
    x = Molecule(325, fill=True)
    print(x.__dict__)