/******************************************************************************
  This source file is part of the Avogadro project.
  This source code is released under the 3-Clause BSD License, (see "LICENSE").
******************************************************************************/

#include "nwchemjson.h"

#include <avogadro/core/crystaltools.h>
#include <avogadro/core/elements.h>
#include <avogadro/core/gaussianset.h>
#include <avogadro/core/molecule.h>
#include <avogadro/core/unitcell.h>
#include <avogadro/core/utilities.h>

#include <nlohmann/json.hpp>

#include <iostream>

namespace Avogadro::QuantumIO {

using std::string;
using std::vector;

using nlohmann::json;

using Core::Array;
using Core::Atom;
using Core::GaussianSet;
using Core::Molecule;
using Core::split;

NWChemJson::NWChemJson() {}

NWChemJson::~NWChemJson() {}

bool NWChemJson::read(std::istream& file, Molecule& molecule)
{
  json root;
  try {
    file >> root;
  } catch (json::parse_error& e) {
    appendError("Error parsing JSON: " + string(e.what()));
    return false;
  }

  if (!root.is_object()) {
    appendError("Error: Input is not a JSON object.");
    return false;
  }

  if (root.find("simulation") == root.end()) {
    appendError("Error: no \"simulation\" key found.");
    return false;
  }

  json simulation = root["simulation"];

  if (simulation.find("calculations") == simulation.end() ||
      !simulation["calculations"].is_array()) {
    appendError("Error: no \"calculations\" array found.");
    return false;
  }

  // Scan the calculations array for calculationSetup.molecule objects.
  json calculations = simulation["calculations"];

  // Iterate through the objects in the array, and print out any molecules.
  json moleculeArray = json::array();
  json basisSetArray = json::array();
  json calculationVib;
  json molecularOrbitals;
  int numberOfElectrons = 0;
  string theory;
  string xcFunctional;
  for (const auto& calcObj : calculations) {
    if (calcObj.is_object()) {
      string calcType = calcObj.value("calculationType", "");
      // Store the last vibrational frequencies calculation object.
      if (calcType == "vibrationalModes")
        calculationVib = calcObj;
      json calcSetup = calcObj.value("calculationSetup", json());
      json calcMol = calcSetup.value("molecule", json());
      numberOfElectrons = calcSetup.value("numberOfElectrons", -1);
      if (calcSetup.count("exchangeCorrelationFunctional") &&
          calcSetup["exchangeCorrelationFunctional"].is_array() &&
          !calcSetup["exchangeCorrelationFunctional"].empty()) {
        json functional = calcSetup["exchangeCorrelationFunctional"][0];
        if (functional.is_object()) {
          xcFunctional = functional.value("xcName", "");
        }
        if (xcFunctional == "B3LYP Method XC Potential") {
          xcFunctional = "b3lyp";
        }
      }
      if (calcSetup.count("waveFunctionTheory")) {
        theory = calcSetup["waveFunctionTheory"].get<std::string>();
        if (theory == "Density Functional Theory") {
          theory = "dft";
        }
      }
      if (!calcMol.is_null() && calcMol.is_object())
        moleculeArray.push_back(calcMol);
      json basisSet = calcSetup.value("basisSet", json());
      if (!basisSet.is_null() && basisSet.is_object())
        basisSetArray.push_back(basisSet);

      json calcResults = calcObj.value("calculationResults", json());
      calcMol = calcResults.value("molecule", json());
      if (!calcMol.is_null() && calcMol.is_object())
        moleculeArray.push_back(calcMol);
      // There is currently one id for all, just get the last one we find.
      if (calcResults.count("molecularOrbitals") &&
          calcResults["molecularOrbitals"].is_object()) {
        molecularOrbitals = calcResults["molecularOrbitals"];
      }
    }
  }

  // For now, we are just grabbing the "last" molecule, and using that. This
  // needs more complex logic to step through the file and do it properly.
  json atoms;
  if (!moleculeArray.empty()) {
    json finalMol = moleculeArray.back();
    atoms = finalMol.value("atoms", json());
  }
  if (atoms.is_array()) {
    for (auto jsonAtom : atoms) {
      if (jsonAtom.is_null() || !jsonAtom.is_object())
        continue;
      Atom a = molecule.addAtom(
        static_cast<unsigned char>(jsonAtom.value("elementNumber", 0)));
      json pos = jsonAtom["cartesianCoordinates"]["value"];
      Vector3 position;
      if (pos.is_array() && pos.size() >= 3)
        position = Vector3(pos[0], pos[1], pos[2]);
      string units = jsonAtom["cartesianCoordinates"]["units"];
      if (units == "bohr")
        position *= BOHR_TO_ANGSTROM_D;
      a.setPosition3d(position);
    }
  }
  // Perceive bonds for the molecule.
  molecule.perceiveBondsSimple();
  molecule.perceiveBondOrders();

  // Add in the electronic structure information if available.
  if (molecularOrbitals.is_object() &&
      molecularOrbitals["atomicOrbitalDescriptions"].is_array()) {
    json basisSet;
    if (!basisSetArray.empty())
      basisSet = basisSetArray.back();
    json orbDesc = molecularOrbitals.value("atomicOrbitalDescriptions", json());

    // Figure out the mapping of basis set to molecular orbitals.
    Array<int> atomNumber;
    Array<string> atomSymbol;
    for (auto& i : orbDesc) {
      string desc = i;
      vector<string> parts = split(desc, ' ');
      assert(parts.size() == 3);
      int num = Core::lexicalCast<int>(parts[0]).value_or(0);
      if (atomNumber.size() > 0 && atomNumber.back() == num)
        continue;
      atomNumber.push_back(num);
      atomSymbol.push_back(parts[1]);
    }

    // Now create the structure, and expand out the orbitals.
    auto* basis = new GaussianSet;
    basis->setMolecule(&molecule);
    string basisSetName;
    for (size_t i = 0; i < atomSymbol.size(); ++i) {
      string symbol = atomSymbol[i];
      json basisFunctions = basisSet.value("basisFunctions", json());
      json currentFunction;
      for (auto& basisFunction : basisFunctions) {
        currentFunction = basisFunction;

        string elementType;
        if (currentFunction.count("elementLabel"))
          elementType = currentFunction["elementLabel"].get<std::string>();
        else if (currentFunction.count("elementType"))
          elementType = currentFunction["elementType"].get<std::string>();

        if (elementType == symbol)
          break;

        currentFunction = json();
      }

      if (currentFunction.is_null())
        break;

      if (currentFunction.count("basisSetName")) {
        if (basisSetName.empty()) {
          basisSetName = currentFunction["basisSetName"].get<std::string>();
        } else if (basisSetName != currentFunction["basisSetName"]) {
          basisSetName = "Custom";
        }
      }

      json contraction = currentFunction.value("basisSetContraction", json());
      bool spherical =
        currentFunction.value("basisSetHarmonicType", "") == "spherical";
      for (auto contractionShell : contraction) {
        string shellType;
        if (contractionShell.count("basisSetShell"))
          shellType = contractionShell["basisSetShell"].get<std::string>();
        else if (contractionShell.count("basisSetShellType"))
          shellType = contractionShell["basisSetShellType"].get<std::string>();
        json exponent = contractionShell.value("basisSetExponent", json());
        json coefficient =
          contractionShell.value("basisSetCoefficient", json());
        assert(exponent.size() == coefficient.size());
        GaussianSet::orbital type = GaussianSet::UU;
        if (shellType == "s")
          type = GaussianSet::S;
        else if (shellType == "p")
          type = GaussianSet::P;
        else if (shellType == "d" && spherical)
          type = GaussianSet::D5;
        else if (shellType == "d")
          type = GaussianSet::D;
        else if (shellType == "f" && spherical)
          type = GaussianSet::F7;
        else if (shellType == "f")
          type = GaussianSet::F;

        if (type != GaussianSet::UU) {
          int b = basis->addBasis(i, type);
          for (size_t k = 0; k < exponent.size() && k < coefficient.size();
               ++k) {
            basis->addGto(b, coefficient[k], exponent[k]);
          }
        }
      }
    }
    // Now to add the molecular orbital coefficients.
    json moCoeffs = molecularOrbitals.value("molecularOrbital", json());
    vector<double> coeffArray;
    vector<double> energyArray;
    vector<unsigned char> occArray;
    vector<unsigned int> numArray;
    for (auto currentMO : moCoeffs) {
      json coeff = currentMO.value("moCoefficients", json());
      for (auto& j : coeff)
        coeffArray.push_back(j);
      if (currentMO.count("orbitalEnergy")) {
        energyArray.push_back(currentMO["orbitalEnergy"].value("value", 0.0));
      }
      if (currentMO.count("orbitalOccupancy")) {
        occArray.push_back(
          static_cast<unsigned char>(currentMO["orbitalOccupancy"]));
      }
      if (currentMO.count("orbitalNumber")) {
        numArray.push_back(
          static_cast<unsigned int>(currentMO["orbitalNumber"]));
      }
    }
    basis->setMolecularOrbitals(coeffArray);
    basis->setMolecularOrbitalEnergy(energyArray);
    basis->setMolecularOrbitalOccupancy(occArray);
    basis->setMolecularOrbitalNumber(numArray);
    basis->setElectronCount(numberOfElectrons);
    basis->setFunctionalName(xcFunctional);
    basis->setName(basisSetName);
    basis->setTheoryName(theory);
    molecule.setBasisSet(basis);
  }

  // Now to see if there was a vibrational frequencies calculation.
  if (!calculationVib.is_null() && calculationVib.is_object()) {
    json normalModes = calculationVib.value("calculationResults", json())
                         .value("normalModes", json());
    if (!normalModes.is_null() && normalModes.is_array()) {
      Array<double> frequencies;
      Array<double> intensities;
      Array<Array<Vector3>> Lx;
      for (const auto& mode : normalModes) {
        frequencies.push_back(
          mode.value("normalModeFrequency", json()).value("value", 0.0));
        intensities.push_back(mode.value("normalModeInfraRedIntensity", json())
                                .value("value", 0.0));
        json lx = mode.value("normalModeVector", json()).value("value", json());
        if (!lx.empty() && lx.is_array()) {
          Array<Vector3> modeLx;
          modeLx.resize(lx.size() / 3);
          for (size_t k = 0; k < lx.size(); ++k)
            modeLx[k / 3][k % 3] = lx[k];
          Lx.push_back(modeLx);
        }
      }
      molecule.setVibrationFrequencies(frequencies);
      molecule.setVibrationIRIntensities(intensities);
      molecule.setVibrationLx(Lx);
    }
  }

  return true;
}

bool NWChemJson::write(std::ostream&, const Molecule&)
{
  return false;
}

vector<std::string> NWChemJson::fileExtensions() const
{
  vector<std::string> ext;
  ext.emplace_back("nwjson");
  return ext;
}

vector<std::string> NWChemJson::mimeTypes() const
{
  vector<std::string> mime;
  mime.emplace_back("chemical/x-nwjson");
  return mime;
}

} // namespace Avogadro::QuantumIO
