/*
 * 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.
 *
 */
#include "xtp_libint2.h"
#define BOOST_TEST_MAIN

#define BOOST_TEST_MODULE test_hdf5

// Standard includes
#include <cassert>

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

// Local VOTCA includes
#include "votca/xtp/checkpoint.h"
#include "votca/xtp/checkpointreader.h"
#include "votca/xtp/checkpointwriter.h"
#include "votca/xtp/orbitals.h"
#include "votca/xtp/qmatom.h"

BOOST_AUTO_TEST_SUITE(test_hdf5)

using namespace votca::xtp;
using namespace votca;
BOOST_AUTO_TEST_CASE(checkpoint_file_test) {
  libint2::initialize();

  Index occupiedLevels = 4;
  Index unoccupiedLevels = 13;
  Index numElectrons = 12;

  Eigen::VectorXd moeTest = Eigen::VectorXd::Random(17);
  Eigen::MatrixXd mocTest = Eigen::MatrixXd::Random(17, 17);

  QMMolecule atoms = QMMolecule(" ", 0);
  atoms.LoadFromFile(std::string(XTP_TEST_DATA_FOLDER) + "/hdf5/molecule.xyz");

  double qmEnergy = -2.1025e-3;

  std::string qmPackage = "NOPE";

  Index rpaMin = '?';
  Index rpaMax = 1e3;

  Index bseVmin = -6019386;
  Index bseCmax = 42;

  double scaHfx = 3.14159;

  bool useTDA = true;

  Eigen::MatrixXd vxcTest = Eigen::MatrixXd::Random(200, 200);
  std::string someECP = "aye aye Cap'n";

  Eigen::VectorXd QPpertEnergiesTest = Eigen::VectorXd::Random(31);
  Eigen::MatrixXd QPdiagEnergiesTest = Eigen::VectorXd::Random(21);
  Eigen::MatrixXd QPdiagCoefficientsTest = Eigen::MatrixXd::Identity(42, 42);

  Eigen::VectorXd BSESingletEnergiesTest = Eigen::VectorXd::Random(25);
  Eigen::MatrixXd BSESingletCoefficientsTest = Eigen::MatrixXd::Random(25, 38);
  Eigen::MatrixXd BSESingletCoefficientsARTest =
      Eigen::MatrixXd::Random(42, 42);

  Eigen::VectorXd BSETripletEnergiesTest = Eigen::VectorXd::Random(33);
  Eigen::MatrixXd BSETripletCoefficientsTest = Eigen::MatrixXd::Random(33, 31);

  Index basisSetSize = 0;

  // Write orbitals
  Orbitals orbWrite;
  orbWrite.setNumberOfOccupiedLevels(occupiedLevels);
  orbWrite.setNumberOfAlphaElectrons(numElectrons);
  orbWrite.MOs().eigenvalues() = moeTest;
  orbWrite.MOs().eigenvectors() = mocTest;

  orbWrite.QMAtoms() = atoms;
  orbWrite.SetupDftBasis(std::string(XTP_TEST_DATA_FOLDER) + "/hdf5/3-21G.xml");
  basisSetSize = orbWrite.getBasisSetSize();

  orbWrite.setQMEnergy(qmEnergy);
  orbWrite.setQMpackage(qmPackage);
  orbWrite.setRPAindices(rpaMin, rpaMax);
  // no need to write qpmin, qpmax
  orbWrite.setBSEindices(bseVmin, bseCmax);
  orbWrite.setScaHFX(scaHfx);
  orbWrite.setTDAApprox(useTDA);
  orbWrite.setECPName(someECP);
  orbWrite.QPpertEnergies() = QPpertEnergiesTest;
  orbWrite.QPdiag().eigenvalues() = QPdiagEnergiesTest;
  orbWrite.QPdiag().eigenvectors() = QPdiagCoefficientsTest;
  orbWrite.BSESinglets().eigenvalues() = BSESingletEnergiesTest;
  orbWrite.BSESinglets().eigenvectors() = BSESingletCoefficientsTest;
  orbWrite.BSESinglets().eigenvectors2() = BSESingletCoefficientsARTest;
  orbWrite.BSETriplets().eigenvalues() = BSETripletEnergiesTest;
  orbWrite.BSETriplets().eigenvectors() = BSETripletCoefficientsTest;

  orbWrite.WriteToCpt("xtp_testing.hdf5");

  // Read Orbitals
  Orbitals orbRead;
  orbRead.ReadFromCpt("xtp_testing.hdf5");

  double tol = 1e-6;

  // Test the read values
  BOOST_CHECK_EQUAL(orbRead.getBasisSetSize(), basisSetSize);
  BOOST_CHECK_EQUAL(orbRead.getBasisSetSize(),
                    occupiedLevels + unoccupiedLevels);
  BOOST_CHECK_EQUAL(orbRead.getNumberOfAlphaElectrons(), numElectrons);
  BOOST_CHECK(orbRead.MOs().eigenvalues().isApprox(moeTest, tol));

  BOOST_CHECK(orbRead.MOs().eigenvectors().isApprox(mocTest, tol));
  BOOST_CHECK_CLOSE(orbRead.getDFTTotalEnergy(), qmEnergy, tol);
  BOOST_CHECK_EQUAL(orbRead.getQMpackage(), qmPackage);
  BOOST_CHECK_EQUAL(orbRead.getRPAmin(), rpaMin);
  BOOST_CHECK_EQUAL(orbRead.getRPAmax(), rpaMax);

  BOOST_CHECK_EQUAL(orbRead.getBSEvmin(), bseVmin);
  BOOST_CHECK_EQUAL(orbRead.getBSEcmax(), bseCmax);

  BOOST_CHECK_CLOSE(orbRead.getScaHFX(), scaHfx, tol);
  BOOST_CHECK_EQUAL(orbRead.getTDAApprox(), useTDA);
  BOOST_CHECK_EQUAL(orbRead.getECPName(), someECP);
  BOOST_CHECK(orbRead.QPpertEnergies().isApprox(QPpertEnergiesTest, tol));
  BOOST_CHECK(orbRead.QPdiag().eigenvalues().isApprox(QPdiagEnergiesTest, tol));
  BOOST_CHECK(orbRead.QPdiag().eigenvectors().isApprox(QPdiagCoefficientsTest));
  BOOST_CHECK(orbRead.BSESinglets().eigenvalues().isApprox(
      BSESingletEnergiesTest, tol));
  BOOST_CHECK(orbRead.BSESinglets().eigenvectors().isApprox(
      BSESingletCoefficientsTest, tol));
  BOOST_CHECK(orbRead.BSESinglets().eigenvectors2().isApprox(
      BSESingletCoefficientsARTest, tol));
  BOOST_CHECK(orbRead.BSETriplets().eigenvalues().isApprox(
      BSETripletEnergiesTest, tol));
  BOOST_CHECK(orbRead.BSETriplets().eigenvectors().isApprox(
      BSETripletCoefficientsTest, tol));

  BOOST_REQUIRE_EQUAL(orbRead.QMAtoms().size(), atoms.size());

  for (Index i = 0; i < atoms.size(); ++i) {
    const auto& atomRead = orbRead.QMAtoms()[i];
    const auto& atomTest = atoms[i];
    BOOST_CHECK_EQUAL(atomRead.getId(), atomTest.getId());
    BOOST_CHECK(atomRead.getPos().isApprox(atomTest.getPos(), tol));
    BOOST_CHECK_EQUAL(atomRead.getNuccharge(), atomTest.getNuccharge());
    BOOST_CHECK_EQUAL(atomRead.getElement(), atomTest.getElement());
  }

  // Check if basissets are equal after a write read cycle
  std::stringstream ssWrite;
  ssWrite << orbWrite.getDftBasis();
  std::stringstream ssRead;
  ssRead << orbRead.getDftBasis();

  bool basissets_equal = (ssWrite.str() == ssRead.str());
  BOOST_CHECK(basissets_equal);

  if (!basissets_equal) {
    std::cout << "Thing we wrote:" << std::endl;
    std::cout << ssWrite.str();
    std::cout << "Thing we read:" << std::endl;
    std::cout << ssRead.str();
  }

  libint2::finalize();
}

