import uproot
from itertools import chain, product
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm as tqdm
import src.mip_eff.tb_data_handling as dh
import src.mip_eff.simulation_analysis as mc
import argparse
from datetime import datetime
import os
import warnings
import csv

def eff(mode, Nchb=8, qc=False):
    """ Estimates the MIP detection efficiency for CaloEvents data.
    Parameters:
    -----------
        mode    :   'dt' or 'calo' for trigger-time correlated hits or CaloEvents, respectively.
        Nchb    :   the number of chambers in the setup. Default: 8.

    Returns:
    --------

        eff_tot :    a dictionary of the efficiency value of the different
                    chambers. The chamber numbers serves as keys.
        pool_tot :   a dictionary of the number of MIP tracks used as reference
                    for the efficiency estimation. The chamber numbers serves as keys.
    """

    if mode == 'calo':
        print('Running efficiency estimation for hits in CaloEvents.')
        dataId = 'caloId'
        df_MIPs = pd.DataFrame(columns = ['eventId', 'chbid', 'xpos', 'ypos', 'dt', 'digit', 'chipid', 'hardid',
        'nhits', 'caloId', 'zx0', 'zx1', 'zy0', 'zy1'])
    else:
        print('Running efficiency estimation for hits which are time-correlated with the trigger.')
        dataId = 'eventId'
        df_MIPs = pd.DataFrame(columns = ['eventId', 'chbid', 'xpos', 'ypos', 'dt', 'digit', 'chipid', 'hardid',
        'nhits', 'zx0', 'zx1', 'zy0', 'zy1'])

    if Nchb == 7:
        runList = pd.read_csv('data/eScan8Layers.csv')
    else:
        runList = pd.read_csv('data/eScan{}Layers.csv'.format(Nchb))

    l = runList.shape[0]

    # for selection quality control purpose
    if qc:
        eff_runs = {'run': [], 'eff': [], 'pool': []}

    # Define list of excluded runs
    if Nchb == 8 or Nchb == 7:
        excluded_runs = ['0036', '0650', '0712']
    elif Nchb == 11:
        excluded_runs = ['0018', '0726']
    else:
        excluded_runs = []

    for i in tqdm(range(l)):

        # # debug

        # if i > 0:
        #     print('DEBUG MODE!!! ONLY 1 RUNS!!!')
        #     break

        run = runList.iloc[i]

        if format(run["time"], '04d') not in excluded_runs:
            runId = '{:08d}-{:04d}-{}'.format(run["day"], run["time"], run["index"])
            print(runId)
            df = dh.read_nov18run(run, Nchb=Nchb)
            df = df[df.chbid <= Nchb]

            if mode == 'calo':
                time_wins = dh.calo_time_wins(df)
            else:
                time_wins = [15]

            df_batch, nEvents = dh.cleaning(df, mode, time_wins)

            if nEvents == 0:
                print('Zero nEvents after cleaning!')
                continue

            # eff_events = df_batch.groupby(dataId)['chbid'].agg(lambda x: len(x) > Nchb-2)
            # eff_events = eff_events[eff_events].index.tolist()

            # Keep only effective events with more than Nchb-2 layers
            eff_events = df_batch.groupby(dataId).agg(lambda x: len(set(x)))
            eff_events = eff_events[eff_events['chbid'] > Nchb-2].index.tolist()
            df_batch_eff = df_batch[df_batch[dataId].isin(eff_events)]

            # for selection quality control purpose
            if qc:
                hresx_tot = [0 for i in range(10)]
                hresy_tot = [0 for i in range(10)]

            if df_batch_eff.shape[0] > 0:
                (effective_MIPs, hresx, hresy, edge) = dh.isMIP(df_batch_eff, mode, Nchb=Nchb, res=2)

                print ('{:.2%} of events are valid MIPs'.format(len(effective_MIPs)/nEvents))
                df_batch_eff = df_batch_eff[df_batch_eff[dataId].isin(effective_MIPs)]

                # df_batch_eff[dataId] = df_batch_eff[dataId].agg(lambda x: '{}_{}'.format(runId, x))
                df_MIPs = pd.concat([df_MIPs,df_batch_eff], axis=0)

                # for selection quality control purpose
                if qc:
                    hresx_tot += hresx
                    hresy_tot += hresy
                    eff_run, pool_run, mult_run = dh.efficiency_estimation(df_batch)
                    eff_runs['run'].append(runId)
                    eff_runs['eff'].append(eff_run)
                    eff_runs['pool'].append(pool_run)

    eff_tot, pool_tot, mult_tot = dh.efficiency_estimation(df_MIPs, mode, Nchb)

    dh.exporter(eff_tot, pool_tot, mult_tot, mode, Nchb)

    return eff_tot, pool_tot, mult_tot

def eff_MC(energy, mipeff, sig, particle='pions', Nchb=8):
    """ Estimates the MIP detection efficiency for CaloEvents data.
    Parameters:
    -----------
        energy  :   beam energy in GeV (2-6).
        mipeff  :   simulation mip threshold.
        particle:   beam's type of particle: 'pions' or 'electrons'. Default: 'pions'.
        mode    :   'dt' or 'calo' for trigger-time correlated hits or CaloEvents, respectively.
        Nchb    :   the number of chambers in the setup. Default: 8.

    Returns:
    --------

        eff_tot :    a dictionary of the efficiency value of the different
                    chambers. The chamber numbers serves as keys.
        pool_tot :   a dictionary of the number of MIP tracks used as reference
                    for the efficiency estimation. The chamber numbers serves as keys.
    """

    # filename, digit = mc.MAX_FILES[particle][energy][mipeff]
    MC_HOME = '/Users/dansh/Projects/g4_mpgd_sdhcal/uMSDHCAL/rootfiles/10chb_high_statistics/digitized/'
    prefix = 'digitized_MeasuredEff_0sig_8chbs_uMSDHCAL_ArCO2_2cm'
    suffix = '0T_10k_FTFP_BERT.root'

    filename = prefix + '_pi_{}_'.format(energy) + suffix
    dataId = 'eventId'
    filename = 'digitized_MeasuredEff_20200923_0.01sig_8chbs_high_statistics_100k_10chbs_5GeV_pi_QGSP_BERT.root'
    # filename = 'digitized_noMult_lowMMHighRP_{}sig_8chbs_high_statistics_1M_10chbs_{}_pi.root'.format(sig, energy)
    print("efficiency estimation for: " + filename)
    # read file to dataframe
    file = uproot.open(MC_HOME + filename)
    tree = file['tvec']
    df= (pd.DataFrame(list(chain(*[list(product([x],[y])) \
                        for x, y in zip(tree.array("eventId"),
                                        tree.array("layer"))])),
                                        columns= ['eventId',"chbid"]))
    df['xpos'] = tree.array('xpos').flatten()
    df['ypos'] = tree.array('ypos').flatten()
    # df['Thr'] = tree.array('Thr').flatten()
    df['digit'] = tree.array('digit').flatten()
    df['pad'] = df[['chbid', 'xpos', 'ypos']].agg(tuple, axis=1)

    df_batch, nEvents = dh.cleaning(df, 'MC')
    print('{} clean events'.format(len(df_batch[dataId].unique())))
    if nEvents == 0:
        raise ValueError('nEvent is zero!')
        return

    eff_events = df_batch.groupby(dataId)['chbid'].agg(lambda x: len(set(x)) > Nchb-2)
    eff_events = eff_events[eff_events].index.tolist()
    df_batch_eff = df_batch[df_batch[dataId].isin(eff_events)]

    if df_batch_eff.shape[0] > 0:
        (effective_MIPs, hresx, hresy, edge) = dh.isMIP(df_batch_eff, mode, Nchb=Nchb, res=2)

        print ('{:.2%} of events are valid MIPs'.format(len(effective_MIPs)/len(df_batch[dataId].unique().tolist())))
        df_batch_eff = df_batch_eff[df_batch_eff[dataId].isin(effective_MIPs)]
        df_MIPs = df_batch_eff

    eff_tot, pool_tot, mult_tot = dh.efficiency_estimation(df_MIPs, 'MC', Nchb)
    df_MIPs.to_pickle(MC_HOME + 'df_MIPs_'+ filename[:-5] + '.pkl')
    return eff_tot, pool_tot, mult_tot

