__doc__ = """Symplectic time steppers and concepts for integrating the kinematic and dynamic equations of rod-like objects.  """
import numpy as np

from elastica.timestepper._stepper_interface import (
    _TimeStepper,
    _LinearExponentialIntegratorMixin,
)

"""
Developer Note
--------------

For the reasons why we define Mixin classes here, the developer
is referred to the same section on `explicit_steppers.py`.
"""

class _SystemInstanceStepperMixin:
    def do_step(self, System, time: np.float64, dt: np.float64):
        for (
            kin_prefactor,
            kin_step,
            dyn_prefactor,
            dyn_step,
        ) in self._steps_and_prefactors[:-1]:
            prefac = kin_prefactor(self, dt)
            time = kin_step(self, System, time, prefac)
            prefac = dyn_prefactor(self, dt)
            time = dyn_step(self, System, time, prefac)

        # Peel the last kinematic step and prefactor alone
        last_kin_prefactor = self._steps_and_prefactors[-1][0]
        last_kin_step = self._steps_and_prefactors[-1][1]

        prefac = last_kin_prefactor(self, dt)
        time = last_kin_step(self, System, time, prefac)
        return time

class _SystemCollectionStepperMixin:
    """
    Symplectic stepper mixin collection class
    """

    def do_step(self, SystemCollection, time: np.float64, dt: np.float64):
        """
        Function for doing symplectic stepper over the user defined rods (system).

        Parameters
        ----------
        SystemCollection: rod object
        time: float
        dt: float

        Returns
        -------

        """
        for (
            kin_prefactor,
            kin_step,
            dyn_prefactor,
            dyn_step,
        ) in self._steps_and_prefactors[:-1]:
            prefac = kin_prefactor(self, dt)
            for system in SystemCollection[:-1]:
                _ = kin_step(self, system, time, prefac)
            time = kin_step(self, SystemCollection[-1], time, prefac)

            # TODO: remove below lines and try to find a better call option to compute internal forces and torques
            # We need internal forces and torques because they are used by interaction module.
            update_internal_forces_torques = self._update_internal_forces_torques[
                0
            ]  # here 0 because you have one item in list
            for system in SystemCollection[:-1]:
                _ = update_internal_forces_torques(self, system, time)
            time = update_internal_forces_torques(self, SystemCollection[-1], time)

            # BoCos, External forces, controls etc.
            SystemCollection.synchronize(time)
            # TODO: remove below line, it should be some other function synchronizeBC
            SystemCollection.synchronizeBC(time)
            prefac = dyn_prefactor(self, dt)
            for system in SystemCollection[:-1]:
                _ = dyn_step(self, system, time, prefac)
            time = dyn_step(self, SystemCollection[-1], time, prefac)

            # TODO: remove below line, it should be some other function synchronizeBC
            SystemCollection.synchronizeBC(time)

        # Peel the last kinematic step and prefactor alone
        last_kin_prefactor = self._steps_and_prefactors[-1][0]
        last_kin_step = self._steps_and_prefactors[-1][1]

        prefac = last_kin_prefactor(self, dt)
        for system in SystemCollection[:-1]:
            _ = last_kin_step(self, system, time, prefac)
        time = last_kin_step(self, SystemCollection[-1], time, prefac)

        # Call back function, will call the user defined call back functions and store data
        SystemCollection.callBack(time, int(time / dt))

        return time

class SymplecticStepper(_TimeStepper):
    """
    Symplectic stepper constructor.

    Attributes
    ----------
    _steps: list
        List containing methods of symplectic time stepper.
    _prefactors: list
        List containing prefactors of symplectic time stepper.
    _update_internal_forces_torques: list
        List containing methods for computing internal forces and torques.
    """

    def __init__(self, cls=None):
        super(SymplecticStepper, self).__init__()
        take_methods_from = self if cls is None else cls()
        # Let the total number of steps for the Symplectic method
        # be (2*n + 1) (for time-symmetry). What we do is collect
        # the first n + 1 entries down in _steps and _prefac below, and then
        # reverse and append it to itself.
        self._steps = [
            v
            for (k, v) in take_methods_from.__class__.__dict__.items()
            if k.endswith("step")
        ]
        # Prefac here is necessary because the linear-exponential integrator
        # needs only the prefactor and not the dt.
        self._prefactors = [
            v
            for (k, v) in take_methods_from.__class__.__dict__.items()
            if k.endswith("prefactor")
        ]

        # We are getting function named as _update_internal_forces_torques from dictionary,
        # it turns a list.
        self._update_internal_forces_torques = [
            v
            for (k, v) in take_methods_from.__class__.__dict__.items()
            if k.endswith("forces_torques")
        ]

        def mirror(in_list):
            """ Mirrors an input list ignoring the last element
            If steps = [A, B, C]
            then this call makes it [A, B, C, B, A]

            Parameters
            ----------
            in_list : input list to be mirrored, modified in-place

            Returns
            -------

            """
            #  syntax is very ugly
            in_list.extend(in_list[-2::-1])

        mirror(self._steps)
        mirror(self._prefactors)

        assert len(self._steps) == len(
            self._prefactors
        ), "Size mismatch in the number of steps and prefactors provided for a Symplectic Stepper!"

        self._kinematic_steps = self._steps[::2]
        self._dynamic_steps = self._steps[1::2]
        self._kinematic_prefactors = self._prefactors[::2]
        self._dynamic_prefactors = self._prefactors[1::2]

        # Avoid this check for MockClasses
        if len(self._kinematic_steps) > 0:
            assert (
                len(self._kinematic_steps) == len(self._dynamic_steps) + 1
            ), "Size mismatch in the number of kinematic and dynamic steps provided for a Symplectic Stepper!"
            assert (
                len(self._kinematic_prefactors) == len(self._dynamic_prefactors) + 1
            ), "Size mismatch in the number of kinematic and dynamic prefactors provided for a Symplectic Stepper!"

        from itertools import zip_longest

        self._steps_and_prefactors = tuple(
            zip_longest(
                self._kinematic_prefactors,
                self._kinematic_steps,
                self._dynamic_prefactors,
                self._dynamic_steps,
            )
        )

    @property
    def n_stages(self):
        return len(self._steps_and_prefactors)

class PositionVerlet(SymplecticStepper):
    """
    Position Verlet symplectic time stepper class, which
    includes methods for second-order position Verlet.
    """

    def __init__(self):
        super(PositionVerlet, self).__init__()

    def _first_kinematic_prefactor(self, dt):
        return 0.5 * dt

    def _first_kinematic_step(self, System, time: np.float64, prefac: np.float64):
        System.kinematic_states += prefac * System.kinematic_rates(time, prefac)
        return time + prefac

    def _first_dynamic_prefactor(self, dt):
        return dt

    def _first_dynamic_step(self, System, time: np.float64, prefac: np.float64):
        System.dynamic_states += prefac * System.dynamic_rates(
            time, prefac
        )  # TODO : Why should we pass dt into System again?
        return time

    # TODO: find a better place for this or a better call option. We need to compute internal forces and torques before external because interaction uses it!
    def _update_internal_forces_torques(self, System, time: np.float64):
        System.update_internal_forces_and_torques(time)
        return time

    # Note : we don't need the second half of the calls as it simply forwards
    # to its equivalent first half. This is taken care in the base class

    # def _second_kinematic_step(self, System, time: np.float64, dt: np.float64):
    #     return self._first_kinematic_step(System, time, dt)

class PEFRL(SymplecticStepper):
    """
    Position Extended Forest-Ruth Like Algorithm of
    I.M. Omelyan, I.M. Mryglod, and R. Folk, Computer Physics Communications 146, 188 (2002),
    http://arxiv.org/abs/cond-mat/0110585
    """

    # xi and chi are confusing, but be careful!
    ξ = np.float64(0.1786178958448091e0)  # ξ
    λ = -np.float64(0.2123418310626054e0)  # λ
    χ = -np.float64(0.6626458266981849e-1)  # χ

    # Pre-calculate other coefficients
    lambda_dash_coeff = 0.5 * (1.0 - 2.0 * λ)
    xi_chi_dash_coeff = 1.0 - 2.0 * (ξ + χ)

    def __init__(self):
        super(PEFRL, self).__init__()

    def _first_kinematic_prefactor(self, dt):
        return self.ξ * dt

    def _first_kinematic_step(self, System, time: np.float64, prefac: np.float64):
        System.kinematic_states += prefac * System.kinematic_rates(time, prefac)
        return time + prefac

    def _first_dynamic_prefactor(self, dt):
        return self.lambda_dash_coeff * dt

    def _first_dynamic_step(self, System, time: np.float64, prefac: np.float64):
        System.dynamic_states += prefac * System.dynamic_rates(time, prefac)
        return time

    def _second_kinematic_prefactor(self, dt):
        return self.χ * dt

    def _second_kinematic_step(self, System, time: np.float64, prefac: np.float64):
        System.kinematic_states += prefac * System.kinematic_rates(time, prefac)
        return time + prefac

    def _second_dynamic_prefactor(self, dt):
        return self.λ * dt

    def _second_dynamic_step(self, System, time: np.float64, prefac: np.float64):
        System.dynamic_states += prefac * System.dynamic_rates(time, prefac)
        return time

    def _third_kinematic_prefactor(self, dt):
        return self.xi_chi_dash_coeff * dt

    def _third_kinematic_step(self, System, time: np.float64, prefac: np.float64):
        # Need to fill in
        System.kinematic_states += prefac * System.kinematic_rates(time, prefac)
        return time + prefac

    # TODO: find a better place for this or a better call option. We need to compute internal forces and torques before external because interaction uses it!
    def _update_internal_forces_torques(self, System, time: np.float64):
        System.update_internal_forces_and_torques(time)
        return time

    # Note : we don't need the second half of the calls as it simply forwards
    # to its equivalent first half. This is taken care in the base class

    # def _third_dynamic_step(self, System, time: np.float64, dt: np.float64):
    #     return self._second_dynamic_step(System, time, dt)
    #
    # def _fourth_kinematic_step(self, System, time: np.float64, dt: np.float64):
    #     return self._second_kinematic_step(System, time, dt)
    #
    # def _fourth_dynamic_step(self, System, time: np.float64, dt: np.float64):
    #     return self._first_dynamic_step(System, time, dt)
    #
    # def _fifth_kinematic_step(self, System, time: np.float64, dt: np.float64):
    #     return self._first_kinematic_step(System, time, dt)
    #     # return time + dt # To avoid numerical precision errors

class SymplecticLinearExponentialIntegrator(
    _LinearExponentialIntegratorMixin, SymplecticStepper
):
    def __init__(self):
        _LinearExponentialIntegratorMixin.__init__(self)
        SymplecticStepper.__init__(self, _LinearExponentialIntegratorMixin)


""" @package forcebalance.lipid Matching of lipid bulk properties.  Under development.

author Lee-Ping Wang
@date 04/2012
"""
from __future__ import division
from __future__ import print_function

from builtins import str
from builtins import zip
from builtins import map
from builtins import range
import abc
import os
import shutil
from forcebalance.finite_difference import *
from forcebalance.nifty import *
from forcebalance.nifty import _exec
from forcebalance.target import Target
import numpy as np
from forcebalance.molecule import Molecule
from re import match, sub
import subprocess
from subprocess import PIPE
try:
    from lxml import etree
except: pass
from pymbar import pymbar
import itertools
from collections import defaultdict, namedtuple, OrderedDict
import csv
import copy

from forcebalance.output import getLogger
logger = getLogger(__name__)

def weight_info(W, PT, N_k, verbose=True):
    C = []
    N = 0
    W += 1.0e-300
    I = np.exp(-1*np.sum((W*np.log(W))))
    for ns in N_k:
        C.append(sum(W[N:N+ns]))
        N += ns
    C = np.array(C)
    if verbose:
        logger.info("MBAR Results for Phase Point %s, Box, Contributions:\n" % str(PT))
        logger.info(str(C) + '\n')
        logger.info("InfoContent: % .2f snapshots (%.2f %%)\n" % (I, 100*I/len(W)))
    return C

