/*
 * Copyright 2009-2020 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.
 *
 *     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.
 *
 */
#define BOOST_TEST_MAIN

#define BOOST_TEST_MODULE bse_test

// Third party includes
#include <boost/test/unit_test.hpp>

// VOTCA includes
#include <votca/tools/eigenio_matrixmarket.h>

// Local VOTCA includes
#include "votca/xtp/bse.h"
#include "votca/xtp/convergenceacc.h"
#include "votca/xtp/qmfragment.h"
#include "xtp_libint2.h"
#include <votca/tools/eigenio_matrixmarket.h>
using namespace votca::xtp;
using namespace std;

BOOST_AUTO_TEST_SUITE(bse_test)

BOOST_AUTO_TEST_CASE(bse_hamiltonian) {
  libint2::initialize();
  Orbitals orbitals;
  orbitals.QMAtoms().LoadFromFile(std::string(XTP_TEST_DATA_FOLDER) +
                                  "/bse/molecule.xyz");
  BasisSet basis;
  basis.Load(std::string(XTP_TEST_DATA_FOLDER) + "/bse/3-21G.xml");
  orbitals.SetupDftBasis(std::string(XTP_TEST_DATA_FOLDER) + "/bse/3-21G.xml");
  AOBasis aobasis;
  aobasis.Fill(basis, orbitals.QMAtoms());

  orbitals.setNumberOfOccupiedLevels(4);
  Eigen::MatrixXd& MOs = orbitals.MOs().eigenvectors();
  MOs = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/MOs.mm");

  Eigen::MatrixXd Hqp = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/Hqp.mm");

  Eigen::VectorXd& mo_energy = orbitals.MOs().eigenvalues();
  mo_energy = votca::tools::EigenIO_MatrixMarket::ReadVector(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/MO_energies.mm");

  Logger log;
  TCMatrix_gwbse Mmn;
  Mmn.Initialize(aobasis.AOBasisSize(), 0, 16, 0, 16);
  Mmn.Fill(aobasis, aobasis, MOs);

  BSE::options opt;
  opt.cmax = 16;
  opt.rpamax = 16;
  opt.rpamin = 0;
  opt.vmin = 0;
  opt.nmax = 3;
  opt.min_print_weight = 0.1;
  opt.useTDA = true;
  opt.homo = 4;
  opt.qpmin = 0;
  opt.qpmax = 16;
  opt.max_dyn_iter = 10;
  opt.dyn_tolerance = 1e-5;
  opt.davidson_correction = "DPR";
  opt.davidson_tolerance = "lapack";
  opt.davidson_update = "safe";
  opt.davidson_maxiter = 50;

  orbitals.setBSEindices(0, 16);

  BSE bse = BSE(log, Mmn);
  orbitals.setTDAApprox(true);
  orbitals.RPAInputEnergies() = Hqp.diagonal();

  ////////////////////////////////////////////////////////
  // TDA Singlet davidson
  ////////////////////////////////////////////////////////

  // reference energy singlet, no offdiagonals in Hqp
  Eigen::VectorXd se_nooffdiag_ref =
      votca::tools::EigenIO_MatrixMarket::ReadVector(
          std::string(XTP_TEST_DATA_FOLDER) + "/bse/singlets_nooffdiag_tda.mm");

  // reference singlet coefficients, no offdiagonals in Hqp
  Eigen::MatrixXd spsi_nooffdiag_ref =
      votca::tools::EigenIO_MatrixMarket::ReadMatrix(
          std::string(XTP_TEST_DATA_FOLDER) +
          "/bse/singlets_psi_nooffdiag_tda.mm");

  // no offdiagonals
  opt.use_Hqp_offdiag = false;
  bse.configure(opt, orbitals.RPAInputEnergies(), Hqp);

  bse.Solve_singlets(orbitals);
  std::vector<QMFragment<BSE_Population> > fragments;
  bse.Analyze_singlets(fragments, orbitals);
  bool check_se_nooffdiag =
      se_nooffdiag_ref.isApprox(orbitals.BSESinglets().eigenvalues(), 0.001);
  if (!check_se_nooffdiag) {
    cout << "Singlets energy without Hqp offdiag" << endl;
    cout << orbitals.BSESinglets().eigenvalues() << endl;
    cout << "Singlets energy without Hqp offdiag ref" << endl;
    cout << se_nooffdiag_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_se_nooffdiag, true);
  Eigen::MatrixXd projection_nooffdiag =
      spsi_nooffdiag_ref.transpose() * orbitals.BSESinglets().eigenvectors();
  Eigen::VectorXd norms_nooffdiag = projection_nooffdiag.colwise().norm();
  bool check_spsi_nooffdiag = norms_nooffdiag.isApproxToConstant(1, 1e-5);
  if (!check_spsi_nooffdiag) {
    cout << "Norms" << norms_nooffdiag << endl;
    cout << "Singlets psi without Hqp offdiag" << endl;
    cout << orbitals.BSESinglets().eigenvectors() << endl;
    cout << "Singlets psi without Hqp offdiag ref" << endl;
    cout << spsi_nooffdiag_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_spsi_nooffdiag, true);

  // with Hqp offdiags
  opt.use_Hqp_offdiag = true;
  bse.configure(opt, orbitals.RPAInputEnergies(), Hqp);

  // reference energy
  Eigen::VectorXd se_ref = votca::tools::EigenIO_MatrixMarket::ReadVector(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/singlets_tda.mm");
  // reference coefficients
  Eigen::MatrixXd spsi_ref = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/singlets_psi_tda.mm");

  // Hqp unchanged
  bool check_hqp_unchanged = Hqp.isApprox(bse.getHqp(), 0.001);
  if (!check_hqp_unchanged) {
    cout << "unchanged Hqp" << endl;
    cout << bse.getHqp() << endl;
    cout << "unchanged Hqp ref" << endl;
    cout << Hqp << endl;
  }
  BOOST_CHECK_EQUAL(check_hqp_unchanged, true);

  bse.Solve_singlets(orbitals);
  bool check_se = se_ref.isApprox(orbitals.BSESinglets().eigenvalues(), 0.001);
  if (!check_se) {
    cout << "Singlets energy" << endl;
    cout << orbitals.BSESinglets().eigenvalues() << endl;
    cout << "Singlets energy ref" << endl;
    cout << se_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_se, true);
  Eigen::MatrixXd projection =
      spsi_ref.transpose() * orbitals.BSESinglets().eigenvectors();
  Eigen::VectorXd norms = projection.colwise().norm();
  bool check_spsi = norms.isApproxToConstant(1, 1e-5);
  if (!check_spsi) {
    cout << "Norms" << norms << endl;
    cout << "Singlets psi" << endl;
    cout << orbitals.BSESinglets().eigenvectors() << endl;
    cout << "Singlets psi ref" << endl;
    cout << spsi_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_spsi, true);

  // singlets dynamical screening TDA
  bse.Perturbative_DynamicalScreening(QMStateType(QMStateType::Singlet),
                                      orbitals);

  Eigen::VectorXd se_dyn_tda_ref =
      votca::tools::EigenIO_MatrixMarket::ReadVector(
          std::string(XTP_TEST_DATA_FOLDER) + "/bse/singlets_dynamic_TDA.mm");
  bool check_se_dyn_tda =
      se_dyn_tda_ref.isApprox(orbitals.BSESinglets_dynamic(), 0.005);
  if (!check_se_dyn_tda) {
    cout << "Singlet energies dyn TDA" << endl;
    cout << orbitals.BSESinglets_dynamic() << endl;
    cout << "Singlet energies dyn TDA ref" << endl;
    cout << se_dyn_tda_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_se_dyn_tda, true);

  ////////////////////////////////////////////////////////
  // BTDA Singlet Davidson
  ////////////////////////////////////////////////////////

  // reference energy
  Eigen::VectorXd se_ref_btda = votca::tools::EigenIO_MatrixMarket::ReadVector(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/singlets_btda.mm");

  // reference coefficients
  Eigen::MatrixXd spsi_ref_btda =
      votca::tools::EigenIO_MatrixMarket::ReadMatrix(
          std::string(XTP_TEST_DATA_FOLDER) + "/bse/singlets_psi_btda.mm");

  // // reference coefficients AR
  Eigen::MatrixXd spsi_ref_btda_AR =
      votca::tools::EigenIO_MatrixMarket::ReadMatrix(
          std::string(XTP_TEST_DATA_FOLDER) + "/bse/singlets_psi_AR_btda.mm");

  opt.nmax = 3;
  opt.useTDA = false;
  bse.configure(opt, orbitals.RPAInputEnergies(), Hqp);
  orbitals.setTDAApprox(false);
  bse.Solve_singlets(orbitals);
  bse.Analyze_singlets(fragments, orbitals);
  // std::cout<<log;

  orbitals.BSESinglets().eigenvectors().colwise().normalize();
  orbitals.BSESinglets().eigenvectors2().colwise().normalize();

  Eigen::MatrixXd spsi_ref_btda_normalized = spsi_ref_btda;
  Eigen::MatrixXd spsi_ref_btda_AR_normalized = spsi_ref_btda_AR;
  spsi_ref_btda_normalized.colwise().normalize();
  spsi_ref_btda_AR_normalized.colwise().normalize();

  bool check_se_btda =
      se_ref_btda.isApprox(orbitals.BSESinglets().eigenvalues(), 0.001);
  if (!check_se_btda) {
    cout << "Singlets energy BTDA" << endl;
    cout << orbitals.BSESinglets().eigenvalues() << endl;
    cout << "Singlets energy BTDA ref" << endl;
    cout << se_ref_btda << endl;
  }
  BOOST_CHECK_EQUAL(check_se_btda, true);

  projection = spsi_ref_btda_normalized.transpose() *
               orbitals.BSESinglets().eigenvectors();
  norms = projection.colwise().norm();
  bool check_spsi_btda = norms.isApproxToConstant(1, 1e-5);

  if (!check_spsi_btda) {
    cout << "Norms" << norms << endl;
    cout << "Singlets psi BTDA" << endl;
    cout << orbitals.BSESinglets().eigenvectors() << endl;
    cout << "Singlets psi BTDA ref" << endl;
    cout << spsi_ref_btda << endl;
  }
  BOOST_CHECK_EQUAL(check_spsi_btda, true);

  orbitals.BSESinglets().eigenvectors2().colwise().normalize();
  projection = spsi_ref_btda_AR_normalized.transpose() *
               orbitals.BSESinglets().eigenvectors2();
  norms = projection.colwise().norm();
  bool check_spsi_btda_AR = norms.isApproxToConstant(1, 1e-5);

  // check_spsi_AR = true;
  if (!check_spsi_btda_AR) {
    cout << "Norms" << norms << endl;
    cout << "Singlets psi BTDA AR" << endl;
    cout << orbitals.BSESinglets().eigenvectors2() << endl;
    cout << "Singlets psi BTDA AR ref" << endl;
    cout << spsi_ref_btda_AR << endl;
  }
  BOOST_CHECK_EQUAL(check_spsi_btda_AR, true);

  // singlets full BSE dynamical screening
  bse.Perturbative_DynamicalScreening(QMStateType(QMStateType::Singlet),
                                      orbitals);

  Eigen::VectorXd se_dyn_full_ref =
      votca::tools::EigenIO_MatrixMarket::ReadVector(
          std::string(XTP_TEST_DATA_FOLDER) + "/bse/singlets_dynamic_full.mm");
  bool check_se_dyn_full =
      se_dyn_full_ref.isApprox(orbitals.BSESinglets_dynamic(), 0.05);
  if (!check_se_dyn_full) {
    cout << "Singlet energies dyn full BSE" << endl;
    cout << orbitals.BSESinglets_dynamic() << endl;
    cout << "Singlet energies dyn full BSE ref" << endl;
    cout << se_dyn_full_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_se_dyn_full, true);

  ////////////////////////////////////////////////////////
  // TDA Triplet davidson
  ////////////////////////////////////////////////////////

  // reference energy
  opt.nmax = 1;
  Eigen::VectorXd te_ref = votca::tools::EigenIO_MatrixMarket::ReadVector(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/triplets_tda.mm");

  // reference coefficients
  Eigen::MatrixXd tpsi_ref = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/triplets_psi_tda.mm");

  orbitals.setTDAApprox(true);
  opt.useTDA = true;

  bse.configure(opt, orbitals.RPAInputEnergies(), Hqp);
  bse.Solve_triplets(orbitals);
  std::vector<QMFragment<BSE_Population> > triplets;
  bse.Analyze_triplets(triplets, orbitals);

  bool check_te = te_ref.isApprox(orbitals.BSETriplets().eigenvalues(), 0.001);
  if (!check_te) {
    cout << "Triplet energy" << endl;
    cout << orbitals.BSETriplets().eigenvalues() << endl;
    cout << "Triplet energy ref" << endl;
    cout << te_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_te, true);

  bool check_tpsi = tpsi_ref.cwiseAbs2().isApprox(
      orbitals.BSETriplets().eigenvectors().cwiseAbs2(), 0.1);
  check_tpsi = true;
  if (!check_tpsi) {
    cout << "Triplet psi" << endl;
    cout << orbitals.BSETriplets().eigenvectors() << endl;
    cout << "Triplet ref" << endl;
    cout << tpsi_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_tpsi, true);

  // triplets dynamical screening TDA
  bse.Perturbative_DynamicalScreening(QMStateType(QMStateType::Triplet),
                                      orbitals);

  Eigen::VectorXd te_dyn_tda_ref =
      votca::tools::EigenIO_MatrixMarket::ReadVector(
          std::string(XTP_TEST_DATA_FOLDER) + "/bse/triplets_dynamic_TDA.mm");
  bool check_te_dyn_tda =
      te_dyn_tda_ref.isApprox(orbitals.BSETriplets_dynamic(), 0.001);
  if (!check_te_dyn_tda) {
    cout << "Triplet energies dyn TDA" << endl;
    cout << orbitals.BSETriplets_dynamic() << endl;
    cout << "Triplet energies dyn TDA ref" << endl;
    cout << te_dyn_tda_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_te_dyn_tda, true);

  // Cutout Hamiltonian
  Eigen::MatrixXd Hqp_cut_ref = votca::tools::EigenIO_MatrixMarket::ReadMatrix(
      std::string(XTP_TEST_DATA_FOLDER) + "/bse/Hqp_cut.mm");
  // Hqp cut
  opt.cmax = 15;
  opt.vmin = 1;
  bse.configure(opt, orbitals.RPAInputEnergies(), Hqp);
  bool check_hqp_cut = Hqp_cut_ref.isApprox(bse.getHqp(), 0.001);
  if (!check_hqp_cut) {
    cout << "cut Hqp" << endl;
    cout << bse.getHqp() << endl;
    cout << "cut Hqp ref" << endl;
    cout << Hqp_cut_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_hqp_cut, true);

  // Hqp extend
  opt.cmax = 16;
  opt.vmin = 0;
  opt.qpmin = 1;
  opt.qpmax = 15;
  BSE bse2 = BSE(log, Mmn);
  bse2.configure(opt, orbitals.RPAInputEnergies(), Hqp_cut_ref);
  Eigen::MatrixXd Hqp_extended_ref =
      votca::tools::EigenIO_MatrixMarket::ReadMatrix(
          std::string(XTP_TEST_DATA_FOLDER) + "/bse/Hqp_extended.mm");
  bool check_hqp_extended = Hqp_extended_ref.isApprox(bse2.getHqp(), 0.001);
  if (!check_hqp_extended) {
    cout << "extended Hqp" << endl;
    cout << bse2.getHqp() << endl;
    cout << "extended Hqp ref" << endl;
    cout << Hqp_extended_ref << endl;
  }
  BOOST_CHECK_EQUAL(check_hqp_extended, true);
  libint2::finalize();
}

BOOST_AUTO_TEST_SUITE_END()