def plot_eff(eff_tot, pool_tot, mode, Nchb=8):
    """Generates plots for efficiency estimation and save them.
    Parameters:
    -----------
        eff_tot :    a dictionary of the efficiency value of the different
                    chambers. The chamber numbers serve as keys.
        pool_tot :   a dictionary of the number of MIP tracks used as reference
                    for the efficiency estimation. The chamber numbers serve as keys.
        mode     :  'dt', 'calo', or 'MC' for trigger-time correlated hits, CaloEvents, or simulation, respectively.
        Nchb :      the number of chambers in the setup. Default: 8.

    Return:
    -------
        None

    """
    fig = plt.figure(figsize=(15, 20), constrained_layout=False)
    gs = fig.add_gridspec(nrows=3, ncols=1, left=0.05, right=0.48, wspace=0.05)
    ax1 = fig.add_subplot(gs[1, :])
    ax2 = fig.add_subplot(gs[-1, :])

    # 16x16 cm^2 uM chambers
    yerr_sm = [1/np.sqrt(pool_tot[i]) for i in range(2, 4)]
    ax1.errorbar(range(2, 4), [eff_tot[i] for i in range(2, 4)], yerr=yerr_sm, fmt='o',
                label='Small MM')

    # 48x48 cm^2 uM chambers
    yerr_lm = [1/np.sqrt(pool_tot[i]) for i in range(4, 7)]
    ax1.errorbar(range(4, 7), [eff_tot[i] for i in range(4, 7)], yerr=yerr_lm, fmt='s',
                label='large MM')
    yerr = [1/np.sqrt(pool_tot[i]) for i in range(7, Nchb+1)]

    # 48x48 cm^2 RPWELL chambers
    if Nchb == 11:
        dict_RPWELL = {7: 'ASU 61', 8: 'ASU 60', 9: 'ASU 51', 10: 'ASU 57', 11: 'ASU 52'}
    elif Nchb in [7,8]:
        dict_RPWELL = {7: 'ASU 61', 8: 'ASU 51'}

    # saving dato to csv
    if not os.path.isdir('./results'):
        os.mkdir("./results")

    ax1.errorbar(range(7, Nchb+1), [eff_tot[i] for i in range(7, Nchb+1)], yerr=yerr, fmt='^',
                label='RPWELL')

    ax1.set_xlabel('layer number')
    ax1.set_ylim(0.4,1)
    ax1.legend()
    ax1.set_ylabel('MIP detection efficiency')

    # Add a table at the bottom of the axes
    # ax2.axis('tight')
    rowlabels = ['Small MM 2',
                 'Small MM 3',
                 'Large MM 1',
                 'Large MM 2',
                 'Large MM 3']
    ax2.axis('off')
    cell_text = []
    for i in range(7, Nchb+1):
        rowlabels.append('RPWELL {}'.format(dict_RPWELL[i]))

    cell_text.append(['{:.2}'.format(eff_tot[i] * 100) for i in range(2, Nchb+1)])
    cell_text.append(['{:.2}'.format(1/np.sqrt(pool_tot[i]) * 100) for i in range(2, Nchb+1)])
    cell_text.append([pool_tot[i] for i in range(2, Nchb+1)])
    table = ax2.table(cellText=np.array(cell_text).T,
                        colLabels=('efficiency[%]', 'error[%]', 'tested tracks'),
                        rowLabels= rowlabels,
                        loc='center')
    table.set_fontsize(18)
    table.auto_set_column_width([0, 1, 2])
    table.scale(10, 3)

    # saving figure
    if not os.path.isdir('./figures'):
        os.mkdir("./figures")

    fig.savefig('figures/mip_eff_{}layers_{}_{}.png'.format(Nchb, mode, datetime.now().strftime('%Y%m%d_%H%M')))

    # plt.show()
    del(fig)
    print('| ASU | Effieicny [%] | number of tested tracks|')
    print('|------|----------------|-----------------------|')
    for i in range(2, 4):
        print("|layer {} | {:.2}+\-{:.2} \t| {} |".format(i,
                                                  eff_tot[i]*100,
                                                  yerr_sm[i-2]*100,
                                                  pool_tot[i]))
    for i in range(4, 7):
        print("|layer {} | {:.2}+\-{:.2} \t| {} |".format(i,
                                                  eff_tot[i]*100,
                                                  yerr_lm[i-4]*100,
                                                  pool_tot[i]))
    for i in range(7,Nchb+1):
        print("|{} | {:.2}+\-{:.2} \t| {} |".format(dict_RPWELL[i],
                                                  eff_tot[i]*100,
                                                  yerr[i-7]*100,
                                                  pool_tot[i]))

def plot_mult(mult_tot, mode, Nchb=8):
    """Generates plots for efficiency estimation and save them.
    Parameters:
    -----------
        mult_tot :  a dictionary of the pad multiplicity value of the different
                    chambers. The chamber numbers serve as keys.
        mode     :  'dt', 'calo', or 'MC' for trigger-time correlated hits, CaloEvents, or simulation, respectively.
        Nchb :      the number of chambers in the setup. Default: 8.

    Return:
    -------
        None

    """

    fig = plt.figure(figsize=(10, 7))
    # 16x16 cm^2 uM chambers
    plt.scatter(range(2, 4), [mult_tot[i] for i in range(2, 4)],
                label='Small MM')

    # 48x48 cm^2 uM chambers
    plt.scatter(range(4, 7), [mult_tot[i] for i in range(4, 7)],
                label='large MM')

    # 48x48 cm^2 RPWELL chambers
    if Nchb == 11:
        dict_RPWELL = {7: 'ASU 61', 8: 'ASU 60', 9: 'ASU 51', 10: 'ASU 57', 11: 'ASU 52'}
    elif Nchb == 8:
        dict_RPWELL = {7: 'ASU 61', 8: 'ASU 51'}

    plt.scatter(range(7, Nchb+1), [mult_tot[i] for i in range(7, Nchb+1)],
                label='RPWELL')

    plt.xlabel('layer number')
    plt.ylim(0.4,1)
    plt.legend()
    plt.ylabel('MIP detection efficiency')

    # saving figure
    if not os.path.isdir('./figures'):
        os.mkdir("./figures")

    fig.savefig('figures/mip_mult_{}layers_{}_{}.png'.format(Nchb, mode, datetime.now().strftime('%Y%m%d_%H%M')))

    # plt.show()
    del(fig)
    print('| \# of Chamber | Multiplicity |')
    print('|------|----------------|-----------------------|')
    for i in range(7,Nchb+1):
        print("|{} | {:.2} \t|".format(i, mult_tot[i]))

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument('-m', '--mode', help='analysis mode (dt, calo or MC)',
                        default='dt', choices=['dt', 'calo', 'MC'])
    parser.add_argument('-n', '--Nchb', help='number of layers in the setup',
                        default='8', choices=['8', '11', '7'])
    parser.add_argument('-e', '--energy', type=int, help='beam energy in GeV.',
                        choices=range(2,7))
    parser.add_argument('-s', '--sigma', type=float, help="width of the residual's fluctuation.",
                        default=1.)
    parser.add_argument('--mipeff', help='mip efficiency')
    parser.add_argument('-p', '--particle', help="The type of the beam's particle.",
                        default='pions', choices=('pions', 'electrons'))

    args = parser.parse_args()
    mode = args.mode
    if mode == 'MC':
        energy = '{}GeV'.format(args.energy)
        particle = args.particle
        mipeff = args.mipeff
        # if mipeff not in list(mc.MAX_FILES[particle][energy].keys()):
        #     raise ValueError('Out of range for specific energy: mipeff.')
    Nchb = int(args.Nchb)

    warnings.simplefilter('ignore', np.RankWarning)

    if mode not in ['calo', 'dt', 'MC']:
        raise ValueError('Not a valid type of analysis mode.')

    if mode == 'MC':
        eff_tot, pool_tot, mult_tot = eff_MC(energy, mipeff, args.sigma, particle=particle, Nchb=Nchb)
        plot_eff(eff_tot, pool_tot, 'MC', Nchb)
        plot_mult(mult_tot, 'MC', Nchb)
    else:
        eff_tot, pool_tot, mult_tot = eff(mode, Nchb=Nchb)
        plot_eff(eff_tot, pool_tot, mode, Nchb)
        plot_mult(mult_tot, mode, Nchb)


"""Defines the Instrument class"""
#***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License.  You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************
import collections as _collections
import numpy as _np
import warnings as _warnings

from ..tools import matrixtools as _mt
from .label import Label as _Label

#from . import labeldicts as _ld
from . import modelmember as _gm
from . import operation as _op
from . import spamvec as _sv

def convert(instrument, toType, basis, extra=None):
    """
    Convert intrument to a new type of parameterization, potentially
    creating a new object.  Raises ValueError for invalid conversions.

    Parameters
    ----------
    instrument : Instrument
        Instrument to convert

    toType : {"full","TP","static","static unitary"}
        The type of parameterizaton to convert to.  See
        :method:`Model.set_all_parameterizations` for more details.

    basis : {'std', 'gm', 'pp', 'qt'} or Basis object
        The basis for `povm`.  Allowed values are Matrix-unit (std),
        Gell-Mann (gm), Pauli-product (pp), and Qutrit (qt)
        (or a custom basis object).

    extra : object, optional
        Additional information for conversion.

    Returns
    -------
    Instrument
       The converted instrument, usually a distinct
       object from the object passed as input.
    """

    if toType == "TP":
        if isinstance(instrument, TPInstrument):
            return instrument
        else:
            return TPInstrument(list(instrument.items()))
    elif toType in ("full", "static", "static unitary"):
        gate_list = [(k, _op.convert(g, toType, basis)) for k, g in instrument.items()]
        return Instrument(gate_list)
    else:
        raise ValueError("Cannot convert an instrument to type %s" % toType)

