/*
 *            Copyright 2009-2023 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.
 *
 */

#include <filesystem>

// Third party includes
#include <boost/format.hpp>
#include <boost/math/constants/constants.hpp>

// Local VOTCA includes
#include "votca/xtp/esp2multipole.h"
#include "votca/xtp/segmentmapper.h"

// Local private VOTCA includes
#include "eqm.h"

namespace votca {
namespace xtp {

void EQM::ParseSpecificOptions(const tools::Property& options) {

  QMPackageFactory{};

  // job tasks
  std::string tasks_string = options.get(".tasks").as<std::string>();

  // We split either on a space or a comma
  tools::Tokenizer tokenizedTasks(tasks_string, " ,");
  std::vector<std::string> tasks = tokenizedTasks.ToVector();

  do_dft_input_ = std::find(tasks.begin(), tasks.end(), "input") != tasks.end();
  do_dft_run_ = std::find(tasks.begin(), tasks.end(), "dft") != tasks.end();
  do_dft_parse_ = std::find(tasks.begin(), tasks.end(), "parse") != tasks.end();
  do_gwbse_ = std::find(tasks.begin(), tasks.end(), "gwbse") != tasks.end();
  do_esp_ = std::find(tasks.begin(), tasks.end(), "esp") != tasks.end();

  // set the basis sets and functional for DFT and GWBSE
  gwbse_options_ = options.get("gwbse");
  package_options_ = options.get("dftpackage");
  esp_options_ = options.get(".esp_options");
}

void EQM::WriteJobFile(const Topology& top) {

  std::cout << "\n... ... Writing job file: " << jobfile_ << std::flush;
  std::ofstream ofs;
  ofs.open(jobfile_, std::ofstream::out);
  if (!ofs.is_open()) {
    throw std::runtime_error("\nERROR: bad file handle: " + jobfile_);
  }
  ofs << "<jobs>" << std::endl;
  Index jobCount = 0;

  const std::vector<Segment>& segments = top.Segments();
  for (const Segment& segment : segments) {
    Index id = segment.getId();
    std::string tag = "";
    tools::Property Input;
    tools::Property& pInput = Input.add("input", "");
    tools::Property& pSegment =
        pInput.add("segment", (boost::format("%1$s") % segment.getId()).str());
    pSegment.setAttribute<std::string>("type", segment.getType());
    pSegment.setAttribute<Index>("id", segment.getId());
    Job job(id, tag, Input, Job::AVAILABLE);
    job.ToStream(ofs);
    jobCount++;
  }

  ofs << "</jobs>" << std::endl;
  ofs.close();

  std::cout << " with " << jobCount << " jobs" << std::flush;
}

void EQM::SetJobToFailed(Job::JobResult& jres, Logger& pLog,
                         const std::string& errormessage) {
  XTP_LOG(Log::error, pLog) << errormessage << std::flush;
  std::cout << pLog;
  jres.setError(errormessage);
  jres.setStatus(Job::FAILED);
}

void EQM::WriteLoggerToFile(const std::string& logfile, Logger& logger) {
  std::ofstream ofs;
  ofs.open(logfile, std::ofstream::out);
  if (!ofs.is_open()) {
    throw std::runtime_error("Bad file handle: " + logfile);
  }
  ofs << logger << std::endl;
  ofs.close();
}

Job::JobResult EQM::EvalJob(const Topology& top, Job& job, QMThread& opThread) {
  Orbitals orbitals;
  Job::JobResult jres = Job::JobResult();
  tools::Property job_input_ = job.getInput();
  std::vector<tools::Property*> lSegments = job_input_.Select("segment");
  Index segId = lSegments.front()->getAttribute<Index>("id");
  std::string segType = lSegments.front()->getAttribute<std::string>("type");
  std::string qmgeo_state = "n";
  if (lSegments.front()->exists("qm_geometry")) {
    qmgeo_state = lSegments.front()->getAttribute<std::string>("qm_geometry");
  }

  QMState state(qmgeo_state);
  const Segment& seg = top.getSegment(segId);

  Logger& pLog = opThread.getLogger();
  QMMapper mapper(pLog);
  mapper.LoadMappingFile(mapfile_);
  orbitals.QMAtoms() = mapper.map(seg, state);
  XTP_LOG(Log::error, pLog)
      << TimeStamp() << " Evaluating site " << seg.getId() << std::flush;

  // directories and files
  std::filesystem::path arg_path;
  std::string eqm_work_dir = "OR_FILES";
  std::string frame_dir =
      "frame_" + boost::lexical_cast<std::string>(top.getStep());
  std::string orb_file =
      (boost::format("%1%_%2%%3%") % "molecule" % segId % ".orb").str();
  std::string mol_dir = (boost::format("%1%%2%%3%") % "molecule" % "_" % segId).str();
  std::string package_append = "workdir_" + Identify();
  std::string work_dir =
      (arg_path / eqm_work_dir / package_append / frame_dir / mol_dir)
          .generic_string();

  tools::Property job_summary;
  tools::Property& output_summary = job_summary.add("output", "");
  tools::Property& segment_summary = output_summary.add("segment", "");
  std::string segName = seg.getType();
  segId = seg.getId();
  segment_summary.setAttribute("id", segId);
  segment_summary.setAttribute("type", segName);
  if (do_dft_input_ || do_dft_run_ || do_dft_parse_) {
    XTP_LOG(Log::error, pLog) << "Running DFT" << std::flush;
    Logger dft_logger(votca::Log::current_level);
    dft_logger.setMultithreading(false);
    dft_logger.setPreface(Log::info, (boost::format("\nDFT INF ...")).str());
    dft_logger.setPreface(Log::error, (boost::format("\nDFT ERR ...")).str());
    dft_logger.setPreface(Log::warning, (boost::format("\nDFT WAR ...")).str());
    dft_logger.setPreface(Log::debug, (boost::format("\nDFT DBG ...")).str());
    std::string package = package_options_.get(".name").as<std::string>();
    std::unique_ptr<QMPackage> qmpackage = QMPackageFactory().Create(package);
    qmpackage->setLog(&dft_logger);
    qmpackage->setRunDir(work_dir);
    qmpackage->Initialize(package_options_);

    // create input for DFT
    if (do_dft_input_) {
      std::filesystem::create_directories(work_dir);
      qmpackage->WriteInputFile(orbitals);
    }

    if (do_dft_run_) {
      bool run_dft_status = qmpackage->Run();
      if (!run_dft_status) {
        std::string output = "DFT run failed";
        SetJobToFailed(jres, pLog, output);
        return jres;
      }
    }

    // parse the log/orbitals files
    if (do_dft_parse_) {
      bool parse_log_status = qmpackage->ParseLogFile(orbitals);
      if (!parse_log_status) {
        std::string output = "log incomplete; ";
        SetJobToFailed(jres, pLog, output);
        return jres;
      }
      bool parse_orbitals_status = qmpackage->ParseMOsFile(orbitals);
      if (!parse_orbitals_status) {
        std::string output = "orbfile failed; ";
        SetJobToFailed(jres, pLog, output);
        return jres;
      }
      // additionally copy *.gbw files for orca (-> iqm guess)
      if (qmpackage->getPackageName() == "orca") {
        std::string DIR = eqm_work_dir + "/molecules/" + frame_dir;
        std::filesystem::create_directories(DIR);
        std::string gbw_file =
            (boost::format("%1%_%2%%3%") % "molecule" % segId % ".gbw").str();
        std::string GBWFILE = DIR + "/" + gbw_file;
        XTP_LOG(Log::error, pLog)
            << "Copying MO data to " << gbw_file << std::flush;
        std::string GBWFILE_workdir =
            work_dir + "/" +
            package_options_.get("temporary_file").as<std::string>() + ".gbw";
        std::filesystem::copy_file(
            GBWFILE_workdir, GBWFILE,
            std::filesystem::copy_options::overwrite_existing);
      }

    }  // end of the parse orbitals/log
    qmpackage->CleanUp();
    WriteLoggerToFile(work_dir + "/dft.log", dft_logger);
  }

  if (!do_dft_parse_ && (do_gwbse_ || do_esp_)) {
    // load the DFT data from serialized orbitals object
    std::string ORB_FILE =
        eqm_work_dir + "/molecules/" + frame_dir + "/" + orb_file;
    XTP_LOG(Log::error, pLog)
        << TimeStamp() << " Loading DFT data from " << ORB_FILE << std::flush;
    orbitals.ReadFromCpt(ORB_FILE);
  }

  if (do_gwbse_) {
    XTP_LOG(Log::error, pLog) << "Running GWBSE" << std::flush;
    try {
      GWBSE gwbse = GWBSE(orbitals);
      Logger gwbse_logger(votca::Log::current_level);
      gwbse_logger.setMultithreading(false);
      gwbse_logger.setPreface(Log::info, (boost::format("\nGWBSE INF ...")).str());
      gwbse_logger.setPreface(Log::error, (boost::format("\nGWBSE ERR ...")).str());
      gwbse_logger.setPreface(Log::warning, (boost::format("\nGWBSE WAR ...")).str());
      gwbse_logger.setPreface(Log::debug, (boost::format("\nGWBSE DBG ...")).str());
      gwbse.setLogger(&gwbse_logger);
      gwbse.Initialize(gwbse_options_);
      gwbse.Evaluate();
      gwbse.addoutput(segment_summary);
      WriteLoggerToFile(work_dir + "/gwbse.log", gwbse_logger);
    } catch (std::runtime_error& error) {
      std::string errormessage(error.what());
      SetJobToFailed(jres, pLog, "GWBSE:" + errormessage);
      return jres;
    }
  }

  if (do_esp_) {
    XTP_LOG(Log::error, pLog) << "Running ESPFIT" << std::flush;
    try {
      Esp2multipole esp2multipole = Esp2multipole(pLog);
      esp2multipole.Initialize(esp_options_);
      std::string ESPDIR =
          "MP_FILES/" + frame_dir + "/" + esp2multipole.GetStateString();
      StaticSegment seg2 = esp2multipole.Extractingcharges(orbitals);
      std::string mps_file = (boost::format("%1%_%2%_%3%.mps") % segType % segId %
                              esp2multipole.GetStateString())
                                 .str();
      std::filesystem::create_directories(ESPDIR);
      seg2.WriteMPS(ESPDIR + "/" + mps_file,
                    "Generated by eqm:" + esp2multipole.GetStateString());
      XTP_LOG(Log::error, pLog)
          << "Written charges to " << (ESPDIR + "/" + mps_file) << std::flush;
      segment_summary.add("partialcharges", (ESPDIR + "/" + mps_file));
    } catch (std::runtime_error& error) {
      std::string errormessage(error.what());
      SetJobToFailed(jres, pLog, "ESPFIT:" + errormessage);
      return jres;
    }
  }
  XTP_LOG(Log::error, pLog) << TimeStamp() << " Finished evaluating site "
                            << seg.getId() << std::flush;

  if (do_dft_parse_ || do_gwbse_) {
    XTP_LOG(Log::error, pLog) << "Saving data to " << orb_file << std::flush;
    std::string DIR = eqm_work_dir + "/molecules/" + frame_dir;
    std::filesystem::create_directories(DIR);
    std::string ORBFILE = DIR + "/" + orb_file;
    orbitals.WriteToCpt(ORBFILE);
  }

  // output of the JOB
  jres.setOutput(job_summary);
  jres.setStatus(Job::COMPLETE);

  return jres;
}
}  // namespace xtp
}  // namespace votca
