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

#include <QAction>
#include <QDialog>
#include <QMessageBox>
#include <QString>

#include <avogadro/core/crystaltools.h>
#include <avogadro/core/unitcell.h>
#include <avogadro/qtgui/molecule.h>
#include <avogadro/qtgui/chartdialog.h>
#include <avogadro/qtgui/chartwidget.h>

#include "pdfoptionsdialog.h"
#include "plotpdf.h"

using Avogadro::Core::CrystalTools;
using Avogadro::Core::UnitCell;
using Avogadro::QtGui::Molecule;

using std::map;

namespace Avogadro::QtPlugins {

using Core::Array;

PlotPdf::PlotPdf(QObject* parent_)
  : Avogadro::QtGui::ExtensionPlugin(parent_), m_actions(QList<QAction*>()),
    m_molecule(nullptr),
    m_pdfOptionsDialog(new PdfOptionsDialog(qobject_cast<QWidget*>(parent()))),
    m_displayDialogAction(new QAction(this))
{
  m_displayDialogAction->setText(tr("Plot Pair Distribution Function…"));
  connect(m_displayDialogAction.data(), &QAction::triggered, this,
          &PlotPdf::displayDialog);
  m_actions.push_back(m_displayDialogAction.data());
  m_displayDialogAction->setProperty("menu priority", 70);

  updateActions();
}

PlotPdf::~PlotPdf() = default;

QList<QAction*> PlotPdf::actions() const
{
  return m_actions;
}

QStringList PlotPdf::menuPath(QAction*) const
{
  return QStringList() << tr("&Crystal");
}

void PlotPdf::setMolecule(QtGui::Molecule* mol)
{
  if (m_molecule == mol)
    return;

  if (m_molecule)
    m_molecule->disconnect(this);

  m_molecule = mol;

  if (m_molecule)
    connect(m_molecule, SIGNAL(changed(uint)), SLOT(moleculeChanged(uint)));

  updateActions();
}

void PlotPdf::moleculeChanged(unsigned int c)
{
  Q_ASSERT(m_molecule == qobject_cast<Molecule*>(sender()));

  auto changes = static_cast<Molecule::MoleculeChanges>(c);

  if (changes & Molecule::UnitCell) {
    if (changes & Molecule::Added || changes & Molecule::Removed)
      updateActions();
  }
}

void PlotPdf::updateActions()
{
  // Disable everything for nullptr molecules.
  if (!m_molecule) {
    foreach (QAction* action, m_actions)
      action->setEnabled(false);
    return;
  }

  // Only display the actions if there is a unit cell
  if (m_molecule->unitCell()) {
    foreach (QAction* action, m_actions)
      action->setEnabled(true);
  } else {
    foreach (QAction* action, m_actions)
      action->setEnabled(false);
  }
}

void PlotPdf::displayDialog()
{
  // Do nothing if the user cancels
  if (m_pdfOptionsDialog->exec() != QDialog::Accepted)
    return;

  // Otherwise, fetch the options and perform the run
  double maxRadius = m_pdfOptionsDialog->maxRadius();
  double step = m_pdfOptionsDialog->step();
  // size_t numpoints = m_pdfOptionsDialog->numDataPoints();
  // double max2theta = m_pdfOptionsDialog->max2Theta();

  PdfData results;
  QString err;
  if (!generatePdfPattern(*m_molecule, results, err, maxRadius, step)) {
    QMessageBox::critical(qobject_cast<QWidget*>(parent()),
                          tr("Failed to generate PDF pattern"),
                          tr("Error message: ") + err);
    return;
  }

  // Now generate a plot with the data
  std::vector<float> xData;
  std::vector<float> yData;
  for (const auto& item : results) {
    xData.push_back(item.first);
    yData.push_back(item.second);
  }

  const char* xTitle = "r (Å)";
  const char* yTitle = "g(r)";
  const char* windowName = "Pair Distribution Function";

  if (!m_chartDialog)
    m_chartDialog.reset(
      new QtGui::ChartDialog(qobject_cast<QWidget*>(this->parent())));

  m_chartDialog->setWindowTitle(windowName);
  auto* chart = m_chartDialog->chartWidget();
  chart->clearPlots();
  chart->addPlot(xData, yData, QtGui::color4ub{ 255, 0, 0, 255 });
  chart->setXAxisTitle(xTitle);
  chart->setYAxisTitle(yTitle);
  m_chartDialog->show();
}

bool PlotPdf::generatePdfPattern(QtGui::Molecule& mol, PdfData& results,
                                 QString& err, double maxRadius, double step)
{
  Array<Vector3> refAtomCoords = mol.atomPositions3d();

  size_t i, j;

  UnitCell* uc = mol.unitCell();
  if (!uc) {
    err = "No unit cell found.";
    return false;
  }

  size_t a = static_cast<size_t>(maxRadius / uc->a()) + 1;
  size_t b = static_cast<size_t>(maxRadius / uc->b()) + 1;
  size_t c = static_cast<size_t>(maxRadius / uc->c()) + 1;

  Vector3 disp = a * uc->aVector() + b * uc->bVector() + c * uc->cVector();
  for (i = 0; i < refAtomCoords.size(); ++i) {
    refAtomCoords[i] += disp;
  }

  Molecule newMolecule = mol;
  CrystalTools::buildSupercell(newMolecule, 2 * a + 1, 2 * b + 1, 2 * c + 1);

  Array<Vector3> newAtomCoords = newMolecule.atomPositions3d();

  map<size_t, size_t> pdfCount;
  double dist, rStep = step;
  size_t k, binIdx;

  for (i = 0; i < refAtomCoords.size(); ++i) {
    for (j = 0; j < newAtomCoords.size(); ++j) {
      dist = (refAtomCoords.at(i) - newAtomCoords.at(j)).norm();
      binIdx = static_cast<size_t>(dist / rStep);
      if (pdfCount.find(binIdx) == pdfCount.end()) {
        pdfCount.insert(std::make_pair(binIdx, 1));
      } else {
        pdfCount[binIdx]++;
      }
    }
  }

  for (k = 0; k < static_cast<size_t>(maxRadius / rStep); k++) {
    if (pdfCount.find(k) == pdfCount.end()) {
      results.push_back(std::make_pair(k * rStep, 0.0));
    } else {
      results.push_back(std::make_pair(
        k * rStep, pdfCount[k] * newMolecule.unitCell()->volume() /
                     (4 * M_PI * pow(k * rStep, 2) * rStep *
                      refAtomCoords.size() * newAtomCoords.size())));
    }
  }

  return true;
}

} // namespace Avogadro::QtPlugins