class Instrument(_gm.ModelMember, _collections.OrderedDict):
    """
    Meant to correspond to a quantum instrument in theory, this class
    generalizes that notion slightly to include a collection of gates that may
    or may not have all of the properties associated by a mathematical quantum
    instrument.
    """

    def __init__(self, op_matrices, items=[]):
        """
        Creates a new Instrument object.

        Parameters
        ----------
        op_matrices : dict of LinearOperator objects
            A dict (or list of key,value pairs) of the gates.
        """
        self._readonly = False  # until init is done
        if len(items) > 0:
            assert(op_matrices is None), "`items` was given when op_matrices != None"

        dim = None
        evotype = None

        if op_matrices is not None:
            if isinstance(op_matrices, dict):
                matrix_list = [(k, v) for k, v in op_matrices.items()]  # gives definite ordering
            elif isinstance(op_matrices, list):
                matrix_list = op_matrices  # assume it's is already an ordered (key,value) list
            else:
                raise ValueError("Invalid `op_matrices` arg of type %s" % type(op_matrices))

            items = []
            for k, v in matrix_list:
                gate = v if isinstance(v, _op.LinearOperator) else \
                    _op.FullDenseOp(v)

                if evotype is None: evotype = gate._evotype
                else: assert(evotype == gate._evotype), \
                    "All instrument gates must have the same evolution type"

                if dim is None: dim = gate.dim
                assert(dim == gate.dim), "All instrument gates must have the same dimension!"
                items.append((k, gate))

        if evotype is None:
            evotype = "densitymx"  # default (if no instrument gates)

        _collections.OrderedDict.__init__(self, items)
        _gm.ModelMember.__init__(self, dim, evotype)
        self._paramvec = self._build_paramvec()
        self._readonly = True

    #No good way to update Instrument on the fly yet...
    #def _update_paramvec(self, modified_obj=None):
    #    """Updates self._paramvec after a member of this Model is modified"""
    #    for obj in self.values():
    #        assert(obj.gpindices is self), "Cannot add/adjust parameter vector!"
    #
    #    #update parameters changed by modified_obj
    #    self._paramvec[modified_obj.gpindices] = modified_obj.to_vector()
    #
    #    #re-initialze any members that also depend on the updated parameters
    #    modified_indices = set(modified_obj.gpindices_as_array())
    #    for obj in self.values()
    #        if obj is modified_obj: continue
    #        if modified_indices.intersection(obj.gpindices_as_array()):
    #            obj.from_vector(self._paramvec[obj.gpindices])

    def _build_paramvec(self):
        """ Resizes self._paramvec and updates gpindices & parent members as needed,
            and will initialize new elements of _paramvec, but does NOT change
            existing elements of _paramvec (use _update_paramvec for this)"""
        v = _np.empty(0, 'd'); off = 0

        # Step 2: add parameters that don't exist yet
        for obj in self.values():
            if obj.gpindices is None or obj.parent is not self:
                #Assume all parameters of obj are new independent parameters
                v = _np.insert(v, off, obj.to_vector())
                num_new_params = obj.allocate_gpindices(off, self)
                off += num_new_params
            else:
                inds = obj.gpindices_as_array()
                M = max(inds) if len(inds) > 0 else -1; L = len(v)
                if M >= L:
                    #Some indices specified by obj are absent, and must be created.
                    w = obj.to_vector()
                    v = _np.concatenate((v, _np.empty(M + 1 - L, 'd')), axis=0)  # [v.resize(M+1) doesn't work]
                    for ii, i in enumerate(inds):
                        if i >= L: v[i] = w[ii]
                off = M + 1
        return v

    def __setitem__(self, key, value):
        if self._readonly: raise ValueError("Cannot alter Instrument elements")
        else: return _collections.OrderedDict.__setitem__(self, key, value)

    def __reduce__(self):
        """ Needed for OrderedDict-derived classes (to set dict items) """
        #need to *not* pickle parent, as __reduce__ bypasses ModelMember.__getstate__
        dict_to_pickle = self.__dict__.copy()
        dict_to_pickle['_parent'] = None

        #Note: must *copy* elements for pickling/copying
        return (Instrument, (None, [(key, gate.copy()) for key, gate in self.items()]), dict_to_pickle)

    def __pygsti_reduce__(self):
        return self.__reduce__()

    def simplify_operations(self, prefix=""):
        """
        Returns a dictionary of gates that belong to the Instrument's parent
        `Model` - that is, whose `gpindices` are set to all or a subset of
        this instruments's gpindices.  These are used internally within
        computations involving the parent `Model`.

        Parameters
        ----------
        prefix : str
            A string, usually identitying this instrument, which may be used
            to prefix the simplified gate keys.

        Returns
        -------
        OrderedDict of Gates
        """
        #Create a "simplified" (Model-referencing) set of element gates
        simplified = _collections.OrderedDict()
        if isinstance(prefix, _Label):  # Deal with case when prefix isn't just a string
            for k, g in self.items():
                comp = g.copy()
                comp.set_gpindices(_gm._compose_gpindices(self.gpindices,
                                                          g.gpindices), self.parent)
                simplified[_Label(prefix.name + "_" + k, prefix.sslbls)] = comp
        else:
            if prefix: prefix += "_"
            for k, g in self.items():
                comp = g.copy()
                comp.set_gpindices(_gm._compose_gpindices(self.gpindices,
                                                          g.gpindices), self.parent)
                simplified[prefix + k] = comp
        return simplified

    def num_elements(self):
        """
        Return the number of total gate elements in this instrument.
        This is in general different from the number of *parameters*,
        which are the number of free variables used to generate all of
        the matrix *elements*.

        Returns
        -------
        int
        """
        return sum([g.size for g in self.values()])

    def num_params(self):
        """
        Get the number of independent parameters which specify this Instrument.

        Returns
        -------
        int
           the number of independent parameters.
        """
        return len(self._paramvec)

    def to_vector(self):
        """
        Extract a vector of the underlying gate parameters from this Instrument.

        Returns
        -------
        numpy array
            a 1D numpy array with length == num_params().
        """
        return self._paramvec

    def from_vector(self, v, close=False, nodirty=False):
        """
        Initialize the Instrument using a vector of its parameters.

        Parameters
        ----------
        v : numpy array
            The 1D vector of gate parameters.  Length
            must == num_params().

        Returns
        -------
        None
        """
        assert(len(v) == self.num_params())
        for gate in self.values():
            gate.from_vector(v[gate.gpindices], close, nodirty)
        self._paramvec = v

    def transform(self, S):
        """
        Update Instrument element matrix G with inv(S) * G * S.

        Parameters
        ----------
        S : GaugeGroupElement
            A gauge group element which specifies the "S" matrix
            (and it's inverse) used in the above similarity transform.
        """
        #Note: since each Mi is a linear function of MT and the Di, we can just
        # transform the MT and Di (self.param_ops) and re-init the elements.
        for gate in self.values():
            gate.transform(S)
            self._paramvec[gate.gpindices] = gate.to_vector()
        self.dirty = True

    def depolarize(self, amount):
        """
        Depolarize this Instrument by the given `amount`.

        Parameters
        ----------
        amount : float or tuple
            The amount to depolarize by.  If a tuple, it must have length
            equal to one less than the dimension of the gate. All but the
            first element of each spam vector (often corresponding to the
            identity element) are multiplied by `amount` (if a float) or
            the corresponding `amount[i]` (if a tuple).

        Returns
        -------
        None
        """
        #Note: since each Mi is a linear function of MT and the Di, we can just
        # depolarize the MT and Di (self.param_ops) and re-init the elements.
        for gate in self.values():
            gate.depolarize(amount)
            self._paramvec[gate.gpindices] = gate.to_vector()
        self.dirty = True

    def rotate(self, amount, mxBasis='gm'):
        """
        Rotate this instrument by the given `amount`.

        Parameters
        ----------
        amount : tuple of floats, optional
            Specifies the rotation "coefficients" along each of the non-identity
            Pauli-product axes.  The gate's matrix `G` is composed with a
            rotation operation `R`  (so `G` -> `dot(R, G)` ) where `R` is the
            unitary superoperator corresponding to the unitary operator
            `U = exp( sum_k( i * rotate[k] / 2.0 * Pauli_k ) )`.  Here `Pauli_k`
            ranges over all of the non-identity un-normalized Pauli operators.

        mxBasis : {'std', 'gm', 'pp', 'qt'} or Basis object
            The source and destination basis, respectively.  Allowed
            values are Matrix-unit (std), Gell-Mann (gm), Pauli-product (pp),
            and Qutrit (qt) (or a custom basis object).

        Returns
        -------
        None
        """
        #Note: since each Mi is a linear function of MT and the Di, we can just
        # rotate the MT and Di (self.param_ops) and re-init the elements.
        for gate in self.values():
            gate.rotate(amount, mxBasis)
            self._paramvec[gate.gpindices] = gate.to_vector()
        self.dirty = True

    def acton(self, state):
        """
        Act with this instrument upon `state`

        Parameters
        ----------
        state : SPAMVec
            The state to act on

        Returns
        -------
        OrderedDict
            A dictionary whose keys are the outcome labels (strings)
            and whose values are `(prob, normalized_state)` tuples
            giving the probability of seeing the given outcome and
            the resulting state that would be obtained if and when
            that outcome is observed.
        """
        # Note: no 'stabilizer' or 'statevec' support yet (how renormalize sframe or how does state vec work?)
        assert(self._evotype in ('densitymx',)), \
            "acton(...) cannot be used with the %s evolution type!" % self._evotype
        assert(state._evotype == self._evotype), "Evolution type mismatch: %s != %s" % (self._evotype, state._evotype)

        staterep = state._rep
        outcome_probs_and_states = _collections.OrderedDict()
        for lbl, element in self.items():
            output_rep = element._rep.acton(staterep)
            output_unnormalized_state = output_rep.todense()
            prob = output_unnormalized_state[0] * state.dim**0.25
            output_normalized_state = output_unnormalized_state / prob  # so [0]th == 1/state_dim**0.25
            outcome_probs_and_states[lbl] = (prob, _sv.StaticSPAMVec(output_normalized_state, self._evotype, 'prep'))

        return outcome_probs_and_states

    def __str__(self):
        s = "Instrument with elements:\n"
        for lbl, element in self.items():
            s += "%s:\n%s\n" % (lbl, _mt.mx_to_string(element.base, width=4, prec=2))
        return s