# NPT_Trajectory = namedtuple('NPT_Trajectory', ['fnm', 'Rhos', 'pVs', 'Energies', 'Grads', 'mEnergies', 'mGrads', 'Rho_errs', 'Hvap_errs'])

class Lipid(Target):

    """ Subclass of Target for lipid property matching."""

    def __init__(self,options,tgt_opts,forcefield):
        # Initialize base class
        super(Lipid,self).__init__(options,tgt_opts,forcefield)
        # Weight of the density
        self.set_option(tgt_opts,'w_rho',forceprint=True)
        # Weight of the thermal expansion coefficient
        self.set_option(tgt_opts,'w_alpha',forceprint=True)
        # Weight of the isothermal compressibility
        self.set_option(tgt_opts,'w_kappa',forceprint=True)
        # Weight of the isobaric heat capacity
        self.set_option(tgt_opts,'w_cp',forceprint=True)
        # Weight of the dielectric constant
        self.set_option(tgt_opts,'w_eps0',forceprint=True)
        # Weight of the area per lipid
        self.set_option(tgt_opts,'w_al',forceprint=True)
        # Weight of the bilayer isothermal compressibility
        self.set_option(tgt_opts,'w_lkappa',forceprint=True)
        # Weight of the deuterium order parameter
        self.set_option(tgt_opts,'w_scd',forceprint=True)
        # Normalize the property contributions to the objective function
        self.set_option(tgt_opts,'w_normalize',forceprint=True)
        # Optionally pause on the zeroth step
        self.set_option(tgt_opts,'manual')
        # Number of time steps in the lipid "equilibration" run
        self.set_option(tgt_opts,'lipid_eq_steps',forceprint=True)
        # Number of time steps in the lipid "production" run
        self.set_option(tgt_opts,'lipid_md_steps',forceprint=True)
        # Number of time steps in the gas "equilibration" run
        self.set_option(tgt_opts,'gas_eq_steps',forceprint=False)
        # Number of time steps in the gas "production" run
        self.set_option(tgt_opts,'gas_md_steps',forceprint=False)
        # Cutoff for nonbonded interactions in the liquid
        if tgt_opts['nonbonded_cutoff'] is not None:
            self.set_option(tgt_opts,'nonbonded_cutoff')
        # Cutoff for vdW interactions if different from other nonbonded interactions
        if tgt_opts['vdw_cutoff'] is not None:
            self.set_option(tgt_opts,'vdw_cutoff')
        # Time step length (in fs) for the lipid production run
        self.set_option(tgt_opts,'lipid_timestep',forceprint=True)
        # Time interval (in ps) for writing coordinates
        self.set_option(tgt_opts,'lipid_interval',forceprint=True)
        # Time step length (in fs) for the gas production run
        self.set_option(tgt_opts,'gas_timestep',forceprint=True)
        # Time interval (in ps) for writing coordinates
        self.set_option(tgt_opts,'gas_interval',forceprint=True)
        # Minimize the energy prior to running any dynamics
        self.set_option(tgt_opts,'minimize_energy',forceprint=True)
        # Isolated dipole (debye) for analytic self-polarization correction.
        self.set_option(tgt_opts,'self_pol_mu0',forceprint=True)
        # Molecular polarizability (ang**3) for analytic self-polarization correction.
        self.set_option(tgt_opts,'self_pol_alpha',forceprint=True)
        # Set up the simulation object for self-polarization correction.
        self.do_self_pol = (self.self_pol_mu0 > 0.0 and self.self_pol_alpha > 0.0)
        # Enable anisotropic periodic box
        self.set_option(tgt_opts,'anisotropic_box',forceprint=True)
        # Whether to save trajectories (0 = never, 1 = delete after good step, 2 = keep all)
        self.set_option(tgt_opts,'save_traj')

        #======================================#
        #     Variables which are set here     #
        #======================================#
        ## LPW 2018-02-11: This is set to True if the target calculates
        ## a single-point property over several existing snapshots.
        self.loop_over_snapshots = False
        # List of trajectory files that may be deleted if self.save_traj == 1.
        self.last_traj = []
        # Extra files to be copied back at the end of a run.
        self.extra_output = []
        # Read the reference data
        self.read_data()
        # Read in lipid starting coordinates.
        if 'n_ic' in self.RefData:
            # Linked IC folder into the temp-directory.
            self.nptfiles += ["IC"]
            # Store IC frames in a dictionary.
            self.lipid_mols = OrderedDict()
            self.lipid_mols_new = OrderedDict()
            for pt in self.PhasePoints:
                pt_label = "IC/%sK-%s%s" % (pt[0], pt[1], pt[2])
                if not os.path.exists(os.path.join(self.root, self.tgtdir, pt_label, self.lipid_coords)):
                    raise RuntimeError("Initial condition files don't exist; please provide IC directory")
                # Create molecule object for each IC.
                all_ic = Molecule(os.path.join(self.root, self.tgtdir, pt_label, self.lipid_coords))
                self.lipid_mols[pt] = []
                n_uniq_ic = int(self.RefData['n_ic'][pt])
                if n_uniq_ic > len(all_ic):
                    raise RuntimeError("Number of frames in initial conditions .gro file is less than the number of parallel simulations requested in data.csv")
                # Index ICs by pressure and temperature in a dictionary.
                for ic in range(n_uniq_ic):
                    self.lipid_mols[pt].append(all_ic[ic])
        else:
            # Read in lipid starting coordinates.
            if not os.path.exists(os.path.join(self.root, self.tgtdir, self.lipid_coords)):
                logger.error("%s doesn't exist; please provide lipid_coords option\n" % self.lipid_coords)
                raise RuntimeError
            self.lipid_mol = Molecule(os.path.join(self.root, self.tgtdir, self.lipid_coords), toppbc=True)
            # Extra files to be linked into the temp-directory.
            self.nptfiles += [self.lipid_coords]
        # Scripts to be copied from the ForceBalance installation directory.
        self.scripts += ['npt_lipid.py']
        # Prepare the temporary directory.
        self.prepare_temp_directory()
        # Build keyword dictionary to pass to engine.
        if self.do_self_pol:
            self.gas_engine_args.update(self.OptionDict)
            self.gas_engine_args.update(options)
            del self.gas_engine_args['name']
            # Create engine object for gas molecule to do the polarization correction.
            self.gas_engine = self.engine_(target=self, mol=self.gas_mol, name="selfpol", **self.gas_engine_args)
        # Don't read indicate.log when calling meta_indicate()
        self.read_indicate = False
        self.write_indicate = False
        # Don't read objective.p when calling meta_get()
        self.read_objective = False
        #======================================#
        #          UNDER DEVELOPMENT           #
        #======================================#
        # Put stuff here that I'm not sure about. :)
        np.set_printoptions(precision=4, linewidth=100)
        np.seterr(under='ignore')
        ## Saved force field mvals for all iterations
        self.SavedMVal = {}
        ## Saved trajectories for all iterations and all temperatures
        self.SavedTraj = defaultdict(dict)
        ## Evaluated energies for all trajectories (i.e. all iterations and all temperatures), using all mvals
        self.MBarEnergy = defaultdict(lambda:defaultdict(dict))

    def prepare_temp_directory(self):
        """ Prepare the temporary directory by copying in important files. """
        abstempdir = os.path.join(self.root,self.tempdir)
        for f in self.nptfiles:
            LinkFile(os.path.join(self.root, self.tgtdir, f), os.path.join(abstempdir, f))
        for f in self.scripts:
            LinkFile(os.path.join(os.path.split(__file__)[0],"data",f),os.path.join(abstempdir,f))

    def read_data(self):
        # Read the 'data.csv' file. The file should contain guidelines.
        with open(os.path.join(self.tgtdir,'data.csv'),'r') as f: R0 = list(csv.reader(f))
        # All comments are erased.
        R1 = [[sub('#.*$','',word) for word in line] for line in R0 if len(line[0]) > 0 and line[0][0] != "#"]
        # All empty lines are deleted and words are converted to lowercase.
        R = [[wrd.lower() for wrd in line] for line in R1 if any([len(wrd) for wrd in line]) > 0]
        global_opts = OrderedDict()
        found_headings = False
        known_vars = ['mbar','rho','hvap','alpha','kappa','cp','eps0','cvib_intra',
                      'cvib_inter','cni','devib_intra','devib_inter', 'al', 'scd', 'n_ic', 'lkappa']
        self.RefData = OrderedDict()
        for line in R:
            if line[0] == "global":
                # Global options are mainly denominators for the different observables.
                if isfloat(line[2]):
                    global_opts[line[1]] = float(line[2])
                elif line[2].lower() == 'false':
                    global_opts[line[1]] = False
                elif line[2].lower() == 'true':
                    global_opts[line[1]] = True
            elif not found_headings:
                found_headings = True
                headings = line
                if len(set(headings)) != len(headings):
                    logger.error('Column headings in data.csv must be unique\n')
                    raise RuntimeError
                if 'p' not in headings:
                    logger.error('There must be a pressure column heading labeled by "p" in data.csv\n')
                    raise RuntimeError
                if 't' not in headings:
                    logger.error('There must be a temperature column heading labeled by "t" in data.csv\n')
                    raise RuntimeError
            elif found_headings:
                try:
                    # Temperatures are in kelvin.
                    t     = [float(val) for head, val in zip(headings,line) if head == 't'][0]
                    # For convenience, users may input the pressure in atmosphere or bar.
                    pval  = [float(val.split()[0]) for head, val in zip(headings,line) if head == 'p'][0]
                    punit = [val.split()[1] if len(val.split()) >= 1 else "atm" for head, val in zip(headings,line) if head == 'p'][0]
                    unrec = set([punit]).difference(['atm','bar'])
                    if len(unrec) > 0:
                        logger.error('The pressure unit %s is not recognized, please use bar or atm\n' % unrec[0])
                        raise RuntimeError
                    # This line actually reads the reference data and inserts it into the RefData dictionary of dictionaries.
                    for head, val in zip(headings,line):
                        if head == 't' or head == 'p' : continue
                        if isfloat(val):
                            self.RefData.setdefault(head,OrderedDict([]))[(t,pval,punit)] = float(val.strip())
                        elif val.lower() == 'true':
                            self.RefData.setdefault(head,OrderedDict([]))[(t,pval,punit)] = True
                        elif val.lower() == 'false':
                            self.RefData.setdefault(head,OrderedDict([]))[(t,pval,punit)] = False
                        elif head == 'scd':
                            self.RefData.setdefault(head,OrderedDict([]))[(t,pval,punit)] = np.array(list(map(float, val.split())))
                except:
                    logger.error(line + '\n')
                    logger.error('Encountered an error reading this line!\n')
                    raise RuntimeError
            else:
                logger.error(line + '\n')
                logger.error('I did not recognize this line!\n')
                raise RuntimeError
        # Check the reference data table for validity.
        default_denoms = defaultdict(int)
        PhasePoints = None
        RefData_copy = copy.deepcopy(self.RefData)
        for head in self.RefData:
            if head == 'n_ic':
                continue
            if head not in known_vars+[i+"_wt" for i in known_vars]:
                # Only hard-coded properties may be recognized.
                logger.error("The column heading %s is not recognized in data.csv\n" % head)
                raise RuntimeError
            if head in known_vars:
                if head+"_wt" not in self.RefData:
                    # If the phase-point weights are not specified in the reference data file, initialize them all to one.
                    RefData_copy[head+"_wt"] = OrderedDict([(key, 1.0) for key in self.RefData[head]])
                wts = np.array(list(RefData_copy[head+"_wt"].values()))
                dat = np.array(list(self.RefData[head].values()))
                # S_cd specifies an array of averages (one for each tail node).  Find avg over axis 0.
                avg = np.average(dat, weights=wts, axis=0)
                if len(wts) > 1:
                    # If there is more than one data point, then the default denominator is the
                    # standard deviation of the experimental values.
                    if head == 'scd':
                        default_denoms[head+"_denom"] = np.average(np.sqrt(np.dot(wts, (dat-avg)**2)/wts.sum()))
                    else:
                        default_denoms[head+"_denom"] = np.sqrt(np.dot(wts, (dat-avg)**2)/wts.sum())
                else:
                    # If there is only one data point, then the denominator is just the single
                    # data point itself.
                    if head == 'scd':
                        default_denoms[head+"_denom"] = np.average(np.sqrt(np.abs(dat[0])))
                    else:
                        default_denoms[head+"_denom"] = np.sqrt(np.abs(dat[0]))
            self.PhasePoints = list(self.RefData[head].keys())
            # This prints out all of the reference data.
            # printcool_dictionary(self.RefData[head],head)
        self.RefData = RefData_copy
        # Create labels for the directories.
        self.Labels = ["%.2fK-%.1f%s" % i for i in self.PhasePoints]
        logger.debug("global_opts:\n%s\n" % str(global_opts))
        logger.debug("default_denoms:\n%s\n" % str(default_denoms))
        for opt in global_opts:
            if "_denom" in opt:
                # Record entries from the global_opts dictionary so they can be retrieved from other methods.
                self.set_option(global_opts,opt,default=default_denoms[opt])
            else:
                self.set_option(global_opts,opt)

    def check_files(self, there):
        there = os.path.abspath(there)
        havepts = 0
        if all([i in os.listdir(there) for i in self.Labels]):
            for d in os.listdir(there):
                if d in self.Labels:
                    if os.path.exists(os.path.join(there, d, 'npt_result.p')):
                        havepts += 1
        if (float(havepts)/len(self.Labels)) > 0.75:
            return 1
        else:
            return 0
    def npt_simulation(self, temperature, pressure, simnum):
        """ Submit a NPT simulation to the Work Queue. """
        wq = getWorkQueue()
        if not os.path.exists('npt_result.p'):
            link_dir_contents(os.path.join(self.root,self.rundir),os.getcwd())
            self.last_traj += [os.path.join(os.getcwd(), i) for i in self.extra_output]
            self.lipid_mol[simnum%len(self.lipid_mol)].write(self.lipid_coords, ftype='tinker' if self.engname == 'tinker' else None)
            cmdstr = '%s python npt_lipid.py %s %.3f %.3f' % (self.nptpfx, self.engname, temperature, pressure)
            if wq is None:
                logger.info("Running condensed phase simulation locally.\n")
                logger.info("You may tail -f %s/npt.out in another terminal window\n" % os.getcwd())
                _exec(cmdstr, copy_stderr=True, outfnm='npt.out')
            else:
                queue_up(wq, command = cmdstr+' &> npt.out',
                         input_files = self.nptfiles + self.scripts + ['forcebalance.p'],
                         output_files = ['npt_result.p', 'npt.out'] + self.extra_output, tgt=self)

    def polarization_correction(self,mvals):
        d = self.gas_engine.get_multipole_moments(optimize=True)['dipole']
        if not in_fd():
            logger.info("The molecular dipole moment is % .3f debye\n" % np.linalg.norm(d))
        # Taken from the original OpenMM interface code, this is how we calculate the conversion factor.
        # dd2 = ((np.linalg.norm(d)-self.self_pol_mu0)*debye)**2
        # eps0 = 8.854187817620e-12 * coulomb**2 / newton / meter**2
        # epol = 0.5*dd2/(self.self_pol_alpha*angstrom**3*4*np.pi*eps0)/(kilojoule_per_mole/AVOGADRO_CONSTANT_NA)
        # In [2]: eps0 = 8.854187817620e-12 * coulomb**2 / newton / meter**2
        # In [7]: 1.0 * debye ** 2 / (1.0 * angstrom**3*4*np.pi*eps0) / (kilojoule_per_mole/AVOGADRO_CONSTANT_NA)
        # Out[7]: 60.240179789402056
        convert = 60.240179789402056
        dd2 = (np.linalg.norm(d)-self.self_pol_mu0)**2
        epol = 0.5*convert*dd2/self.self_pol_alpha
        return epol

    def indicate(self):
        AGrad = hasattr(self, 'Gp')
        PrintDict = OrderedDict()
        def print_item(key, heading, physunit):
            if self.Xp[key] > 0:
                printcool_dictionary(self.Pp[key], title='%s %s%s\nTemperature  Pressure  Reference  Calculated +- Stdev     Delta    Weight    Term   ' %
                                     (self.name, heading, " (%s) " % physunit if physunit else ""), bold=True, color=4, keywidth=15)
                bar = printcool("%s objective function: % .3f%s" % (heading, self.Xp[key], ", Derivative:" if AGrad else ""))
                if AGrad:
                    self.FF.print_map(vals=self.Gp[key])
                    logger.info(bar)
                PrintDict[heading] = "% 10.5f % 8.3f % 14.5e" % (self.Xp[key], self.Wp[key], self.Xp[key]*self.Wp[key])

        print_item("Rho", "Density", "kg m^-3")
        print_item("Alpha", "Thermal Expansion Coefficient", "10^-4 K^-1")
        print_item("Kappa", "Isothermal Compressibility", "10^-6 bar^-1")
        print_item("Cp", "Isobaric Heat Capacity", "cal mol^-1 K^-1")
        print_item("Eps0", "Dielectric Constant", None)
        print_item("Al", "Average Area per Lipid", "nm^2")
        print_item("Scd", "Deuterium Order Parameter", None)
        print_item("LKappa", "Bilayer Isothermal Compressibility", "mN/m")

        PrintDict['Total'] = "% 10s % 8s % 14.5e" % ("","",self.Objective)

        Title = "%s Condensed Phase Properties:\n %-20s %40s" % (self.name, "Property Name", "Residual x Weight = Contribution")
        printcool_dictionary(PrintDict,color=4,title=Title,keywidth=31)
        return

    def objective_term(self, points, expname, calc, err, grad, name="Quantity", SubAverage=False):
        if expname in self.RefData:
            exp = self.RefData[expname]
            Weights = self.RefData[expname+"_wt"]
            Denom = getattr(self,expname+"_denom")
        else:
            # If the reference data doesn't exist then return nothing.
            return 0.0, np.zeros(self.FF.np), np.zeros((self.FF.np,self.FF.np)), None

        Sum = sum(Weights.values())
        for i in Weights:
            Weights[i] /= Sum
        logger.info("Weights have been renormalized to " + str(sum(Weights.values())) + "\n")
        # Use least-squares or hyperbolic (experimental) objective.
        LeastSquares = True

        logger.info("Physical quantity %s uses denominator = % .4f\n" % (name, Denom))
        if not LeastSquares:
            # If using a hyperbolic functional form
            # we still want the contribution to the
            # objective function to be the same when
            # Delta = Denom.
            Denom /= 3 ** 0.5

        Objective = 0.0
        Gradient = np.zeros(self.FF.np)
        Hessian = np.zeros((self.FF.np,self.FF.np))
        Objs = {}
        GradMap = []
        avgCalc = 0.0
        avgExp  = 0.0
        avgGrad = np.zeros(self.FF.np)
        for i, PT in enumerate(points):
            avgCalc += Weights[PT]*calc[PT]
            avgExp  += Weights[PT]*exp[PT]
            avgGrad += Weights[PT]*grad[PT]
        for i, PT in enumerate(points):
            if SubAverage:
                G = grad[PT]-avgGrad
                Delta = calc[PT] - exp[PT] - avgCalc + avgExp
            else:
                G = grad[PT]
                Delta = calc[PT] - exp[PT]
            if hasattr(Delta, "__len__"):
                Delta = np.average(Delta)
            if LeastSquares:
                # Least-squares objective function.
                ThisObj = Weights[PT] * Delta ** 2 / Denom**2
                Objs[PT] = ThisObj
                ThisGrad = 2.0 * Weights[PT] * Delta * G / Denom**2
                GradMap.append(G)
                Objective += ThisObj
                Gradient += ThisGrad
                # Gauss-Newton approximation to the Hessian.
                Hessian += 2.0 * Weights[PT] * (np.outer(G, G)) / Denom**2
            else:
                # L1-like objective function.
                D = Denom
                S = Delta**2 + D**2
                ThisObj  = Weights[PT] * (S**0.5-D) / Denom
                ThisGrad = Weights[PT] * (Delta/S**0.5) * G / Denom
                ThisHess = Weights[PT] * (1/S**0.5-Delta**2/S**1.5) * np.outer(G,G) / Denom
                Objs[PT] = ThisObj
                GradMap.append(G)
                Objective += ThisObj
                Gradient += ThisGrad
                Hessian += ThisHess
        GradMapPrint = [["#PhasePoint"] + self.FF.plist]
        for PT, g in zip(points,GradMap):
            GradMapPrint.append([' %8.2f %8.1f %3s' % PT] + ["% 9.3e" % i for i in g])
        o = wopen('gradient_%s.dat' % name)
        for line in GradMapPrint:
            print(' '.join(line), file=o)
        o.close()

        Delta = np.array([calc[PT] - exp[PT] for PT in points])
        delt = {PT : r for PT, r in zip(points,Delta)}
        if expname == 'scd':
            print_out = OrderedDict([('    %8.2f %8.1f %3s' % PT, '\n %s' % (' '.join('\t \t \t %9.6f    %9.6f +- %-7.6f % 7.6f \n' % F for F in zip(exp[PT], calc[PT], flat(err[PT]), delt[PT])))) for PT in calc])
        else:
            print_out = OrderedDict([('    %8.2f %8.1f %3s' % PT, "%9.3f    %9.3f +- %-7.3f % 7.3f % 9.5f % 9.5f" % (exp[PT],calc[PT],err[PT],delt[PT],Weights[PT],Objs[PT])) for PT in calc])

        return Objective, Gradient, Hessian, print_out

    def submit_jobs(self, mvals, AGrad=True, AHess=True):
        # This routine is called by Objective.stage() will run before "get".
        # It submits the jobs to the Work Queue and the stage() function will wait for jobs to complete.
        #
        # First dump the force field to a pickle file
        lp_dump((self.FF,mvals,self.OptionDict,AGrad),'forcebalance.p')

        # Give the user an opportunity to copy over data from a previous (perhaps failed) run.
        if (not self.evaluated) and self.manual:
            warn_press_key("Now's our chance to fill the temp directory up with data!\n(Considering using 'read' or 'continue' for better checkpointing)", timeout=7200)

        # If self.save_traj == 1, delete the trajectory files from a previous good optimization step.
        if self.evaluated and self.goodstep and self.save_traj < 2:
            for fn in self.last_traj:
                if os.path.exists(fn):
                    os.remove(fn)
        self.last_traj = []

        # Set up and run the NPT simulations.
        snum = 0
        for label, pt in zip(self.Labels, self.PhasePoints):
            T = pt[0]
            P = pt[1]
            Punit = pt[2]
            if Punit == 'bar':
                P *= 1.0 / 1.01325
            if not os.path.exists(label):
                os.makedirs(label)
                os.chdir(label)
                if 'n_ic' in self.RefData:
                    n_uniq_ic = int(self.RefData['n_ic'][pt])
                    # Loop over parallel trajectories.
                    for trj in range(n_uniq_ic):
                        rel_trj = "trj_%i" % trj
                        # Create directories for each parallel simulation.
                        if not os.path.exists(rel_trj):
                            os.makedirs(rel_trj)
                            os.chdir(rel_trj)
                            # Pull each simulation molecule from the lipid_mols dictionary.
                            # lipid_mols is a dictionary of paths to either the initial
                            # geometry files, or the geometries from the final frame of the
                            # previous iteration.
                            self.lipid_mol = self.lipid_mols[pt][trj]
                            self.lipid_mol.write(self.lipid_coords)
                            if not self.lipid_coords in self.nptfiles:
                                self.nptfiles += [self.lipid_coords]
                            self.npt_simulation(T,P,snum)
                        os.chdir('..')
                else:
                    self.npt_simulation(T,P,snum)
                os.chdir('..')
                snum += 1

    def get(self, mvals, AGrad=True, AHess=True):

        """
        Fitting of lipid bulk properties.  This is the current major
        direction of development for ForceBalance.  Basically, fitting
        the QM energies / forces alone does not always give us the
        best simulation behavior.  In many cases it makes more sense
        to try and reproduce some experimentally known data as well.

        In order to reproduce experimentally known data, we need to
        run a simulation and compare the simulation result to
        experiment.  The main challenge here is that the simulations
        are computationally intensive (i.e. they require energy and
        force evaluations), and furthermore the results are noisy.  We
        need to run the simulations automatically and remotely
        (i.e. on clusters) and a good way to calculate the derivatives
        of the simulation results with respect to the parameter values.

        This function contains some experimentally known values of the
        density and enthalpy of vaporization (Hvap) of lipid water.
        It launches the density and Hvap calculations on the cluster,
        and gathers the results / derivatives.  The actual calculation
        of results / derivatives is done in a separate file.

        After the results come back, they are gathered together to form
        an objective function.

        @param[in] mvals Mathematical parameter values
        @param[in] AGrad Switch to turn on analytic gradient
        @param[in] AHess Switch to turn on analytic Hessian
        @return Answer Contribution to the objective function

        """

        mbar_verbose = False

        Answer = {}

        Results = {}
        Points = []  # These are the phase points for which data exists.
        BPoints = [] # These are the phase points for which we are doing MBAR for the condensed phase.
        tt = 0
        for label, PT in zip(self.Labels, self.PhasePoints):
            if 'n_ic' in self.RefData:
                self.lipid_mols[PT] = [Molecule(last_frame) for last_frame in self.lipid_mols[PT]]
                n_uniq_ic = int(self.RefData['n_ic'][PT])
                for ic in range(n_uniq_ic):
                    if os.path.exists('./%s/trj_%s/npt_result.p' % (label, ic)):
                        # Read in each each parallel simulation's data, and concatenate each property time series.
                        ts = lp_load('./%s/trj_%s/npt_result.p' % (label, ic))
                        if ic == 0:
                            ts_concat = list(ts)
                        else:
                            for d_arr in range(len(ts)):
                                if isinstance(ts[d_arr], np.ndarray):
                                    # Gradients need a unique append format.
                                    if d_arr == 5:
                                        ts_concat[d_arr] = np.append(ts_concat[d_arr], ts[d_arr], axis = 1)
                                    else:
                                        ts_concat[d_arr] = np.append(ts_concat[d_arr], ts[d_arr], axis = 0)
                                if isinstance(ts_concat[d_arr], list):
                                    ts_concat[d_arr] = [np.append(ts_concat[d_arr][i], ts[d_arr][i], axis = 1) for i in range(len(ts_concat[d_arr]))]
                        # Write concatenated time series to a pickle file.
                        if ic == (int(n_uniq_ic) - 1):
                            lp_dump((ts_concat), './%s/npt_result.p' % label)
            if os.path.exists('./%s/npt_result.p' % label):
                logger.info('Reading information from ./%s/npt_result.p\n' % label)
                Points.append(PT)
                Results[tt] = lp_load('./%s/npt_result.p' % label)
                tt += 1
            else:
                logger.warning('The file ./%s/npt_result.p does not exist so we cannot read it\n' % label)
                pass
                # for obs in self.RefData:
                #     del self.RefData[obs][PT]
        if len(Points) == 0:
            logger.error('The lipid simulations have terminated with \x1b[1;91mno readable data\x1b[0m - this is a problem!\n')
            raise RuntimeError

        # Assign variable names to all the stuff in npt_result.p
        Rhos, Vols, Potentials, Energies, Dips, Grads, GDips, \
            Rho_errs, Alpha_errs, Kappa_errs, Cp_errs, Eps0_errs, NMols, Als, Al_errs, Scds, Scd_errs, LKappa_errs = ([Results[t][i] for t in range(len(Points))] for i in range(18))
        # Determine the number of molecules
        if len(set(NMols)) != 1:
            logger.error(str(NMols))
            logger.error('The above list should only contain one number - the number of molecules\n')
            raise RuntimeError
        else:
            NMol = list(set(NMols))[0]

        R  = np.array(list(itertools.chain(*list(Rhos))))
        V  = np.array(list(itertools.chain(*list(Vols))))
        E  = np.array(list(itertools.chain(*list(Energies))))
        Dx = np.array(list(itertools.chain(*list(d[:,0] for d in Dips))))
        Dy = np.array(list(itertools.chain(*list(d[:,1] for d in Dips))))
        Dz = np.array(list(itertools.chain(*list(d[:,2] for d in Dips))))
        G  = np.hstack(tuple(Grads))
        GDx = np.hstack(tuple(gd[0] for gd in GDips))
        GDy = np.hstack(tuple(gd[1] for gd in GDips))
        GDz = np.hstack(tuple(gd[2] for gd in GDips))
        A  = np.array(list(itertools.chain(*list(Als))))
        S  = np.array(list(itertools.chain(*list(Scds))))

        Rho_calc = OrderedDict([])
        Rho_grad = OrderedDict([])
        Rho_std  = OrderedDict([])
        Alpha_calc = OrderedDict([])
        Alpha_grad = OrderedDict([])
        Alpha_std  = OrderedDict([])
        Kappa_calc = OrderedDict([])
        Kappa_grad = OrderedDict([])
        Kappa_std  = OrderedDict([])
        Cp_calc = OrderedDict([])
        Cp_grad = OrderedDict([])
        Cp_std  = OrderedDict([])
        Eps0_calc = OrderedDict([])
        Eps0_grad = OrderedDict([])
        Eps0_std  = OrderedDict([])
        Al_calc = OrderedDict([])
        Al_grad = OrderedDict([])
        Al_std  = OrderedDict([])
        LKappa_calc = OrderedDict([])
        LKappa_grad = OrderedDict([])
        LKappa_std  = OrderedDict([])
        Scd_calc = OrderedDict([])
        Scd_grad = OrderedDict([])
        Scd_std  = OrderedDict([])

        # The unit that converts atmospheres * nm**3 into kj/mol :)
        pvkj=0.061019351687175

        # Run MBAR using the total energies. Required for estimates that use the kinetic energy.
        BSims = len(BPoints)
        Shots = len(Energies[0])
        Shots_m = [len(i) for i in Energies]
        N_k = np.ones(BSims)*Shots
        # Use the value of the energy for snapshot t from simulation k at potential m
        U_kln = np.zeros([BSims,BSims,Shots])
        for m, PT in enumerate(BPoints):
            T = PT[0]
            P = PT[1] / 1.01325 if PT[2] == 'bar' else PT[1]
            beta = 1. / (kb * T)
            for k in range(BSims):
                # The correct Boltzmann factors include PV.
                # Note that because the Boltzmann factors are computed from the conditions at simulation "m",
                # the pV terms must be rescaled to the pressure at simulation "m".
                kk = Points.index(BPoints[k])
                U_kln[k, m, :]   = Energies[kk] + P*Vols[kk]*pvkj
                U_kln[k, m, :]  *= beta
        W1 = None
        if len(BPoints) > 1:
            logger.info("Running MBAR analysis on %i states...\n" % len(BPoints))
            mbar = pymbar.MBAR(U_kln, N_k, verbose=mbar_verbose, relative_tolerance=5.0e-8)
            W1 = mbar.getWeights()
            logger.info("Done\n")
        elif len(BPoints) == 1:
            W1 = np.ones((BPoints*Shots,BPoints))
            W1 /= BPoints*Shots

        def fill_weights(weights, phase_points, mbar_points, snapshots):
            """ Fill in the weight matrix with MBAR weights where MBAR was run,
            and equal weights otherwise. """
            new_weights = np.zeros([len(phase_points)*snapshots,len(phase_points)])
            for m, PT in enumerate(phase_points):
                if PT in mbar_points:
                    mm = mbar_points.index(PT)
                    for kk, PT1 in enumerate(mbar_points):
                        k = phase_points.index(PT1)
                        logger.debug("Will fill W2[%i:%i,%i] with W1[%i:%i,%i]\n" % (k*snapshots,k*snapshots+snapshots,m,kk*snapshots,kk*snapshots+snapshots,mm))
                        new_weights[k*snapshots:(k+1)*snapshots,m] = weights[kk*snapshots:(kk+1)*snapshots,mm]
                else:
                    logger.debug("Will fill W2[%i:%i,%i] with equal weights\n" % (m*snapshots,(m+1)*snapshots,m))
                    new_weights[m*snapshots:(m+1)*snapshots,m] = 1.0/snapshots
            return new_weights

        W2 = fill_weights(W1, Points, BPoints, Shots)

        if self.do_self_pol:
            EPol = self.polarization_correction(mvals)
            GEPol = np.array([(f12d3p(fdwrap(self.polarization_correction, mvals, p), h = self.h, f0 = EPol)[0] if p in self.pgrad else 0.0) for p in range(self.FF.np)])
            bar = printcool("Self-polarization correction to \nenthalpy of vaporization is % .3f kJ/mol%s" % (EPol, ", Derivative:" if AGrad else ""))
            if AGrad:
                self.FF.print_map(vals=GEPol)
                logger.info(bar)

        for i, PT in enumerate(Points):
            T = PT[0]
            P = PT[1] / 1.01325 if PT[2] == 'bar' else PT[1]
            PV = P*V*pvkj
            H = E + PV
            # The weights that we want are the last ones.
            W = flat(W2[:,i])
            C = weight_info(W, PT, np.ones(len(Points), dtype=int)*Shots, verbose=mbar_verbose)
            Gbar = flat(np.dot(G,col(W)))
            mBeta = -1/kb/T
            Beta  = 1/kb/T
            kT    = kb*T
            # Define some things to make the analytic derivatives easier.
            def avg(vec):
                return np.dot(W,vec)
            def covde(vec):
                return flat(np.dot(G,col(W*vec))) - avg(vec)*Gbar
            def deprod(vec):
                return flat(np.dot(G,col(W*vec)))
            ## Density.
            Rho_calc[PT]   = np.dot(W,R)
            Rho_grad[PT]   = mBeta*(flat(np.dot(G,col(W*R))) - np.dot(W,R)*Gbar)
            ## Ignore enthalpy.
            ## Thermal expansion coefficient.
            Alpha_calc[PT] = 1e4 * (avg(H*V)-avg(H)*avg(V))/avg(V)/(kT*T)
            GAlpha1 = -1 * Beta * deprod(H*V) * avg(V) / avg(V)**2
            GAlpha2 = +1 * Beta * avg(H*V) * deprod(V) / avg(V)**2
            GAlpha3 = deprod(V)/avg(V) - Gbar
            GAlpha4 = Beta * covde(H)
            Alpha_grad[PT] = 1e4 * (GAlpha1 + GAlpha2 + GAlpha3 + GAlpha4)/(kT*T)
            ## Isothermal compressibility.
            bar_unit = 0.06022141793 * 1e6
            Kappa_calc[PT] = bar_unit / kT * (avg(V**2)-avg(V)**2)/avg(V)
            GKappa1 = +1 * Beta**2 * avg(V**2) * deprod(V) / avg(V)**2
            GKappa2 = -1 * Beta**2 * avg(V) * deprod(V**2) / avg(V)**2
            GKappa3 = +1 * Beta**2 * covde(V)
            Kappa_grad[PT] = bar_unit*(GKappa1 + GKappa2 + GKappa3)
            ## Isobaric heat capacity.
            Cp_calc[PT] = 1000/(4.184*NMol*kT*T) * (avg(H**2) - avg(H)**2)
            if hasattr(self,'use_cvib_intra') and self.use_cvib_intra:
                logger.debug("Adding " + str(self.RefData['devib_intra'][PT]) + " to the heat capacity\n")
                Cp_calc[PT] += self.RefData['devib_intra'][PT]
            if hasattr(self,'use_cvib_inter') and self.use_cvib_inter:
                logger.debug("Adding " + str(self.RefData['devib_inter'][PT]) + " to the heat capacity\n")
                Cp_calc[PT] += self.RefData['devib_inter'][PT]
            GCp1 = 2*covde(H) * 1000 / 4.184 / (NMol*kT*T)
            GCp2 = mBeta*covde(H**2) * 1000 / 4.184 / (NMol*kT*T)
            GCp3 = 2*Beta*avg(H)*covde(H) * 1000 / 4.184 / (NMol*kT*T)
            Cp_grad[PT] = GCp1 + GCp2 + GCp3
            ## Static dielectric constant.
            prefactor = 30.348705333964077
            D2 = avg(Dx**2)+avg(Dy**2)+avg(Dz**2)-avg(Dx)**2-avg(Dy)**2-avg(Dz)**2
            Eps0_calc[PT] = 1.0 + prefactor*(D2/avg(V))/T
            GD2  = 2*(flat(np.dot(GDx,col(W*Dx))) - avg(Dx)*flat(np.dot(GDx,col(W)))) - Beta*(covde(Dx**2) - 2*avg(Dx)*covde(Dx))
            GD2 += 2*(flat(np.dot(GDy,col(W*Dy))) - avg(Dy)*flat(np.dot(GDy,col(W)))) - Beta*(covde(Dy**2) - 2*avg(Dy)*covde(Dy))
            GD2 += 2*(flat(np.dot(GDz,col(W*Dz))) - avg(Dz)*flat(np.dot(GDz,col(W)))) - Beta*(covde(Dz**2) - 2*avg(Dz)*covde(Dz))
            Eps0_grad[PT] = prefactor*(GD2/avg(V) - mBeta*covde(V)*D2/avg(V)**2)/T
            ## Average area per lipid
            Al_calc[PT]   = np.dot(W,A)
            Al_grad[PT]   = mBeta*(flat(np.dot(G,col(W*A))) - np.dot(W,A)*Gbar)
            ## Bilayer Isothermal compressibility.
            A_m2 = A * 1e-18
            kbT = 1.3806488e-23 * T
            LKappa_calc[PT] = (1e3 * 2 * kbT / 128) * (avg(A_m2) / (avg(A_m2**2)-avg(A_m2)**2))
            al_avg = avg(A_m2)
            al_sq_avg = avg(A_m2**2)
            al_avg_sq = al_avg**2
            al_var = al_sq_avg - al_avg_sq
            GLKappa1 = covde(A_m2) / al_var
            GLKappa2 = (al_avg / al_var**2) * (covde(A_m2**2) - (2 * al_avg * covde(A)))
            LKappa_grad[PT] = (1e3 * 2 * kbT / 128) * (GLKappa1 - GLKappa2)
            ## Deuterium order parameter
            Scd_calc[PT]   = np.dot(W,S)
            # LPW: In case I did not do the conversion correctly, the line of code previously here was:
            # Scd_grad[PT]   = mBeta * (flat(np.average(np.mat(G) * (S * W[:, np.newaxis]), axis = 1)) - np.average(np.average(S * W[:, np.newaxis], axis = 0), axis = 0) * Gbar)
            Scd_grad[PT]   = mBeta * (flat(np.average(np.dot(G, (S * W[:, np.newaxis])), axis = 1)) - np.average(np.average(S * W[:, np.newaxis], axis = 0), axis = 0) * Gbar)
            ## Estimation of errors.
            Rho_std[PT]    = np.sqrt(sum(C**2 * np.array(Rho_errs)**2))
            Alpha_std[PT]   = np.sqrt(sum(C**2 * np.array(Alpha_errs)**2)) * 1e4
            Kappa_std[PT]   = np.sqrt(sum(C**2 * np.array(Kappa_errs)**2)) * 1e6
            Cp_std[PT]   = np.sqrt(sum(C**2 * np.array(Cp_errs)**2))
            Eps0_std[PT]   = np.sqrt(sum(C**2 * np.array(Eps0_errs)**2))
            Al_std[PT]    = np.sqrt(sum(C**2 * np.array(Al_errs)**2))
            # LPW: In case I did not do the conversion correctly, the line of code previously here was:
            # Scd_std[PT]    = np.sqrt(sum(np.mat(C**2) * np.array(Scd_errs)**2))
            Scd_std[PT]    = np.sqrt(sum(np.dot(row(C**2), np.array(Scd_errs)**2)))
            LKappa_std[PT]   = np.sqrt(sum(C**2 * np.array(LKappa_errs)**2)) * 1e6

        # Get contributions to the objective function
        X_Rho, G_Rho, H_Rho, RhoPrint = self.objective_term(Points, 'rho', Rho_calc, Rho_std, Rho_grad, name="Density")
        X_Alpha, G_Alpha, H_Alpha, AlphaPrint = self.objective_term(Points, 'alpha', Alpha_calc, Alpha_std, Alpha_grad, name="Thermal Expansion")
        X_Kappa, G_Kappa, H_Kappa, KappaPrint = self.objective_term(Points, 'kappa', Kappa_calc, Kappa_std, Kappa_grad, name="Compressibility")
        X_Cp, G_Cp, H_Cp, CpPrint = self.objective_term(Points, 'cp', Cp_calc, Cp_std, Cp_grad, name="Heat Capacity")
        X_Eps0, G_Eps0, H_Eps0, Eps0Print = self.objective_term(Points, 'eps0', Eps0_calc, Eps0_std, Eps0_grad, name="Dielectric Constant")
        X_Al, G_Al, H_Al, AlPrint = self.objective_term(Points, 'al', Al_calc, Al_std, Al_grad, name="Avg Area per Lipid")
        X_Scd, G_Scd, H_Scd, ScdPrint = self.objective_term(Points, 'scd', Scd_calc, Scd_std, Scd_grad, name="Deuterium Order Parameter")
        X_LKappa, G_LKappa, H_LKappa, LKappaPrint = self.objective_term(Points, 'lkappa', LKappa_calc, LKappa_std, LKappa_grad, name="Bilayer Compressibility")

        Gradient = np.zeros(self.FF.np)
        Hessian = np.zeros((self.FF.np,self.FF.np))

        if X_Rho == 0: self.w_rho = 0.0
        if X_Alpha == 0: self.w_alpha = 0.0
        if X_Kappa == 0: self.w_kappa = 0.0
        if X_Cp == 0: self.w_cp = 0.0
        if X_Eps0 == 0: self.w_eps0 = 0.0
        if X_Al == 0: self.w_al = 0.0
        if X_Scd == 0: self.w_scd = 0.0
        if X_LKappa == 0: self.w_lkappa = 0.0

        if self.w_normalize:
            w_tot = self.w_rho + self.w_alpha + self.w_kappa + self.w_cp + self.w_eps0 + self.w_al + self.w_scd + self.w_lkappa
        else:
            w_tot = 1.0
        w_1 = self.w_rho / w_tot
        w_3 = self.w_alpha / w_tot
        w_4 = self.w_kappa / w_tot
        w_5 = self.w_cp / w_tot
        w_6 = self.w_eps0 / w_tot
        w_7 = self.w_al / w_tot
        w_8 = self.w_scd / w_tot
        w_9 = self.w_lkappa / w_tot

        Objective    = w_1 * X_Rho + w_3 * X_Alpha + w_4 * X_Kappa + w_5 * X_Cp + w_6 * X_Eps0 + w_7 * X_Al + w_8 * X_Scd + w_9 * X_LKappa
        if AGrad:
            Gradient = w_1 * G_Rho + w_3 * G_Alpha + w_4 * G_Kappa + w_5 * G_Cp + w_6 * G_Eps0 + w_7 * G_Al + w_8 * G_Scd + w_9 * G_LKappa
        if AHess:
            Hessian  = w_1 * H_Rho + w_3 * H_Alpha + w_4 * H_Kappa + w_5 * H_Cp + w_6 * H_Eps0 + w_7 * H_Al + w_8 * H_Scd + w_9 * H_LKappa

        if not in_fd():
            self.Xp = {"Rho" : X_Rho, "Alpha" : X_Alpha,
                           "Kappa" : X_Kappa, "Cp" : X_Cp, "Eps0" : X_Eps0, "Al" : X_Al, "Scd" : X_Scd, "LKappa" : X_LKappa}
            self.Wp = {"Rho" : w_1, "Alpha" : w_3,
                           "Kappa" : w_4, "Cp" : w_5, "Eps0" : w_6, "Al" : w_7, "Scd" : w_8, "LKappa" : w_9}
            self.Pp = {"Rho" : RhoPrint, "Alpha" : AlphaPrint,
                           "Kappa" : KappaPrint, "Cp" : CpPrint, "Eps0" : Eps0Print, "Al" : AlPrint, "Scd" : ScdPrint, "LKappa": LKappaPrint}
            if AGrad:
                self.Gp = {"Rho" : G_Rho, "Alpha" : G_Alpha,
                               "Kappa" : G_Kappa, "Cp" : G_Cp, "Eps0" : G_Eps0, "Al" : G_Al, "Scd" : G_Scd, "LKappa" : G_LKappa}
            self.Objective = Objective

        Answer = {'X':Objective, 'G':Gradient, 'H':Hessian}
        return Answer