BOOST_AUTO_TEST_CASE(open_file_error) {
  BOOST_REQUIRE_THROW(
      CheckpointFile cpf("/bin/mr/root/man.pls", CheckpointAccessLevel::READ),
      std::runtime_error);
}

BOOST_AUTO_TEST_CASE(checkpoint_open_non_existing_loc) {
  CheckpointFile cpf("testin_yo.ab", CheckpointAccessLevel::MODIFY);
  BOOST_REQUIRE_THROW(CheckpointReader r = cpf.getReader("/some/bulshit"),
                      std::runtime_error);
}

BOOST_AUTO_TEST_CASE(read_non_exisiting_matrix) {

  CheckpointFile cpf("xtp_testing.hdf5", CheckpointAccessLevel::READ);
  CheckpointReader r = cpf.getReader("/QMdata");

  Eigen::MatrixXd someMatrix;

  BOOST_REQUIRE_THROW(r(someMatrix, "someMatrix012'5915.jb"),
                      std::runtime_error);
}

BOOST_AUTO_TEST_CASE(read_non_existing_scalar) {
  CheckpointFile cpf("xtp_testing.hdf5", CheckpointAccessLevel::READ);
  CheckpointReader r = cpf.getReader("/QMdata");

  float someThing = 0;
  BOOST_REQUIRE_THROW(r(someThing, "someThing"), std::runtime_error);
}