class TPInstrument(_gm.ModelMember, _collections.OrderedDict):
    """
    A trace-preservng quantum instrument which is a collection of gates whose
    sum is a trace-preserving map.  The instrument's elements may or may not
    have all of the properties associated by a mathematical quantum instrument.

    If M1,M2,...Mn are the elements of the instrument, then we parameterize
    1. MT = (M1+M2+...Mn) as a TPParmeterizedGate
    2. Di = Mi - MT for i = 1..(n-1) as FullyParameterizedGates

    So to recover M1...Mn we compute:
    Mi = Di + MT for i = 1...(n-1)
       = -(n-2)*MT-sum(Di) = -(n-2)*MT-[(MT-Mi)-n*MT] for i == (n-1)
    """
    #Scratch:
    #    Scratch
    # M1+M2+M3+M4  MT
    #   -M2-M3-M4  M1-MT
    #-M1   -M3-M4  M2-MT
    #-M1-M2   -M4  M3-MT
    #
    #(M1-MT) + (M2-MT) + (M3-MT) = (MT-M4) - 3*MT = -2*MT-M4
    # M4 = -(sum(Di)+(4-2=2)*MT) = -(sum(all)+(4-3=1)*MT)
    #n=2 case: (M1-MT) = (MT-M2)-MT = -M2, so M2 = -sum(Di)

    def __init__(self, op_matrices, items=[]):
        """
        Creates a new Instrument object.

        Parameters
        ----------
        gates : dict of numpy arrays
            A dict (or list of key,value pairs) of the operation matrices whose sum
            must be a trace-preserving (TP) map.
        """
        self._readonly = False  # until init is done
        if len(items) > 0:
            assert(op_matrices is None), "`items` was given when op_matrices != None"

        dim = None
        self.param_ops = []  # first element is TP sum (MT), following
        #elements are fully-param'd (Mi-Mt) for i=0...n-2

        #Note: when un-pickling using items arg, these members will
        # remain the above values, but *will* be set when state dict is copied
        # in (so unpickling works as desired)

        if op_matrices is not None:
            if isinstance(op_matrices, dict):
                matrix_list = [(k, v) for k, v in op_matrices.items()]  # gives definite ordering
            elif isinstance(op_matrices, list):
                matrix_list = op_matrices  # assume it's is already an ordered (key,value) list
            else:
                raise ValueError("Invalid `op_matrices` arg of type %s" % type(op_matrices))

            # Create gate objects that are used to parameterize this instrument
            MT = _op.TPDenseOp(sum([v for k, v in matrix_list]))
            MT.set_gpindices(slice(0, MT.num_params()), self)
            self.param_ops.append(MT)

            dim = MT.dim; off = MT.num_params()
            for k, v in matrix_list[:-1]:
                Di = _op.FullDenseOp(v - MT)
                Di.set_gpindices(slice(off, off + Di.num_params()), self)
                assert(Di.dim == dim)
                self.param_ops.append(Di); off += Di.num_params()

            #Create a TPInstrumentOp for each operation matrix
            # Note: TPInstrumentOp sets it's own parent and gpindices
            items = [(k, _op.TPInstrumentOp(self.param_ops, i))
                     for i, (k, v) in enumerate(matrix_list)]

            #DEBUG
            #print("POST INIT PARAM GATES:")
            #for i,v in enumerate(self.param_ops):
            #    print(i,":\n",v)
            #
            #print("POST INIT ITEMS:")
            #for k,v in items:
            #    print(k,":\n",v)

        _collections.OrderedDict.__init__(self, items)
        _gm.ModelMember.__init__(self, dim, "densitymx")
        self._readonly = True

    def __setitem__(self, key, value):
        if self._readonly: raise ValueError("Cannot alter POVM elements")
        else: return _collections.OrderedDict.__setitem__(self, key, value)

    def __reduce__(self):
        """ Needed for OrderedDict-derived classes (to set dict items) """
        #Don't pickle TPInstrumentGates b/c they'll each pickle the same
        # param_ops and I don't this will unpickle correctly.  Instead, just
        # strip the numpy array from each element and call __init__ again when
        # unpickling:
        op_matrices = [(lbl, _np.asarray(val)) for lbl, val in self.items()]
        return (TPInstrument, (op_matrices, []), {'_gpindices': self._gpindices})

    def __pygsti_reduce__(self):
        return self.__reduce__()

    def simplify_operations(self, prefix=""):
        """
        Returns a dictionary of gates that belong to the Instrument's parent
        `Model` - that is, whose `gpindices` are set to all or a subset of
        this instruments's gpindices.  These are used internally within
        computations involving the parent `Model`.

        Parameters
        ----------
        prefix : str
            A string, usually identitying this instrument, which may be used
            to prefix the simplified gate keys.

        Returns
        -------
        OrderedDict of Gates
        """
        #Create a "simplified" (Model-referencing) set of param gates
        param_simplified = []
        for g in self.param_ops:
            comp = g.copy()
            comp.set_gpindices(_gm._compose_gpindices(self.gpindices,
                                                      g.gpindices), self.parent)
            param_simplified.append(comp)

        # Create "simplified" elements, which infer their parent and
        # gpindices from the set of "param-gates" they're constructed with.
        if isinstance(prefix, _Label):  # Deal with case when prefix isn't just a string
            simplified = _collections.OrderedDict(
                [(_Label(prefix.name + "_" + k, prefix.sslbls), _op.TPInstrumentOp(param_simplified, i))
                 for i, k in enumerate(self.keys())])
        else:
            if prefix: prefix += "_"
            simplified = _collections.OrderedDict(
                [(prefix + k, _op.TPInstrumentOp(param_simplified, i))
                 for i, k in enumerate(self.keys())])
        return simplified

    def num_elements(self):
        """
        Return the number of total gate elements in this instrument.
        This is in general different from the number of *parameters*,
        which are the number of free variables used to generate all of
        the matrix *elements*.

        Returns
        -------
        int
        """
        return sum([g.size for g in self.values()])

    def num_params(self):
        """
        Get the number of independent parameters which specify this Instrument.

        Returns
        -------
        int
           the number of independent parameters.
        """
        return sum([g.num_params() for g in self.param_ops])

    def to_vector(self):
        """
        Extract a vector of the underlying gate parameters from this Instrument.

        Returns
        -------
        numpy array
            a 1D numpy array with length == num_params().
        """
        v = _np.empty(self.num_params(), 'd')
        for gate in self.param_ops:
            v[gate.gpindices] = gate.to_vector()
        return v

    def from_vector(self, v, close=False, nodirty=False):
        """
        Initialize the Instrument using a vector of its parameters.

        Parameters
        ----------
        v : numpy array
            The 1D vector of gate parameters.  Length
            must == num_params().

        Returns
        -------
        None
        """
        for gate in self.param_ops:
            gate.from_vector(v[gate.gpindices], close, nodirty)
        for instGate in self.values():
            instGate._construct_matrix()

    def transform(self, S):
        """
        Update Instrument element matrix G with inv(S) * G * S.

        Parameters
        ----------
        S : GaugeGroupElement
            A gauge group element which specifies the "S" matrix
            (and it's inverse) used in the above similarity transform.
        """
        #Note: since each Mi is a linear function of MT and the Di, we can just
        # transform the MT and Di (self.param_ops) and re-init the elements.
        for gate in self.param_ops:
            gate.transform(S)

        for element in self.values():
            element._construct_matrix()  # construct from param gates
        self.dirty = True

    def depolarize(self, amount):
        """
        Depolarize this Instrument by the given `amount`.

        Parameters
        ----------
        amount : float or tuple
            The amount to depolarize by.  If a tuple, it must have length
            equal to one less than the dimension of the gate. All but the
            first element of each spam vector (often corresponding to the
            identity element) are multiplied by `amount` (if a float) or
            the corresponding `amount[i]` (if a tuple).

        Returns
        -------
        None
        """
        #Note: since each Mi is a linear function of MT and the Di, we can just
        # depolarize the MT and Di (self.param_ops) and re-init the elements.
        for gate in self.param_ops:
            gate.depolarize(amount)

        for element in self.values():
            element._construct_matrix()  # construct from param gates
        self.dirty = True

    def rotate(self, amount, mxBasis='gm'):
        """
        Rotate this instrument by the given `amount`.

        Parameters
        ----------
        amount : tuple of floats, optional
            Specifies the rotation "coefficients" along each of the non-identity
            Pauli-product axes.  The gate's matrix `G` is composed with a
            rotation operation `R`  (so `G` -> `dot(R, G)` ) where `R` is the
            unitary superoperator corresponding to the unitary operator
            `U = exp( sum_k( i * rotate[k] / 2.0 * Pauli_k ) )`.  Here `Pauli_k`
            ranges over all of the non-identity un-normalized Pauli operators.

        mxBasis : {'std', 'gm', 'pp', 'qt'} or Basis object
            The source and destination basis, respectively.  Allowed
            values are Matrix-unit (std), Gell-Mann (gm), Pauli-product (pp),
            and Qutrit (qt) (or a custom basis object).

        Returns
        -------
        None
        """
        #Note: since each Mi is a linear function of MT and the Di, we can just
        # rotate the MT and Di (self.param_ops) and re-init the elements.
        for gate in self.param_ops:
            gate.rotate(amount, mxBasis)

        for element in self.values():
            element._construct_matrix()  # construct from param gates
        self.dirty = True

    def __str__(self):
        s = "TPInstrument with elements:\n"
        for lbl, element in self.items():
            s += "%s:\n%s\n" % (lbl, _mt.mx_to_string(element.base, width=4, prec=2))
        return s


#!/usr/bin/env python3
import copy
import logging
import os
import sys
import tempfile

import numpy as np
import offsb.op.geometry
import simtk.unit
import simtk.unit as unit
import offsb.treedi
import offsb.treedi.tree
from rdkit import Chem
from simtk import openmm

import openforcefield as oFF
from openforcefield.topology import Molecule
from openforcefield.typing.engines.smirnoff.parameters import \
    UnassignedProperTorsionParameterException

from .. import qcarchive as qca
from .. import rdutil
from ..tools import const

logger = logging.getLogger(__name__)
oFFlogger = logging.getLogger("openforcefield")
oFFlogger.setLevel(logging.ERROR)

def load_geometric_opt_trj_xyz(fnm, dims=3):
    N = 0
    xyz = []
    with open(fnm, "r") as fd:
        N = int(fd.readline().split()[0])
        total_N = 1
        for line in fd:
            total_N += 1
        n_frames = total_N // (N + 2)
    ene_in_au = []
    sym = list([""] * N)
    with open(fnm, "r") as fd:
        xyz = np.empty((n_frames, N, dims))
        for j in range(n_frames):
            N = int(fd.readline().split()[0])
            energy = float(fd.readline().split()[-1])
            ene_in_au.append(energy)
            for i in range(N):
                line = fd.readline().split()
                pos = line[1:]
                sym[i] = line[0]
                xyz[j][i][:] = list(map(float, pos))

    # xyz in angstrom
    return sym, xyz, np.array(ene_in_au)