# Copyright (c) 2021 Institute for Quantum Computing, Baidu Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Quantum chemistry module
"""

import os
import re
import numpy as np
import psi4
import openfermion
from openfermion import MolecularData, transforms
from openfermion.ops import general_basis_change
from paddle_quantum.utils import Hamiltonian

__all__ = [
    "geometry",
    "get_molecular_data",
    "active_space",
    "fermionic_hamiltonian",
    "spin_hamiltonian"
]

def _hamiltonian_transformation(spin_h, tol=1e-8):
    r"""将哈密顿量从 openfermion 格式转换成 Paddle Quantum 格式。

    Warning:
        输入的哈密顿量必须为埃尔米特的，输入的哈密顿中虚数的系数会和实数一起转换成他们的范数 (norm)。

    Args:
        spin_h (openfermion.ops.operators.qubit_operator.QubitOperator): openfermion 格式的哈密顿量
        tol (float, optional): 系数小于 tol 的值将被忽略掉，默认为 1e-8

    Returns:
        paddle_quantum.Hamiltonian object: Paddle Quantum 格式的哈密顿量
    """
    terms = spin_h.__str__().split('+\n')
    spin_h.compress(abs_tol=tol)
    pauli_str = []
    for term in terms:
        decomposed_term = re.match(r"(.*) \[(.*)\].*", term).groups()
        if decomposed_term[1] == '':
            try:
                pauli_str.append([float(decomposed_term[0]), 'I'])
            except ValueError:
                if complex(decomposed_term[0]).real > 0:
                    pauli_str.append([abs(complex(decomposed_term[0])), 'I'])
                else:
                    pauli_str.append([-abs(complex(decomposed_term[0])), 'I'])
        else:
            term_str = ', '.join(re.split(r' ', decomposed_term[1]))
            try:
                pauli_str.append([float(decomposed_term[0]), term_str])
            except ValueError:
                if complex(decomposed_term[0]).real > 0:
                    pauli_str.append([abs(complex(decomposed_term[0])), term_str])
                else:
                    pauli_str.append([-abs(complex(decomposed_term[0])), term_str])
    return Hamiltonian(pauli_str)

def _geo_str(geometry):
    r"""创建分子几何信息的字符串。

    Args:
        geometry (list): 包含了分子的几何信息，以 H2 分子为例
        [['H', [-1.68666, 1.79811, 0.0]], ['H', [-1.12017, 1.37343, 0.0]]]。

    Returns:
        str: 分子几何信息的字符串。
    """
    geo_str = ''
    for item in geometry:
        atom_symbol = item[0]
        position = item[1]
        line = '{} {} {} {}'.format(atom_symbol,
                                    position[0],
                                    position[1],
                                    position[2])
        if len(geo_str) > 0:
            geo_str += '\n'
        geo_str += line
    geo_str += '\nsymmetry c1'
    return geo_str

def _run_psi4(
    molecule,
    charge,
    multiplicity,
    method,
    basis,
    if_print,
    if_save
):
    r"""计算分子的必要信息，包括单体积分 (one-body integrations) 和双体积分 (two-body integrations)，
    以及用 scf 和 fci 的方法计算基态的能量。

    Args:
        molecule (MolecularData object): 包含分子所有信息的类 (class)。
        charge (int): 分子的电荷。
        multiplicity (int): 分子的多重度。
        method (str): 用于计算基态能量的方法，包括 'scf'和 'fci'。
        basis (str): 常用的基组是 'sto-3g', '6-31g'等。更多的基组选择可以参考网站。
        https://psicode.org/psi4manual/master/basissets_byelement.html#apdx-basiselement。
        if_print (Boolean): 是否需要打印出选定方法 (method) 计算出的分子基态能量。
        if_save (Boolean): 是否需要将分子信息存储成 .hdf5 文件。
    """
    psi4.set_memory('500 MB')
    psi4.set_options({'soscf': 'false',
                      'scf_type': 'pk'})
    geo = molecule.geometry
    mol = psi4.geometry(_geo_str(geo))
    mol.set_multiplicity(multiplicity)
    mol.set_molecular_charge(charge)

    if molecule.multiplicity == 1:
        psi4.set_options({'reference': 'rhf',
                          'guess': 'sad'})
    else:
        psi4.set_options({'reference': 'rohf',
                          'guess': 'gwh'})

    # HF calculation
    hf_energy, hf_wfn = psi4.energy('scf/' + basis, molecule=mol, return_wfn='on')
    # Get orbitals and Fock matrix.
    molecule.hf_energy = hf_energy
    molecule.nuclear_repulsion = mol.nuclear_repulsion_energy()
    molecule.canonical_orbitals = np.asarray(hf_wfn.Ca())
    molecule.overlap_integrals = np.asarray(hf_wfn.S())
    molecule.n_orbitals = molecule.canonical_orbitals.shape[0]
    molecule.n_qubits = 2 * molecule.n_orbitals
    molecule.orbital_energies = np.asarray(hf_wfn.epsilon_a())
    molecule.fock_matrix = np.asarray(hf_wfn.Fa())

    # Get integrals using MintsHelper.
    mints = psi4.core.MintsHelper(hf_wfn.basisset())

    molecule.one_body_integrals = general_basis_change(
        np.asarray(mints.ao_kinetic()), molecule.canonical_orbitals, (1, 0))
    molecule.one_body_integrals += general_basis_change(
        np.asarray(mints.ao_potential()), molecule.canonical_orbitals, (1, 0))
    two_body_integrals = np.asarray(mints.ao_eri())
    two_body_integrals.reshape((molecule.n_orbitals, molecule.n_orbitals,
                                molecule.n_orbitals, molecule.n_orbitals))
    two_body_integrals = np.einsum('psqr', two_body_integrals)
    two_body_integrals = general_basis_change(
        two_body_integrals, molecule.canonical_orbitals, (1, 1, 0, 0))
    molecule.two_body_integrals = two_body_integrals

    # FCI calculation
    psi4.set_options({'qc_module': 'detci'})
    fci_energy, fci_wfn = psi4.energy('fci/' + basis, molecule=mol, return_wfn='on')
    molecule.fci_energy = fci_energy

    if if_save is True:
        molecule.save()

    if if_print is True:
        if method == 'scf':
            print('Hartree-Fock energy for {} ({} electrons) is {}.'.format(
                molecule.name, molecule.n_electrons, hf_energy))

        elif method == 'fci':
            print('FCI energy for {} ({} electrons) is {}.'.format(
                molecule.name, molecule.n_electrons, fci_energy))
        elif method == '':
            print('Calculation is done')

def geometry(structure=None, file=None):
    r"""读取分子的几何信息。

    Args:
        structure (string, optional): 分子几何信息的字符串形式，以 H2 分子为例
            ``[['H', [-1.68666, 1.79811, 0.0]], ['H', [-1.12017, 1.37343, 0.0]]]``
        file (string, optional): .xyz 文件的路径

    Returns:
        str: 分子的几何信息

    Raises:
        AssertionError: 两个输入参数不可以同时为 ``None`` 。
    """
    if ((structure is None) and (file is None)):
        raise AssertionError('Input must be structure or .xyz file')
    elif file is None:
        shape = np.array(structure).shape
        assert shape[1] == 2, 'The shape of structure must be (n, 2)'
        for i in range(shape[0]):
            assert type(np.array(structure)[:, 0][i]) == str, 'The first position must be element symbol'
            assert len(np.array(structure)[:, 1][i]) == 3, 'The second position represents coordinate ' \
                                                           'of particle: x, y, z'
        geo = structure
    elif structure is None:
        assert file[-4:] == '.xyz', 'The file is supposed to be .xyz'
        geo = []
        with open(file) as f:
            for line in f.readlines()[2:]:
                one_geo = []

                symbol, x, y, z = line.split()
                one_geo.append(symbol)
                one_geo.append([float(x), float(y), float(z)])
                geo.append(one_geo)

    return geo

def get_molecular_data(
    geometry,
    charge=0,
    multiplicity=1,
    basis='sto-3g',
    method='scf',
    if_save=True,
    if_print=True,
    name="",
    file_path="."
):
    r"""计算分子的必要信息，包括单体积分（one-body integrations）和双体积分（two-body integrations），
    以及用选定的方法计算基态的能量。

    Args:
        geometry (str): 分子的几何信息
        charge (int, optional): 分子的电荷，默认值为 0
        multiplicity (int, optional): 分子的多重度，默认值为 1
        basis (str, optional): 常用的基组是 ``'sto-3g'`` 、 ``'6-31g'`` 等，默认的基组是 ``'sto-3g'``，更多的基组选择可以参考网站
            https://psicode.org/psi4manual/master/basissets_byelement.html#apdx-basiselement
        method (str, optional): 用于计算基态能量的方法，包括 ``'scf'`` 和 ``'fci'`` ，默认方法为 ``'scf'``
        if_save (bool, optional): 是否需要将分子信息存储成 .hdf5 文件，默认为 ``True``
        if_print (bool, optional): 是否需要打印出选定方法 (method) 计算出的分子基态能量，默认为 ``True``
        name (str, optional): 命名储存的文件，默认为 ``""``
        file_path (str, optional): 文件的储存路径，默认为 ``"."``

    Returns:
        MolecularData: 包含分子所有信息的类
    """
    methods = ['scf', 'fci']
    assert method in methods, 'We provide 2 methods: scf and fci'

    if if_save is True:
        path = file_path + '/qchem_data/'
        folder = os.path.exists(path)
        if not folder:
            os.makedirs(path)
        if name == "":
            elements = np.array(geometry)[:, 0]
            symbol, counts = np.unique(elements, return_counts=True)
            filename = path
            for i in range(len(symbol)):
                filename += symbol[i]+str(counts[i])+'_'
            filename += basis + '_' + method + '.hdf5'
        else:
            if name[-5:] == '.hdf5':
                filename = name
            else:
                filename = name + '.hdf5'

    molecule = MolecularData(geometry,
                             basis=basis,
                             multiplicity=multiplicity,
                             charge=charge,
                             filename=filename)

    _run_psi4(molecule,
              charge,
              multiplicity,
              method,
              basis,
              if_print,
              if_save)

    return molecule

def active_space(electrons,
                 orbitals,
                 multiplicity=1,
                 active_electrons=None,
                 active_orbitals=None):
    r"""对于给定的活跃电子和活跃轨道计算相应的活跃空间（active space）。

    Args:
        electrons (int): 电子数
        orbitals (int): 轨道数
        multiplicity (int, optional): 自旋多重度
        active_electrons (int, optional): 活跃 (active) 电子数，默认情况为所有电子均为活跃电子
        active_orbitals (int, optional): 活跃 (active) 轨道数，默认情况为所有轨道均为活跃轨道

    Returns:
        tuple: 核心轨道和活跃轨道的索引
    """
    assert type(electrons) == int and electrons > 0, 'Number of electrons must be positive integer.'
    assert type(orbitals) == int and orbitals > 0, 'Number of orbitals must be positive integer.'
    assert type(multiplicity) == int and multiplicity >= 0, 'The multiplicity must be non-negative integer.'

    if active_electrons is None:
        no_core_orbitals = 0
        core_orbitals = []
    else:
        assert type(active_electrons) == int, 'The number of active electrons must be integer.'
        assert active_electrons > 0, 'The number of active electrons must be greater than 0.'
        assert electrons >= active_electrons, 'The number of electrons should more than or equal ' \
                                              'to the number of active electrons.'
        assert active_electrons >= multiplicity - 1, 'The number of active electrons should greater than ' \
                                                     'or equal to multiplicity - 1.'
        assert multiplicity % 2 != active_electrons % 2, 'Mulitiplicity and active electrons should be one odd ' \
                                                         'and the other one even.'

        no_core_orbitals = (electrons - active_electrons) // 2
        core_orbitals = list(np.arange(0, no_core_orbitals))

    if active_orbitals is None:
        active_orbitals = list(np.arange(no_core_orbitals, orbitals))
    else:
        assert type(active_orbitals) == int, 'The number of active orbitals must be integer.'
        assert active_orbitals > 0, 'The number of active orbitals must be greater than 0.'
        assert no_core_orbitals + active_orbitals <= orbitals, 'The summation of core orbitals and active ' \
                                                               'orbitals should be smaller than orbitals.'
        assert no_core_orbitals + active_orbitals > (electrons + multiplicity - 1) / 2, \
            'The summation of core orbitals and active orbitals should be greater than ' \
            '(electrons + multiplicity - 1)/2.'

        active_orbitals = list(np.arange(no_core_orbitals, no_core_orbitals + active_orbitals))

    return core_orbitals, active_orbitals

def fermionic_hamiltonian(molecule,
                          filename=None,
                          multiplicity=1,
                          active_electrons=None,
                          active_orbitals=None):
    r"""计算给定分子的费米哈密顿量。

    Args:
        molecule (MolecularData): 包含分子所有信息的类
        filename (str, optional): 分子的 .hdf5 文件的路径
        multiplicity (int, optional): 自旋多重度
        active_electrons (int, optional): 活跃 (active) 电子数，默认情况为所有电子均为活跃电子
        active_orbitals (int, optional): 活跃 (active) 轨道数，默认情况为所有轨道均为活跃轨道

    Returns:
        openfermion.ops.operators.qubit_operator.QubitOperator: openfermion 格式的哈密顿量
    """
    if molecule is None:
        assert type(filename) == str, 'Please provide the path of .hdf5 file.'
        assert filename[-5:] == '.hdf5', 'The filename is supposed to be .hdf5 file'
        molecule = MolecularData(filename=filename)
    core_orbitals, active_orbitals = active_space(
        molecule.n_electrons,
        molecule.n_orbitals,
        multiplicity,
        active_electrons,
        active_orbitals)
    terms_molecular_hamiltonian = molecule.get_molecular_hamiltonian(
        occupied_indices=core_orbitals, active_indices=active_orbitals
    )
    fermionic_hamiltonian = openfermion.transforms.get_fermion_operator(terms_molecular_hamiltonian)

    return fermionic_hamiltonian

def spin_hamiltonian(molecule,
                     filename=None,
                     multiplicity=1,
                     mapping_method='jordan_wigner',
                     active_electrons=None,
                     active_orbitals=None):
    r"""生成 Paddle Quantum 格式的哈密顿量

    Args:
        molecule (openfermion.ops.operators.qubit_operator.QubitOperator): openfermion 格式的哈密顿量
        filename (str, optional): 分子的 .hdf5 文件的路径
        multiplicity (int, optional): 自旋多重度
        mapping_method (str, optional): 映射方法，这里默认为 ``'jordan_wigner'`` ，此外还提供 ``'bravyi_kitaev'``
        active_electrons (int, optional): 活跃 (active) 电子数，默认情况为所有电子均为活跃电子
        active_orbitals (int, optional): 活跃 (active) 轨道数默认情况为所有轨道均为活跃轨道

    Returns:
        paddle_quantum.utils.Hamiltonian: Paddle Quantum 格式的哈密顿量
    """
    assert mapping_method in ['jordan_wigner', 'bravyi_kitaev'], "Please choose the mapping " \
                                                                 "in ['jordan_wigner', 'bravyi_kitaev']."
    fermionic_h = fermionic_hamiltonian(molecule,
                                        filename,
                                        multiplicity,
                                        active_electrons,
                                        active_orbitals)

    if mapping_method == 'jordan_wigner':
        spin_h = transforms.jordan_wigner(fermionic_h)
    elif mapping_method == 'bravyi_kitaev':
        spin_h = transforms.bravyi_kitaev(fermionic_h)
    return _hamiltonian_transformation(spin_h, tol=1e-8)


from warnings import warn
from pathlib import Path

import numpy as np
import pandas as pd

from astropy.nddata import CCDData, StdDevUncertainty
from astropy.table import Table
from astropy.io import fits
import ccdproc
from ccdproc import combine, trim_image

from .filemgmt import load_if_exists, make_summary
from .ccdutil import CCDData_astype, load_ccd
from .misc import chk_keyval

__all__ = ["sstd", "weighted_mean", "stack_FITS", "combine_ccd"]

def sstd(a, **kwargs):
    ''' Sample standard deviation function
    '''
    return np.std(a, ddof=1, **kwargs)

# FIXME: Add this to Ccdproc esp. for mem_limit
def weighted_mean(ccds, unit='adu'):
    datas = []
    ws = []  # weights = 1 / sigma**2
    for ccd in ccds:
        datas.append(ccd.data)
        ws.append(1 / ccd.uncertainty.array**2)
    wmean = np.average(np.array(datas), axis=0, weights=ws)
    wuncert = np.sqrt(1 / np.sum(np.array(ws), axis=0))
    nccd = CCDData(data=wmean, header=ccds[0].header, unit=unit)
    nccd.uncertainty = StdDevUncertainty(wuncert)
    return nccd

def group_FITS(summary_table, type_key=None, type_val=None, group_key=None):
    ''' Organize the group_by and type_key for stack_FITS
    Parameters
    ----------
    summary_table: pandas.DataFrame or astropy.table.Table
        The table which contains the metadata (header) of files. If it
        is in the astropy table format, it will be converted to
        `~pandas.DataFrame` object.

    type_key, type_val: None, str, list of str, optional
        The header keyword for the ccd type, and the value you want to
        match.

    group_key : None, str, list of str, optional
        The header keyword which will be used to make groups for the
        CCDs that have selected from ``type_key`` and ``type_val``. If
        ``None`` (default), no grouping will occur, but it will return
        the `~pandas.DataFrameGroupBy` object will be returned for the
        sake of consistency.

    Return
    ------
    grouped : ~pandas.DataFrameGroupBy
        The table after the grouping process.

    group_type_key : list of str
        The ``type_key`` that can directly be used for ``stack_FITS``
        for each element of ``grouped.groups``.
        Basically this is ``type_key + group_key``.

    Example
    -------
    >>> allfits = list(Path('.').glob("*.fits"))
    >>> summary_table = make_summary(allfits)
    >>> type_key = ["OBJECT"]
    >>> type_val = ["dark"]
    >>> group_key = ["EXPTIME"]
    >>> gs, g_key = group_FITS(summary_table,
    ...                        type_key,
    ...                        type_val,
    ...                        group_key)
    >>> for g_val, group in gs:
    >>>     _ = combine_ccd(group["file"],
    ...                     type_key=g_key,
    ...                     type_val=g_val)
    '''
    if ((not isinstance(summary_table, Table))
            and (not isinstance(summary_table, pd.DataFrame))):
        raise TypeError("summary_table must be an astropy Table or Pandas "
                        + f"DataFrame. It's now {type(summary_table)}.")
    elif isinstance(summary_table, Table):
        st = summary_table.to_pandas()
    else:
        st = summary_table.copy()

    type_key, type_val, group_key = chk_keyval(type_key=type_key,
                                               type_val=type_val,
                                               group_key=group_key)

    if len(group_key + type_key) == 0:
        raise ValueError("At least one of type_key and group_key should not "
                         + "be empty!")

    # For simplicity, crop the original data by type_key and type_val first.
    for k, v in zip(type_key, type_val):
        st = st[st[k] == v]

    group_type_key = type_key + group_key
    grouped = st.groupby(group_key)

    return grouped, group_type_key

def stack_FITS(fitslist=None, summary_table=None, extension=0,
               unit='adu', table_filecol="file", trim_fits_section=None,
               loadccd=True, type_key=None, type_val=None):
    ''' Stacks the FITS files specified in fitslist
    Parameters
    ----------
    fitslist: None, list of path-like, or list of CCDData
        The list of path to FITS files or the list of CCDData to be
        stacked. It is useful to give list of CCDData if you have
        already stacked/loaded FITS file into a list by your own
        criteria. If ``None`` (default), you must give ``fitslist`` or
        ``summary_table``. If it is not ``None``, this function will do
        very similar job to that of ``ccdproc.combine``. Although it is
        not a good idea, a mixed list of CCDData and paths to the files
        is also acceptable.

    summary_table: None, pandas.DataFrame or astropy.table.Table
        The table which contains the metadata of files. If there are
        many FITS files and you want to use stacking many times, it is
        better to make a summary table by ``filemgmt.make_summary`` and
        use that instead of opening FITS files' headers every time you
        call this function. If you want to use ``summary_table`` instead
        of ``fitslist`` and have set ``loadccd=True``, you must not have
        ``None`` or ``NaN`` value in the
        ``summary_table[table_filecol]``.

    extension: int or str
        The extension of FITS to be stacked. For single extension, set
        it as 0.

    unit: Unit or str, optional
        The unit of the CCDs to be loaded.
        Used only when ``fitslist`` is not a list of ``CCDData`` and
        ``loadccd`` is ``True``.

    table_filecol: str
        The column name of the ``summary_table`` which contains the path
        to the FITS files.

    trim_fits_section : str or None, optional
        The ``fits_section`` of ``ccdproc.trim_image``. Region of
        ``ccd`` from which the overscan is extracted; see
        `~ccdproc.subtract_overscan` for details.
        Default is ``None``.

    loadccd: bool, optional
        Whether to return file paths or loaded CCDData. If ``False``, it
        is a function to select FITS files using ``type_key`` and
        ``type_val`` without using much memory.
        This is ignored if ``fitslist`` is given and composed of
        ``CCDData`` objects.

    type_key, type_val: str, list of str
        The header keyword for the ccd type, and the value you want to
        match.

    Return
    ------
    matched: list of Path or list of CCDData
        list containing Path to files if ``loadccd`` is ``False``.
        Otherwise it is a list containing loaded CCDData after loading
        the files. If ``ccdlist`` is given a priori, list of CCDData
        will be returned regardless of ``loadccd``.
    '''
    def _parse_val(value):
        val = str(value)
        if val.lstrip('+-').isdigit():  # if int
            result = int(val)
        else:
            try:
                result = float(val)
            except ValueError:
                result = str(val)
        return result

    def _check_mismatch(row):
        mismatch = False
        for k, v in zip(type_key, type_val):
            hdr_val = _parse_val(row[k])
            parse_v = _parse_val(v)
            if (hdr_val != parse_v):
                mismatch = True
                break
        return mismatch

    if ((fitslist is not None) + (summary_table is not None) != 1):
        raise ValueError("One and only one of fitslist or summary_table must "
                         + "be not None.")

    # If fitslist
    if fitslist is not None:
        table_mode = False
        try:
            fitslist = list(fitslist)
        except TypeError:
            raise TypeError("fitslist must be convertable to list. "
                            + f"It's now {type(fitslist)}.")

    # If summary_table
    if summary_table is not None:
        table_mode = True
        if ((not isinstance(summary_table, Table))
                and (not isinstance(summary_table, pd.DataFrame))):
            raise TypeError("summary_table must be an astropy Table or Pandas "
                            + f"DataFrame. It's now {type(summary_table)}.")

    # Check for type_key and type_val
    type_key, type_val, _ = chk_keyval(type_key=type_key,
                                       type_val=type_val,
                                       group_key=None)

    # Setting whether to group
    grouping = False
    if len(type_key) > 0:
        grouping = True

    print("Analyzing FITS... ", end='')
    # Set fitslist and summary_table based on the given input and grouping.
    if table_mode:
        if isinstance(summary_table, Table):
            summary_table = summary_table.to_pandas()
        fitslist = summary_table[table_filecol].tolist()
    else:
        if grouping:
            summary_table = make_summary(fitslist,
                                         extension=extension,
                                         verbose=True,
                                         fname_option='relative',
                                         keywords=type_key,
                                         sort_by=None,
                                         pandas=True)
        # else: no need to make summary_table.

    print("Done", end='')

    if loadccd:
        print(" and loading FITS... ")
    else:
        print(".")

    matched = []

    # Append appropriate CCDs or filepaths to matched
    if grouping:  # summary_table is used.
        for i, row in summary_table.iterrows():
            mismatch = _check_mismatch(row)
            if mismatch:  # skip this row (file)
                continue

            # if not skipped:
            # TODO: Is it better to remove Path here?
            if isinstance(fitslist[i], CCDData):
                matched.append(fitslist[i])
            else:  # it must be a path to the file
                fpath = Path(fitslist[i])
                if loadccd:
                    ccd_i = load_ccd(fpath, extension=extension, unit=unit)
                    if trim_fits_section is not None:
                        ccd_i = trim_image(ccd_i,
                                           fits_section=trim_fits_section)
                    matched.append(ccd_i)
                else:
                    matched.append(fpath)
    else:  # summary_table is not used.
        for item in fitslist:
            if isinstance(item, CCDData):
                matched.append(item)
            else:
                if loadccd:
                    ccd_i = load_ccd(item, extension=extension, unit=unit)
                    if trim_fits_section is not None:
                        ccd_i = trim_image(
                            ccd_i, fits_section=trim_fits_section)
                    matched.append(ccd_i)
                else:  # TODO: Is is better to remove Path here?
                    matched.append(Path(item))

    # Generate warning OR information messages
    if len(matched) == 0:
        if grouping:
            warn('No FITS file had "{:s} = {:s}"'.format(str(type_key),
                                                         str(type_val))
                 + "Maybe int/float/str confusing?")
        else:
            warn('No FITS file found')
    else:
        if grouping:
            N = len(matched)
            ks = str(type_key)
            vs = str(type_val)
            if loadccd:
                print(f'{N} FITS files with "{ks} = {vs}" are loaded.')
            else:
                print(f'{N} FITS files with "{ks} = {vs}" are selected.')
        else:
            if loadccd:
                print('{:d} FITS files are loaded.'.format(len(matched)))

    return matched

def combine_ccd(fitslist=None, summary_table=None, table_filecol="file",
                trim_fits_section=None, output=None, unit='adu',
                subtract_frame=None, combine_method='median',
                reject_method=None, normalize_exposure=False,
                normalize_average=False,
                exposure_key='EXPTIME', mem_limit=2e9,
                combine_uncertainty_function=None,
                extension=0, type_key=None, type_val=None,
                dtype="float32", uncertainty_dtype="float32",
                output_verify='fix', overwrite=False,
                verbose=True, **kwargs):
    ''' Combining images
    Slight variant from ccdproc.
    # TODO: accept the input like ``sigma_clip_func='median'``, etc.
    Parameters
    ----------
    fitslist: path-like, list of path-like, or list of CCDData
        The list of path to FITS files or the list of CCDData to be
        stacked. It is useful to give list of CCDData if you have
        already stacked/loaded FITS file into a list by your own
        criteria. If ``None`` (default), you must give ``fitslist`` or
        ``summary_table``. If it is not ``None``, this function will do
        very similar job to that of ``ccdproc.combine``. Although it is
        not a good idea, a mixed list of CCDData and paths to the files
        is also acceptable.

    summary_table: pandas.DataFrame or astropy.table.Table
        The table which contains the metadata of files. If there are
        many FITS files and you want to use stacking many times, it is
        better to make a summary table by ``filemgmt.make_summary`` and
        use that instead of opening FITS files' headers every time you
        call this function. If you want to use ``summary_table`` instead
        of ``fitslist`` and have set ``loadccd=True``, you must not have
        ``None`` or ``NaN`` value in the
        ``summary_table[table_filecol]``.

    table_filecol: str
        The column name of the ``summary_table`` which contains the path
        to the FITS files.

    trim_fits_section : str or None, optional
        The ``fits_section`` of ``ccdproc.trim_image``. Region of
        ``ccd`` from which the overscan is extracted; see
        `~ccdproc.subtract_overscan` for details.
        Default is ``None``.

    output : path-like or None, optional.
        The path if you want to save the resulting ``ccd`` object.
        Default is ``None``.

    unit : `~astropy.units.Unit` or str, optional.
        The units of the data.
        Default is ``'adu'``.

    subtract_frame : array-like, optional.
        The frame you want to subtract from the image after the
        combination. It can be, e.g., dark frame, because it is easier
        to calculate Poisson error before the dark subtraction and
        subtract the dark later.
        TODO: This maybe unnecessary.
        Default is ``None``.

    combine_method : str or None, optinal.
        The ``method`` for ``ccdproc.combine``, i.e., {'average',
        'median', 'sum'}
        Default is ``None``.

    reject_method : str
        Made for simple use of ``ccdproc.combine``, [None, 'minmax',
        'sigclip' == 'sigma_clip', 'extrema' == 'ext']. Automatically
        turns on the option, e.g., ``clip_extrema = True`` or
        ``sigma_clip = True``. Leave it blank for no rejection.
        Default is ``None``.

    normalize_exposure : bool, optional.
        Whether to normalize the values by the exposure time of each
        frame before combining.
        Default is ``False``.

    normalize_average : bool, optional.
        Whether to normalize the values by the average value of each
        frame before combining.
        Default is ``False``.

    exposure_key : str, optional
        The header keyword for the exposure time.
        Default is ``"EXPTIME"``.

    combine_uncertainty_function : callable, None, optional
        The uncertainty calculation function of ``ccdproc.combine``. If
        ``None`` use the default uncertainty func when using average,
        median or sum combine, otherwise use the function provided.
        Default is ``None``.

    extension: int or str, optional
        The extension to be used.
        Default is ``0``.

    dtype : str or `numpy.dtype` or None, optional
        Allows user to set dtype. See `numpy.array` ``dtype`` parameter
        description. If ``None`` it uses ``np.float64``.
        Default is ``None``.

    type_key, type_val: str, list of str
        The header keyword for the ccd type, and the value you want to
        match. For an open HDU named ``hdu``, e.g., only the files which
        satisfies ``hdu[extension].header[type_key] == type_val`` among
        all the ``fitslist`` will be used.

    output_verify : str
        Output verification option.  Must be one of ``"fix"``,
        ``"silentfix"``, ``"ignore"``, ``"warn"``, or ``"exception"``.
        May also be any combination of ``"fix"`` or ``"silentfix"`` with
        ``"+ignore"``, ``+warn``, or ``+exception" (e.g.
        ``"fix+warn"``).  See the astropy documentation below:
        http://docs.astropy.org/en/stable/io/fits/api/verification.html#verify

    mem_limit : float, optional
        Maximum memory which should be used while combining (in bytes).
        Default is ``2.e9``.

    **kwarg:
        kwargs for the ``ccdproc.combine``. See its documentation. This
        includes (RHS are the default values)
        ```
        weights=None,
        scale=None,
        mem_limit=16000000000.0,
        clip_extrema=False,
        nlow=1,
        nhigh=1,
        minmax_clip=False,
        minmax_clip_min=None,
        minmax_clip_max=None,
        sigma_clip=False,
        sigma_clip_low_thresh=3,
        sigma_clip_high_thresh=3,
        sigma_clip_func=<numpy.ma.core._frommethod instance>,
        sigma_clip_dev_func=<numpy.ma.core._frommethod instance>,
        combine_uncertainty_function=None, **ccdkwargs
        ```

    Returns
    -------
    master: astropy.nddata.CCDData
        Resulting combined ccd.
    '''
    def _set_reject_method(reject_method):
        ''' Convenience function for ccdproc.combine reject switches
        '''
        clip_extrema, minmax_clip, sigma_clip = False, False, False

        if reject_method in ['extrema', 'ext']:
            clip_extrema = True
        elif reject_method in ['minmax']:
            minmax_clip = True
        elif reject_method in ['sigma_clip', 'sigclip']:
            sigma_clip = True
        else:
            if reject_method not in [None, 'no']:
                raise KeyError("reject must be one in [None, 'minmax', "
                               + "'sigclip'=='sigma_clip', 'extrema'=='ext']")

        return clip_extrema, minmax_clip, sigma_clip

    def _print_info(combine_method, Nccd, reject_method, **kwargs):
        if reject_method is None:
            reject_method = 'no'

        info_str = ('"{:s}" combine {:d} images by "{:s}" rejection')

        print(info_str.format(combine_method, Nccd, reject_method))
        print(dict(**kwargs))
        return

    # def _normalize_exptime(ccdlist, exposure_key):
    #     _ccdlist = ccdlist.copy()
    #     exptimes = []

    #     for i in range(len(_ccdlist)):
    #         exptime = _ccdlist[i].header[exposure_key]
    #         exptimes.append(exptime)
    #         _ccdlist[i] = _ccdlist[i].divide(exptime)

    #     if verbose:
    #         if len(np.unique(exptimes)) != 1:
    #             print('There are more than one exposure times:\n\t', end=' ')
    #             print(np.unique(exptimes), end=' ')
    #             print('seconds')
    #         print(f'Normalized images by exposure time ("{exposure_key}").')

    #     return _ccdlist

    def _add_and_print(s, header, verbose):
        header.add_history(s)
        if verbose:
            print(s)
    # Give only one
    if ((fitslist is not None) + (summary_table is not None) != 1):
        raise ValueError(
            "One and only one of [fitslist, summary_table] must be given.")

    # If fitslist
    if fitslist is not None:
        try:
            fitslist = list(fitslist)
        except TypeError:
            raise TypeError("fitslist must be convertable to list. "
                            + f"It's now {type(fitslist)}.")

    # If summary_table
    if summary_table is not None:
        if ((not isinstance(summary_table, Table))
                and (not isinstance(summary_table, pd.DataFrame))):
            raise TypeError("summary_table must be an astropy Table or Pandas "
                            + f"DataFrame. It's now {type(summary_table)}.")

    # Check for type_key and type_val
    if ((type_key is None) ^ (type_val is None)):
        raise ValueError(
            "type_key and type_val must be both specified or both None.")

    if (output is not None) and (Path(output).exists()):
        if overwrite:
            print(f"{output} already exists:\n\tBut will be overridden.")
        else:
            print(f"{output} already exists:")
            return load_if_exists(output, loader=CCDData.read, if_not=None)

    # Do we really need to accept all three of normalize & scale?
    # if scale is None:
    #     scale = np.ones(len(ccdlist))
    if (((normalize_average) + (normalize_exposure)) > 1):
        raise ValueError("Only up to one of [normalize_average, "
                         + "normalize_exposure] must be not None.")

    # Set history messages
    str_history = ('{:d} images with {:s} = {:s} are "{:s}" combined '
                   + 'using "{:s}" rejection (additional kwargs: {})')
    str_nexp = "Each frame normalized by exposure time before combination."
    str_navg = "Each frame normalized by average value before combination."
    str_subt = "Subtracted a user-provided frame"
    str_trim = "Trim by FITS section {}"

    if reject_method is None:
        reject_method = 'no'

    # Select CCDs by
    ccdlist = stack_FITS(fitslist=fitslist,
                         summary_table=summary_table,
                         table_filecol=table_filecol,
                         extension=extension,
                         unit=unit,
                         type_key=type_key,
                         type_val=type_val,
                         loadccd=False)
    #  trim_fits_section=trim_fits_section,
    # loadccd=False: Loading CCD here may cause memory blast...

    try:
        header = ccdlist[0].header
    except AttributeError:
        header = fits.getheader(ccdlist[0])

    if verbose:
        _print_info(combine_method=combine_method,
                    Nccd=len(ccdlist),
                    reject_method=reject_method,
                    dtype=dtype,
                    **kwargs)

    scale = None
    # Normalize by exposure
    if normalize_exposure:
        tmp = make_summary(fitslist=fitslist,
                           keywords=[exposure_key],
                           verbose=False,
                           sort_by=None)
        exptimes = tmp[exposure_key].tolist()
        scale = 1 / np.array(exptimes)
        _add_and_print(str_nexp, header, verbose)

    # Normalize by pixel average
    if normalize_average:
        def invavg(a):
            return 1 / np.mean(a)
        scale = invavg
        _add_and_print(str_navg, header, verbose)

    # Set rejection switches
    clip_extrema, minmax_clip, sigma_clip = _set_reject_method(reject_method)

    if len(ccdlist) == 1:
        if isinstance(ccdlist[0], CCDData):
            master = ccdlist[0]
        else:
            master = load_ccd(ccdlist[0], extension=extension, unit=unit)
    else:
        master = combine(img_list=ccdlist,
                         method=combine_method,
                         clip_extrema=clip_extrema,
                         minmax_clip=minmax_clip,
                         sigma_clip=sigma_clip,
                         mem_limit=mem_limit,
                         combine_uncertainty_function=combine_uncertainty_function,
                         unit=unit,
                         hdu=extension,
                         scale=scale,
                         dtype=dtype,
                         **kwargs)

    ncombine = len(ccdlist)
    header["COMBVER"] = (ccdproc.__version__,
                         "ccdproc version used for combine.")
    header["NCOMBINE"] = (ncombine, "Number of combined images")
    header["COMBMETH"] = (combine_method, "Combining method")

    header.add_history(str_history.format(ncombine,
                                          str(type_key),
                                          str(type_val),
                                          str(combine_method),
                                          str(reject_method),
                                          kwargs))

    if subtract_frame is not None:
        subtract = CCDData(subtract_frame.copy())
        master.data = master.subtract(subtract).data
        _add_and_print(str_subt, header, verbose)

    if trim_fits_section is not None:
        master = trim_image(master, fits_section=trim_fits_section)
        _add_and_print(str_trim.format(trim_fits_section), header, verbose)

    master.header = header
    master = CCDData_astype(master, dtype=dtype,
                            uncertainty_dtype=uncertainty_dtype)

    if output is not None:
        if verbose:
            print(f"Writing FITS to {output}... ", end='')
        master.write(output, output_verify=output_verify, overwrite=overwrite)
        if verbose:
            print("Saved.")

    return master