/*
 *            Copyright 2009-2017 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      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.
 *
 */

// Overload of uBLAS prod function with MKL/GSL implementations

#include "votca/xtp/aobasis.h"
#include "votca/xtp/qminterface.h"
#include <votca/xtp/dftengine.h>

#include <boost/format.hpp>
#include <boost/filesystem.hpp>
#include <boost/numeric/ublas/operation.hpp>
#include <votca/xtp/aomatrix.h>
#include <votca/xtp/threecenters.h>

#include <votca/xtp/qmpackagefactory.h>
#include <boost/math/constants/constants.hpp>
#include <boost/numeric/ublas/symmetric.hpp>
#include <votca/tools/linalg.h>
#include <votca/tools/constants.h>

#include <votca/xtp/elements.h>
#include <votca/xtp/diis.h>

#include <votca/ctp/xinteractor.h>
#include <votca/ctp/logger.h>

using boost::format;
using namespace boost::filesystem;

namespace votca {
    namespace xtp {
        namespace ub = boost::numeric::ublas;
        // +++++++++++++++++++++++++++++ //
        // DFTENGINE MEMBER FUNCTIONS        //
        // +++++++++++++++++++++++++++++ //

        void DFTENGINE::CleanUp() {

        }

        void DFTENGINE::Initialize(Property* options) {

            string key = Identify();

            // get OpenMP thread number

            _openmp_threads = options->ifExistsReturnElseReturnDefault<int>(key + ".openmp", 0);

            // basis sets
            _dftbasis_name = options->ifExistsReturnElseThrowRuntimeError<string>(key + ".dftbasis");

            if (options->exists(key + ".auxbasis")) {
                _auxbasis_name = options->get(key + ".auxbasis").as<string>();
                _with_RI = true;
            } else {
                _with_RI = false;
            }

            if (options->exists(key + ".ecp")) {
                _ecp_name = options->get(key + ".ecp").as<string>();
                _with_ecp = true;
            } else {
                _with_ecp = false;
            }
            _with_guess = options->ifExistsReturnElseReturnDefault<bool>(key + ".read_guess", false);
            _initial_guess = options->ifExistsReturnElseReturnDefault<string>(key + ".initial_guess", "atom");

            // numerical integrations
            _grid_name = options->ifExistsReturnElseReturnDefault<string>(key + ".integration_grid", "medium");
            _use_small_grid = options->ifExistsReturnElseReturnDefault<bool>(key + ".integration_grid_small", true);
            _grid_name_small = Choosesmallgrid(_grid_name);

            // exchange and correlation as in libXC

            _xc_functional_name = options->ifExistsReturnElseThrowRuntimeError<string>(key + ".xc_functional");

            _numofelectrons = 0;

            if (options->exists(key + ".convergence")) {

                _Econverged = options->ifExistsReturnElseReturnDefault<double>(key + ".convergence.energy", 1e-7);
                _error_converged = options->ifExistsReturnElseReturnDefault<double>(key + ".convergence.error", 1e-7);
                _max_iter = options->ifExistsReturnElseReturnDefault<int>(key + ".convergence.max_iterations", 100);

                if (options->exists(key + ".convergence.method")) {
                    string method = options->get(key + ".convergence.method").as<string>();
                    if (method == "DIIS") {
                        _usediis = true;
                    } else if (method == "mixing") {
                        _usediis = false;
                    } else {
                        cout << "WARNING method not known. Using Mixing" << endl;
                        _usediis = false;
                    }
                } else {
                    _usediis = true;
                }
                if (!_usediis) {
                    _histlength = 1;
                    _maxout = false;
                }

                if (options->exists(key + ".convergence.mixing")) {
                    _useautomaticmixing = false;
                    _mixingparameter = options->get(key + ".convergence.mixing").as<double>();
                } else {
                    _useautomaticmixing = true;
                    _mixingparameter = -10;
                }

                _levelshift = options->ifExistsReturnElseReturnDefault<double>(key + ".convergence.levelshift", 0.0);
                _levelshiftend = options->ifExistsReturnElseReturnDefault<double>(key + ".convergence.levelshift_end", 0.8);
                _maxout = options->ifExistsReturnElseReturnDefault<bool>(key + ".convergence.DIIS_maxout", false);
                _histlength = options->ifExistsReturnElseReturnDefault<int>(key + ".convergence.DIIS_length", 10);
                _diis_start = options->ifExistsReturnElseReturnDefault<double>(key + ".convergence.DIIS_start", 0.01);
                _adiis_start = options->ifExistsReturnElseReturnDefault<double>(key + ".convergence.ADIIS_start", 2);

            } else {
                _Econverged = 1e-7;
                _error_converged = 1e-7;
                _maxout = false;
                _diis_start = 0.01;
                _adiis_start = 2;
                _histlength = 10;
                _useautomaticmixing = true;
                _mixingparameter = -10;
                _usediis = true;
                _max_iter = 100;
                _levelshift = 0.25;
                _levelshiftend = 0.8;
            }

            return;
        }

        /*
         *    Density Functional theory implementation
         *
         */

        bool DFTENGINE::Evaluate(Orbitals* _orbitals) {

            // set the parallelization
#ifdef _OPENMP

            omp_set_num_threads(_openmp_threads);

#endif

            /**** END OF PREPARATION ****/

            /**** Density-independent matrices ****/

            ub::vector<double>& MOEnergies = _orbitals->MOEnergies();
            ub::matrix<double>& MOCoeff = _orbitals->MOCoefficients();
            if (MOEnergies.size() != _dftbasis.AOBasisSize()) {
                MOEnergies.resize(_dftbasis.AOBasisSize());
            }
            if (MOCoeff.size1() != _dftbasis.AOBasisSize() || MOCoeff.size2() != _dftbasis.AOBasisSize()) {
                MOCoeff.resize(_dftbasis.AOBasisSize(), _dftbasis.AOBasisSize());
            }

            /**** Construct initial density  ****/

            ub::matrix<double> H0 = _dftAOkinetic.Matrix() + _dftAOESP.getNuclearpotential();

            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Constructed inital density " << flush;

            NuclearRepulsion();

            if (_addexternalsites) {
                H0 += _dftAOESP.getExternalpotential();

                H0 += _dftAODipole_Potential.getExternalpotential();
                H0 += _dftAOQuadrupole_Potential.getExternalpotential();

                double estat = ExternalRepulsion();
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " E_electrostatic " << estat << flush;
                E_nucnuc += estat;

            }

            if (_do_externalfield) {
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Integrated external potential on grid " << flush;
                double extneralgrid_nucint = ExternalGridRepulsion(_externalgrid_nuc);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Nuclei external potential interaction " << extneralgrid_nucint << " Hartree" << flush;
                H0 += _gridIntegration_ext.IntegrateExternalPotential(_externalgrid);
                E_nucnuc += extneralgrid_nucint;
            }

            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Nuclear Repulsion Energy is " << E_nucnuc << flush;

            // if we have a guess we do not need this.
            if (_with_guess) {
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Reading guess from orbitals object/file" << flush;
                _dftAOdmat = _orbitals->DensityMatrixGroundState();
            } else if (guess_set) {
                ConfigOrbfile(_orbitals);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Using starting guess from last iteration" << flush;
                _dftAOdmat = last_dmat;
            } else {
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Setup Initial Guess using: " << _initial_guess << flush;
                // this temp is necessary because eigenvalues_general returns MO^T and not MO
                if (_initial_guess == "independent") {
                    ub::matrix<double> copy = H0;
                    _diis.SolveFockmatrix(MOEnergies, MOCoeff, copy);
                    _dftAOdmat = _orbitals->DensityMatrixGroundState();

                } else if (_initial_guess == "atom") {

                    _dftAOdmat = AtomicGuess(_orbitals);
                    //cout<<_dftAOdmat<<endl;

                    if (_with_RI) {
                        _ERIs.CalculateERIs(_dftAOdmat);
                    } else {
                        _ERIs.CalculateERIs_4c_small_molecule(_dftAOdmat);
                    }
                    if (_use_small_grid) {
                        _orbitals->AOVxc() = _gridIntegration_small.IntegrateVXC(_dftAOdmat);
                        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled approximate DFT Vxc matrix " << flush;
                    } else {
                        _orbitals->AOVxc() = _gridIntegration.IntegrateVXC(_dftAOdmat);
                        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT Vxc matrix " << flush;
                    }
                    ub::matrix<double> H = H0 + _ERIs.getERIs() + _orbitals->AOVxc();
                    _diis.SolveFockmatrix(MOEnergies, MOCoeff, H);
                    _dftAOdmat = _orbitals->DensityMatrixGroundState();
                    //cout<<_dftAOdmat<<endl;
                    //Have to do one full iteration here, levelshift needs MOs;
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Full atomic density Matrix gives N=" << std::setprecision(9) << linalg_traceofProd(_dftAOdmat, _dftAOoverlap.Matrix()) << " electrons." << flush;
                } else {
                    throw runtime_error("Initial guess method not known/implemented");
                }

            }

            _orbitals->setQMpackage("xtp");

            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " STARTING SCF cycle" << flush;
            CTP_LOG(ctp::logDEBUG, *_pLog) << " --------------------------------------------------------------------------" << flush;

            double energyold = 0;
            double diiserror = 100; //is evolved in DIIs scheme
            Mixing Mixer(_useautomaticmixing, _mixingparameter, &_dftAOoverlap.Matrix(), _pLog);

            for (_this_iter = 0; _this_iter < _max_iter; _this_iter++) {
                CTP_LOG(ctp::logDEBUG, *_pLog) << flush;
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Iteration " << _this_iter + 1 << " of " << _max_iter << flush;

                if (_with_RI) {
                    _ERIs.CalculateERIs(_dftAOdmat);
                } else {
                    _ERIs.CalculateERIs_4c_small_molecule(_dftAOdmat);
                }

                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT Electron repulsion matrix of dimension: " << _ERIs.getSize1() << " x " << _ERIs.getSize2() << flush;
                double vxcenergy = 0.0;
                if (_use_small_grid && diiserror > 1e-5) {
                    _orbitals->AOVxc() = _gridIntegration_small.IntegrateVXC(_dftAOdmat);
                    vxcenergy = _gridIntegration_small.getTotEcontribution();
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled approximate DFT Vxc matrix " << flush;
                } else {
                    _orbitals->AOVxc() = _gridIntegration.IntegrateVXC(_dftAOdmat);
                    vxcenergy = _gridIntegration.getTotEcontribution();
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT Vxc matrix " << flush;
                }
                //cout<<_dftAOdmat<<endl;
                ub::matrix<double> H = H0 + _ERIs.getERIs() + _orbitals->AOVxc();
                //cout<<H0<<endl;
                //exit(0);
                double Eone = linalg_traceofProd(_dftAOdmat, H0);
                double Etwo = 0.5 * _ERIs.getERIsenergy() + vxcenergy;
                double totenergy = Eone + E_nucnuc + Etwo;
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Single particle energy " << std::setprecision(12) << Eone << flush;
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Two particle energy " << std::setprecision(12) << Etwo << flush;
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << std::setprecision(12) << " Exc contribution " << vxcenergy << flush;
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Total Energy " << std::setprecision(12) << totenergy << flush;

                diiserror = _diis.Evolve(_dftAOdmat, H, MOEnergies, MOCoeff, _this_iter, totenergy);
                //cout<<"Energies "<<MOEnergies<<endl;
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " DIIs error " << diiserror << flush;

                ub::matrix<double> dmatin = _dftAOdmat;
                _dftAOdmat = _orbitals->DensityMatrixGroundState();
                if (!(diiserror < _adiis_start && _usediis && _this_iter > 2)) {
                    _dftAOdmat = Mixer.MixDmat(dmatin, _dftAOdmat);

                } else {
                    Mixer.Updatemix(dmatin, _dftAOdmat);
                }

                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Updated Density Matrix " << flush;

                if (tools::globals::verbose) {
                    for (int i = 0; i<int(MOEnergies.size()); i++) {
                        if (i <= _numofelectrons / 2 - 1) {
                            CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t" << i << " occ " << MOEnergies(i) << flush;
                        } else {
                            CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t" << i << " vir " << MOEnergies(i) << flush;

                        }
                    }
                }

                CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\tGAP " << MOEnergies(_numofelectrons / 2) - MOEnergies(_numofelectrons / 2 - 1) << flush;

                if (std::abs(totenergy - energyold) < _Econverged && diiserror < _error_converged) {
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Total Energy has converged to " << std::setprecision(9) << std::abs(totenergy - energyold) << "[Ha] after " << _this_iter + 1 <<
                            " iterations. DIIS error is converged up to " << _error_converged << "[Ha]" << flush;
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Final Single Point Energy " << std::setprecision(12) << totenergy << " Ha" << flush;
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " MO Energies  [Ha]" << flush;
                    for (int i = 0; i<int(MOEnergies.size()); i++) {
                        if (i <= _numofelectrons / 2 - 1) {
                            CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t" << i << " occ " << std::setprecision(12) << MOEnergies(i) << flush;
                        } else {
                            CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t" << i << " vir " << std::setprecision(12) << MOEnergies(i) << flush;
                        }
                    }

                    last_dmat = _dftAOdmat;
                    guess_set = true;
                    // orbitals saves total energies in [eV]
                    _orbitals->setQMEnergy(totenergy * tools::conv::hrt2ev);
                    break;
                } else {
                    energyold = totenergy;
                }
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Density Matrix gives N=" << std::setprecision(9) << linalg_traceofProd(_dftAOdmat, _dftAOoverlap.Matrix()) << " electrons." << flush;

            }
            return true;
        }