class OpenMMEnergy(offsb.treedi.tree.PartitionTree):
    """Creates an openMM system of each entry
    Stores the total energy of molecules
    """

    def __init__(self, filename, source_tree, name):
        super().__init__(source_tree, name)
        import logging

        logger = logging.getLogger()
        level = logger.getEffectiveLevel()
        logger.setLevel(level=logging.ERROR)
        from pkg_resources import iter_entry_points

        self.filename = filename

        self._select = "Molecule"
        self.minimize = False
        self.constrain = False
        self.geometric = True

        # Seems to bottleneck at 4, so mask here to avoid memory waste
        self.processes = 4
        search_pth = list(
            iter_entry_points(group="openforcefield.smirnoff_forcefield_directory")
        )
        abspth = os.path.join(".", filename)
        print("Searching", ".")
        found = False
        if os.path.exists(abspth):
            self.abs_path = abspth
            print("Found")
            found = True
        if not found:
            for entry_point in search_pth:
                pth = entry_point.load()()[0]
                abspth = os.path.join(pth, filename)
                abspth = abspth if abspth.endswith("offxml") else abspth + ".offxml"
                print("Searching", abspth)
                if os.path.exists(abspth):
                    self.abs_path = abspth
                    print("Found")
                    break
                raise Exception("Forcefield could not be found")
        self.forcefield = oFF.typing.engines.smirnoff.ForceField(
            self.abs_path, disable_version_check=True, allow_cosmetic_attributes=True
        )
        logger.setLevel(level=level)
        print("My db is", self.db)

    def to_pickle(self, db=True, name=None):
        super().to_pickle(db=db, name=name)

    def _apply_initialize(self, targets):
        pass

    def _apply_finalize(self, targets):
        pass

    def _unpack_result(self, ret):
        self.db.update(ret)

    def _generate_apply_kwargs(self, i, target, kwargs=None):

        # labels = self.source.db[target.payload]["data"]
        entry = self.source.source.db[target.payload]["data"]

        out_str = ""

        if kwargs is None:
            kwargs = {}

        mol = kwargs.get("mol")

        if mol is None:
            smi = entry.attributes["canonical_isomeric_explicit_hydrogen_mapped_smiles"]
            if "initial_molecule" in entry.dict():
                qcid = entry.dict()["initial_molecule"]
            elif "initial_molecules" in entry.dict():
                qcid = entry.dict()["initial_molecules"]
            else:
                out_str += "{:d} initial mol was empty: {:s}".format(i, str(qcid))
                return {"error": out_str}

            if isinstance(qcid, set):
                qcid = list(qcid)
            if isinstance(qcid, list):
                qcid = str(qcid[0])

            qcmolid = "QCM-" + qcid

            if qcmolid not in self.source.source.db:
                out_str += "{:d} initial mol was empty: {:s}".format(i, str(qcmolid))
                return {"error": out_str}

            if "data" in self.source.source.db.get(qcmolid):
                qcmol = self.source.source.db.get(qcmolid).get("data")
            else:
                out_str += "{:d} initial mol was empty: {:s}".format(i, str(qcmolid))
                return {"error": out_str}

            mol = offsb.rdutil.mol.rdmol_from_smiles_and_qcmol(smi, qcmol)
            kwargs["mol"] = mol
            kwargs["qcmolid"] = qcmolid

    def apply_single(self, i, target, **kwargs):
        def unmap(xyz, map_idx):
            inv = [(map_idx[i] - 1) for i in range(len(xyz))]
            return xyz[inv]

        def remap(xyz, map_idx):
            remap_idx = {v - 1: k for k, v in map_idx.items()}
            inv = [remap_idx[i] for i in range(len(xyz))]
            return xyz[inv]

        # mol = kwargs["mol"]
        # qcmolid = kwargs["qcmolid"]

        ret_str = []
        # if n < 192:
        #    continue
        entry_node = next(self.source.node_iter_to_root(target, select="Entry"))
        attrs = self.source.db[entry_node.payload]["data"].dict()["attributes"]
        # attrs = self.source.db[target.payload]['data'].dict()['attributes']

        # Since we now have spec nodes, need one under the spec
        # Just grab the first.. only need metadata (e.g. mapping)
        try:
            mol_node = next(
                self.source.node_iter_depth_first(entry_node, select="Molecule")
            )
        except Exception:
            ret_str.append(
                "ERROR: This job has no molecules: {} {}\n".format("", target.payload)
            )
            return {target: ret_str}
            # Likely StopIteration (no molecules; all failed)

        qcmol = self.source.db[mol_node.payload]["data"]

        #         if isinstance(qcmolid, list):
        #             qcmolid = qcmolid[0]
        #         qcmolid = "QCM-" + str(qcmolid)
        #         breakpoint()
        #         if qcmolid in self.source.db:
        #             qcmol = self.source.db.get(qcmolid).get("data")
        #         else:
        #             ret_str.append(
        #                 "ERROR: Molcule ID {:s} not in the local database\n".format(qcmolid)
        #             )
        #             return {target: ret_str}

        smiles_pattern = attrs.get("canonical_isomeric_explicit_hydrogen_mapped_smiles")

        mol = rdutil.mol.build_from_smiles(smiles_pattern)
        map_idx = rdutil.mol.atom_map(mol)
        ret = rdutil.mol.embed_qcmol_3d(mol, qcmol)
        if ret < 0:
            ret_str.append(
                "ERROR: Could not generate a conformation in RDKit. {} {}\n".format(
                    "", target.payload
                )
            )
            qca.qcmol_to_xyz(
                qcmol,
                fnm="mol." + "" + "." + target.index + ".rdconfgenfail.xyz",
                comment="" + " rdconfgen fail " + target.payload,
            )
            return {target: ret_str}
        # conf = mol.GetConformer()
        # ids = AllChem.EmbedMultipleConfs( mol, numConfs=1)
        # try:
        #    conf = mol.GetConformer(ids[0])
        # except IndexError:
        #    print("ERROR: Could not generate a conformation in RDKit.", qcmolid, target.payload)
        #    qca.qcmol_to_xyz( qcmol,
        #        fnm="mol."+qcmolid+"."+target.payload+".rdconfgenfail.xyz", comment=qcmolid + " rdconfgen fail " + target.payload)
        #    continue
        # conf = mol.GetConformer(ids[0])

        use_min_mol_for_charge = True
        if use_min_mol_for_charge:
            minidx = None
            minene = None
            minmol = None
            try:
                for opt in self.source.node_iter_depth_first(
                    target, select="Optimization"
                ):
                    opt_rec = self.source.db[opt.payload]["data"]
                    status = opt_rec.status[:]
                    if status != "COMPLETE":
                        ret_str.append("This opt is not complete.. skipping..\n")
                        continue
                    allene = opt_rec.energies
                    if allene is None:
                        ret_str.append(
                            "ERROR: No energies. {} {}\n".format("", target.payload)
                        )
                        # ret_str.append(str(e) + "\n")
                        qca.qcmol_to_xyz(
                            qcmol,
                            fnm="mol." + "" + "." + target.index + ".noenefail.xyz",
                            comment="" + " noene fail " + target.payload,
                        )
                        return {target: ret_str}
                    ene = allene[-1]
                    if minene is None:
                        minene = ene
                        minidx = opt.children[-1]
                    elif ene < minene:
                        minene = ene
                        minidx = opt.children[-1]
            except (TypeError, IndexError) as e:
                # ene is not a list above if TypeError
                # IndexError it is []
                ret_str.append("ERROR: No energies. {} {}\n".format("", target.payload))
                ret_str.append(str(e) + "\n")
                qca.qcmol_to_xyz(
                    qcmol,
                    fnm="mol." + "" + "." + target.index + ".noenefail.xyz",
                    comment="" + " noene fail " + target.payload,
                )
                return {target: ret_str}

            if minidx == None:
                ret_str.append("EMPTY. SKIPPING\n")
                return {target: ret_str}
            grad_node = self.source.node_index[minidx]
            mol_node = self.source.node_index[grad_node.children[0]]
            qcmol = self.source.db[mol_node.payload]["data"]

            ret_str.append("min mol is {} ene is {} au\n".format(mol_node, minene))

        xyz = qcmol.geometry
        sym = qcmol.symbols
        # for i, a in enumerate(mol.GetAtoms()):
        #    conf.SetAtomPosition(i, xyz[ map_idx[ i] - 1] * const.bohr2angstrom)

        # Chem.rdmolops.AssignStereochemistryFrom3D( mol, ids[0], replaceExistingTags=True)
        # Chem.rdmolops.AssignStereochemistryFrom3D( mol, ids[0], replaceExistingTags=True)
        Chem.rdmolops.AssignAtomChiralTagsFromStructure(
            mol, -1, replaceExistingTags=True
        )

        with tempfile.TemporaryDirectory() as tmpdir:
            cwd = os.getcwd()
            os.chdir(tmpdir)
            mmol = oFF.topology.Molecule.from_rdkit(mol, allow_undefined_stereo=True)
            try:
                top = oFF.topology.Topology().from_molecules(mmol)
            except AssertionError:
                ret_str.append(
                    "ERROR: Could not setup molecule in oFF. {} {}\n".format(
                        "", target.payload
                    )
                )
                qca.qcmol_to_xyz(
                    qcmol,
                    fnm="mol." + "" + "." + target.index + ".offmolfail.xyz",
                    comment="" + " oFF fail " + target.payload,
                )
                # pdb.set_trace()
                return {target: ret_str}
            try:
                mmol.compute_partial_charges_am1bcc()

            except Exception as e:
                os.chdir(cwd)
                ret_str.append(
                    "ERROR: Could not compute partial charge. {} {}\n".format(
                        "", target.payload
                    )
                )
                ret_str.append(str(e.__traceback__.__repr__()) + "\n")
                ret_str.append(str(e) + "\n")
                qca.qcmol_to_xyz(
                    qcmol,
                    fnm="mol." + "" + "." + target.index + ".chrgfail.xyz",
                    comment="" + " charge fail " + target.payload,
                )
                # pdb.set_trace()
                return {target: ret_str}
            os.chdir(cwd)

        gen_MM_charge = [mmol.partial_charges]

        fail = True
        nodes = list(self.source.node_iter_depth_first(target, select="Molecule"))
        order = np.arange(len(nodes))
        vals = []
        for mol_node in nodes:
            val = tuple(
                [
                    c.payload[2]
                    for c in self.source.node_iter_to_root(
                        mol_node, select="Constraint"
                    )
                ]
            )
            if len(val) > 0:
                vals.append(val)
        if len(vals) > 0:
            vals = np.array(vals)
            order = np.lexsort(vals.T)
            nodes_in_order = [nodes[i] for i in order]
        else:
            nodes_in_order = nodes

        ret_obj = {}

        for mol_node in nodes_in_order:
            fail = True
            qcmol = self.source.db[mol_node.payload]["data"]
            xyz = qcmol.geometry
            xyz = xyz * const.bohr2angstrom

            # this will take the xyz from qcmol and put them in the order of
            # the cmiles, based on the mapped indices (1-index)
            # sends qcmol to mmmol
            xyz = unmap(xyz, map_idx)

            constraints = None
            if self.constrain:
                # constraints = [ [c.payload[1], c.payload[2]] for c in \
                # self.source.node_iter_to_root( mol_node, select="Constraint")]
                # print("constraints are", constraints)
                swap_map = {v - 1: k for k, v in map_idx.items()}
                constraints = [
                    [[swap_map[i] for i in c.payload[1]], c.payload[2]]
                    for c in self.source.node_iter_to_root(
                        mol_node, select="Constraint"
                    )
                ]
                # qca.qcmol_to_xyz( {"symbols": unmap(qcmol['symbols'],  map_idx), "geometry":xyz*const.angstrom2bohr},
                #     fnm="mol."+qcmolid+"."+target.index+".map.xyz", comment=qcmolid + " map fail " + str(constraints))
                # constraints = [ [unmap(np.array(c.payload[1]), map_idx), c.payload[2]] for c in \
                # self.source.node_iter_to_root( mol_node, select="Constraint")]

                # print("angle before OpenMM is ", self.calculate_dihedral(xyz, *constraints[0][0])/unit.degrees)
                # print("constraints are", constraints)
            try:
                total_ene, pos = self.calc_mm_energy(
                    top, xyz, charge=gen_MM_charge, constraints=constraints
                )
                fail = False
            except UnassignedProperTorsionParameterException as e:
                ret_str.append(
                    "ERROR: oFF could not assign torsion for {} {}\n".format(
                        qcmolid, target.payload
                    )
                )
                ret_str.append(str(e) + "\n")
                break
            except Exception as e:
                ret_str.append(
                    "ERROR: oFF exception for {} {}\n".format(qcmolid, target.payload)
                )
                ret_str.append(str(e.__traceback__.__repr__()) + "\n")
                ret_str.append(str(e) + "\n")
                fail = True
                break

            constraints = [
                c.payload
                for c in self.source.node_iter_to_root(mol_node, select="Constraint")
            ]
            # print("    {} {} {}\n".format( mol_node, constraints, total_ene), end="")
            ret_str.append("    {} {} {}\n".format(mol_node, constraints, total_ene))

            pl = {}
            if self.minimize and pos is not None:
                pl = qcmol.copy()
                pl.geometry = remap(np.array(pos) * const.angstrom2bohr, map_idx)

            pl["energy"] = total_ene

            ret_obj.update({mol_node.payload: {"data": pl}})
            # self.db.__setitem__( )

        if fail:
            qca.qcmol_to_xyz(
                qcmol,
                fnm="mol." + qcmolid + "." + target.index + ".labelfail.xyz",
                comment=qcmolid + " label fail " + target.payload,
            )
        return {target: ret_str, "return": ret_obj}

    def apply(self, targets=None):
        if targets is None:
            targets = list(self.source.iter_entry())
        elif not hasattr(targets, "__iter__"):
            targets = [targets]

        # expand if a generator
        targets = list(targets)

        n_targets = len(targets)

        if self.processes is None or self.processes > 1:
            import concurrent.futures

            exe = concurrent.futures.ProcessPoolExecutor(max_workers=self.processes)

            work = [
                exe.submit(OpenMMEnergy.apply_single, self, n, target)
                for n, target in enumerate(targets, 1)
            ]
            for n, future in enumerate(concurrent.futures.as_completed(work), 1):
                if future.done:
                    try:
                        val = future.result()
                    except RuntimeError:
                        print("RUNTIME ERROR; race condition??")
                if val is None:
                    print("data is None?!?")
                    continue
                for tgt, ret in val.items():
                    if tgt == "return":
                        self.db.update(ret)
                    else:
                        print(n, "/", n_targets, tgt)
                        for line in ret:
                            print(line, end="")

            exe.shutdown()

        # # This version seems to be more stable, but all of the results must
        # # finish before iterating

        # exe = Pool(processes=self.processes)
        # work = [ exe.apply_async( OpenMMEnergy.apply_single, ( self, n, target) ) for n, target in enumerate(targets, 1) ]
        # out = [result.get() for result in work if result is not None]
        # for n, val in enumerate(out, 1):
        #     for tgt, ret in val.items():
        #         if tgt == "return":
        #             self.db.update(ret)
        #         else:
        #             print( n,"/", n_targets, tgt)
        #             for line in ret:
        #                 print(line, end="")
        # exe.close()

        # for i, val in enumerate( out):
        #     for tgt, ret in val.items():
        #         print( n,"/", n_targets, tgt)
        #         print( ret)

        # single process mode; does not launch any new processes
        if self.processes == 1:
            for n, target in enumerate(targets, 1):
                val = self.apply_single(n, target)
                for tgt, ret in val.items():
                    if tgt == "return":
                        self.db.update(ret)
                    else:
                        print(n, "/", n_targets, tgt)
                        for line in ret:
                            print(line, end="")
        ########################################

    def calculate_dihedral(self, coords, idx1, idx2, idx3, idx4):
        # from:
        # https://math.stackexchange.com/questions/47059/how-do-i-calculate-a-dihedral-angle-given-cartesian-coordinates
        # https://stackoverflow.com/questions/20305272/dihedral-torsion-angle-from-four-points-in-cartesian-coordinates-in-python
        b1 = coords[idx1] - coords[idx2]
        b1 /= np.sqrt(np.dot(b1, b1))
        b2 = coords[idx2] - coords[idx3]
        b2 /= np.sqrt(np.dot(b2, b2))
        b3 = coords[idx3] - coords[idx4]
        b3 /= np.sqrt(np.dot(b3, b3))
        # print (b1, b2, b3)
        n1 = np.cross(b1, b2)
        n2 = np.cross(b2, b3)
        m1 = np.cross(n1, b2)
        # print ( n1,n2,m1 )
        x = np.dot(n2, n1)
        y = np.dot(n2, m1)
        theta = np.arctan2(y, x) * unit.radian
        # print ( theta, np.degrees ( theta ) * unit.degree )
        return theta

    def add_harmonic_dihedral_restraints(
        self, system, coords, k, hw, indices_to_restrain, debug=False
    ):
        print("Adding harmonic dihedral restraints: ")
        if len(indices_to_restrain) != 4:
            print(
                "WARNING: Must be exactly 4 atoms specified for a dihedral restraint.  Skipping ..."
            )
            return system
        theta0 = self.calculate_dihedral(coords, *indices_to_restrain)  # in radians
        print(theta0)
        expr = "0.5*k*max(0,( min(d_theta,2*pi-d_theta) - hw ))^2;"
        expr += "d_theta = abs(theta - theta0);"
        expr += f"pi = {np.pi:.10f}"
        force = openmm.CustomTorsionForce(expr)
        force.addGlobalParameter("k", k)
        force.addGlobalParameter("hw", hw)
        force.addPerTorsionParameter("theta0")
        force.addTorsion(*indices_to_restrain, [theta0])
        system.addForce(force)
        return system

    def mm_potential(
        self, forcefield, top, xyz, charge=False, constraints=None, use_geometric=False
    ):

        if isinstance(charge, bool):
            if charge:
                system = forcefield.create_openmm_system(top)
            else:
                mols = [
                    Molecule(mol.reference_molecule) for mol in top.topology_molecules
                ]
                for i, _ in enumerate(mols):
                    mols[i].partial_charges = simtk.unit.Quantity(
                        np.zeros(mols[i].n_atoms), simtk.unit.elementary_charge
                    )
                system = forcefield.create_openmm_system(
                    top, charge_from_molecules=mols
                )
        else:
            mols = [Molecule(mol.reference_molecule) for mol in top.topology_molecules]
            for i, _ in enumerate(mols):
                mols[i].partial_charges = charge[i]
            system = forcefield.create_openmm_system(top, charge_from_molecules=mols)

        integrator = openmm.VerletIntegrator(0.1 * simtk.unit.femtoseconds)
        sim = openmm.app.simulation.Simulation(top, system, integrator)
        sim.context.setPositions(xyz * const.angstrom2nm)

        # params = forcefield.label_molecules(top)[0]
        # for p in params:
        #     print(p)
        #     print(dict(params[p]))

        getPositions = self.minimize
        # system = sim.context.getSystem()
        # forces = system.getForces()
        use_geometric = self.geometric
        if True and use_geometric and self.minimize:

            import geometric.optimize as geoopt

            with tempfile.TemporaryDirectory() as tmpdir:
                cwd = os.getcwd()
                os.chdir(tmpdir)
                system = forcefield.create_openmm_system(
                    top, charge_from_molecules=mols
                )
                xml = openmm.XmlSerializer.serialize(system)
                open("system.xml", "w").write(xml)
                with open("mol.pdb", "w") as fid:
                    openmm.app.pdbfile.PDBFile.writeFile(top.to_openmm(), xyz, fid)

                args = {
                    "input": "system.xml",
                    "pdb": "mol.pdb",
                    "openmm": True,
                    "maxiter": 2000,
                }
                if constraints is not None:

                    with open("constraints", "w") as fid:
                        fid.write("$set\n")
                        for constr in constraints:
                            ids = [x + 1 for x in constr[0]]
                            val = constr[1]
                            out_str = "dihedral {:d} {:d} {:d} {:d} {:12.5f}\n".format(
                                *ids, val
                            )
                            # out_str = "dihedral {:d} {:d} {:d} {:d}\n".format(*ids)
                            fid.write(out_str)
                            # print(out_str, end="")
                    args["constraints"] = "constraints"

                success = False
                with tempfile.TemporaryFile("w") as null:
                    sys.stderr = null
                    try:
                        geoopt.run_optimizer(**args)
                        success = True
                    except Exception as e:
                        print("This optimization failed!")
                        print(e)
                    sys.stderr = sys.__stderr__
                    if success:
                        _, traj, ene = load_geometric_opt_trj_xyz("system_optim.xyz")
                os.chdir(cwd)

            if not success:
                return None, None

            return ene[-1] * const.hartree2kcalmol, traj[-1]

        elif True and self.minimize and constraints is not None:
            # print("Using frozen torsion...")
            for constr in constraints:
                # assume torsion...
                # print("Searching torsion", constr)
                for idx in constr[0]:
                    system.setParticleMass(idx, 0.0)
            if self.minimize:
                sim.minimizeEnergy()
            state = sim.context.getState(getEnergy=True, getPositions=getPositions)
            energy = state.getPotentialEnergy().in_units_of(
                simtk.unit.kilocalories_per_mole
            )
            pos = None
            if self.minimize:
                pos = state.getPositions(asNumpy=True) / simtk.unit.angstroms

        elif False and self.minimize and constraints is not None:
            print("Using custom torsion force...")
            cforce = openmm.CustomTorsionForce(
                "0.5*k*max(0.0,min(dtheta, 2*pi-dtheta))^2; dtheta = abs(theta-theta0); pi = 3.1415926535"
            )
            cforce.addPerTorsionParameter("k")
            cforce.addPerTorsionParameter("theta0")
            for constr in constraints:
                # assume torsion...
                # print("Searching torsion", constr)
                if len(constr[0]) == 4:
                    print("adding torsion", constr[0])
                    cforce.addTorsion(*constr[0], [300.0, constr[1] * np.pi / 180])

            # print("adding force...")
            iforce = system.addForce(cforce)

            integrator = openmm.VerletIntegrator(0.1 * simtk.unit.femtoseconds)
            sim = openmm.app.simulation.Simulation(top, system, integrator)
            sim.context.setPositions(xyz * const.angstrom2nm)
            getPositions = self.minimize
            # system = sim.context.getSystem()
            # forces = system.getForces()
            state = sim.context.getState(getEnergy=True, getPositions=getPositions)
            energy = state.getPotentialEnergy().in_units_of(
                simtk.unit.kilocalories_per_mole
            )
            print("energy is", energy)
            if self.minimize:
                sim.minimizeEnergy()
            state = sim.context.getState(getEnergy=True, getPositions=getPositions)
            energy = state.getPotentialEnergy().in_units_of(
                simtk.unit.kilocalories_per_mole
            )
            pos = state.getPositions(asNumpy=True) / simtk.unit.nanometers
            print("energy is", energy)

            system2 = forcefield.create_openmm_system(top, charge_from_molecules=mols)
            integrator = openmm.VerletIntegrator(0.1 * simtk.unit.femtoseconds)
            sim = openmm.app.simulation.Simulation(top, system2, integrator)
            sim.context.setPositions(pos)
            state = sim.context.getState(getEnergy=True, getPositions=getPositions)
            energy = state.getPotentialEnergy().in_units_of(
                simtk.unit.kilocalories_per_mole
            )
            print("energy is\n", energy)
            pos = None
            if self.minimize:
                pos = state.getPositions(asNumpy=True) / simtk.unit.angstroms
            return energy, pos

        elif self.minimize and constraints is not None:
            forces = system.getForces()
            print("Using iterative torsion force...")
            state = sim.context.getState(getEnergy=True, getPositions=getPositions)
            pos = state.getPositions(asNumpy=True) / simtk.unit.angstroms
            print(
                "angle before OpenMM opt is ",
                self.calculate_dihedral(pos, *constraints[0][0]) / unit.degrees,
            )
            dx = list([np.inf] * len(constraints))
            dk = 1.0 * simtk.unit.kilojoules_per_mole
            dE = 0.0
            energy = 0.0
            eps = 0.001
            maxsteps = 100000
            nstep = 0
            pha = None
            k = None
            per = 0
            ref_k = None
            ref_pha = None
            while nstep < maxsteps and all([abs(dxi) > eps for dxi in dx]):
                for constr in constraints:
                    # assume torsion...
                    # print("Searching torsion", constr)
                    if len(constr[0]) == 4:

                        torsions = [
                            f for f in forces if type(f) == openmm.PeriodicTorsionForce
                        ][0]
                        N = torsions.getNumTorsions()
                        found = False
                        # print("Found N existing torsions:", N)

                        if pha is None:
                            pha = constr[1] * np.pi / 180
                        for idx in range(N):
                            (
                                p1,
                                p2,
                                p3,
                                p4,
                                per,
                                phai,
                                ki,
                            ) = torsions.getTorsionParameters(idx)
                            if k is None:
                                k = ki
                                ref_k = ki
                                ref_pha = phai
                            target = tuple(constr[0])
                            if target[-1] < target[0]:
                                target = tuple(target[::-1])
                            # print("Retreived", p1,p2,p3,p4,per,pha,ki, target)
                            if tuple([p1, p2, p3, p4]) == target:
                                # print("retreived", idx, p1,p2,p3,p4,per,phai,ki)
                                # print("Found torsion", target, constr[1])
                                # print("mass: ", [system.getParticleMass(x) for x in constr[0]])
                                k += dk
                                # per = 1
                                # print("setting", idx, p1,p2,p3,p4,per,pha,k)
                                torsions.setTorsionParameters(
                                    idx, p1, p2, p3, p4, per, pha, k
                                )
                                found = True
                                # print("Set to", idx, p1, p2, p3, p4, per, pha, k)
                        if not found:
                            raise Exception("The constrained torsion was not found!")
                            # p1,p2,p3,p4 = constr[0]
                            # pha = constr[1]*np.pi/180.
                            # k = 9999999.0
                            # per = 1
                            # constr[1],  torsions.addTorsion(p1,p2,p3,p4,per,pha,k)
                        torsions.updateParametersInContext(sim.context)
                sim.minimizeEnergy(tolerance=1e-5, maxIterations=1)
                state = sim.context.getState(getEnergy=True, getPositions=getPositions)
                energy2 = state.getPotentialEnergy() / simtk.unit.kilocalories_per_mole
                pos = state.getPositions(asNumpy=True) / simtk.unit.angstroms

                # now check to see if the constraints moved
                for i, constr in enumerate(constraints):
                    # assume torsion...
                    # print("Calculating torsion", constr)
                    if len(constr[0]) == 4:
                        angle = offsb.op.geometry.TorsionOperation.measure_praxeolitic_single(
                            {"geometry": pos}, list(target)
                        )
                        dx[i] = abs(angle - constr[1])
                        dE = energy2 - energy
                        energy = energy2
                        # print("Dx", i, "is ", dx[i], "k=", k, target, constr[1], angle,"dE=",dE, "pha=", pha * 180/np.pi)
                        angle2 = (
                            self.calculate_dihedral(pos, *list(target)) / unit.degrees
                        )
                        if (
                            nstep == 0
                            or nstep == maxsteps
                            or all([abs(dxi) < eps for dxi in dx])
                        ):
                            print(
                                "C {:d} dt= {:14.8e} k= {:12.3f} {:s} t_0= {:6.2f} t= {:6.2f} t2= {:6.2f} dE= {:16.8e} pha= {:14.8e}".format(
                                    i,
                                    dx[i],
                                    k / unit.kilojoules_per_mole,
                                    str(target),
                                    constr[1],
                                    angle,
                                    angle2,
                                    dE,
                                    pha * 180 / np.pi,
                                )
                            )

                        # angle = constr[1] *np.pi/180
                        # if angle < constr[1] and angle < 0:
                        #     pha *= 1. - (.01 / np.log10(nstep+2)
                        #     # print("decrease1")
                        # elif angle < constr[1] and angle > 0:
                        #     pha *= 1. + (.1 /np.log10(nstep+2))
                        #     # print("increase1")
                        # elif angle > constr[1] and angle > 0:
                        #     pha *= 1. - (.1 / np.log10(nstep+2)
                        #     # print("decrease2")
                        # elif angle > constr[1] and angle < 0:
                        #     pha *= 1. + (.1 /np.log10(nstep+2))
                        #     # print("increase2")
                        pha -= (0.75) * ((angle - constr[1]) * np.pi / 180)
                        if pha < -2 * np.pi:
                            pha += 2 * np.pi
                        elif pha > 2 * np.pi:
                            pha -= 2 * np.pi

                nstep += 1

            for constr in constraints:
                # assume torsion...
                # print("Searching torsion", constr)
                if len(constr[0]) == 4:

                    torsions = [
                        f for f in forces if type(f) == openmm.PeriodicTorsionForce
                    ][0]
                    N = torsions.getNumTorsions()
                    found = False
                    # print("Found N existing torsions:", N)

                    for idx in range(N):
                        p1, p2, p3, p4, per, phai, ki = torsions.getTorsionParameters(
                            idx
                        )
                        target = tuple(constr[0])
                        if target[-1] < target[0]:
                            target = tuple(target[::-1])
                        if tuple([p1, p2, p3, p4]) == target:
                            # print("Found torsion", target, constr[1])
                            # per = 1
                            torsions.setTorsionParameters(
                                idx, p1, p2, p3, p4, per, ref_pha, ref_k
                            )

            state = sim.context.getState(getEnergy=True, getPositions=getPositions)
            energy = state.getPotentialEnergy().in_units_of(
                simtk.unit.kilocalories_per_mole
            )
            # print("energy (remove)", energy)
            torsions.updateParametersInContext(sim.context)
            state = sim.context.getState(getEnergy=True, getPositions=getPositions)
            energy = state.getPotentialEnergy().in_units_of(
                simtk.unit.kilocalories_per_mole
            )
            # print("energy (remove)", energy)
            pos = None
            if self.minimize:
                pos = state.getPositions(asNumpy=True) / simtk.unit.angstroms

        else:
            if self.minimize:
                sim.minimizeEnergy()
            state = sim.context.getState(getEnergy=True, getPositions=getPositions)
            energy = state.getPotentialEnergy().in_units_of(
                simtk.unit.kilocalories_per_mole
            )
            pos = None
            if self.minimize:
                pos = state.getPositions(asNumpy=True) / simtk.unit.angstroms
        return energy, pos

    def calc_mm_energy(self, top, xyz, component=None, charge=False, constraints=None):
        forcefield = self.forcefield
        if component is None:
            ene, pos = self.mm_potential(
                forcefield, top, xyz, charge=charge, constraints=constraints
            )
            return ene, pos

        modff = copy.deepcopy(forcefield)

        force = modff.get_parameter_handler(component)
        for term in force.parameters:
            if component == "vdW":
                term.epsilon *= 0.0
            if component in ["Bonds", "Angles"]:
                term.k *= 0.0
            if component in ["ProperTorsions", "ImproperTorsions"]:
                for i, _ in enumerate(term.k):
                    term.k[i] *= 0.0

        ene, pos = self.mm_potential(
            modff, top, xyz, charge=charge, constraints=constraints
        )
        return ene, pos

