/* Copyright (c) 2015  Gerald Knizia
 * 
 * This file is part of the IboView program (see: http://www.iboview.org)
 * 
 * IboView is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 3.
 * 
 * IboView is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with bfint (LICENSE). If not, see http://www.gnu.org/licenses/
 * 
 * Please see IboView documentation in README.txt for:
 * -- A list of included external software and their licenses. The included
 *    external software's copyright is not touched by this agreement.
 * -- Notes on re-distribution and contributions to/further development of
 *    the IboView software
 */

#include "Iv.h"
#include <boost/format.hpp>
#include <iostream>
#include <math.h> // for sqrt.

#include "CtIo.h"
#include "CtMatrix.h"
#include "CxColor.h"

#include "IvDocument.h"
#include "IvScript.h"
#include "IvIrc.h"
// #include "IvOrbitalFile.h"

#include "CtAtomData.h"
#include "CtBasisLibrary.h"
#include "CxPodArray.h"
using namespace ct;

static const bool g_UseOccupationNumbers = false;

// make the difference in charge for IBOs along all frames.
bool MakeIboChangeCurve(TArray<float> &CurveData, uint iRow, FDocument *document)
{
   // WARNING: due to the way the document loading is done, this function is
   // invoking at least O(N^2) cost (any maybe O(N^3)) when loading lots of frames,
   // since it is called after every frame.
   // Should be fixed... might get expensive?
   FOrbital
      *pRefOrb = dynamic_cast<FOrbital*>(document->GetRowCol(iRow, 0, false));
   if (!pRefOrb)
      // data set is not there or is not an orbital.
      return false;
   uint
      nCols = document->GetNumFrames();
   CurveData.clear();
   CurveData.resize(nCols, 0.);
   TArray<double>
      RefChg = pRefOrb->MakeIaoCharges(g_UseOccupationNumbers);
//    PrintChargeArray(std::cout, "ref charge", RefChg);
   if (RefChg.empty())
      return false; // IAO coeffs not made.
   for (uint iCol = 0; iCol < nCols; ++ iCol) {
      FOrbital
         *pCurOrb = dynamic_cast<FOrbital*>(document->GetRowCol(iRow, iCol, false));
      TArray<double>
         CurChg = pCurOrb->MakeIaoCharges(g_UseOccupationNumbers);
//       PrintChargeArray(std::cout, "cur charge", RefChg);
      double f = 0;
      assert(RefChg.size() == CurChg.size());
      for (uint iAt = 0; iAt < RefChg.size(); ++ iAt)
         f += sqr(RefChg[iAt] - CurChg[iAt]);
      CurveData[iCol] = 1.0 * std::sqrt(f);
   }
   return true;
}

// root mean square deviation between two vectors (i.e., 2-norm of their difference)
static double Rmsd(double const *pA, double const *pB, size_t N) {
   double RmsdSq = 0;
   for (size_t i = 0; i < N; ++ i)
      RmsdSq += sqr(pA[i] - pB[i]);
   return sqrt(RmsdSq);
}

// estimate the IRC arc positions for all loaded frames. Optionally do not do the mass weighting.
void MakeIrcArcLengths(TArray<float> &ArcLengths, FDocument *document, uint Flags)
{
   ArcLengths.clear();

   // get a pointer to the initial geometry.
   IFrame
      *pFrame0 = document->GetFrame(0);
   if (pFrame0 == 0)
      return;
   FAtomSet
      *pAtoms0 = pFrame0->pGetAtoms();
   if (pAtoms0 == 0)
      return;

   // make a list of the atomic masses
   FAtomicMassType
      MassType = ATMASS_StandardAtomicWeight;
   if (Flags & ARCLENGTH_UseIsotopeMasses)
      MassType = ATMASS_MostCommonIsotope;
   uint
      nAt = pAtoms0->size();
   TArray<double>
      Masses(nAt, 1.);
   if ((Flags & ARCLENGTH_NoMassWeighting) == 0) {
      for (uint iAt = 0; iAt < nAt; ++ iAt)
         Masses[iAt] = GetAtomicMass((*pAtoms0)[iAt].AtomicNumber, MassType);
   }

   // assemble a matrix of the mass weighted coordinates for all frames.
   // The numbers and types of atoms in all frames must be equal.
   uint
      nFrames = (unsigned)document->GetNumFrames();
   FMemoryStack2
      Mem(3 * nAt * nFrames * sizeof(double) + 2000000);
   FStackMatrix
      Mwcs(3 * nAt, nFrames, &Mem);
   for (uint iFrame = 0; iFrame < nFrames; ++ iFrame) {
      IFrame
         *pFrame = document->GetFrame(iFrame);
      assert(pFrame != 0);
      FAtomSet
         *pAtoms = pFrame->pGetAtoms();
      if (pAtoms == 0 || pAtoms->size() != nAt)
         IvNotify(NOTIFY_Error, "IrcArcLength: Number of atoms not consistent across different frames.");
      for (uint iAt = 0; iAt < nAt; ++ iAt) {
         if ((*pAtoms)[iAt].AtomicNumber != (*pAtoms0)[iAt].AtomicNumber)
            IvNotify(NOTIFY_Error, "IrcArcLength: Type of atoms not consistent across different frames.");
         for (uint ixyz = 0; ixyz < 3; ++ ixyz)
            Mwcs(3*iAt + ixyz, iFrame) = sqrt(Masses[iAt]) * (*pAtoms)[iAt].vPos[ixyz];
      }
   }

   // approximate the IRC arc lengths via the differences of the 2-norm of the mass-weighted
   // coordinates (note: if we have gradients we could calculate them more accurately).
   ArcLengths.resize(nFrames, 0.);
   double ArcLength = 0.;
   ArcLengths[0] = ArcLength;
   for (uint iFrame = 1; iFrame < nFrames; ++ iFrame) {
      double Delta = Rmsd(&Mwcs(0, iFrame), &Mwcs(0, iFrame-1), Mwcs.nRows);
      ArcLength += Delta;
      ArcLengths[iFrame] = float(ArcLength);
   }

   // that's it already.
}