        // SETUP INVARIANT AOMatrices

        void DFTENGINE::SetupInvariantMatrices() {

            // local variables for checks
            // check eigenvalues of overlap matrix, if too small basis might have linear dependencies
            ub::vector<double> _eigenvalues;
            ub::matrix<double> _eigenvectors;

            {
                // DFT AOOverlap matrix
                _dftAOoverlap.Fill(_dftbasis);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT Overlap matrix of dimension: " << _dftAOoverlap.Dimension() << flush;
                //cout<<"overlap"<<_dftAOoverlap.Matrix()<<endl;
                // check DFT basis for linear dependence
                linalg_eigenvalues(_dftAOoverlap.Matrix(), _eigenvalues, _eigenvectors);

                //Not brilliant but we need S-1/2 for DIIS and I do not want to calculate it each time
                ub::matrix<double> _diagS = ub::zero_matrix<double>(_eigenvectors.size1(), _eigenvectors.size1());
                for (unsigned _i = 0; _i < _eigenvalues.size(); _i++) {

                    _diagS(_i, _i) = 1.0 / sqrt(_eigenvalues[_i]);
                }
                ub::matrix<double> _temp = ub::prod(_diagS, ub::trans(_eigenvectors));
                _Sminusonehalf = ub::prod(_eigenvectors, _temp);

                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Smallest eigenvalue of DFT Overlap matrix : " << _eigenvalues[0] << flush;
            }

            _dftAOkinetic.Fill(_dftbasis);
            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT Kinetic energy matrix of dimension: " << _dftAOkinetic.Dimension() << flush;

            _dftAOESP.Fillnucpotential(_dftbasis, _atoms, _with_ecp);
            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT nuclear potential matrix of dimension: " << _dftAOESP.Dimension() << flush;

            if (_addexternalsites) {
                _dftAOESP.Fillextpotential(_dftbasis, _externalsites);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT external pointcharge potential matrix of dimension: " << _dftAOESP.Dimension() << flush;

                _dftAODipole_Potential.Fillextpotential(_dftbasis, _externalsites);
                if (_dftAODipole_Potential.Dimension() > 0) {
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT external dipole potential matrix of dimension: " << _dftAODipole_Potential.Dimension() << flush;
                }
                _dftAOQuadrupole_Potential.Fillextpotential(_dftbasis, _externalsites);
                if (_dftAOQuadrupole_Potential.Dimension()) {
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT external quadrupole potential matrix of dimension: " << _dftAOQuadrupole_Potential.Dimension() << flush;
                }
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " External sites\t Name \t Coordinates \t charge \t dipole \t quadrupole" << flush;

                for (unsigned i = 0; i < _externalsites.size(); i++) {

                    vector<ctp::APolarSite*> ::iterator pit;
                    for (pit = _externalsites[i]->begin(); pit < _externalsites[i]->end(); ++pit) {

                        CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t " << (*pit)->getName() << " | " << (*pit)->getPos().getX()
                                << " " << (*pit)->getPos().getY() << " " << (*pit)->getPos().getZ() << " | " << (*pit)->getQ00();
                        if ((*pit)->getRank() > 0) {
                            tools::vec dipole = (*pit)->getQ1();
                            CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t " << " | " << dipole.getX()
                                    << " " << dipole.getY() << " " << dipole.getZ();
                        }
                        if ((*pit)->getRank() > 1) {
                            std::vector<double> quadrupole = (*pit)->getQ2();
                            CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t " << " | " << quadrupole[0] << " " << quadrupole[1] << " " << quadrupole[2] << " "
                                    << quadrupole[3] << " " << quadrupole[4];
                        }
                        CTP_LOG(ctp::logDEBUG, *_pLog)  << flush;
                    }
                }
            }

            if (_with_ecp) {

                _dftAOECP.Fill(_dftbasis, vec(0, 0, 0), &_ecp);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT ECP matrix of dimension: " << _dftAOECP.Dimension() << flush;
                _dftAOESP.getNuclearpotential() += _dftAOECP.Matrix();
            }

            _diis.Configure(_usediis, true, _histlength, _maxout, _diismethod, _adiis_start, _diis_start, _levelshift, _levelshiftend, _numofelectrons / 2);
            _diis.setLogger(_pLog);
            _diis.setOverlap(&_dftAOoverlap.Matrix());
            _diis.setSqrtOverlap(&_Sminusonehalf);

            if (_with_RI) {

                AOCoulomb _auxAOcoulomb;

                _auxAOcoulomb.Fill(_auxbasis);

                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled AUX Coulomb matrix of dimension: " << _auxAOcoulomb.Dimension() << flush;

                int dimensions = _auxAOcoulomb.Invert_DFT();

                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Inverted AUX Coulomb matrix, removed " << dimensions << " functions from aux basis" << flush;

                // prepare invariant part of electron repulsion integrals
                _ERIs.Initialize(_dftbasis, _auxbasis, _auxAOcoulomb.Matrix());
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Setup invariant parts of Electron Repulsion integrals " << flush;
            } else {
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Calculating 4c integrals. " << flush;
                _ERIs.Initialize_4c_small_molecule(_dftbasis);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Calculated 4c integrals. " << flush;
            }

            return;
        }