#    def calc_vdw_direct(xyz, labels):
#        """ doesn't work yet"""
#        na = len(xyz)
#        atoms = range(na)
#        ene = 0.0
#        r = distance.cdist(xyz,xyz)
#        for i in atoms:
#            ii = (i,)
#            for j in atoms:
#                if(j >= i):
#                    break
#                jj = (j,)
#                eps = np.sqrt(labels[ii].epsilon * labels[jj].epsilon) / labels[jj].epsilon.unit
#                rmin = (labels[ii].rmin_half + labels[jj].rmin_half)/2.0
#                rmin = rmin / rmin.unit
#                rij = r[i,j]
#                a = rmin/rij
#                a = a**6
#                ene += eps * (a**2 - 2*a)
#                #print(i,j,"r", rij, "ene", ene, "rmin", rmin, "eps", eps, "a", a)
#        return ene * simtk.unit.kilocalorie_per_mole

#    def apply( self, targets=None):
#        gen_MM_charge = True
#        try:
#            total_ene = self.calc_mm_energy(forcefield, top, xyz, charge=gen_MM_charge)
#        except Exception as e:
#            log.write(str(e) + '\n')
#            success = False
#            break
# log.write("  Conformation energy {:4d}/{:4d}\n".format(qcmol_i+1,n_qcmol))
# log.flush()
# if(get_qm_energy):
#    ene_str = "{:13.8f} a.u.".format(ene[qcmol_i])
#    log.write("    QM {:18s}= {:s} \n".format("Energy", ene_str))
#    log.flush()

