/*
 *            Copyright 2009-2018 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/qmatom.h"
#include <votca/xtp/dftengine.h>

#include <boost/format.hpp>
#include <boost/filesystem.hpp>
#include <votca/xtp/aomatrix.h>
#include <votca/xtp/orbitals.h>
#include <votca/xtp/qmpackagefactory.h>
#include <boost/math/constants/constants.hpp>
#include <votca/tools/constants.h>
#include <votca/tools/elements.h>
#include <votca/ctp/xinteractor.h>

using boost::format;
using namespace boost::filesystem;
using std::flush;
using namespace votca::tools;

namespace votca {
  namespace xtp {

    void DFTEngine::Initialize(Property& options) {

      string key = "package";

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

      _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 (!_with_RI) {
        std::vector<std::string> choices = {"direct", "cache"};
        _four_center_method = options.ifExistsAndinListReturnElseThrowRuntimeError<std::string>(key + ".four_center_method", choices);

        _with_screening = options.ifExistsReturnElseReturnDefault<bool>(key + ".with_screening", true);
        _screening_eps = options.ifExistsReturnElseReturnDefault<double>(key + ".screening_eps", 1e-9);
      }

      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");

      _grid_name = options.ifExistsReturnElseReturnDefault<string>(key + ".integration_grid", "medium");
      _use_small_grid = options.ifExistsReturnElseReturnDefault<bool>(key + ".integration_grid_small", true);
      _grid_name_small = ReturnSmallGrid(_grid_name);
      _xc_functional_name = options.ifExistsReturnElseThrowRuntimeError<string>(key + ".xc_functional");

      if (options.exists(key + ".externaldensity")){
        _integrate_ext_density=true;
          _orbfilename=options.ifExistsReturnElseThrowRuntimeError<string>(key + ".externaldensity.orbfile");
          _gridquality=options.ifExistsReturnElseThrowRuntimeError<string>(key + ".externaldensity.gridquality");
          _state=options.ifExistsReturnElseThrowRuntimeError<string>(key + ".externaldensity.state");
      }

      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;
        }
        _mixingparameter = options.ifExistsReturnElseReturnDefault<double>(key + ".convergence.mixing", 0.7);
        _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;
        _mixingparameter = 0.7;
        _usediis = true;
        _max_iter = 100;
        _levelshift = 0.25;
        _levelshiftend = 0.8;
      }

      return;
    }

void DFTEngine::PrintMOs(const Eigen::VectorXd& MOEnergies){
      CTP_LOG(ctp::logDEBUG, *_pLog) << "  Orbital energies: " << flush;
      CTP_LOG(ctp::logDEBUG, *_pLog) << "  index occupation energy(Hartree) " << flush;
      for (int i = 0; i<MOEnergies.size(); i++) {
        int occupancy=0;
        if (i < _numofelectrons / 2 ) {
          occupancy=2;
        }
        CTP_LOG(ctp::logDEBUG, *_pLog) <<(boost::format(" %1$5d      %2$1d   %3$+1.10f")
                                     %i %occupancy %MOEnergies(i)).str()<<flush;

      }
      return;
    }

void DFTEngine::CalcElDipole(){
  AODipole dipole;
  dipole.Fill(_dftbasis);
  CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Electric Dipole is[e*bohr]:\n\t\t dx="
                          << _dftAOdmat.cwiseProduct(dipole.Matrix()[0]).sum()
          << "\n\t\t dy=" << _dftAOdmat.cwiseProduct(dipole.Matrix()[1]).sum()
          << "\n\t\t dz=" << _dftAOdmat.cwiseProduct(dipole.Matrix()[2]).sum()<<flush;
  return;
}

   /*
     *    Density Functional theory implementation
     *
     */
    bool DFTEngine::Evaluate(Orbitals& orbitals) {
      // set the parallelization
#ifdef _OPENMP
      omp_set_num_threads(_openmp_threads);
#endif

      Eigen::VectorXd& MOEnergies = orbitals.MOEnergies();
      Eigen::MatrixXd& MOCoeff = orbitals.MOCoefficients();
      if (MOEnergies.size() != _dftbasis.AOBasisSize()) {
        MOEnergies.resize(_dftbasis.AOBasisSize());
      }
      if (MOCoeff.rows() != _dftbasis.AOBasisSize() || MOCoeff.cols() != _dftbasis.AOBasisSize()) {
        MOCoeff.conservativeResize(_dftbasis.AOBasisSize(), _dftbasis.AOBasisSize());
      }

      Eigen::MatrixXd H0 = _dftAOkinetic.Matrix() + _dftAOESP.getNuclearpotential();
      CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Constructed independent particle hamiltonian " << flush;
      NuclearRepulsion();
      if(_with_ecp){
          H0+=_dftAOECP.Matrix();
      }
      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(_integrate_ext_density){
        Orbitals extdensity;
        extdensity.ReadFromCpt(_orbfilename);
        H0+=IntegrateExternalDensity(extdensity);
      }

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

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

      if (_with_guess) {
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Reading guess from orbitals object/file" << flush;
        orbitals.MOCoefficients()=OrthogonalizeGuess(orbitals.MOCoefficients() );
        _dftAOdmat = _conv_accelerator.DensityMatrix(MOCoeff,MOEnergies);
      } else {
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Setup Initial Guess using: " << _initial_guess << flush;
        if (_initial_guess == "independent") {
          _conv_accelerator.SolveFockmatrix(MOEnergies, MOCoeff, H0);
          _dftAOdmat = _conv_accelerator.DensityMatrix(MOCoeff,MOEnergies);

        } else if (_initial_guess == "atom") {
          _dftAOdmat = AtomicGuess(orbitals);
          CalculateERIs(_dftbasis, _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;
          }
          Eigen::MatrixXd H = H0 + _ERIs.getERIs() + orbitals.AOVxc();
          if(_ScaHFX>0){
              if (_with_RI) {
             _ERIs.CalculateEXX(_dftAOdmat);
           } else {
             _ERIs.CalculateEXX_4c_small_molecule(_dftAOdmat);
           }
            H-=0.5*_ScaHFX*_ERIs.getEXX();
          }
          _conv_accelerator.SolveFockmatrix(MOEnergies, MOCoeff, H);
          _dftAOdmat = _conv_accelerator.DensityMatrix(MOCoeff,MOEnergies);

          CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Full atomic density Matrix gives N="
                  << std::setprecision(9) << _dftAOdmat.cwiseProduct(_dftAOoverlap.Matrix()).sum() << " electrons." << flush;
        } else {
          throw std::runtime_error("Initial guess method not known/implemented");
        }

      }

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

      double energyold = std::numeric_limits<double>::max();

      for (int 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;

        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT Electron repulsion matrix" << flush;
        double vxcenergy = 0.0;
        if (_use_small_grid && _conv_accelerator.getDIIsError() > 1e-3) {
          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;
        }
        CalculateERIs(_dftbasis, _dftAOdmat);
        Eigen::MatrixXd H = H0 + _ERIs.getERIs() + orbitals.AOVxc();
        if(_ScaHFX>0){
              if (_with_RI) {
                if(_conv_accelerator.getUseMixing()){
                  _ERIs.CalculateEXX(_dftAOdmat);
                }else{
                  Eigen::Block<Eigen::MatrixXd> occblock=MOCoeff.block(0,0,MOEnergies.rows(), _numofelectrons / 2);
                  _ERIs.CalculateEXX(occblock,_dftAOdmat);
                }
           } else {
             _ERIs.CalculateEXX_4c_small_molecule(_dftAOdmat);
           }
            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Filled DFT Electron exchange matrix"<< flush;
            H-=0.5*_ScaHFX*_ERIs.getEXX();
          }

        double Eone = _dftAOdmat.cwiseProduct(H0).sum();
        double Etwo = 0.5 * _ERIs.getERIsenergy() + vxcenergy;
        if(_ScaHFX>0){
          Etwo-=_ScaHFX/4*_ERIs.getEXXsenergy();
        }
        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) << " Local Exc contribution " << vxcenergy << flush;
        if(_ScaHFX>0){
          CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                  << std::setprecision(12) << " Non local Ex contribution " << -_ScaHFX/4*_ERIs.getEXXsenergy() << flush;
        }
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Total Energy " << std::setprecision(12) << totenergy << flush;

        _dftAOdmat = _conv_accelerator.Iterate(_dftAOdmat, H, MOEnergies, MOCoeff, totenergy);

        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " DIIs error " << _conv_accelerator.getDIIsError() << std::flush;
        double deltaE=totenergy - energyold;
        if(this_iter == 0){
          deltaE=0;
        }
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Delta Etot " <<deltaE<< std::flush;
        if (tools::globals::verbose) {
          PrintMOs(MOEnergies);
        }

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

        if (std::abs(totenergy - energyold) < _Econverged && _conv_accelerator.getDIIsError() < _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;

          PrintMOs(MOEnergies);
           CalcElDipole();
          // 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="
                << _dftAOdmat.cwiseProduct(_dftAOoverlap.Matrix()).sum() << " electrons." << flush;

      }
      return true;
    }

    // SETUP INVARIANT AOMatrices

    void DFTEngine::SetupInvariantMatrices() {

        _dftAOoverlap.Fill(_dftbasis);
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Filled DFT Overlap matrix."<< flush;

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

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

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

        _dftAODipole_Potential.Fillextpotential(_dftbasis, _externalsites);
        if (_dftAODipole_Potential.Dimension() > 0) {
          CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                  << " Filled DFT external dipole potential matrix" << flush;
        }
        _dftAOQuadrupole_Potential.Fillextpotential(_dftbasis, _externalsites);
        if (_dftAOQuadrupole_Potential.Dimension()) {
          CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                  << " Filled DFT external quadrupole potential matrix."<< flush;
        }
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()<< " External sites"<<flush;
        CTP_LOG(ctp::logDEBUG, *_pLog)<<" Name      Coordinates[nm]     charge[e]         dipole[e*nm]    "
                                        "              quadrupole[e*nm^2]         " << flush;

        for (auto segment:_externalsites) {
          for (ctp::APolarSite* site:*segment){
            std::string output=(boost::format("  %1$s"
                                            "   %2$+1.4f %3$+1.4f %4$+1.4f"
                                            "   %5$+1.4f")
                                            %site->getName()
                                            %site->getPos().getX() %site->getPos().getY() %site->getPos().getZ()
                                            %site->getQ00()).str();
            if (site->getRank() > 0) {
              tools::vec dipole = site->getQ1();
              output+=(boost::format("   %1$+1.4f %2$+1.4f %3$+1.4f")
                                     %dipole.getX() %dipole.getY() %dipole.getZ()).str();
            }
            if (site->getRank() > 1) {
              std::vector<double> quadrupole = site->getQ2();
              output+=(boost::format("   %1$+1.4f %2$+1.4f %3$+1.4f %4$+1.4f %5$+1.4f")
                                     %quadrupole[0] %quadrupole[1] %quadrupole[2]
                                     %quadrupole[3] %quadrupole[4]).str();
            }
            CTP_LOG(ctp::logDEBUG, *_pLog) <<output<< flush;
          }
        }
      }

      if (_with_ecp) {
        _dftAOECP.setECP(&_ecp);
        _dftAOECP.Fill(_dftbasis);
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Filled DFT ECP matrix" << flush;
      }

      _conv_accelerator.Configure(ConvergenceAcc::KSmode::closed, _usediis,
              true, _histlength, _maxout, _adiis_start, _diis_start,
              _levelshift, _levelshiftend, _numofelectrons, _mixingparameter);
      _conv_accelerator.setLogger(_pLog);
      _conv_accelerator.setOverlap(&_dftAOoverlap, 1e-8);

      if (_with_RI) {

        AOCoulomb auxAOcoulomb;
        auxAOcoulomb.Fill(_auxbasis);
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Filled auxiliary Coulomb matrix"<< flush;

        Eigen::MatrixXd Inverse=auxAOcoulomb.Pseudo_InvSqrt(1e-8);

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

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

        if (_four_center_method=="cache") {

          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;
        }

        if (_with_screening && _four_center_method=="direct") {
          CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Calculating 4c diagonals. " << flush;
          _ERIs.Initialize_4c_screening(_dftbasis, _screening_eps);
          CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Calculated 4c diagonals. " << flush;
        }
      }

      return;
    }

    Eigen::MatrixXd DFTEngine::RunAtomicDFT_unrestricted(QMAtom* uniqueAtom){
      bool with_ecp = _with_ecp;
      if (uniqueAtom->getType() == "H" || uniqueAtom->getType() == "He") {
        with_ecp = false;
      }

       std::vector<QMAtom*> atom;
        atom.push_back(uniqueAtom);

        AOBasis dftbasis;
        NumericalIntegration gridIntegration;
        dftbasis.AOBasisFill(_dftbasisset, atom);
        AOBasis ecp;
        if (with_ecp) {
          ecp.ECPFill(_ecpbasisset, atom);
        }
        gridIntegration.GridSetup(_grid_name, atom, dftbasis);
        gridIntegration.setXCfunctional(_xc_functional_name);
        int numofelectrons = uniqueAtom->getNuccharge();
        int alpha_e = 0;
        int beta_e = 0;

        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;
        ERIs ERIs_atom;

        // DFT AOOverlap matrix

        dftAOoverlap.Fill(dftbasis);
        dftAOkinetic.Fill(dftbasis);

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

        Eigen::VectorXd MOEnergies_alpha;
        Eigen::MatrixXd MOCoeff_alpha;
        Eigen::VectorXd MOEnergies_beta;
        Eigen::MatrixXd MOCoeff_beta;
        ConvergenceAcc Convergence_alpha;

        ConvergenceAcc Convergence_beta;
        double adiisstart = 0;
        double diisstart = 0;

        Convergence_alpha.Configure(ConvergenceAcc::KSmode::open, true, false,
                20, 0, adiisstart, diisstart, 0.1, diisstart, alpha_e, _mixingparameter);
        Convergence_alpha.setLogger(_pLog);
        Convergence_alpha.setOverlap(&dftAOoverlap, 1e-8);
        Convergence_beta.Configure(ConvergenceAcc::KSmode::open, true, false,
                20, 0, adiisstart, diisstart, 0.1, diisstart, beta_e, _mixingparameter);
        Convergence_beta.setLogger(_pLog);
        Convergence_beta.setOverlap(&dftAOoverlap, 1e-8);
        /**** Construct initial density  ****/

        Eigen::MatrixXd H0 = dftAOkinetic.Matrix() + dftAOESP.getNuclearpotential();
        if (with_ecp) {
          dftAOECP.setECP(&ecp);
          dftAOECP.Fill(dftbasis);
          H0 += dftAOECP.Matrix();
        }
        Convergence_alpha.SolveFockmatrix(MOEnergies_alpha, MOCoeff_alpha, H0);

        Eigen::MatrixXd dftAOdmat_alpha = Convergence_alpha.DensityMatrix(MOCoeff_alpha, MOEnergies_alpha);
        if (uniqueAtom->getType() == "H") {
          return dftAOdmat_alpha;
        }
        Convergence_beta.SolveFockmatrix(MOEnergies_beta, MOCoeff_beta, H0);
        Eigen::MatrixXd dftAOdmat_beta = Convergence_beta.DensityMatrix(MOCoeff_beta, MOEnergies_beta);

        double energyold = 0;
        int maxiter = 80;
        for (int this_iter = 0; this_iter < maxiter; this_iter++) {
          ERIs_atom.CalculateERIs_4c_small_molecule(dftAOdmat_alpha + dftAOdmat_beta);
          double E_two_alpha = ERIs_atom.getERIs().cwiseProduct(dftAOdmat_alpha).sum();
          double E_two_beta = ERIs_atom.getERIs().cwiseProduct(dftAOdmat_beta).sum();
          Eigen::MatrixXd H_alpha = H0 + ERIs_atom.getERIs();
          Eigen::MatrixXd H_beta = H0 + ERIs_atom.getERIs();

            Eigen::MatrixXd AOVxc_alpha = gridIntegration.IntegrateVXC(dftAOdmat_alpha);
            double E_vxc_alpha = gridIntegration.getTotEcontribution();

            Eigen::MatrixXd 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;

            if(_ScaHFX>0){
              ERIs_atom.CalculateEXX_4c_small_molecule(dftAOdmat_alpha);
              double E_exx_alpha =-0.5*_ScaHFX*ERIs_atom.getEXX().cwiseProduct(dftAOdmat_alpha).sum();
              H_alpha -=_ScaHFX* ERIs_atom.getEXX();
              E_two_alpha += E_exx_alpha;
              ERIs_atom.CalculateEXX_4c_small_molecule(dftAOdmat_beta);
              double E_exx_beta =-0.5*_ScaHFX*ERIs_atom.getEXX().cwiseProduct(dftAOdmat_beta).sum();
              H_beta -=_ScaHFX*ERIs_atom.getEXX();
              E_two_beta += E_exx_beta;
            }

          double E_one_alpha = dftAOdmat_alpha.cwiseProduct(H0).sum();
          double E_one_beta = dftAOdmat_beta.cwiseProduct(H0).sum();
          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
          dftAOdmat_alpha = Convergence_alpha.Iterate(dftAOdmat_alpha, H_alpha,
                  MOEnergies_alpha, MOCoeff_alpha, E_alpha);
          //evolve beta
          dftAOdmat_beta = Convergence_beta.Iterate(dftAOdmat_beta, H_beta,
                  MOEnergies_beta, MOCoeff_beta, E_beta);

          if (tools::globals::verbose) {
            CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Iter " << this_iter << " of " << maxiter
                    << " Etot " << totenergy
                    << " diise_a " << Convergence_alpha.getDIIsError()
                    << " diise_b " << Convergence_beta.getDIIsError()
                    << "\n\t\t a_gap " << MOEnergies_alpha(alpha_e) - MOEnergies_alpha(alpha_e - 1)
                    << " b_gap " << MOEnergies_beta(beta_e) - MOEnergies_beta(beta_e - 1)
                    <<" Nalpha="<<dftAOoverlap.Matrix().cwiseProduct(dftAOdmat_alpha).sum()
                    <<" Nbeta="<<dftAOoverlap.Matrix().cwiseProduct(dftAOdmat_beta).sum()<<flush;
          }
          bool converged = (std::abs(totenergy - energyold) < _Econverged &&
                  Convergence_alpha.getDIIsError() < _error_converged &&
                  Convergence_beta.getDIIsError() < _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. Unconverged density.\n\t\t\t"
                      <<" DIIsError_alpha=" << Convergence_alpha.getDIIsError()
                      << " DIIsError_beta=" << Convergence_beta.getDIIsError() << flush;
            }
            break;

          } else {
            energyold = totenergy;
          }
        }
        Eigen::MatrixXd avgmatrix=SphericalAverageShells(dftAOdmat_alpha + dftAOdmat_beta,dftbasis);
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Atomic density Matrix for " << uniqueAtom->getType() << " gives N="
                << std::setprecision(9) << avgmatrix.cwiseProduct(dftAOoverlap.Matrix()).sum() << " electrons." << flush;
        return avgmatrix;
    }

    Eigen::MatrixXd DFTEngine::AtomicGuess(Orbitals& orbitals) {

      std::vector<QMAtom*> uniqueelements;

      CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Scanning molecule of size " << _atoms.size() << " for unique elements" << flush;
      for (QMAtom* atom:_atoms) {
        bool exists = false;
        if (uniqueelements.size() == 0) {
          exists = false;
        } else {
          for (const QMAtom* unique_atom:uniqueelements){
            if (atom->getType() == unique_atom->getType()) {
              exists = true;
              break;
            }
          }
        }
        if (!exists) {
          uniqueelements.push_back(atom);
        }
      }

      CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " " << uniqueelements.size() << " unique elements found" << flush;
      std::vector< Eigen::MatrixXd > uniqueatom_guesses;
      for ( QMAtom* unique_atom:uniqueelements) {
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Calculating atom density for " << unique_atom->getType() << flush;
        Eigen::MatrixXd dmat_unrestricted=RunAtomicDFT_unrestricted(unique_atom);
        uniqueatom_guesses.push_back(dmat_unrestricted);
      }

      Eigen::MatrixXd guess = Eigen::MatrixXd::Zero(_dftbasis.AOBasisSize(), _dftbasis.AOBasisSize());
      int start = 0;
      for (QMAtom* atom:_atoms) {
        unsigned index = 0;
        for (index = 0; index < uniqueelements.size(); index++) {
          if (atom->getType() == uniqueelements[index]->getType()) {
            break;
          }
        }
        Eigen::MatrixXd& dmat_unrestricted= uniqueatom_guesses[index];
        guess.block(start, start, dmat_unrestricted.rows(),dmat_unrestricted.cols())=dmat_unrestricted;
        start += dmat_unrestricted.rows();
      }

      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.setQMpackage("xtp");
      orbitals.setDFTbasis(_dftbasis_name);
      orbitals.setBasisSetSize(_dftbasis.AOBasisSize());
      orbitals.setScaHFX(_ScaHFX);
      if (_with_ecp) {
        orbitals.setECP(_ecp_name);
      }
      if (_with_RI) {
        orbitals.setAuxbasis(_auxbasis_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(tools::globals::VOTCA_MKL){
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Using MKL overload for Eigen " << flush;
      } else {
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Using native Eigen implementation, no BLAS overload " << flush;
      }

      if (_atoms.size() == 0) {
        _atoms = orbitals.QMAtoms();
      }

      CTP_LOG(ctp::logDEBUG, *_pLog) << " Molecule Coordinates [A] " << flush;
      for (const QMAtom* atom:_atoms) {

      std::string output=(boost::format("  %1$s"
                                         "   %2$+1.4f %3$+1.4f %4$+1.4f")
                         %atom->getType() %(atom->getPos().getX()*tools::conv::bohr2ang)
                         %(atom->getPos().getY()*tools::conv::bohr2ang)
                         %(atom->getPos().getZ()*tools::conv::bohr2ang) ).str();

        CTP_LOG(ctp::logDEBUG, *_pLog) <<output<< flush;
      }

      _dftbasisset.LoadBasisSet(_dftbasis_name);

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

      if (_with_RI) {
        _auxbasisset.LoadBasisSet(_auxbasis_name);
        _auxbasis.AOBasisFill(_auxbasisset, _atoms);
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Loaded AUX Basis Set " << _auxbasis_name
                << " with "<<_auxbasis.AOBasisSize()<<" functions"<< flush;
      }
      if (_with_ecp) {
        _ecpbasisset.LoadPseudopotentialSet(_ecp_name);
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Loaded ECP library " << _ecp_name << flush;

        _ecp.ECPFill(_ecpbasisset, _atoms);
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Filled ECP Basis of size " << _ecp.getNumofShells() << flush;
      }

      _gridIntegration.GridSetup(_grid_name, _atoms, _dftbasis);
      _gridIntegration.setXCfunctional(_xc_functional_name);

      _ScaHFX = _gridIntegration.getExactExchange(_xc_functional_name);
      if (_ScaHFX > 0) {
        CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp()
                << " Using hybrid functional with alpha=" << _ScaHFX << flush;
      }

      CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() << " Setup numerical integration grid "
              << _grid_name << " for vxc functional "
              << _xc_functional_name <<  flush;
      CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t "<<" with " << _gridIntegration.getGridSize() << " points"
              << " divided into "<< _gridIntegration.getBoxesSize() << " boxes" << flush;
      if (_use_small_grid) {
        _gridIntegration_small.GridSetup(_grid_name_small, _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  << flush;
        CTP_LOG(ctp::logDEBUG, *_pLog) << "\t\t " << " with " << _gridIntegration_small.getGridSize() << " points"
                << " divided into "<< _gridIntegration_small.getBoxesSize() << " boxes" << flush;
      }

      if (_do_externalfield) {
        _gridIntegration_ext.GridSetup(_grid_name_ext, _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;
      }

      for (auto& atom : _atoms) {
        _numofelectrons += atom->getNuccharge();
      }

      // 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;
    }

    void DFTEngine::NuclearRepulsion() {
      _E_nucnuc = 0.0;

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

    double DFTEngine::ExternalRepulsion(ctp::Topology* top) {

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

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

      for (unsigned i = 0; i < nuclei.size(); ++i) {
        ctp::APolarSite* nucleus = nuclei[i];
        nucleus->setIsoP(0.0);
        double Q = _atoms[i]->getNuccharge();
        nucleus->setQ00(Q, 0);
      }
      ctp::XInteractor actor;
      actor.ResetEnergy();
      nuclei.CalcPos();
      double E_ext = 0.0;
      for (std::shared_ptr<ctp::PolarSeg>  seg : _externalsites) {
        seg->CalcPos();
        tools::vec s = tools::vec(0.0);
        if (top != NULL) {
          s = top->PbShortestConnect(nuclei.getPos(),
                  seg->getPos())+nuclei.getPos()-seg->getPos();
        }

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

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

    double DFTEngine::ExternalGridRepulsion(std::vector<double> externalpotential_nuc) {
      double E_ext = 0.0;
      if (!_do_externalfield) {
        return 0;
      }
      for (unsigned i = 0; i < _atoms.size(); i++) {
        double Q = _atoms[i]->getNuccharge();
        E_ext += Q * externalpotential_nuc[i];
      }
      return E_ext;
    }

    string DFTEngine::ReturnSmallGrid(const 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.

    Eigen::MatrixXd DFTEngine::SphericalAverageShells(const Eigen::MatrixXd& dmat, AOBasis& dftbasis) {
      Eigen::MatrixXd avdmat = Eigen::MatrixXd::Zero(dmat.rows(), dmat.cols());
      int start = 0.0;
      std::vector<int> starts;
      std::vector<int> ends;
      for (AOShell* shell:dftbasis) {
        int end = shell->getNumFunc() + start;

        if (shell->isCombined()) {
          std::vector<int> temp = NumFuncSubShell(shell->getType());
          int numfunc = start;
          for (int & SubshellFunc:temp) {
            starts.push_back(numfunc);
            numfunc +=SubshellFunc;
            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;
    }

    Eigen::MatrixXd DFTEngine::IntegrateExternalDensity(Orbitals& extdensity){
      BasisSet basis;
      basis.LoadBasisSet(extdensity.getDFTbasis());
      AOBasis aobasis;
      aobasis.AOBasisFill(basis,extdensity.QMAtoms());
      NumericalIntegration numint;
      numint.GridSetup(_gridquality,extdensity.QMAtoms(),aobasis);
      Eigen::MatrixXd dmat=extdensity.DensityMatrixGroundState();

      numint.IntegrateDensity(dmat);
      CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() <<" Calculated external density"<<flush;
      Eigen::MatrixXd e_contrib=numint.IntegratePotential(_dftbasis);
      CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() <<" Calculated potential from electron density"<<flush;
      AOESP esp;
      esp.Fillnucpotential(_dftbasis,extdensity.QMAtoms());

      double nuc_energy=0.0;
      for(const QMAtom* atom:_atoms){
        nuc_energy+=numint.IntegratePotential(atom->getPos())*atom->getNuccharge();
        for(const QMAtom* extatom:extdensity.QMAtoms()){
          const double dist=tools::abs(atom->getPos()-extatom->getPos());
          nuc_energy+=atom->getNuccharge()*extatom->getNuccharge()/dist;
        }
      }
      CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() <<" Calculated potential from nuclei"<<flush;
      CTP_LOG(ctp::logDEBUG, *_pLog) << ctp::TimeStamp() <<" Elelctrostatic: "<<nuc_energy<<flush;
      _E_nucnuc+=nuc_energy;
      return e_contrib+esp.getNuclearpotential();
    }

    void DFTEngine::CalculateERIs(const AOBasis& dftbasis, const Eigen::MatrixXd& DMAT) {

      if (_with_RI)
        _ERIs.CalculateERIs(_dftAOdmat);
      else if (_four_center_method.compare("cache") == 0)
        _ERIs.CalculateERIs_4c_small_molecule(_dftAOdmat);
      else if (_four_center_method.compare("direct") == 0)
        _ERIs.CalculateERIs_4c_direct(_dftbasis, _dftAOdmat);
    }

    Eigen::MatrixXd DFTEngine::OrthogonalizeGuess(const Eigen::MatrixXd& GuessMOs ){
      Eigen::MatrixXd nonortho=GuessMOs.transpose()*_dftAOoverlap.Matrix()*GuessMOs;
      Eigen::SelfAdjointEigenSolver<Eigen::MatrixXd> es(nonortho);
      Eigen::MatrixXd result=GuessMOs*es.operatorInverseSqrt();
      return result;
    }

  }
}


/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: t; c-basic-offset: 4 -*- vim:set ts=4 sw=4 sts=4 noet: */
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include "compiler/bcc/parser.h"

#include <cassert>
#include <cctype>
#include <cerrno>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iostream>

#include <boost/math/constants/constants.hpp>

#include "compiler/bcc/token.h"
#include "compiler/bcc/tokenizer.h"
#include "flint/bc.h"

namespace flint {
namespace compiler {
namespace bcc {

namespace {

void ReportUnexpectedToken(Token &t)
{
	std::cerr << "unexpected token: ";
	t.Write(std::cerr) << std::endl;
}

int ConvertToInteger(const char *p, const char *q)
{
	assert(p < q);
	int n = 0;
	while (p < q) {
		n *= 10;
		char c = *p++;
		assert('0' <= c && c <= '9');
		n += c - '0';
	}
	return n;
}

int GetIntegerFromAddress(const Token &t)
{
	assert(t.type == Token::Type::kAddress);
	const char *p = t.lexeme + 1; // add offset for "$"
	const char *q = t.lexeme + t.size;
	return ConvertToInteger(p, q);
}

int GetIntegerFromLabel(const Token &t)
{
	assert(t.type == Token::Type::kLabel);
	const char *p = t.lexeme + 1; // add offset for "L"
	const char *q = t.lexeme + t.size;
	return ConvertToInteger(p, q);
}

int GetIntegerFromNumber(const Token &t)
{
	assert(t.type == Token::Type::kNumber);
	const char *p = t.lexeme;
	const char *q = t.lexeme + t.size;
	return ConvertToInteger(p, q);
}

int GetIntegerFromIntReg(const Token &t)
{
	assert(t.type == Token::Type::kIntReg);
	const char *p = t.lexeme + 2; // offset for "$i"
	const char *q = t.lexeme + t.size; // subtract offset too
	return ConvertToInteger(p, q);
}

}

class ParserImpl {
public:
	explicit ParserImpl(const char *input);

	bool Parse(Body *body);

private:
	bool ParseInst(bc::Code *code);
	bool ParseLabelColon(int n, Body *body);
	bool ParseInstOp(int a, bc::Code *code);
	bool ParseCall1(int a, bc::Call1::Op op, bc::Code *code);
	bool ParseCall2(int a, bc::Call2::Op op, bc::Code *code);
	bool ParseGen1(int a, bc::Gen1::Type type, bc::Code *code);
	bool ParseGen2(int a, bc::Gen2::Type type, bc::Code *code);
	bool ParseFr(int *a);
	bool ParseInstBr(bc::Code *code);
	bool ParseLabel(int *n);
	bool ParseInstJmp(bc::Code *code);
	bool ParseInstLb(bc::Code *code);
	bool ParseInstLc(bc::Code *code);
	bool ParseInstLd(bc::Code *code);
	bool ParseInstLoad(bc::Code *code);
	bool ParseInstLoadi(bc::Code *code);
	bool ParseImm(double *d);
	bool ParseInstStore(bc::Code *code);
	bool ParseInstRefer(bc::Code *code);
	bool ParseIr(int *i);
	bool ParseInstDeref(bc::Code *code);
	bool ParseInstAlloc(bc::Code *code);
	bool ParseInstSave(bc::Code *code);
	bool ParseInstMove(bc::Code *code);
	bool ParseInstTranspose(bc::Code *code);
	bool ParseInstOuterproduct(bc::Code *code);
	bool ParseInstScalarproduct(bc::Code *code);
	bool ParseInstVectorproduct(bc::Code *code);
	bool ParseInstDeterminant(bc::Code *code);
	bool ParseInstSelect2(bc::Code *code);
	bool ParseInstSelect3(bc::Code *code);
	bool ParseInstSelrow(bc::Code *code);
	bool ParseInstMult(bc::Code *code);
	bool ParseInstMmul(bc::Code *code);

	bool ReadIdentifier(Token *t);
	bool ReadInteger(int *i);
	bool ReadToken(Token::Type type);
	bool ReadToken(Token::Type type, Token *t);
	bool ReadToken(Token *t);
	bool SkipParenClose();
	bool SkipSpace();

	template<bc::Call1::Op op1, bc::Call2::Op op2>
	bool ParseCall(int a, bc::Code *code) {
		int a1;
		if (!ParseFr(&a1))
			return false;
		Token t;
		if (!ReadToken(&t))
			return false;
		switch (t.type) {
		case Token::Type::kParenClose:
			{
				code->set_type(bc::Code::kCall1);
				auto *call1 = code->mutable_call1();
				call1->set_a(a);
				call1->set_op(op1);
				call1->set_a1(a1);
				return true;
			}
		case Token::Type::kSpace:
			// one more argument
			break;
		default:
			ReportUnexpectedToken(t);
			return false;
		}
		if ( !ReadToken(Token::Type::kAddress, &t) ||
			 !SkipParenClose() )
			return false;
		code->set_type(bc::Code::kCall2);
		auto *call2 = code->mutable_call2();
		call2->set_a(a);
		call2->set_op(op2);
		call2->set_a1(a1);
		call2->set_a2(GetIntegerFromAddress(t));
		return true;
	}

	void ReportTruncatedLine();

	Tokenizer lexer_;
};

ParserImpl::ParserImpl(const char *input)
	: lexer_(input)
{
	assert(input);
}

bool ParserImpl::Parse(Body *body)
{
	Token t;
	while ( int r = lexer_(&t) ) {
		if (r < 0)
			return false;
		if (t.type != Token::Type::kSpace) {
			std::cerr << "line starting with non-space: ";
			t.Write(std::cerr) << std::endl;
			return false;
		}
		switch (lexer_(&t)) {
		case 1:
			break;
		case 0:
			ReportTruncatedLine();
			return false;
		default:
			return false;
		}
		switch (t.type) {
		case Token::Type::kSpace:
			{
				bc::Code code;
				if (!ParseInst(&code))
					return false;
				body->code.push_back(code);
			}
			break;
		case Token::Type::kLabel:
			if (!ParseLabelColon(GetIntegerFromLabel(t), body))
				return false;
			break;
		default:
			ReportUnexpectedToken(t);
			return false;
		}
		if (!ReadToken(Token::Type::kEol))
			return false;
	}
	return true;
}

bool ParserImpl::ParseInst(bc::Code *code)
{
	Token t;
	switch (lexer_(&t)) {
	case 1:
		break;
	case 0:
		ReportTruncatedLine();
		return false;
	default:
		return false;
	}
	switch (t.type) {
	case Token::Type::kAddress:
		return ParseInstOp(GetIntegerFromAddress(t), code);
	case Token::Type::kBr:
		return ParseInstBr(code);
	case Token::Type::kJmp:
		return ParseInstJmp(code);
	case Token::Type::kLb:
		return ParseInstLb(code);
	case Token::Type::kLc:
		return ParseInstLc(code);
	case Token::Type::kLd:
		return ParseInstLd(code);
	case Token::Type::kLoad:
		return ParseInstLoad(code);
	case Token::Type::kLoadi:
		return ParseInstLoadi(code);
	case Token::Type::kRet:
		code->set_type(bc::Code::kRet);
		return true;
	case Token::Type::kStore:
		return ParseInstStore(code);
	case Token::Type::kRefer:
		return ParseInstRefer(code);
	case Token::Type::kDeref:
		return ParseInstDeref(code);
	case Token::Type::kAlloc:
		return ParseInstAlloc(code);
	case Token::Type::kSave:
		return ParseInstSave(code);
	case Token::Type::kMove:
		return ParseInstMove(code);
	case Token::Type::kTranspose:
		return ParseInstTranspose(code);
	case Token::Type::kOuterproduct:
		return ParseInstOuterproduct(code);
	case Token::Type::kScalarproduct:
		return ParseInstScalarproduct(code);
	case Token::Type::kVectorproduct:
		return ParseInstVectorproduct(code);
	case Token::Type::kDeterminant:
		return ParseInstDeterminant(code);
	case Token::Type::kSelect2:
		return ParseInstSelect2(code);
	case Token::Type::kSelect3:
		return ParseInstSelect3(code);
	case Token::Type::kSelrow:
		return ParseInstSelrow(code);
	case Token::Type::kMult:
		return ParseInstMult(code);
	case Token::Type::kMmul:
		return ParseInstMmul(code);
	default:
		break;
	}
	ReportUnexpectedToken(t);
	return false;
}

bool ParserImpl::ParseLabelColon(int n, Body *body)
{
	assert(n >= 0);
	if (body->labels.size() <= static_cast<size_t>(n))
		body->labels.resize(n + 1);
	body->labels.at(n) = static_cast<int>(body->code.size());
	return ReadToken(Token::Type::kColon);
}

bool ParserImpl::ParseInstOp(int a, bc::Code *code)
{
	if ( !SkipSpace() ||
		 !ReadToken(Token::Type::kEqualSign) ||
		 !SkipSpace() ||
		 !ReadToken(Token::Type::kParenOpen) )
		return false;
	Token t;
	switch (lexer_(&t)) {
	case 1:
		break;
	case 0:
		ReportTruncatedLine();
		return false;
	default:
		return false;
	}
	switch (t.type) {
	case Token::Type::kAbs:
		return ParseCall1(a, bc::Call1::kAbs, code);
	case Token::Type::kArccos:
		return ParseCall1(a, bc::Call1::kArccos, code);
	case Token::Type::kArccosh:
		return ParseCall1(a, bc::Call1::kArccosh, code);
	case Token::Type::kArccot:
		return ParseCall1(a, bc::Call1::kArccot, code);
	case Token::Type::kArccoth:
		return ParseCall1(a, bc::Call1::kArccoth, code);
	case Token::Type::kArccsc:
		return ParseCall1(a, bc::Call1::kArccsc, code);
	case Token::Type::kArccsch:
		return ParseCall1(a, bc::Call1::kArccsch, code);
	case Token::Type::kArcsec:
		return ParseCall1(a, bc::Call1::kArcsec, code);
	case Token::Type::kArcsech:
		return ParseCall1(a, bc::Call1::kArcsech, code);
	case Token::Type::kArcsin:
		return ParseCall1(a, bc::Call1::kArcsin, code);
	case Token::Type::kArcsinh:
		return ParseCall1(a, bc::Call1::kArcsinh, code);
	case Token::Type::kArctan:
		return ParseCall1(a, bc::Call1::kArctan, code);
	case Token::Type::kArctanh:
		return ParseCall1(a, bc::Call1::kArctanh, code);
	case Token::Type::kCeiling:
		return ParseCall1(a, bc::Call1::kCeiling, code);
	case Token::Type::kCos:
		return ParseCall1(a, bc::Call1::kCos, code);
	case Token::Type::kCosh:
		return ParseCall1(a, bc::Call1::kCosh, code);
	case Token::Type::kCot:
		return ParseCall1(a, bc::Call1::kCot, code);
	case Token::Type::kCoth:
		return ParseCall1(a, bc::Call1::kCoth, code);
	case Token::Type::kCsc:
		return ParseCall1(a, bc::Call1::kCsc, code);
	case Token::Type::kCsch:
		return ParseCall1(a, bc::Call1::kCsch, code);
	case Token::Type::kExp:
		return ParseCall1(a, bc::Call1::kExp, code);
	case Token::Type::kFloor:
		return ParseCall1(a, bc::Call1::kFloor, code);
	case Token::Type::kFactorial:
		return ParseCall1(a, bc::Call1::kFactorial, code);
	case Token::Type::kLn:
		return ParseCall1(a, bc::Call1::kLn, code);
	case Token::Type::kLog10:
		return ParseCall1(a, bc::Call1::kLog10, code);
	case Token::Type::kMinus: // 1 or 2 args
		return ParseCall<bc::Call1::kMinus1, bc::Call2::kMinus2>(a, code);
	case Token::Type::kRoot: // 1 or 2 args
		return ParseCall<bc::Call1::kRoot1, bc::Call2::kRoot2>(a, code);
	case Token::Type::kSec:
		return ParseCall1(a, bc::Call1::kSec, code);
	case Token::Type::kSech:
		return ParseCall1(a, bc::Call1::kSech, code);
	case Token::Type::kSin:
		return ParseCall1(a, bc::Call1::kSin, code);
	case Token::Type::kSinh:
		return ParseCall1(a, bc::Call1::kSinh, code);
	case Token::Type::kTan:
		return ParseCall1(a, bc::Call1::kTan, code);
	case Token::Type::kTanh:
		return ParseCall1(a, bc::Call1::kTanh, code);
	// Call2
	case Token::Type::kDivide:
		return ParseCall2(a, bc::Call2::kDivide, code);
	case Token::Type::kEq:
		return ParseCall2(a, bc::Call2::kEq, code);
	case Token::Type::kGeq:
		return ParseCall2(a, bc::Call2::kGeq, code);
	case Token::Type::kGt:
		return ParseCall2(a, bc::Call2::kGt, code);
	case Token::Type::kLeq:
		return ParseCall2(a, bc::Call2::kLeq, code);
	case Token::Type::kLog:
		return ParseCall2(a, bc::Call2::kLog, code);
	case Token::Type::kLt:
		return ParseCall2(a, bc::Call2::kLt, code);
	case Token::Type::kMax:
		return ParseCall2(a, bc::Call2::kMax, code);
	case Token::Type::kMin:
		return ParseCall2(a, bc::Call2::kMin, code);
	case Token::Type::kNeq:
		return ParseCall2(a, bc::Call2::kNeq, code);
	case Token::Type::kPlus:
		return ParseCall2(a, bc::Call2::kPlus, code);
	case Token::Type::kPower:
		return ParseCall2(a, bc::Call2::kPower, code);
	case Token::Type::kRem: // remainder
		return ParseCall2(a, bc::Call2::kRemainder, code);
	case Token::Type::kTimes:
		return ParseCall2(a, bc::Call2::kTimes, code);
	case Token::Type::kMod: // Modulo
		return ParseCall2(a, bc::Call2::kModulo, code);
	// Gen1
	case Token::Type::kExponentialVariate:
		return ParseGen1(a, bc::Gen1::kExponentialVariate, code);
	case Token::Type::kPoissonVariate:
		return ParseGen1(a, bc::Gen1::kPoissonVariate, code);
	// Gen2
	case Token::Type::kGammaVariate:
		return ParseGen2(a, bc::Gen2::kGammaVariate, code);
	case Token::Type::kGaussVariate:
		return ParseGen2(a, bc::Gen2::kGaussVariate, code);
	case Token::Type::kLognormalVariate:
		return ParseGen2(a, bc::Gen2::kLognormalVariate, code);
	case Token::Type::kUniformVariate:
		return ParseGen2(a, bc::Gen2::kUniformVariate, code);
	case Token::Type::kWeibullVariate:
		return ParseGen2(a, bc::Gen2::kWeibullVariate, code);
	default:
		std::cerr << "unknown instruction: ";
		t.Write(std::cerr) << std::endl;
		return false;
	}
}

bool ParserImpl::ParseCall1(int a, bc::Call1::Op op, bc::Code *code)
{
	int a1;
	if ( ParseFr(&a1) && SkipParenClose() ) {
		code->set_type(bc::Code::kCall1);
		auto *call1 = code->mutable_call1();
		call1->set_a(a);
		call1->set_op(op);
		call1->set_a1(a1);
		return true;
	}
	return false;
}

bool ParserImpl::ParseCall2(int a, bc::Call2::Op op, bc::Code *code)
{
	int a1, a2;
	if ( ParseFr(&a1) && ParseFr(&a2) && SkipParenClose() ) {
		code->set_type(bc::Code::kCall2);
		auto *call2 = code->mutable_call2();
		call2->set_a(a);
		call2->set_op(op);
		call2->set_a1(a1);
		call2->set_a2(a2);
		return true;
	}
	return false;
}

bool ParserImpl::ParseGen1(int a, bc::Gen1::Type type, bc::Code *code)
{
	int a1;
	if ( ParseFr(&a1) && SkipParenClose() ) {
		code->set_type(bc::Code::kGen1);
		auto *gen1 = code->mutable_gen1();
		gen1->set_a(a);
		gen1->set_type(type);
		gen1->set_a1(a1);
		return true;
	}
	return false;
}

bool ParserImpl::ParseGen2(int a, bc::Gen2::Type type, bc::Code *code)
{
	int a1, a2;
	if ( ParseFr(&a1) && ParseFr(&a2) && SkipParenClose() ) {
		code->set_type(bc::Code::kGen2);
		auto *gen2 = code->mutable_gen2();
		gen2->set_a(a);
		gen2->set_type(type);
		gen2->set_a1(a1);
		gen2->set_a2(a2);
		return true;
	}
	return false;
}

bool ParserImpl::ParseFr(int *a)
{
	Token t;
	if ( SkipSpace() && ReadToken(Token::Type::kAddress, &t) ) {
		*a = GetIntegerFromAddress(t);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstBr(bc::Code *code)
{
	int a, n;
	if ( ParseFr(&a) && ParseLabel(&n) ) {
		code->set_type(bc::Code::kBr);
		auto *br = code->mutable_br();
		br->set_a(a);
		br->set_l(n);
		return true;
	}
	return false;
}

bool ParserImpl::ParseLabel(int *n)
{
	Token t;
	if ( SkipSpace() && ReadToken(Token::Type::kLabel, &t) ) {
		*n = GetIntegerFromLabel(t);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstJmp(bc::Code *code)
{
	int n;
	if (ParseLabel(&n)) {
		code->set_type(bc::Code::kJmp);
		auto *jmp = code->mutable_jmp();
		jmp->set_l(n);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstLb(bc::Code *code)
{
	int a, d;
	Token t;
	if ( ParseFr(&a) &&
		 SkipSpace() && ReadIdentifier(&t) &&
		 ParseFr(&d) ) {
		code->set_type(bc::Code::kLb);
		auto *lb = code->mutable_lb();
		lb->set_a(a);
		lb->set_v(t.lexeme, t.size);
		lb->set_d(d);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstLc(bc::Code *code)
{
	int a, i0, d;
	if ( ParseFr(&a) &&
		 SkipSpace() && ReadInteger(&i0) &&
		 ParseFr(&d) ) {
		code->set_type(bc::Code::kLc);
		auto *lc = code->mutable_lc();
		lc->set_a(a);
		lc->set_i0(i0);
		lc->set_d(d);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstLd(bc::Code *code)
{
	int a, i0, i1, d;
	if ( ParseFr(&a) &&
		 SkipSpace() && ReadInteger(&i0) &&
		 SkipSpace() && ReadInteger(&i1) &&
		 ParseFr(&d) ) {
		code->set_type(bc::Code::kLd);
		auto *ld = code->mutable_ld();
		ld->set_a(a);
		ld->set_i0(i0);
		ld->set_i1(i1);
		ld->set_d(d);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstLoad(bc::Code *code)
{
	int a;
	Token t;
	if ( ParseFr(&a) &&
		 SkipSpace() && ReadIdentifier(&t) ) {
		code->set_type(bc::Code::kLoad);
		auto *load = code->mutable_load();
		load->set_a(a);
		load->set_v(t.lexeme, t.size); // TODO
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstLoadi(bc::Code *code)
{
	int a;
	double v;
	if ( ParseFr(&a) && ParseImm(&v) ) {
		code->set_type(bc::Code::kLoadi);
		auto *loadi = code->mutable_loadi();
		loadi->set_a(a);
		loadi->set_v(v);
		return true;
	}
	return false;
}

bool ParserImpl::ParseImm(double *d)
{
	if (!SkipSpace())
		return false;
	Token t;
	if (!ReadToken(&t))
		return false;
	switch (t.type) {
	case Token::Type::kEulergamma:
		*d = boost::math::constants::euler<double>();
		return true;
	case Token::Type::kExponentiale:
		*d = boost::math::constants::e<double>();
		return true;
	case Token::Type::kPi:
		*d = boost::math::constants::pi<double>();
		return true;
	case Token::Type::kTrue:
		*d = 1;
		return true;
	case Token::Type::kFalse:
		*d = 0;
		return true;
	case Token::Type::kNumber:
		{
			const char *c = std::strpbrk(t.lexeme, "/\n");
			if ( c && *c == '/' ) { // rationals
				double num = std::atof(t.lexeme);
				if (num == HUGE_VAL) {
					std::cerr << "numerator is out of range: ";
					t.Write(std::cerr) << std::endl;
					return false;
				}
				double den = std::atof(++c);
				if (den == HUGE_VAL) {
					std::cerr << "denominator is out of range: ";
					t.Write(std::cerr) << std::endl;
					return false;
				}
				if (den == 0) {
					std::cerr << "invalid denominator: ";
					t.Write(std::cerr) << std::endl;
					return false;
				}
				*d = num/den;
			} else {
				*d = std::atof(t.lexeme);
			}
			return true;
		}
	default:
		std::cerr << "immediate value expected, but got: ";
		t.Write(std::cerr) << std::endl;
		return false;
	}
}

bool ParserImpl::ParseInstStore(bc::Code *code)
{
	Token t;
	int a;
	if ( SkipSpace() && ReadIdentifier(&t) &&
		 ParseFr(&a) ) {
		code->set_type(bc::Code::kStore);
		auto *store = code->mutable_store();
		store->set_v(t.lexeme, t.size);
		store->set_a(a);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstRefer(bc::Code *code)
{
	int i0;
	Token t;
	if ( ParseIr(&i0) &&
		 SkipSpace() && ReadIdentifier(&t) ) {
		code->set_type(bc::Code::kRefer);
		auto *x = code->mutable_refer();
		x->set_i0(i0);
		x->set_v(t.lexeme, t.size);
		return true;
	}
	return false;
}

bool ParserImpl::ParseIr(int *i)
{
	Token t;
	if ( SkipSpace() && ReadToken(Token::Type::kIntReg, &t) ) {
		*i = GetIntegerFromIntReg(t);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstDeref(bc::Code *code)
{
	int f0, i1, k;
	if ( ParseFr(&f0) &&
		 ParseIr(&i1) &&
		 SkipSpace() && ReadInteger(&k) ) {
		code->set_type(bc::Code::kDeref);
		auto *x = code->mutable_deref();
		x->set_f0(f0);
		x->set_i1(i1);
		x->set_k(k);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstAlloc(bc::Code *code)
{
	int i0, k;
	if ( ParseIr(&i0) &&
		 SkipSpace() && ReadInteger(&k) ) {
		code->set_type(bc::Code::kAlloc);
		auto *x = code->mutable_alloc();
		x->set_i0(i0);
		x->set_k(k);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstSave(bc::Code *code)
{
	int i1, k;
	Token t;
	if ( SkipSpace() && ReadIdentifier(&t) &&
		 ParseIr(&i1) &&
		 SkipSpace() && ReadInteger(&k) ) {
		code->set_type(bc::Code::kSave);
		auto *x = code->mutable_save();
		x->set_v(t.lexeme, t.size);
		x->set_i1(i1);
		x->set_k(k);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstMove(bc::Code *code)
{
	int i0, f1, k;
	if ( ParseIr(&i0) &&
		 ParseFr(&f1) &&
		 SkipSpace() && ReadInteger(&k) ) {
		code->set_type(bc::Code::kMove);
		auto *x = code->mutable_move();
		x->set_i0(i0);
		x->set_f1(f1);
		x->set_k(k);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstTranspose(bc::Code *code)
{
	int i0, i1, kr, kc;
	if ( ParseIr(&i0) &&
		 ParseIr(&i1) &&
		 SkipSpace() && ReadInteger(&kr) &&
		 SkipSpace() && ReadInteger(&kc) ) {
		code->set_type(bc::Code::kTranspose);
		auto *x = code->mutable_transpose();
		x->set_i0(i0);
		x->set_i1(i1);
		x->set_kc(kc);
		x->set_kr(kr);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstOuterproduct(bc::Code *code)
{
	int i0, k1, i1, k2, i2;
	if ( ParseIr(&i0) &&
		 SkipSpace() && ReadInteger(&k1) &&
		 ParseIr(&i1) &&
		 SkipSpace() && ReadInteger(&k2) &&
		 ParseIr(&i2) ) {
		code->set_type(bc::Code::kOuterproduct);
		auto *x = code->mutable_outerproduct();
		x->set_i0(i0);
		x->set_k1(k1);
		x->set_i1(i1);
		x->set_k2(k2);
		x->set_i2(i2);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstScalarproduct(bc::Code *code)
{
	int f0, k, i1, i2;
	if ( ParseFr(&f0) &&
		 SkipSpace() && ReadInteger(&k) &&
		 ParseIr(&i1) &&
		 ParseIr(&i2) ) {
		code->set_type(bc::Code::kScalarproduct);
		auto *x = code->mutable_scalarproduct();
		x->set_f0(f0);
		x->set_k(k);
		x->set_i1(i1);
		x->set_i2(i2);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstVectorproduct(bc::Code *code)
{
	int i0, i1, i2;
	if ( ParseIr(&i0) && ParseIr(&i1) && ParseIr(&i2) ) {
		code->set_type(bc::Code::kVectorproduct);
		auto *x = code->mutable_vectorproduct();
		x->set_i0(i0);
		x->set_i1(i1);
		x->set_i2(i2);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstDeterminant(bc::Code *code)
{
	int f0, k, i1;
	if ( ParseFr(&f0) &&
		 SkipSpace() && ReadInteger(&k) &&
		 ParseIr(&i1) ) {
		code->set_type(bc::Code::kDeterminant);
		auto *x = code->mutable_determinant();
		x->set_f0(f0);
		x->set_k(k);
		x->set_i1(i1);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstSelect2(bc::Code *code)
{
	int f0, i1, f2;
	if ( ParseFr(&f0) && ParseIr(&i1) && ParseFr(&f2) ) {
		code->set_type(bc::Code::kSelect2);
		auto *x = code->mutable_select2();
		x->set_f0(f0);
		x->set_i1(i1);
		x->set_f2(f2);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstSelect3(bc::Code *code)
{
	int f0, kr, kc, i1, f2, f3;
	if ( ParseFr(&f0) &&
		 SkipSpace() && ReadInteger(&kr) &&
		 SkipSpace() && ReadInteger(&kc) &&
		 ParseIr(&i1) &&
		 ParseFr(&f2) &&
		 ParseFr(&f3) ) {
		code->set_type(bc::Code::kSelect3);
		auto *x = code->mutable_select3();
		x->set_f0(f0);
		x->set_kc(kc);
		x->set_kr(kr);
		x->set_i1(i1);
		x->set_f2(f2);
		x->set_f3(f3);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstSelrow(bc::Code *code)
{
	int i0, kr, kc, i1, f2;
	if ( ParseIr(&i0) &&
		 SkipSpace() && ReadInteger(&kr) &&
		 SkipSpace() && ReadInteger(&kc) &&
		 ParseIr(&i1) &&
		 ParseFr(&f2) ) {
		code->set_type(bc::Code::kSelrow);
		auto *x = code->mutable_selrow();
		x->set_i0(i0);
		x->set_kc(kc);
		x->set_kr(kr);
		x->set_i1(i1);
		x->set_f2(f2);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstMult(bc::Code *code)
{
	int i0, k, f1, i2;
	if ( ParseIr(&i0) &&
		 SkipSpace() && ReadInteger(&k) &&
		 ParseFr(&f1) &&
		 ParseIr(&i2) ) {
		code->set_type(bc::Code::kMult);
		auto *x = code->mutable_mult();
		x->set_i0(i0);
		x->set_k(k);
		x->set_f1(f1);
		x->set_i2(i2);
		return true;
	}
	return false;
}

bool ParserImpl::ParseInstMmul(bc::Code *code)
{
	int i0, kr, kx, kc, i1, i2;
	if ( ParseIr(&i0) &&
		 SkipSpace() && ReadInteger(&kr) &&
		 SkipSpace() && ReadInteger(&kx) &&
		 SkipSpace() && ReadInteger(&kc) &&
		 ParseIr(&i1) &&
		 ParseIr(&i2) ) {
		code->set_type(bc::Code::kMmul);
		auto *x = code->mutable_mmul();
		x->set_i0(i0);
		x->set_kr(kr);
		x->set_kx(kx);
		x->set_kc(kc);
		x->set_i1(i1);
		x->set_i2(i2);
		return true;
	}
	return false;
}

bool ParserImpl::ReadIdentifier(Token *t)
{
	return ReadToken(Token::Type::kIdentifier, t);
}

bool ParserImpl::ReadInteger(int *i)
{
	Token t;
	if (ReadToken(Token::Type::kNumber, &t)) {
		*i = GetIntegerFromNumber(t);
		return true;
	}
	return false;
}

bool ParserImpl::ReadToken(Token::Type type)
{
	Token t;
	return ReadToken(type, &t);
}

bool ParserImpl::ReadToken(Token::Type type, Token *t)
{
	if (!ReadToken(t))
		return false;
	if (t->type == type)
		return true;
	std::cerr << "expected token of type "
			  << static_cast<int>(type)
			  << ", but got: ";
	t->Write(std::cerr) << std::endl;
	return false;
}

bool ParserImpl::ReadToken(Token *t)
{
	switch (lexer_(t)) {
	case 1:
		return true;
	case 0:
		ReportTruncatedLine();
		return false;
	default:
		return false;
	}
}

bool ParserImpl::SkipParenClose()
{
	return ReadToken(Token::Type::kParenClose);
}

bool ParserImpl::SkipSpace()
{
	return ReadToken(Token::Type::kSpace);
}

void ParserImpl::ReportTruncatedLine()
{
	std::cerr << "truncated line: ";
	lexer_.Dump(std::cerr) << std::endl;
}

Parser::Parser(const char *input)
	: impl_(new ParserImpl(input))
{
	assert(input);
}

Parser::~Parser() = default;

bool Parser::operator()(Body *body)
{
	assert(body);
	return impl_->Parse(body);
}

}
}
}


// This file is part of sibilla : inference in epidemics with Belief Propagation
// Author: Alfredo Braunstein

#include <pybind11/pybind11.h>
#include <pybind11/stl_bind.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <pybind11/pytypes.h>
#include <pybind11/iostream.h>
#include <string>
#include <sstream>
#include <numeric>
#include <boost/lexical_cast.hpp>
#include <iterator>
#include <exception>
#include "bp.h"
#include "drop.h"

PYBIND11_MAKE_OPAQUE(std::valarray<real_t>);
PYBIND11_MAKE_OPAQUE(std::vector<real_t>);
PYBIND11_MAKE_OPAQUE(std::vector<int>);
PYBIND11_MAKE_OPAQUE(std::vector<Node>);

namespace py = pybind11;
using namespace std;
using boost::lexical_cast;

vector<real_t> make_vector(py::list & l)
{
    vector<real_t> v(l.size());
    int i = 0;
    for (py::handle o : l) {
        v[i++] = py::cast<real_t>(o);
    }
    return v;
}

map<int, vector<times_t> >
get_times(FactorGraph const & f) {
        map<int, vector<times_t> > times;
        for (int i = 0; i < int(f.nodes.size()); ++i) {
            times[i] = f.nodes[i].times;
            times[i].pop_back();
        }
        return times;
}

vector<tuple<real_t, real_t, real_t>>
get_marginal(Node const & n)
{
        vector<real_t> rbt(n.bt.size());
        vector<real_t> lbg(n.bg.size());
        int const T = n.bt.size() - 1;
        lbg[0] = n.bg[0];
        rbt[T] = n.bt[T];
        for (int t = 1; t <= T; ++t) {
            lbg[t] = lbg[t-1] + n.bg[t];
            rbt[T - t] = rbt[T - t + 1] + n.bt[T - t];
        }
        auto marg = vector<tuple<real_t, real_t, real_t>>(T-1);
        for (int t = 1; t < T; ++t)
            marg[t-1] = make_tuple(rbt[t], 1-rbt[t]-lbg[t-1], lbg[t-1]);
        return marg;
}

tuple<real_t, real_t, real_t> get_marginal_index(Node const & n, int t)
{
        if (t < 0 || t >= int(n.bt.size()) - 2)
            throw py::key_error("Time out of range");
        return get_marginal(n)[t];
}

void check_index(FactorGraph const & G, int i)
{
        if (i < 0 || i >= int(G.nodes.size()))
                throw invalid_argument("unexistent index");

}

template<int i>
real_t mygetter(Proba & p)
{
    return p.theta[i];
}

template<int i>
void mysetter(Proba & p, real_t x)
{
    p.theta[i] = x;
}

PYBIND11_MODULE(_sib, m) {
    py::class_<RealParams>(m, "RealParams", py::buffer_protocol())
        .def(py::init([](py::buffer const b) {
                py::buffer_info info = b.request();
                if (info.format != py::format_descriptor<real_t>::format() || info.ndim != 1)
                throw std::runtime_error("Incompatible buffer format!");

                auto v = new RealParams(info.shape[0]);
                memcpy(&(*v)[0], info.ptr, sizeof(real_t) * (size_t) (v->size()));
                return v;
                }))
        .def(py::init([](vector<real_t> const & p)->RealParams {return RealParams(&p[0], p.size());}))
        .def(py::init([](py::list & l)->RealParams {auto v = make_vector(l); return RealParams(&v[0], v.size());}))
        .def_buffer([](RealParams &p) -> py::buffer_info {
            return py::buffer_info(
                &p[0],                               /* Pointer to buffer */
                sizeof(real_t),                          /* Size of one scalar */
                py::format_descriptor<real_t>::format(), /* Python struct-style format descriptor */
                1,                                      /* Number of dimensions */
                { p.size() },                 /* Buffer dimensions */
                { sizeof(real_t) }             /* Strides (in bytes) for each index */
                );
        })
        .def("__add__", [](RealParams & p, RealParams & q)->RealParams { return p + q; })
        .def("__getitem__", [](const RealParams &p, ssize_t i) {
                if (i > int(p.size()))
                    throw py::index_error();
                return p[i];
                })
        .def("__setitem__", [](RealParams &p, ssize_t i, real_t v) {
                if (i > int(p.size()))
                    throw py::index_error();
                p[i] = v;
                })
        .def("__repr__", [](RealParams &p) {
                    string s = "RealParams([";
                    for (size_t i = 0; i < p.size(); ++i)
                        s += (i ? ",":"") + lexical_cast<string>(p[i]);
                    s+="])";
                    return s;
                });
    // py::add_ostream_redirect(m, "ostream_redirect");
    py::bind_vector<std::vector<int>>(m, "VectorInt");
    py::bind_vector<std::vector<real_t>>(m, "VectorReal");
    py::bind_vector<std::vector<Node>>(m, "VectorNode");

    py::class_<Test, shared_ptr<Test>>(m, "Test")
        .def(py::init<real_t,real_t,real_t>(), py::arg("ps")=0.0, py::arg("pi")=0.0, py::arg("pr")=0.0)
        .def(py::init([](int s)->Test{return Test(s==0, s==1, s==2);}))
        .def("__repr__", &lexical_cast<string, Test>)
        .def_readwrite("ps", &Test::ps, "probability of S")
        .def_readwrite("pi", &Test::pi, "probability of I")
        .def_readwrite("pr", &Test::pr, "probability of R");

    py::class_<Proba, shared_ptr<Proba>>(m, "Proba")
        .def("__call__", [](Proba const & p, real_t d) { return p(d); } )
        .def("grad", [](Proba const & p, real_t d) { RealParams dtheta(0.0, p.theta.size()); p.grad(dtheta, d); return dtheta;} )
        .def("__repr__", &lexical_cast<string, Proba>)
        .def_readwrite("theta", &Proba::theta);

    py::class_<Uniform, Proba, shared_ptr<Uniform>>(m, "Uniform")
        .def(py::init<real_t>(), py::arg("p") = 1.0)
        .def_property("p", &mygetter<0>, &mysetter<0>);

    py::class_<Exponential, Proba, shared_ptr<Exponential>>(m, "Exponential")
        .def(py::init<real_t>(), py::arg("mu") = 0.1)
        .def_property("mu", &mygetter<0>, &mysetter<0>);

    py::class_<Gamma, Proba, shared_ptr<Gamma>>(m, "Gamma")
        .def(py::init<real_t, real_t>(), py::arg("k") = 1.0, py::arg("mu") = 0.1)
        .def_property("k", &mygetter<0>, &mysetter<0>)
        .def_property("mu", &mygetter<1>, &mysetter<1>);

    py::class_<ConstantRate, Proba, shared_ptr<ConstantRate>>(m, "ConstantRate")
        .def(py::init<real_t, real_t>(), py::arg("gamma") = 1.0, py::arg("Dt") = 1.0)
        .def_property("gamma", &mygetter<0>, &mysetter<0>)
        .def_property("Dt", [](ConstantRate & p){ return p.Dt; }, [](ConstantRate & p, real_t Dt) { p.Dt = Dt; });

    py::class_<UnnormalizedGammaPDF, Proba, shared_ptr<UnnormalizedGammaPDF>>(m, "UnnormalizedGammaPDF")
        .def(py::init<real_t, real_t>(), py::arg("k") = 1.0, py::arg("mu") = 0.1)
        .def_property("k", &mygetter<0>, &mysetter<0>)
        .def_property("mu", &mygetter<1>, &mysetter<1>);
    py::class_<PiecewiseLinear, Proba, shared_ptr<PiecewiseLinear>>(m, "PiecewiseLinear")
        .def(py::init<RealParams const &, real_t>(), py::arg("theta"), py::arg("step") = 1.0)
        .def(py::init<Proba const &, int, real_t>(), py::arg("prob"), py::arg("num"), py::arg("step") = 1.0);

    py::class_<Cached, Proba, shared_ptr<Cached>>(m, "Cached")
        .def(py::init<std::shared_ptr<Proba> const &, int>(), py::arg("prob"), py::arg("T"))
        .def_property("theta", &Cached::get_theta, &Cached::set_theta);

    py::class_<Scaled, Proba, shared_ptr<Scaled>>(m, "Scaled")
        .def(py::init<std::shared_ptr<Proba> const &, real_t>(), py::arg("prob"), py::arg("scale") = 1.0);

    py::class_<PDF, Proba, shared_ptr<PDF>>(m, "PDF")
        .def(py::init<std::shared_ptr<Proba> const &>());

    py::class_<Params>(m, "Params")
        .def(py::init<shared_ptr<Proba> const &, shared_ptr<Proba> const &, real_t, real_t, real_t, real_t>(),
                "SIB Params class. prob_i and prob_r parameters are defaults.",
                py::arg("prob_i") = *new Uniform(1.0),
                py::arg("prob_r") = *new Exponential(0.1),
                py::arg("pseed") = 0.01,
                py::arg("psus") = 0.5,
                py::arg("pautoinf") = 0.0,
                py::arg("learn_rate") = 0.0)

        .def_readwrite("prob_r", &Params::prob_r)
        .def_readwrite("prob_i", &Params::prob_i)
        .def_readonly("fakeobs", &Params::fakeobs)
        .def_readwrite("pseed", &Params::pseed)
        .def_readwrite("psus", &Params::psus)
        .def_readwrite("pautoinf", &Params::pautoinf)
        .def_readwrite("learn_rate", &Params::learn_rate)
        .def("__repr__", &lexical_cast<string, Params>);

    py::class_<FactorGraph>(m, "FactorGraph", "SIB class representing the graphical model of the epidemics")
        .def(py::init<Params const &,
                vector<tuple<int,int,times_t,real_t>>,
                vector<tuple<int,shared_ptr<Test>,times_t>>,
                vector<tuple<int,shared_ptr<Proba>,shared_ptr<Proba>,shared_ptr<Proba>,shared_ptr<Proba>>>
                >(),
                py::arg("params") = Params(shared_ptr<Proba>(new Uniform(1.0)), shared_ptr<Proba>(new Exponential(0.5)), 0.1, 0.45, 0.0, 0.0),
                py::arg("contacts") = vector<tuple<int,int,times_t,real_t>>(),
                py::arg("tests") = vector<tuple<int,shared_ptr<Test>,times_t>>(),
                py::arg("individuals") = vector<tuple<int,shared_ptr<Proba>,shared_ptr<Proba>,shared_ptr<Proba>,shared_ptr<Proba>>>())
        .def("update", &FactorGraph::iteration,
                py::arg("damping") = 0.0,
                py::arg("learn") = false,
                "perform one iteration")
        .def("loglikelihood", &FactorGraph::loglikelihood, "compute the bethe log-likelihood")
        .def("__repr__", &lexical_cast<string, FactorGraph>)
        .def("append_contact", (void (FactorGraph::*)(int,int,times_t,real_t,real_t)) &FactorGraph::append_contact,
                py::arg("i"),
                py::arg("j"),
                py::arg("t"),
                py::arg("lambdaij"),
                py::arg("lambdaji") = real_t(FactorGraph::DO_NOT_OVERWRITE),
                "appends a new contact from i to j at time t with transmission probabilities lambdaij, lambdaji")
        .def("reset_observations", &FactorGraph::reset_observations,
                py::arg("obs"),
                "resets all observations")
        .def("append_observation", (void (FactorGraph::*)(int,int,times_t)) &FactorGraph::append_observation,
                py::arg("i"),
                py::arg("s"),
                py::arg("t"),
                "adds a new observation state s to node i at time t")
        .def("append_observation", (void (FactorGraph::*)(int,shared_ptr<Test> const &,times_t)) &FactorGraph::append_observation,
                py::arg("i"),
                py::arg("o"),
                py::arg("t"),
                "adds a new test o to node i at time t")
        .def("show", &FactorGraph::show_graph)
        .def("drop_contacts", &FactorGraph::drop_contacts, "drop contacts at time t (first time)")
        .def("drop_time", &drop_time, "drop time t (first time)")
        .def("drop_sc", &drop_sc,
                py::arg("t"),
                py::arg("maxit_bp") = 1,
                py::arg("tol_bp") = 1e-3,
                py::arg("damping_bp") = 0.0,
                py::arg("maxit_sc") = 20,
                py::arg("tol_sc") = 1e-3,
                py::arg("damping_sc") = 0.1,
                "drop contacts at time t (first time), adjusting fields")

        .def("showmsg", [](FactorGraph & f){f.show_msg(std::cout);}, "show messages for debugging")
        .def_readonly("nodes", &FactorGraph::nodes, "all nodes in this FactorGraph")
        .def_readonly("params", &FactorGraph::params, "parameters");

    py::class_<Node>(m, "Node", "SIB class representing an individual")
        .def("marginal", &get_marginal, "compute marginal probabilities (pS,pI,pR) corresponding to times n.times[1:]")
        .def("marginal_index", &get_marginal_index, "marginal at a given time (excluding time -1)")
        .def_readwrite("ht", &Node::ht, "external prior on ti")
        .def_readwrite("hg", &Node::hg, "external prior on gi")
        .def_readonly("bt", &Node::bt, "belief on ti")
        .def_readonly("bg", &Node::bg, "belief on gi")
        .def_readonly("err", &Node::err_, "error on update")
        .def_readonly("df_i", &Node::df_i, "gradient on prior_i params")
        .def_readonly("df_r", &Node::df_r, "gradient on prior_r params")
        .def_readonly("times", &Node::times, "event times of this node")
        .def_readonly("index", &Node::index, "node index (deprecated, do not use)")
        .def_readonly("prob_i", &Node::prob_i, "probability of infection as function of t-ti")
        .def_readonly("prob_r", &Node::prob_r, "cumulative probability of recovery P(tr>t)")
        .def_readonly("prob_i0", &Node::prob_i0, "probability of infection as function of t-ti for ti=0")
        .def_readonly("prob_r0", &Node::prob_r0, "cumulative probability of recovery P(tr>t) for ti=0");

    m.def("set_num_threads", &omp_set_num_threads, "sets the maximum number of simultaneous cpu threads");
    m.def("version", [](){return VERSION;}, "compiled version of sib");
}