        ub::matrix<double> DFTENGINE::AtomicGuess(Orbitals* _orbitals) {
            ub::matrix<double> guess = ub::zero_matrix<double>(_dftbasis.AOBasisSize());

            std::vector<ctp::QMAtom*> uniqueelements;
            std::vector<ctp::QMAtom*>::const_iterator at;
            std::vector<ctp::QMAtom*>::iterator st;
            std::vector< ub::matrix<double> > uniqueatom_guesses;
            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Scanning molecule of size " << _atoms.size() << " for unique elements" << flush;
            for (at = _atoms.begin(); at < _atoms.end(); ++at) {
                bool exists = false;
                if (uniqueelements.size() == 0) {
                    exists = false;
                } else {
                    for (st = uniqueelements.begin(); st < uniqueelements.end(); ++st) {
                        if ((*at)->type == (*st)->type) {
                            exists = true;
                            break;
                        }
                    }
                }
                if (!exists) {
                    uniqueelements.push_back((*at));
                }
            }
            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " " << uniqueelements.size() << " unique elements found" << flush;
            Elements _elements;
            for (st = uniqueelements.begin(); st < uniqueelements.end(); ++st) {

                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Calculating atom density for " << (*st)->type << flush;
                bool with_ecp = _with_ecp;
                if ((*st)->type == "H" || (*st)->type == "He") {
                    with_ecp = false;
                }
                std::vector<ctp::QMAtom*> atom;
                atom.push_back(*st);

                AOBasis dftbasis;
                AOBasis ecp;
                NumericalIntegration gridIntegration;
                dftbasis.AOBasisFill(&_dftbasisset, atom);
                if (with_ecp) {
                    ecp.ECPFill(&_ecpbasisset, atom);
                }
                gridIntegration.GridSetup(_grid_name, &_dftbasisset, atom, &dftbasis);
                gridIntegration.setXCfunctional(_xc_functional_name);
                int numofelectrons = int(_elements.getNucCrg((*st)->type));
                int alpha_e = 0;
                int beta_e = 0;
                if (with_ecp) {
                    numofelectrons -= int(_ecpbasisset.getElement((*st)->type)->getNcore());
                }

                if ((numofelectrons % 2) != 0) {
                    alpha_e = numofelectrons / 2 + numofelectrons % 2;
                    beta_e = numofelectrons / 2;
                } else {
                    alpha_e = numofelectrons / 2;
                    beta_e = alpha_e;
                }

                AOOverlap dftAOoverlap;
                AOKinetic dftAOkinetic;
                AOESP dftAOESP;
                AOECP dftAOECP;
                ub::vector<double> eigenvalues;
                ub::matrix<double> eigenvectors;
                ub::matrix<double> Sminusonehalf;
                ERIs ERIs_atom;

                // DFT AOOverlap matrix

                dftAOoverlap.Fill(dftbasis);
                linalg_eigenvalues(dftAOoverlap.Matrix(), eigenvalues, eigenvectors);

                //Not brilliant but we need S-1/2 for DIIS and I do not want to calculate it each time
                ub::matrix<double> _diagS = ub::zero_matrix<double>(eigenvectors.size1(), eigenvectors.size1());
                for (unsigned _i = 0; _i < eigenvalues.size(); _i++) {

                    _diagS(_i, _i) = 1.0 / sqrt(eigenvalues[_i]);
                }
                ub::matrix<double> _temp = ub::prod(_diagS, ub::trans(eigenvectors));
                Sminusonehalf = ub::prod(eigenvectors, _temp);

                dftAOkinetic.Fill(dftbasis);

                dftAOESP.Fillnucpotential(dftbasis, atom, with_ecp);
                ERIs_atom.Initialize_4c_small_molecule(dftbasis);

                ub::vector<double>MOEnergies_alpha;
                ub::matrix<double>MOCoeff_alpha;
                ub::vector<double>MOEnergies_beta;
                ub::matrix<double>MOCoeff_beta;
                Diis diis_alpha;
                Diis diis_beta;
                double adiisstart = 0;
                double diisstart = 0.0001;
                Mixing Mix_alpha(true, 0.7, &dftAOoverlap.Matrix(), _pLog);
                Mixing Mix_beta(true, 0.7, &dftAOoverlap.Matrix(), _pLog);
                diis_alpha.Configure(true, false, 20, 0, "", adiisstart, diisstart, 0.1, 0.000, alpha_e);
                diis_alpha.setLogger(_pLog);
                diis_alpha.setOverlap(&dftAOoverlap.Matrix());
                diis_alpha.setSqrtOverlap(&Sminusonehalf);
                diis_beta.Configure(true, false, 20, 0, "", adiisstart, diisstart, 0.1, 0.000, beta_e);
                diis_beta.setLogger(_pLog);
                diis_beta.setOverlap(&dftAOoverlap.Matrix());
                diis_beta.setSqrtOverlap(&Sminusonehalf);
                /**** Construct initial density  ****/

                ub::matrix<double> H0 = dftAOkinetic.Matrix() + dftAOESP.getNuclearpotential();
                if (with_ecp) {
                    dftAOECP.Fill(dftbasis, vec(0, 0, 0), &ecp);
                    H0 += dftAOECP.Matrix();
                }
                ub::matrix<double> copy = H0;
                diis_alpha.SolveFockmatrix(MOEnergies_alpha, MOCoeff_alpha, copy);

                MOEnergies_beta = MOEnergies_alpha;
                MOCoeff_beta = MOCoeff_alpha;

                //ub::matrix<double>dftAOdmat_alpha = DensityMatrix_frac(MOCoeff_alpha,MOEnergies_alpha,alpha_e);
                //ub::matrix<double>dftAOdmat_beta = DensityMatrix_frac(MOCoeff_beta,MOEnergies_beta,beta_e);
                ub::matrix<double>dftAOdmat_alpha = DensityMatrix_unres(MOCoeff_alpha, alpha_e);
                if ((*st)->type == "H") {
                    uniqueatom_guesses.push_back(dftAOdmat_alpha);
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Atomic density Matrix for " << (*st)->type <<
                            " gives N=" << std::setprecision(9) << linalg_traceofProd(dftAOdmat_alpha, dftAOoverlap.Matrix()) << " electrons." << flush;
                    continue;
                }

                ub::matrix<double>dftAOdmat_beta = DensityMatrix_unres(MOCoeff_beta, beta_e);
                bool _HF = false;
                double energyold = 0;
                int maxiter = 50;
                for (int this_iter = 0; this_iter < maxiter; this_iter++) {

                    ERIs_atom.CalculateERIs_4c_small_molecule(dftAOdmat_alpha + dftAOdmat_beta);
                    double E_two_alpha = linalg_traceofProd(ERIs_atom.getERIs(), dftAOdmat_alpha);
                    double E_two_beta = linalg_traceofProd(ERIs_atom.getERIs(), dftAOdmat_beta);
                    ub::matrix<double> H_alpha = H0 + ERIs_atom.getERIs();
                    ub::matrix<double> H_beta = H0 + ERIs_atom.getERIs();
                    if (_HF) {
                        ERIs_atom.CalculateEXX_4c_small_molecule(dftAOdmat_alpha);
                        double E_exx_alpha = -linalg_traceofProd(ERIs_atom.getEXX(), dftAOdmat_alpha);
                        H_alpha -= ERIs_atom.getEXX();
                        E_two_alpha += E_exx_alpha;
                        ERIs_atom.CalculateEXX_4c_small_molecule(dftAOdmat_beta);
                        double E_exx_beta = -linalg_traceofProd(ERIs_atom.getEXX(), dftAOdmat_beta);
                        H_beta -= ERIs_atom.getEXX();
                        E_two_beta += E_exx_beta;

                    } else {
                        ub::matrix<double> AOVxc_alpha = gridIntegration.IntegrateVXC(dftAOdmat_alpha);
                        double E_vxc_alpha = gridIntegration.getTotEcontribution();
                        ub::matrix<double> AOVxc_beta = gridIntegration.IntegrateVXC(dftAOdmat_beta);
                        double E_vxc_beta = gridIntegration.getTotEcontribution();
                        H_alpha += AOVxc_alpha;
                        H_beta += AOVxc_beta;
                        E_two_alpha += E_vxc_alpha;
                        E_two_beta += E_vxc_beta;

                    }
                    double E_one_alpha = linalg_traceofProd(dftAOdmat_alpha, H0);
                    double E_one_beta = linalg_traceofProd(dftAOdmat_beta, H0);

                    double E_alpha = E_one_alpha + E_two_alpha;
                    double E_beta = E_one_beta + E_two_beta;

                    double totenergy = E_alpha + E_beta;
                    //evolve alpha
                    double diiserror_alpha = diis_alpha.Evolve(dftAOdmat_alpha, H_alpha, MOEnergies_alpha, MOCoeff_alpha, this_iter, E_alpha);

                    ub::matrix<double> dmatin_alpha = dftAOdmat_alpha;
                    dftAOdmat_alpha = DensityMatrix_unres(MOCoeff_alpha, alpha_e);

                    if (!(diiserror_alpha < 0.005 && this_iter > 2)) {
                        dftAOdmat_alpha = Mix_alpha.MixDmat(dmatin_alpha, dftAOdmat_alpha, false);
                        //cout<<"mixing_alpha"<<endl;
                    } else {
                        Mix_alpha.Updatemix(dmatin_alpha, dftAOdmat_alpha);
                    }
                    //evolve beta
                    double diiserror_beta = diis_beta.Evolve(dftAOdmat_beta, H_beta, MOEnergies_beta, MOCoeff_beta, this_iter, E_beta);
                    ub::matrix<double> dmatin_beta = dftAOdmat_beta;
                    dftAOdmat_beta = DensityMatrix_unres(MOCoeff_beta, beta_e);
                    //dftAOdmat_beta=DensityMatrix_frac(MOCoeff_beta,MOEnergies_beta,beta_e);
                    if (!(diiserror_beta < 0.005 && this_iter > 2)) {
                        dftAOdmat_beta = Mix_beta.MixDmat(dmatin_beta, dftAOdmat_beta, false);

                    } else {
                        Mix_beta.Updatemix(dmatin_beta, dftAOdmat_beta);
                    }

                    if (tools::globals::verbose) {
                        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Iter " << this_iter << " of " << maxiter
                                << " Etot " << totenergy << " diise_a " << diiserror_alpha << " diise_b " << diiserror_beta
                                << " a_gap " << MOEnergies_alpha(alpha_e) - MOEnergies_alpha(alpha_e - 1) << " b_gap " << MOEnergies_beta(beta_e) - MOEnergies_beta(beta_e - 1) << flush;
                    }
                    bool converged = (std::abs(totenergy - energyold) < _Econverged && diiserror_alpha < _error_converged && diiserror_beta < _error_converged);
                    if (converged || this_iter == maxiter - 1) {

                        if (converged) {
                            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Converged after " << this_iter + 1 << " iterations" << flush;
                        } else {
                            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Not converged after " << this_iter + 1 <<
                                    " iterations. Using unconverged density. DIIsError_alpha=" << diiserror_alpha << " DIIsError_beta=" << diiserror_beta << flush;
                        }

                        ub::matrix<double> avdmat = AverageShells(dftAOdmat_alpha + dftAOdmat_beta, dftbasis);
                        uniqueatom_guesses.push_back(avdmat);
                        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Atomic density Matrix for " << (*st)->type << " gives N="
                                << std::setprecision(9) << linalg_traceofProd(avdmat, dftAOoverlap.Matrix()) << " electrons." << flush;
                        break;

                    } else {
                        energyold = totenergy;
                    }
                }
            }
            unsigned start = 0;
            unsigned end = 0;
            for (at = _atoms.begin(); at < _atoms.end(); ++at) {
                unsigned index = 0;
                for (unsigned i = 0; i < uniqueelements.size(); i++) {
                    if ((*at)->type == uniqueelements[i]->type) {
                        index = i;
                        break;
                    }
                }

                Element* element = _dftbasisset.getElement((*at)->type);
                for (Element::ShellIterator its = element->firstShell(); its != element->lastShell(); its++) {
                    end += (*its)->getnumofFunc();
                }
                ub::project(guess, ub::range(start, end), ub::range(start, end)) = uniqueatom_guesses[index];

                start = end;
            }

            return guess;
        }