# ene_name_str = "Energy" if gen_MM_charge else "EnergyNoElec"
# log.write("    MM {:18s}= {:10.5f} {:s}\n".format(ene_name_str,
#    (total_ene/total_ene.unit ),
#    str(total_ene.unit)))
# log.flush()

# sim.context.setPositions(xyz * angstrom2nm)
# state = sim.context.getState(getEnergy = True, getPositions=True)
# sim_ene = state.getPotentialEnergy()
# reporter.report(sim, state)
# log.write("    MM {:18s}= {:9.4f} {:s}\n\n".format("App Energy",
#    (sim_ene.value_in_unit(simtk.unit.kilocalorie_per_mole)),
#    str(simtk.unit.kilocalorie_per_mole)))
# log.flush()
# sim.minimizeEnergy()
# state = sim.context.getState(getEnergy = True, getPositions=True)
# sim_min_ene = state.getPotentialEnergy()
# reporter.report(sim, state)
# log.write("    MM {:18s}= {:9.4f} {:s}\n\n".format("App Min Energy",
#    (sim_min_ene.value_in_unit(simtk.unit.kilocalorie_per_mole)),
#    str(simtk.unit.kilocalorie_per_mole)))

# log.flush()
#        ene_sum = total_ene
# if('direct_vdw' not in mol_data['energy']['oFF']):
#    mol_data['energy']['oFF']['direct_vdw'] = []
# mol_data['energy']['oFF']['direct_vdw'].append(calc_vdw_direct(xyz, labels['vdW']))
# print(mol_data['energy']['oFF']['direct_vdw'])