BOOST_AUTO_TEST_CASE(read_vector_strings) {
  CheckpointFile cpf("xtp_vector_string.hdf5");
  CheckpointWriter w = cpf.getWriter();

  std::vector<std::string> test_vec = {
      "a", "b", "advhsbsavc", "dumuunnadnjdsahjads", "DASASDFAFNADDH blasndd"};
  w(test_vec, "vec");

  std::vector<std::string> test_vec2;
  CheckpointReader r = cpf.getReader();
  r(test_vec2, "vec");

  BOOST_CHECK_EQUAL(test_vec.size(), test_vec2.size());

  for (Index i = 0; i > Index(test_vec.size()); i++) {
    BOOST_CHECK_EQUAL(test_vec[i], test_vec2[i]);
  }
}

BOOST_AUTO_TEST_CASE(staticsegment) {

  CheckpointFile cpf("xtp_staticsegment.hdf5");
  CheckpointWriter w = cpf.getWriter();
  StaticSegment seg = StaticSegment("test", 0);
  for (Index i = 0; i < 10; ++i) {
    seg.push_back(StaticSite(0, "O", Eigen::Vector3d::Random()));
    seg.push_back(StaticSite(25, "O", Eigen::Vector3d::Random()));
    seg.push_back(StaticSite(32, "O", Eigen::Vector3d::Random()));
    seg.push_back(StaticSite(100, "O", Eigen::Vector3d::Random()));
    seg.push_back(StaticSite(2, "Si", Eigen::Vector3d::Random()));
    seg.push_back(StaticSite(3145, "N", Eigen::Vector3d::Random()));
  }

  seg.WriteToCpt(w);
  CheckpointReader r = cpf.getReader();
  StaticSegment seg2 = StaticSegment(r);

  BOOST_REQUIRE_EQUAL(seg2.size(), seg.size());
  for (Index i = 0; i < seg.size(); ++i) {
    const auto& atomRead = seg2[i];
    const auto& atomTest = seg[i];
    BOOST_CHECK_EQUAL(atomRead.getId(), atomTest.getId());
    BOOST_CHECK(atomRead.getPos().isApprox(atomTest.getPos(), 1e-7));
    BOOST_CHECK_EQUAL(atomRead.getCharge(), atomTest.getCharge());
    BOOST_CHECK_EQUAL(atomRead.getElement(), atomTest.getElement());
  }
}

BOOST_AUTO_TEST_SUITE_END()