        void DFTENGINE::ConfigOrbfile(Orbitals* _orbitals) {
            if (_with_guess) {

                if (_orbitals->hasDFTbasis()) {
                    if (_orbitals->getDFTbasis() != _dftbasis_name) {
                        throw runtime_error((boost::format("Basisset Name in guess orb file and in dftengine option file differ %1% vs %2%") % _orbitals->getDFTbasis() % _dftbasis_name).str());
                    }
                } else {
                    CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " WARNING: Orbital file has no basisset information,using it as a guess might work or not for calculation with " << _dftbasis_name << flush;
                }
            }
            _orbitals->setDFTbasis(_dftbasis_name);
            _orbitals->setBasisSetSize(_dftbasis.AOBasisSize());
            if (_with_ecp) {
                _orbitals->setECP(_ecp_name);
            }

            if (_with_guess) {
                if(_orbitals->hasECP() || _with_ecp){
                    if(_orbitals->getECP()!=_ecp_name){
                        throw runtime_error((boost::format("ECPs in orb file: %1% and options %2% differ") % _orbitals->getECP() % _ecp_name).str());
                    }
                }
                if (_orbitals->getNumberOfElectrons() != _numofelectrons / 2) {
                    throw runtime_error((boost::format("Number of electron in guess orb file: %1% and in dftengine: %2% differ.") % _orbitals->getNumberOfElectrons() % (_numofelectrons / 2)).str());
                }
                if (_orbitals->getNumberOfLevels() != _dftbasis.AOBasisSize()) {
                    throw runtime_error((boost::format("Number of levels in guess orb file: %1% and in dftengine: %2% differ.") % _orbitals->getNumberOfLevels() % _dftbasis.AOBasisSize()).str());
                }
            } else {
                _orbitals->setNumberOfElectrons(_numofelectrons / 2);
                _orbitals->setNumberOfLevels(_numofelectrons / 2, _dftbasis.AOBasisSize() - _numofelectrons / 2);
            }
            return;
        }

        // PREPARATION

        void DFTENGINE::Prepare(Orbitals* _orbitals) {
            #ifdef _OPENMP

            omp_set_num_threads(_openmp_threads);
            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Using " << omp_get_max_threads() << " threads" << flush;

#endif

            if ( _atoms.size() == 0 ){
            for (const auto& atom : _orbitals->QMAtoms()) {
                if (!atom->from_environment) {
                    _atoms.push_back(atom);
                }
            }

            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Molecule Coordinates [A] " << flush;
            for (unsigned i = 0; i < _atoms.size(); i++) {
                CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t " << _atoms[i]->type << " " << _atoms[i]->x << " " << _atoms[i]->y << " " << _atoms[i]->z << " " << flush;
            }

            // load and fill DFT basis set
            _dftbasisset.LoadBasisSet(_dftbasis_name);

            _dftbasis.AOBasisFill(&_dftbasisset, _atoms);
            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Loaded DFT Basis Set " << _dftbasis_name << flush;

            if (_with_RI) {
                // load and fill AUX basis set
                _auxbasisset.LoadBasisSet(_auxbasis_name);
                //_orbitals->setDFTbasis( _dftbasis_name );
                _auxbasis.AOBasisFill(&_auxbasisset, _atoms);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Loaded AUX Basis Set " << _auxbasis_name << flush;
            }
            if (_with_ecp) {
                // load ECP (element-wise information) from xml file
                _ecpbasisset.LoadPseudopotentialSet(_ecp_name);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Loaded ECP library " << _ecp_name << flush;

                // fill auxiliary ECP basis by going through all atoms
                _ecp.ECPFill(&_ecpbasisset, _atoms);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled ECP Basis of size " << _ecp.getNumofShells() << flush;
            }

            // setup numerical integration grid
            _gridIntegration.GridSetup(_grid_name, &_dftbasisset, _atoms, &_dftbasis);
            _gridIntegration.setXCfunctional(_xc_functional_name);

            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Setup numerical integration grid " << _grid_name << " for vxc functional "
                    << _xc_functional_name << " with " << _gridIntegration.getGridSize() << " points" << flush;
            CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t " << " divided into " << _gridIntegration.getBoxesSize() << " boxes" << flush;
            if (_use_small_grid) {
                _gridIntegration_small.GridSetup(_grid_name_small, &_dftbasisset, _atoms, &_dftbasis);
                _gridIntegration_small.setXCfunctional(_xc_functional_name);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Setup small numerical integration grid " << _grid_name_small << " for vxc functional "
                        << _xc_functional_name << " with " << _gridIntegration_small.getGridpoints().size() << " points" << flush;
                CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t " << " divided into " << _gridIntegration_small.getBoxesSize() << " boxes" << flush;
            }

            if (_do_externalfield) {
                _gridIntegration_ext.GridSetup(_grid_name_ext, &_dftbasisset, _atoms, &_dftbasis);
                CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Setup numerical integration grid " << _grid_name_ext
                        << " for external field with " << _gridIntegration_ext.getGridpoints().size() << " points" << flush;
            }

            Elements _elements;
            //set number of electrons and such

            for (unsigned i = 0; i < _atoms.size(); i++) {
                _numofelectrons += _elements.getNucCrg(_atoms[i]->type);
            }

            // if ECP
            if (_with_ecp) {
                for (unsigned i = 0; i < _atoms.size(); i++) {
                    if (_atoms[i]->type == "H" || _atoms[i]->type == "He") {
                        continue;
                    } else {
                        _numofelectrons -= _ecpbasisset.getElement(_atoms[i]->type)->getNcore();
                    }
                }
            }
            // here number of electrons is actually the total number, everywhere else in votca it is just alpha_electrons

            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Total number of electrons: " << _numofelectrons << flush;

            ConfigOrbfile(_orbitals);
            }
            SetupInvariantMatrices();
            return;
        }

        ub::matrix<double> DFTENGINE::DensityMatrix_unres(const ub::matrix<double>& MOs, int numofelec) {
            if (numofelec == 0) {
                return ub::zero_matrix<double>(MOs.size1());
            }

            ub::matrix<double> _dmatGS = ub::zero_matrix<double>(MOs.size1());
#pragma omp parallel for
            for (unsigned _i = 0; _i < MOs.size1(); _i++) {
                for (unsigned _j = 0; _j < MOs.size1(); _j++) {
                    for (int _level = 0; _level < numofelec; _level++) {

                        _dmatGS(_i, _j) += MOs(_level, _i) * MOs(_level, _j);

                    }
                }
            }
            //}
            // return
            return _dmatGS;
        }

        ub::matrix<double> DFTENGINE::DensityMatrix_frac(const ub::matrix<double>& MOs, const ub::vector<double>& MOEnergies, int numofelec) {
            if (numofelec == 0) {
                return ub::zero_matrix<double>(MOs.size1());
            }

            ub::vector<double>occupation = ub::zero_vector<double>(MOEnergies.size());

            double buffer = 0.0001;
            double homo_energy = MOEnergies(numofelec - 1);
            std::vector<unsigned> degeneracies;

            for (unsigned _level = 0; _level < occupation.size(); _level++) {
                if (MOEnergies(_level)<(homo_energy - buffer)) {
                    occupation(_level) = 1.0;
                    numofelec--;
                } else if (std::abs(MOEnergies(_level) - homo_energy) < buffer) {
                    degeneracies.push_back(_level);
                } else if (MOEnergies(_level)>(homo_energy + buffer)) {
                    occupation(_level) = 0.0;
                }
            }
            double deg_occupation = double(numofelec) / double(degeneracies.size());
            for (unsigned _level = 0; _level < degeneracies.size(); _level++) {
                occupation(degeneracies[_level]) = deg_occupation;
            }
            ub::matrix<double> _dmatGS = ub::zero_matrix<double>(MOs.size1());
#pragma omp parallel for
            for (unsigned _i = 0; _i < MOs.size1(); _i++) {
                for (unsigned _j = 0; _j < MOs.size1(); _j++) {
                    for (unsigned _level = 0; _level < occupation.size(); _level++) {

                        _dmatGS(_i, _j) += occupation(_level) * MOs(_level, _i) * MOs(_level, _j);

                    }
                }
            }
            return _dmatGS;
        }

        void DFTENGINE::NuclearRepulsion() {
            Elements element;
            E_nucnuc = 0.0;

            std::vector<double> charge;
            for (unsigned i = 0; i < _atoms.size(); i++) {
                string name = _atoms[i]->type;
                //cout << " Using atom " << name << "\n" << endl;
                double Q = element.getNucCrg(name);
                bool HorHe = (name == "H" || name == "He");
                if (_with_ecp && !HorHe) {
                    Q -= _ecpbasisset.getElement(name)->getNcore();
                }
                charge.push_back(Q);
            }

            for (unsigned i = 0; i < _atoms.size(); i++) {
                const tools::vec& r1 = _atoms[i]->getPos() * tools::conv::ang2bohr;
                double charge1 = charge[i];
                for (unsigned j = 0; j < i; j++) {
                    const tools::vec& r2 = _atoms[j]->getPos() * tools::conv::ang2bohr;
                    double charge2 = charge[j];
                    E_nucnuc += charge1 * charge2 / (abs(r1 - r2));
                }
            }
            return;
        }

        double DFTENGINE::ExternalRepulsion(ctp::Topology* top) {
            Elements element;

            if (_externalsites.size() == 0) {
                return 0;
            }

            QMMInterface qmminter;
            ctp::PolarSeg nuclei = qmminter.Convert(_atoms);

            ctp::PolarSeg::iterator pes;
            for (ctp::APolarSite* nucleus:nuclei) {
                nucleus->setIsoP(0.0);
                string name = nucleus->getName();
                double Q = element.getNucCrg(name);
                bool HorHe = (name == "H" || name == "He");
                if (_with_ecp && !HorHe) {
                    Q -= _ecpbasisset.getElement(name)->getNcore();
                }
                nucleus->setQ00(Q, 0);
            }
            ctp::XInteractor actor;
            actor.ResetEnergy();
            nuclei.CalcPos();
            double E_ext = 0.0;
            for (ctp::PolarSeg* seg:_externalsites) {
                seg->CalcPos();
                tools::vec s = tools::vec(0.0);
                if (top == NULL) {
                    s = nuclei.getPos() -seg->getPos();
                } else {
                    s = top->PbShortestConnect(nuclei.getPos(),seg->getPos()) + nuclei.getPos() - seg->getPos();
                }

                for (auto nucleus:nuclei) {
                    for (auto site:(*seg)) {
                        actor.BiasIndu(*nucleus, *site, s);
                        nucleus->Depolarize();
                        E_ext += actor.E_f(*nucleus, *site);

                    }
                }
            }
            return E_ext * tools::conv::int2eV * tools::conv::ev2hrt;
        }