#        if('epot' not in mol_data['energy']['oFF']):
#            mol_data['energy']['oFF']['epot'] = []
#        mol_data["energy"]['oFF']['epot'].append(total_ene)
#        for component in ['vdW' ] + valence_params:
#            energy_sans_component = calc_mm_energy(forcefield, top, xyz, component=component, charge=gen_MM_charge)
#            energy_component = total_ene - energy_sans_component
#
#            if(component not in mol_data["energy"]['oFF']):
#                mol_data["energy"]['oFF'][component] = []
#            mol_data["energy"]['oFF'][component].append(energy_component)
#            ene_sum -= energy_component
#            log.write("    MM {:18s}= {:10.5f} {:s}\n".format(component,
#                (energy_component/energy_component.unit ),
#                str(energy_component.unit)))
#            log.flush()
#        if(gen_MM_charge):
#            if('Electrostatics' not in mol_data['energy']['oFF']):
#                mol_data['energy']['oFF']['Electrostatics'] = []
#            mol_data['energy']['oFF']['Electrostatics'].append(ene_sum)
#            log.write("    MM {:18s}= {:10.5f} {:s}\n\n".format("Electrostatics",
#                (ene_sum/ene_sum.unit ),
#                str(ene_sum.unit)))
#
#        #log.write("    MM {:18s}= {:10.5f} {:s}\n\n".format("Direct vdW",
#        #    (mol_data['energy']['oFF']['direct_vdw'][qcmol_i] / mol_data['energy']['oFF']['direct_vdw'][qcmol_i].unit ),
#        #    str(mol_data['energy']['oFF']['direct_vdw'][qcmol_i].unit)))
#        log.flush()

#    if(success == False):
#        continue
#    log.write("\n  Conformation energy summary\n")
#    if(len(ene) > 1):
#        if(get_qm_energy):
#            ene_str = "{:13.8f} +- {:13.8f} a.u.".format(np.mean(ene), np.std(ene))
#            log.write("    QM {:18s}= {:s} \n".format("Energy", ene_str))
#
#        log.write("    MM {:18s}= {:10.5f} +- {:10.5f} {:s}\n\n".format(ene_name_str,
#            np.mean([i/i.unit for i in mol_data["energy"]['oFF']['epot']]),
#            np.std([i/i.unit for i in mol_data["energy"]['oFF']['epot']]),
#            str(mol_data["energy"]['oFF']['epot'][0].unit)))
#        components_to_print = ['vdW' ] + valence_params
#        if(gen_MM_charge):
#            components_to_print += ['Electrostatics']
#        for component in components_to_print:
#            log.write("    MM {:18s}= {:10.5f} +- {:10.5f} {:s}\n".format(component,
#                np.mean([i/i.unit for i in mol_data["energy"]['oFF'][component]]),
#                np.std([i/i.unit for i in mol_data["energy"]['oFF'][component]]),
#                str(mol_data["energy"]['oFF'][component][0].unit)))


import os
import random
import sys
import torch
from skimage.morphology import remove_small_objects,binary_opening
import numpy as np

def average(list):
    s = 0
    for item in list:
        s += item
    return s/len(list)

def sum(list):
    s = 0
    for item in list:
        s += item
    return s

def analysis(x,y):
    '''
    对输入的两个四维张量[B,1,H,W]进行逐图的DSC、PPV、Sensitivity计算
    其中x表示网络输出的预测值
    y表示实际的预想结果mask
    返回为一个batch中DSC、PPV、Sen的平均值及batch大小
    '''

    x = torch.from_numpy(x.astype(np.uint8)).cuda()
    # y = torch.from_numpy(y.astype(np.uint8))
    x=x.type(dtype=torch.uint8)
    y=y.type(dtype=torch.uint8)#保证类型为uint8
    DSC=[]
    PPV=[]
    Sen=[]
    if x.shape==y.shape:
        batch=x.shape[0]
        for i in range(batch):#按第一个维度分开

            tmp = torch.eq(x[i],y[i])

            tp=int(torch.sum(torch.mul(x[i]==1,tmp==1))) #真阳性
            fp=int(torch.sum(torch.mul(x[i]==1,tmp==0))) #假阳性
            fn=int(torch.sum(torch.mul(x[i]==0,tmp==0))) #假阴性

            try:
                DSC.append(2*tp/(fp+2*tp+fn))
            except:
                DSC.append(0)
            try:
                PPV.append(tp/(tp+fp))
            except:
                PPV.append(0)
            try:
                Sen.append(tp/(tp+fn))
            except:
                Sen.append(0)

    else:
        sys.stderr.write('Analysis input dimension error')

    DSC = sum(DSC)/batch
    PPV = sum(PPV)/batch
    Sen = sum(Sen)/batch
    return DSC, PPV, Sen, batch

def post_process(img,min_size=100):
    '''
    图像后处理过程
    包括开运算和去除过小体素
    返回uint16格式numpy二值数组
    '''
    img = img.cpu()
    img = img.numpy().astype(np.bool)
    b,c,w,h = img.shape
    if c==1:
        for i in range(b):
            img_tmp = img[i,0,:,:]
            img_tmp = binary_opening(img_tmp)
            remove_small_objects(img_tmp, min_size=min_size, in_place=True)
            img_tmp = ~remove_small_objects(~img_tmp, min_size=min_size)
            img[i,0,:,:] = img_tmp

    return img.astype(np.uint16)