        double DFTENGINE::ExternalGridRepulsion(std::vector<double> externalpotential_nuc) {
            Elements element;
            double E_ext = 0.0;

            if (!_do_externalfield) {
                return 0;
            }

            for (unsigned i = 0; i < _atoms.size(); i++) {
                string name = _atoms[i]->type;
                double Q = element.getNucCrg(name);
                bool HorHe = (name == "H" || name == "He");
                if (_with_ecp && !HorHe) {
                    Q -= _ecpbasisset.getElement(name)->getNcore();
                }
                E_ext += Q * externalpotential_nuc[i];
            }

            return E_ext;
        }

        string DFTENGINE::Choosesmallgrid(string largegrid) {
            string smallgrid;

            if (largegrid == "xfine") {
                smallgrid = "fine";
            } else if (largegrid == "fine") {
                smallgrid = "medium";
            } else if (largegrid == "medium") {
                _use_small_grid = false;
                smallgrid = "medium";
            } else if (largegrid == "coarse") {
                _use_small_grid = false;
                smallgrid = "coarse";
            } else if (largegrid == "xcoarse") {
                _use_small_grid = false;
                smallgrid = "xcoarse";
            } else {
                throw runtime_error("Grid name for Vxc integration not known.");
            }

            return smallgrid;
        }

        //average atom densities matrices, for SP and other combined shells average each subshell separately. Does not really work yet!!

        ub::matrix<double> DFTENGINE::AverageShells(const ub::matrix<double>& dmat, AOBasis& dftbasis) {
            ub::matrix<double> avdmat = ub::zero_matrix<double>(dmat.size1());
            AOBasis::AOShellIterator it;
            int start = 0.0;
            std::vector<int> starts;
            std::vector<int> ends;
            for (it = dftbasis.firstShell(); it < dftbasis.lastShell(); ++it) {
                const AOShell* shell = dftbasis.getShell(it);
                int end = shell->getNumFunc() + start;

                if (shell->getLmax() != shell->getLmin()) {
                    std::vector<int> temp = NumFuncSubShell(shell->getType());
                    int numfunc = start;
                    for (unsigned i = 0; i < temp.size(); i++) {

                        starts.push_back(numfunc);
                        numfunc += temp[i];
                        ends.push_back(numfunc);
                    }
                } else {
                    starts.push_back(start);
                    ends.push_back(end);
                }
                start = end;
            }
            for (unsigned k = 0; k < starts.size(); k++) {
                int s1 = starts[k];
                int e1 = ends[k];
                int len1 = e1 - s1;
                for (unsigned l = 0; l < starts.size(); l++) {
                    int s2 = starts[l];
                    int e2 = ends[l];
                    int len2 = e2 - s2;
                    double diag = 0.0;
                    double offdiag = 0.0;
                    for (int i = 0; i < len1; ++i) {
                        for (int j = 0; j < len2; ++j) {
                            if (i == j) {
                                diag += dmat(s1 + i, s2 + j);
                            } else {
                                offdiag += dmat(s1 + i, s2 + j);
                            }
                        }
                    }
                    if (len1 == len2) {
                        diag = diag / double(len1);
                        offdiag = offdiag / double(len1 * (len1 - 1));
                    } else {
                        double avg = (diag + offdiag) / double(len1 * len2);
                        diag = avg;
                        offdiag = avg;
                    }
                    for (int i = 0; i < len1; ++i) {
                        for (int j = 0; j < len2; ++j) {
                            if (i == j) {
                                avdmat(s1 + i, s2 + j) = diag;
                            } else {
                                avdmat(s1 + i, s2 + j) = offdiag;
                            }
                        }
                    }
                }
            }

            return avdmat;
        }
    }
}


//   Copyright 2015 Patrick Putnam
//
//   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.
#ifndef CLOTHO_GENERAL_FITNESS_SELECTION_GENERATOR_HPP_
#define CLOTHO_GENERAL_FITNESS_SELECTION_GENERATOR_HPP_

#include <boost/property_tree/ptree.hpp>

#include <boost/random/discrete_distribution.hpp>
#include <boost/random/uniform_int_distribution.hpp>

#include "clotho/data_spaces/selection/fitness_selection_generator.hpp"
#include "clotho/data_spaces/fitness/general_fitness.hpp"

namespace clotho {
namespace genetics {

template < class RNG >
class SelectionGenerator< RNG, fitness_selection< GeneralFitness > > {
public:
    typedef RNG             random_engine_type;
    typedef GeneralFitness  fitness_space_type;

    typedef typename fitness_space_type::individual_id_type                     individual_id_type;

    typedef std::vector< std::pair< individual_id_type, individual_id_type > >  mate_pair_vector;

    typedef typename mate_pair_vector::iterator                                 iterator;
    typedef typename mate_pair_vector::const_iterator                           const_iterator;

    SelectionGenerator( random_engine_type * rng, boost::property_tree::ptree & config ) :
        m_rand( rng )
    {}

    void update( const fitness_space_type & fit, unsigned int N ) {
        typedef typename fitness_space_type::fitness_type   fitness_score_type;
        typedef typename fitness_space_type::const_iterator iterator;

        bool constant_fitness = true;
        unsigned int M = fit.individual_count();

        iterator it = fit.begin(), end = fit.end();

        if( it != end ) {
            fitness_score_type prev = *it++;
            while( constant_fitness && it != end ) {
                fitness_score_type cur = *it++;
                constant_fitness = (prev == cur);
            }
        }

        if( constant_fitness ) {
            boost::random::uniform_int_distribution< individual_id_type > uni( 0, ((M == 0) ? 0 : (M - 1)) );
            generate( uni, N );
        } else {
            boost::random::discrete_distribution< individual_id_type, fitness_score_type > disc( fit.begin(), fit.end() );
            generate( disc, N );
        }
    }

    mate_pair_vector & getMatePairs() const {
        return m_pairs;
    }

    unsigned int individual_count() const {
        return m_pairs.size();
    }

    iterator    begin() {
        return m_pairs.begin();
    }

    iterator    end() {
        return m_pairs.end();
    }

    const_iterator    begin() const {
        return m_pairs.begin();
    }

    const_iterator    end() const {
        return m_pairs.end();
    }

    unsigned int size() const {
        return individual_count();
    }

    virtual ~SelectionGenerator() {}

protected:

    template < class DistributionType >
    void generate( DistributionType & dist, size_t count ) {
        m_pairs.clear();

        while( m_pairs.size() < count ) {
            individual_id_type  id0 = dist( *m_rand );
            individual_id_type  id1 = dist( *m_rand );

            m_pairs.push_back( std::make_pair( id0, id1 ) );
        }
    }

    random_engine_type  * m_rand;
    mate_pair_vector    m_pairs;
};

}   // namespace genetics
}   // namespace clotho

#endif  // CLOTHO_GENERAL_FITNESS_SELECTION_GENERATOR_HPP_


/***************************************************************************************************************************
 * aruco_det_v2.cpp
 * Author: Jario
 * Update Time: 2020.12.11
 *
 * 说明: 单个二维码识别程序，可识别的二维码在Prometheus/Modules/object_detection/config/aruco_images文件夹中
 *      视野里只允许存在一个二维码 且二维码的字典类型要对应
 *      默认二维码的边长为0.2m
 *      1. 【订阅】图像话题 (默认来自web_cam)
 *         /prometheus/camera/rgb/image_raw
 *      2. 【发布】目标位置，发布话题见 Prometheus/Modules/msgs/msg/DetectionInfo.msg
 *         /prometheus/object_detection/aruco_det
 *      3. 【发布】检测结果的可视化图像话题
 *         /prometheus/camera/rgb/image_aruco_det
***************************************************************************************************************************/
#include <pthread.h>
#include <thread>
#include <chrono>
#include <boost/thread/mutex.hpp>
#include <boost/thread/shared_mutex.hpp>

#include <ros/ros.h>
#include <ros/package.h>
#include <yaml-cpp/yaml.h>
#include <image_transport/image_transport.h>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/image_encodings.h>
#include <geometry_msgs/PoseStamped.h>
#include <std_msgs/Bool.h>
#include <prometheus_msgs/DetectionInfo.h>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/core/core.hpp>
#include <opencv2/aruco.hpp>
#include <opencv2/aruco/dictionary.hpp>
#include <opencv2/aruco/charuco.hpp>
#include <opencv2/calib3d.hpp>

#include <Eigen/Dense>
#include <opencv2/core/eigen.hpp>

// #include "message_utils.h"

using namespace std;
using namespace cv;

//【订阅】输入图像
image_transport::Subscriber image_subscriber;
//【发布】检测得到的位置与姿态信息
ros::Publisher pose_pub;
//【发布】输入检测结果图像
image_transport::Publisher aruco_pub;

// 使用cout打印消息
bool local_print = true;

// 相机话题中的图像同步相关变量
int frame_width, frame_height;
std_msgs::Header image_header;
cv::Mat cam_image_copy;
boost::shared_mutex mutex_image_callback;
bool image_status = false;
boost::shared_mutex mutex_image_status;

// 图像接收回调函数，接收web_cam的话题，并将图像保存在cam_image_copy中
void cameraCallback(const sensor_msgs::ImageConstPtr& msg)
{
    if (local_print) ROS_DEBUG("[ArucoDetector] USB image received.");

    cv_bridge::CvImagePtr cam_image;

    try {
        cam_image = cv_bridge::toCvCopy(msg, sensor_msgs::image_encodings::BGR8);
        image_header = msg->header;
    } catch (cv_bridge::Exception& e) {
        if (local_print) ROS_ERROR("cv_bridge exception: %s", e.what());
        return;
    }

    if (cam_image) {
        {
            boost::unique_lock<boost::shared_mutex> lockImageCallback(mutex_image_callback);
            cam_image_copy = cam_image->image.clone();
        }
        {
            boost::unique_lock<boost::shared_mutex> lockImageStatus(mutex_image_status);
            image_status = true;
        }
        frame_width = cam_image->image.size().width;
        frame_height = cam_image->image.size().height;
    }
    return;
}

// 用此函数查看是否收到图像话题
bool getImageStatus(void)
{
    boost::shared_lock<boost::shared_mutex> lock(mutex_image_status);
    return image_status;
}

static bool readCameraParameters(string filename, Mat &camMatrix, Mat &distCoeffs) {
    FileStorage fs(filename, FileStorage::READ);
    if(!fs.isOpened())
        return false;
    fs["camera_matrix"] >> camMatrix;
    fs["distortion_coefficients"] >> distCoeffs;
    return true;
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "aruco_det_v2");
    ros::NodeHandle nh("~");
    image_transport::ImageTransport it(nh);
    // 更新频率为60HZ
    ros::Rate loop_rate(60);
    //【发布】识别
    pose_pub = nh.advertise<geometry_msgs::PoseStamped>("/prometheus/object_detection/aruco_det_v2", 1);

    std::string camera_topic = "/prometheus/camera/rgb/image_raw";
    std::string camera_params_yaml;
    std::string output_topic = "/prometheus/camera/rgb/image_aruco_det";

    if (nh.getParam("camera_topic", camera_topic)) {
        if (local_print) ROS_INFO("camera_topic is %s", camera_topic.c_str());
    } else {
        if (local_print) ROS_WARN("didn't find parameter camera_topic");
    }
    if (nh.getParam("camera_parameters", camera_params_yaml)) {
        if (local_print) ROS_INFO("camera_parameters is %s", camera_params_yaml.c_str());
    } else {
        if (local_print) ROS_WARN("didn't find camera_parameters");
    }
    if (nh.getParam("output_topic", output_topic)) {
        if (local_print) ROS_INFO("output_topic is %s", output_topic.c_str());
    } else {
        if (local_print) ROS_WARN("didn't find parameter output_topic");
    }

    // 接收图像的话题
    image_subscriber = it.subscribe(camera_topic.c_str(), 1, cameraCallback);
    // 发布ArUco检测结果的话题
    aruco_pub = it.advertise(output_topic.c_str(), 1);

    std::string ros_path = ros::package::getPath("prometheus_detection");
    if (local_print) ROS_INFO("DETECTION_PATH: %s", ros_path.c_str());

    cv::Mat camMatrix, distCoeffs;
    bool readOk = readCameraParameters(camera_params_yaml.c_str(), camMatrix, distCoeffs);
    if (!readOk) {
        cerr << "Invalid camera file" << endl;
        return 0;
    }

    if (local_print) {
        cout << "[camMatrix]:" << endl;
        cout << camMatrix << endl;
        cout << "[distCoeffs]:" << endl;
        cout << distCoeffs << endl;
    }

    int dictionaryId(2);
    float markerLength(0.2);
    Ptr<aruco::Dictionary> dictionary =
        aruco::getPredefinedDictionary(aruco::PREDEFINED_DICTIONARY_NAME(dictionaryId));
    Ptr<aruco::DetectorParameters> detectorParams = aruco::DetectorParameters::create();

    cv::Mat frame, frameCopy;
    const auto wait_duration = std::chrono::milliseconds(1000);
    while (ros::ok())
	{
        while (!getImageStatus() && ros::ok())
        {
            if (local_print) cout << "Waiting for image." << endl;
            std::this_thread::sleep_for(wait_duration);
            ros::spinOnce();
        }

        {
            boost::unique_lock<boost::shared_mutex> lockImageCallback(mutex_image_callback);
            frame = cam_image_copy.clone();
        }

        if (!frame.empty())
        {
            vector< int > ids;
            vector< vector< Point2f > > corners, rejected;
            vector< Vec3d > rvecs, tvecs;

            // detect markers and estimate pose
            aruco::detectMarkers(frame, dictionary, corners, ids, detectorParams, rejected);
            if (ids.size() > 0)
                aruco::estimatePoseSingleMarkers(corners, markerLength, camMatrix, distCoeffs, rvecs, tvecs);

            frame.copyTo(frameCopy);
            if(ids.size() > 0) {
                aruco::drawDetectedMarkers(frameCopy, corners, ids);

                for(unsigned int i = 0; i < ids.size(); i++) {
                    aruco::drawAxis(frameCopy, camMatrix, distCoeffs, rvecs[i], tvecs[i], markerLength * 0.5f);
                    cv::Mat rotation_matrix;
                    cv::Rodrigues(rvecs[i], rotation_matrix);
                    Eigen::Matrix3d rotation_matrix_eigen;
                    cv::cv2eigen(rotation_matrix, rotation_matrix_eigen);
                    Eigen::Quaterniond q = Eigen::Quaterniond(rotation_matrix_eigen);
                    q.normalize();

                    geometry_msgs::PoseStamped pose;
                    pose.header.frame_id = "map";
                    pose.pose.position.x = tvecs[i][0];
                    pose.pose.position.y = tvecs[i][1];
                    pose.pose.position.z = tvecs[i][2];
                    pose.pose.orientation.x = q.x();
                    pose.pose.orientation.y = q.y();
                    pose.pose.orientation.z = q.z();
                    pose.pose.orientation.w = q.w();
                    pose_pub.publish(pose);
                }
            }

            sensor_msgs::ImagePtr det_output_msg = cv_bridge::CvImage(std_msgs::Header(), "bgr8", frameCopy).toImageMsg();
            aruco_pub.publish(det_output_msg);
            // cv::imshow("frame", frame);
            // cv::waitKey(10);
        }

        ros::spinOnce();
        loop_rate.sleep();
    }

    return 0;
}


#include <boost/python/def.hpp>
#include <boost/python/class.hpp>
#include <boost/python/args.hpp>
#include <boost/python/return_value_policy.hpp>
#include <boost/python/copy_const_reference.hpp>
#include <boost/python/return_value_policy.hpp>
#include <boost/python/return_internal_reference.hpp>
#include <boost/python/return_by_value.hpp>
#include <cctbx/adp_restraints/aniso_restraints.h>

namespace cctbx {namespace adp_restraints {
namespace {

  void wrap_all ()
  {
    using namespace boost::python;
    using namespace cctbx::geometry_restraints;
    using namespace cctbx::xray;
    typedef eval_adp_aniso_restraints w_t;
    typedef return_value_policy<return_by_value> rbv;
    class_<w_t>("eval_adp_aniso_restraints", no_init)
      .def(init<
        af::const_ref<scatterer<> > const&,
        af::const_ref<scitbx::sym_mat3<double> > const&,
        af::const_ref<double> const&,
        af::const_ref<bond_simple_proxy> const&,
        af::const_ref<bool> const&,
        af::const_ref<bool> const&,
        unsigned,
        bool>((
          arg("scatterers"),
          arg("u_cart"),
          arg("u_iso"),
          arg("bond_proxies"),
          arg("selection"),
          arg("hd_selection"),
          arg("n_grad_u_iso"),
          arg("use_hd"))))
      .def("gradients_iso", &w_t::gradients_iso)
      .def("gradients_aniso_cart", &w_t::gradients_aniso_cart)
      .def_readonly("target", &w_t::target)
      .def_readonly("number_of_restraints", &w_t::number_of_restraints)
    ;
  }
}
namespace boost_python {
  void wrap_aniso_restraints () { wrap_all(); }
}
}}


//头文件
#include <ros/ros.h>

#include <iostream>
#include <Eigen/Eigen>

#include <state_from_mavros.h>
#include <command_to_mavros.h>

#include <px4_command_utils.h>
#include <OptiTrackFeedBackRigidBody.h>
#include <px4_command/ControlCommand.h>
#include <px4_command/DroneState.h>
#include <px4_command/TrajectoryPoint.h>
#include <px4_command/AttitudeReference.h>

#include <px4_command/Trajectory.h>
//msg 头文件
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <mavros_msgs/AttitudeTarget.h>
#include <mavros_msgs/PositionTarget.h>
#include <mavros_msgs/ActuatorControl.h>
#include <sensor_msgs/Imu.h>

#include <bitset>
#include <geometry_msgs/Vector3.h>
#include <geometry_msgs/TwistStamped.h>
#include <geometry_msgs/PoseStamped.h>
#include <std_msgs/Bool.h>
#include <geometry_msgs/Pose.h>
#include <geometry_msgs/Point.h>
#include <px4_command/Topic_for_log.h>

using namespace std;
//---------------------------------------相关参数-----------------------------------------------
px4_command::Topic_for_log _Topic_for_log;

Eigen::Vector3d pos_drone_mocap;                          //无人机当前位置 (vicon)
Eigen::Quaterniond q_mocap;
Eigen::Vector3d Euler_mocap;                              //无人机当前姿态 (vicon)
rigidbody_state UAVstate;
Eigen::Quaterniond q_fcu_target;
Eigen::Vector3d euler_fcu_target;
float Thrust_target;

//>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>函数声明<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
void printf_info();                                                                       //打印函数
//>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>回调函数<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
void log_cb(const px4_command::Topic_for_log::ConstPtr &msg)
{
    _Topic_for_log = *msg;
}

void att_target_cb(const mavros_msgs::AttitudeTarget::ConstPtr& msg)
{
    q_fcu_target = Eigen::Quaterniond(msg->orientation.w, msg->orientation.x, msg->orientation.y, msg->orientation.z);

    //Transform the Quaternion to euler Angles
    euler_fcu_target = quaternion_to_euler(q_fcu_target);

    Thrust_target = msg->thrust;
}

void optitrack_cb(const geometry_msgs::PoseStamped::ConstPtr& msg)
{
    //位置 -- optitrack系 到 ENU系
    int optitrack_frame = 0; //Frame convention 0: Z-up -- 1: Y-up

    if(optitrack_frame == 0)
    {
        // Read the Drone Position from the Vrpn Package [Frame: Vicon]  (Vicon to ENU frame)
        pos_drone_mocap = Eigen::Vector3d(msg->pose.position.x,msg->pose.position.y,msg->pose.position.z);
        // Read the Quaternion from the Vrpn Package [Frame: Vicon[ENU]]
        q_mocap = Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x, msg->pose.orientation.y, msg->pose.orientation.z);
    }
    else
    {
        // Read the Drone Position from the Vrpn Package [Frame: Vicon]  (Vicon to ENU frame)
        pos_drone_mocap = Eigen::Vector3d(-msg->pose.position.x,msg->pose.position.z,msg->pose.position.y);
        // Read the Quaternion from the Vrpn Package [Frame: Vicon[ENU]]
        q_mocap = Eigen::Quaterniond(msg->pose.orientation.w, msg->pose.orientation.x, msg->pose.orientation.z, msg->pose.orientation.y); //Y-up convention, switch the q2 & q3
    }

    // Transform the Quaternion to Euler Angles
    Euler_mocap = quaternion_to_euler(q_mocap);

}

//>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>主 函 数<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
int main(int argc, char **argv)
{
    ros::init(argc, argv, "ground_station");
    ros::NodeHandle nh("~");

    // 【订阅】optitrack估计位置
    ros::Subscriber optitrack_sub = nh.subscribe<geometry_msgs::PoseStamped>("/vrpn_client_node/UAV/pose", 10, optitrack_cb);

    ros::Subscriber log_sub = nh.subscribe<px4_command::Topic_for_log>("/px4_command/topic_for_log", 10, log_cb);

    ros::Subscriber attitude_target_sub = nh.subscribe<mavros_msgs::AttitudeTarget>("/mavros/setpoint_raw/target_attitude", 10,att_target_cb);

    // 频率
    ros::Rate rate(10.0);

    OptiTrackFeedBackRigidBody UAV("/vrpn_client_node/UAV/pose",nh,3,3);

//>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>Main Loop<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
    while(ros::ok())
    {
        //回调一次 更新传感器状态
        ros::spinOnce();

        //利用OptiTrackFeedBackRigidBody类获取optitrack的数据
        //UAV.GetOptiTrackState();

        UAV.RosWhileLoopRun();
        UAV.GetState(UAVstate);

        //打印
        printf_info();
        rate.sleep();
    }

    return 0;

}

void printf_info()
{
    cout <<">>>>>>>>>>>>>>>>>>>>>>>> Ground Station  <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<" <<endl;
    //固定的浮点显示
    cout.setf(ios::fixed);
    //setprecision(n) 设显示小数精度为n位
    cout<<setprecision(2);
    //左对齐
    cout.setf(ios::left);
    // 强制显示小数点
    cout.setf(ios::showpoint);
    // 强制显示符号
    cout.setf(ios::showpos);

    px4_command_utils::prinft_drone_state(_Topic_for_log.Drone_State);

    px4_command_utils::printf_command_control(_Topic_for_log.Control_Command);

    px4_command_utils::prinft_attitude_reference(_Topic_for_log.Attitude_Reference);

    cout <<">>>>>>>>>>>>>>>>>>>>>>>> Control Output  <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<" <<endl;

    cout << "u_l [X Y Z]  : " << _Topic_for_log.Control_Output.u_l[0] << " [ ] "<< _Topic_for_log.Control_Output.u_l[1] <<" [ ] "<< _Topic_for_log.Control_Output.u_l[2] <<" [ ] "<<endl;

    cout << "u_d [X Y Z]  : " << _Topic_for_log.Control_Output.u_d[0] << " [ ] "<< _Topic_for_log.Control_Output.u_d[1] <<" [ ] "<< _Topic_for_log.Control_Output.u_d[2] <<" [ ] "<<endl;
    cout << "NE  [X Y Z]  : " << _Topic_for_log.Control_Output.NE[0] << " [ ] "<< _Topic_for_log.Control_Output.NE[1] <<" [ ] "<< _Topic_for_log.Control_Output.NE[2] <<" [ ] "<<endl;

    cout << "Thrust  [X Y Z]  : " << _Topic_for_log.Control_Output.Thrust[0] << " [ ] "<< _Topic_for_log.Control_Output.Thrust[1] <<" [ ] "<< _Topic_for_log.Control_Output.Thrust[2] <<" [ ] "<<endl;

    cout << "Throttle  [X Y Z]  : " << _Topic_for_log.Control_Output.Throttle[0] << " [ ] "<< _Topic_for_log.Control_Output.Throttle[1] <<" [ ] "<< _Topic_for_log.Control_Output.Throttle[2] <<" [ ] "<<endl;

    cout <<">>>>>>>>>>>>>>>>>>>>>>>> Target Info FCU <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<" <<endl;

    cout << "Pos_vicon [X Y Z]  : " << pos_drone_mocap[0] << " [ m ] "<< pos_drone_mocap[1] <<" [ m ] "<< pos_drone_mocap[2] <<" [ m ] "<<endl;

    cout << "Att_target [R P Y] : " << euler_fcu_target[0] * 180/M_PI <<" [deg]  "<<euler_fcu_target[1] * 180/M_PI << " [deg]  "<< euler_fcu_target[2] * 180/M_PI<<" [deg]  "<<endl;

    cout << "Thr_target [ 0-1 ] : " << Thrust_target <<endl;

    cout <<">>>>>>>>>>>>>>>>>>>>>>>>Error Info [ Longhao ]<<<<<<<<<<<<<<<<<<<<<<<<<" <<endl;
    cout << "Error_pos      : " << UAVstate.Position[0] - _Topic_for_log.Drone_State.position[0] << " [ m ] "<< UAVstate.Position[1] - _Topic_for_log.Drone_State.position[1]<<" [ m ] "<< UAVstate.Position[2] - _Topic_for_log.Drone_State.position[2]<<" [ m ] "<<endl;
    cout << "Error_vel      : " << UAVstate.V_I[0] - _Topic_for_log.Drone_State.velocity[0] << " [m/s] "<< UAVstate.V_I[1] - _Topic_for_log.Drone_State.velocity[1]<<" [m/s] "<< UAVstate.V_I[2] - _Topic_for_log.Drone_State.velocity[2]<<" [m/s] "<<endl;
    cout << "Error_att      : " << UAVstate.Euler[0]*57.3 - _Topic_for_log.Drone_State.attitude[0]*57.3 << " [deg] "<< UAVstate.Euler[1]*57.3 - _Topic_for_log.Drone_State.attitude[1]*57.3<<" [deg] "<< UAVstate.Euler[2]*57.3 - _Topic_for_log.Drone_State.attitude[2]*57.3<<" [deg] "<<endl;
    cout << "Error_att_rate : " << UAVstate.Omega_BI[0]*57.3 - _Topic_for_log.Drone_State.attitude_rate[0]*57.3 << " [deg] "<< UAVstate.Omega_BI[1]*57.3 - _Topic_for_log.Drone_State.attitude_rate[1]*57.3<<" [deg] "<< UAVstate.Omega_BI[2]*57.3 - _Topic_for_log.Drone_State.attitude_rate[2]*57.3<<" [deg] "<<endl;

}


#include "Phase2EndcapRing.h"

#include "FWCore/MessageLogger/interface/MessageLogger.h"

#include "TrackingTools/DetLayers/interface/DetLayerException.h"
#include "TrackingTools/DetLayers/interface/MeasurementEstimator.h"
#include "TrackingTools/GeomPropagators/interface/HelixForwardPlaneCrossing.h"
#include "TrackingTools/DetLayers/interface/rangesIntersect.h"
#include "TrackingTools/DetLayers/interface/ForwardRingDiskBuilderFromDet.h"

#include "LayerCrossingSide.h"
#include "DetGroupMerger.h"
#include "CompatibleDetToGroupAdder.h"

#include "TkDetUtil.h"
#include "DataFormats/GeometryVector/interface/VectorUtil.h"
#include <boost/function.hpp>

using namespace std;

typedef GeometricSearchDet::DetWithState DetWithState;

class DetGroupElementZLess {
public:
  bool operator()(DetGroup a,DetGroup b)
  {
    return (fabs(a.front().det()->position().z()) < fabs(b.front().det()->position().z()));
  }
};

Phase2EndcapRing::Phase2EndcapRing(vector<const GeomDet*>& innerDets,
			       vector<const GeomDet*>& outerDets,
			       const vector<const GeomDet*>& innerDetBrothers,
			       const vector<const GeomDet*>& outerDetBrothers):
  GeometricSearchDet(true),
  theFrontDets(innerDets.begin(),innerDets.end()),
  theBackDets(outerDets.begin(),outerDets.end()),
  theFrontDetBrothers(innerDetBrothers.begin(),innerDetBrothers.end()),
  theBackDetBrothers(outerDetBrothers.begin(),outerDetBrothers.end())
{
  theDets.assign(theFrontDets.begin(),theFrontDets.end());
  theDets.insert(theDets.end(),theBackDets.begin(),theBackDets.end());
  theDets.insert(theDets.end(),theFrontDetBrothers.begin(),theFrontDetBrothers.end());
  theDets.insert(theDets.end(),theBackDetBrothers.begin(),theBackDetBrothers.end());

  // the dets should be already phi-ordered. TO BE CHECKED
  //sort( theFrontDets.begin(), theFrontDets.end(), DetLessPhi() );
  //sort( theBackDets.begin(), theBackDets.end(), DetLessPhi() );

  theDisk = ForwardRingDiskBuilderFromDet()( theDets );

  theFrontDisk = ForwardRingDiskBuilderFromDet()( theFrontDets );
  theBackDisk  = ForwardRingDiskBuilderFromDet()( theBackDets );

  theFrontBinFinder = BinFinderType( theFrontDets.front()->surface().position().phi(),
				     theFrontDets.size());
  theBackBinFinder  = BinFinderType( theBackDets.front()->surface().position().phi(),
				     theBackDets.size());

#ifdef EDM_ML_DEBUG
  LogDebug("TkDetLayers") << "DEBUG INFO for Phase2EndcapRing" ;
  for(vector<const GeomDet*>::const_iterator it=theFrontDets.begin();
      it!=theFrontDets.end(); it++){
    LogDebug("TkDetLayers") << "frontDet detId,phi,z,r: "
                            << (*it)->geographicalId().rawId()  << " , "
			    << (*it)->surface().position().phi()  << " , "
			    << (*it)->surface().position().z()    << " , "
			    << (*it)->surface().position().perp() ;
  }

  if(!theFrontDetBrothers.empty()){
    for(vector<const GeomDet*>::const_iterator it=theFrontDetBrothers.begin();
        it!=theFrontDetBrothers.end(); it++){
      LogDebug("TkDetLayers") << "frontDet brothers detId,phi,z,r: "
                            << (*it)->geographicalId().rawId()  << " , "
  			    << (*it)->surface().position().phi()  << " , "
  			    << (*it)->surface().position().z()    << " , "
  			    << (*it)->surface().position().perp() ;
    }
  }

  for(vector<const GeomDet*>::const_iterator it=theBackDets.begin();
      it!=theBackDets.end(); it++){
    LogDebug("TkDetLayers") << "backDet detId,phi,z,r: "
                            << (*it)->geographicalId().rawId()  << " , "
			    << (*it)->surface().position().phi() << " , "
			    << (*it)->surface().position().z()   << " , "
			    << (*it)->surface().position().perp() ;
  }

  if(!theBackDetBrothers.empty()){
    for(vector<const GeomDet*>::const_iterator it=theBackDetBrothers.begin();
        it!=theBackDetBrothers.end(); it++){
      LogDebug("TkDetLayers") << "backDet brothers detId,phi,z,r: "
                            << (*it)->geographicalId().rawId()  << " , "
  			    << (*it)->surface().position().phi() << " , "
  			    << (*it)->surface().position().z()   << " , "
  			    << (*it)->surface().position().perp() ;
    }
  }
#endif

}

Phase2EndcapRing::~Phase2EndcapRing(){

}

const vector<const GeometricSearchDet*>&
Phase2EndcapRing::components() const
{
  throw DetLayerException("Phase2EndcapRing doesn't have GeometricSearchDet components");
}

pair<bool, TrajectoryStateOnSurface>
Phase2EndcapRing::compatible( const TrajectoryStateOnSurface&, const Propagator&,
		  const MeasurementEstimator&) const{
  edm::LogError("TkDetLayers") << "temporary dummy implementation of Phase2EndcapRing::compatible()!!" ;
  return pair<bool,TrajectoryStateOnSurface>();
}

void
Phase2EndcapRing::groupedCompatibleDetsV( const TrajectoryStateOnSurface& tsos,
				 const Propagator& prop,
				 const MeasurementEstimator& est,
				 std::vector<DetGroup>& result) const
{
  SubLayerCrossings  crossings;
  crossings = computeCrossings( tsos, prop.propagationDirection());
  if(! crossings.isValid()) return;

  std::vector<DetGroup> closestResult;
  std::vector<DetGroup> closestBrotherResult;
  addClosest( tsos, prop, est, crossings.closest(), closestResult,closestBrotherResult);
  if (closestResult.empty())     return;

  DetGroupElement closestGel( closestResult.front().front());
  int crossingSide = LayerCrossingSide().endcapSide( closestGel.trajectoryState(), prop);
  float phiWindow =  tkDetUtil::computeWindowSize( closestGel.det(), closestGel.trajectoryState(), est);
  searchNeighbors( tsos, prop, est, crossings.closest(), phiWindow,
		   closestResult, closestBrotherResult, false);

  vector<DetGroup> closestCompleteResult;
  DetGroupMerger::orderAndMergeTwoLevels(std::move(closestResult),std::move(closestBrotherResult),closestCompleteResult,
					 0, crossingSide);

  vector<DetGroup> nextResult;
  vector<DetGroup> nextBrotherResult;
  searchNeighbors( tsos, prop, est, crossings.other(), phiWindow,
		   nextResult, nextBrotherResult, true);

  vector<DetGroup> nextCompleteResult;
  DetGroupMerger::orderAndMergeTwoLevels(std::move(nextResult),std::move(nextBrotherResult),nextCompleteResult,
					 0, crossingSide);

  DetGroupMerger::orderAndMergeTwoLevels( std::move(closestCompleteResult), std::move(nextCompleteResult), result,
					  crossings.closestIndex(), crossingSide);

  //due to propagator problems, when we add single pt sub modules, we should order them in z (endcap)
  if(!theFrontDetBrothers.empty() && !theBackDetBrothers.empty())
    sort(result.begin(),result.end(),DetGroupElementZLess());

#ifdef EDM_ML_DEBUG
  LogTrace("TkDetLayers") <<"Number of groups : " << result.size() << std::endl;
  for (auto&  grp : result) {
    if ( grp.empty() )  continue;
      LogTrace("TkDetLayers") <<"New group in Phase2EndcapRing made by : " << std::endl;
      for (auto const & det : grp) {
          LogTrace("TkDetLayers") <<" geom det at r: " << det.det()->position().perp() <<" id:" << det.det()->geographicalId().rawId()
                    <<" tsos at:" << det.trajectoryState().globalPosition() << std::endl;
      }
    }
#endif

}

SubLayerCrossings
Phase2EndcapRing::computeCrossings(const TrajectoryStateOnSurface& startingState,
			  PropagationDirection propDir) const
{
  auto rho = startingState.transverseCurvature();

  HelixPlaneCrossing::PositionType startPos( startingState.globalPosition() );
  HelixPlaneCrossing::DirectionType startDir( startingState.globalMomentum() );
  HelixForwardPlaneCrossing crossing(startPos,startDir,rho,propDir);

  pair<bool,double> frontPath = crossing.pathLength( *theFrontDisk);
  if (!frontPath.first) return SubLayerCrossings();

  pair<bool,double> backPath = crossing.pathLength( *theBackDisk);
  if (!backPath.first) return SubLayerCrossings();

  GlobalPoint gFrontPoint(crossing.position(frontPath.second));
  GlobalPoint gBackPoint( crossing.position(backPath.second));

  int frontIndex = theFrontBinFinder.binIndex(gFrontPoint.barePhi());
  SubLayerCrossing frontSLC( 0, frontIndex, gFrontPoint);

  int backIndex = theBackBinFinder.binIndex(gBackPoint.barePhi());
  SubLayerCrossing backSLC( 1, backIndex, gBackPoint);

  // 0ss: frontDisk has index=0, backDisk has index=1
  float frontDist = std::abs(Geom::deltaPhi( gFrontPoint.barePhi(),
					     theFrontDets[frontIndex]->surface().phi()));
  float backDist = std::abs(Geom::deltaPhi( gBackPoint.barePhi(),
					    theBackDets[backIndex]->surface().phi()));

  if (frontDist < backDist) {
    return SubLayerCrossings( frontSLC, backSLC, 0);
  }
  else {
    return SubLayerCrossings( backSLC, frontSLC, 1);
  }
}

bool Phase2EndcapRing::addClosest( const TrajectoryStateOnSurface& tsos,
				 const Propagator& prop,
				 const MeasurementEstimator& est,
				 const SubLayerCrossing& crossing,
				 vector<DetGroup>& result,
				 vector<DetGroup>& brotherresult) const
{
  const vector<const GeomDet*>& sub( subLayer( crossing.subLayerIndex()));
  const GeomDet* det(sub[crossing.closestDetIndex()]);
  bool firstgroup = CompatibleDetToGroupAdder::add( *det, tsos, prop, est, result);
  if(theFrontDetBrothers.empty() && theBackDetBrothers.empty())   return firstgroup;
  // it assumes that the closestDetIndex is ok also for the brother detectors: the crossing is NOT recomputed
  const vector<const GeomDet*>& subBrothers( subLayerBrothers( crossing.subLayerIndex()));
  const GeomDet* detBrother(subBrothers[crossing.closestDetIndex()]);
  bool brothergroup = CompatibleDetToGroupAdder::add( *detBrother, tsos, prop, est, brotherresult);
  return firstgroup || brothergroup;
}

void Phase2EndcapRing::searchNeighbors( const TrajectoryStateOnSurface& tsos,
				     const Propagator& prop,
				     const MeasurementEstimator& est,
				     const SubLayerCrossing& crossing,
				     float window,
				     vector<DetGroup>& result,
				     vector<DetGroup>& brotherresult,
				     bool checkClosest) const
{
  const GlobalPoint& gCrossingPos = crossing.position();

  const vector<const GeomDet*>& sLayer( subLayer( crossing.subLayerIndex()));
  // It assumes that what is ok for the front modules in the pt modules is ok also for the back module
  const vector<const GeomDet*>& sBrotherLayer( subLayerBrothers( crossing.subLayerIndex()));

  int closestIndex = crossing.closestDetIndex();
  int negStartIndex = closestIndex-1;
  int posStartIndex = closestIndex+1;

  if (checkClosest) { // must decide if the closest is on the neg or pos side
    if ( Geom::phiLess( gCrossingPos.barePhi(), sLayer[closestIndex]->surface().phi())) {
      posStartIndex = closestIndex;
    }
    else {
      negStartIndex = closestIndex;
    }
  }

  const BinFinderType& binFinder = (crossing.subLayerIndex()==0 ? theFrontBinFinder : theBackBinFinder);

  typedef CompatibleDetToGroupAdder Adder;
  int half = sLayer.size()/2;  // to check if dets are called twice....
  for (int idet=negStartIndex; idet >= negStartIndex - half; idet--) {
    const GeomDet & neighborDet = *sLayer[binFinder.binIndex(idet)];
    if (!tkDetUtil::overlapInPhi( gCrossingPos, neighborDet, window)) break;
    if (!Adder::add( neighborDet, tsos, prop, est, result)) break;
    if(theFrontDetBrothers.empty() && theBackDetBrothers.empty()) break;
    // If the two above checks are passed also the brother module will be added with no further checks
    const GeomDet & neighborBrotherDet = *sBrotherLayer[binFinder.binIndex(idet)];
    Adder::add( neighborBrotherDet, tsos, prop, est, brotherresult);
    // maybe also add shallow crossing angle test here???
  }
  for (int idet=posStartIndex; idet < posStartIndex + half; idet++) {
    const GeomDet & neighborDet = *sLayer[binFinder.binIndex(idet)];
    if (!tkDetUtil::overlapInPhi( gCrossingPos, neighborDet, window)) break;
    if (!Adder::add( neighborDet, tsos, prop, est, result)) break;
    if(theFrontDetBrothers.empty() && theBackDetBrothers.empty()) break;
    // If the two above checks are passed also the brother module will be added with no further checks
    const GeomDet & neighborBrotherDet = *sBrotherLayer[binFinder.binIndex(idet)];
    Adder::add( neighborBrotherDet, tsos, prop, est, brotherresult);
    // maybe also add shallow crossing angle test here???
  }
}