/******************************************************************************

  Class for surface canonical MC simulation (Fast) 

        Written by k-fleak 

  Note: Input POSCAR for GS-search must be atoms with spin=spinArray[0] only.
        (For surface, all the atoms TO BE MOVED have spin=spinArray[0].)

 ********************************************************************************/

#pragma warning( disable : 4786 )
#include <iostream>
#include <fstream>
#include <vector>
#include <set>
#include <map>
#include <time.h>
#include <unistd.h>

#include "../mt/mt19937ar.h"
#include "../mt/mt19937ar.c"

#include "../vector3D.hpp"
#include "../parsePoscar.cpp"
#include "../atom.cpp"
#include "../mkSupercell.cpp"
#include "../readInput.cpp"
#include "../cluster.cpp"
#include "../equivAtom.cpp"
#include "../multiComponent.cpp"
#include "../getCorrelation.cpp"

#include "cellBank.cpp"
#include "spinBank.cpp"


typedef std::vector<std::vector<std::vector<std::vector<int> > > > vector4;

class GSSURFFAST{

protected:

  std::vector<Cluster> clusterArray;
  std::vector<Atom> atomArrayMC;
  vector4 coordination;

  int step1MC;
  int step2MC;

  /// surface ///
  std::vector<double> slab1ZPos;
  std::vector<double> slab2ZPos;
  double eps;
  ///////////////

  /// GS Fast ///
  int numCell;
  int numBaseAtom;
  std::vector<std::vector<int> > splitCellArrayAll, splitCellArray; // All: all the atoms, splitCellArray: Atoms to be moved only
  std::vector<int> baseMovAtomArray;
  std::vector<int> NatomArray;
  std::vector<int> spinArray;
  double energyProd;
  std::vector<double> compArray;
  ///////////////
  

  std::ofstream cmcFile;
  void setClusterArray(std::vector<Atom>& atomArrayCorre, std::vector<Atom>& atomArrayCorreLattice,
		       double& supercellSize, double & supercellSizeCorre, std::map<int,double>& eci);
  void setAtomArray(std::vector<Atom>& atomArraySupercell, std::vector<Atom>& atomArrayCorreLattice,
		    const std::vector<int>& nCell);
  void setPosition(int& natom_corre, const std::vector<std::vector<int> >& positionOut,
		   std::map<int,double>& eci);

  virtual void gsSearch();
  virtual double calcSpin(std::vector<std::vector<int> > atomVector, int flipAtom, int flipSpin);
  virtual void GSSURFFAST::calcNewSpinConfig(std::vector<std::pair<int, int> > spinFlipArray, long double &energyIn);

public: 

  GSSURFFAST(std::vector<Atom>& atomArray, std::vector<Atom>& atomArrayLattice, 
	     std::vector<Atom>& atomArraySupercell, const std::vector<int>& nCell, std::map<int, double>& eci,
	     const std::vector<std::vector<int> >& positionOut, std::vector<double> MovAtomZPosIn, 
	     std::vector<int>& splitCellIn, std::vector<double>& composition, double energyProductIn);
  ~GSSURFFAST();

  virtual void start();


};

GSSURFFAST::GSSURFFAST(std::vector<Atom>& atomArray, std::vector<Atom>& atomArrayLattice, 
		       std::vector<Atom>& atomArraySupercell, const std::vector<int>& nCell, std::map<int, double>& eci,
		       const std::vector<std::vector<int> >& positionOut, std::vector<double> MovAtomZPosIn, 
		       std::vector<int>& splitCellIn, std::vector<double>& composition, double energyProductIn){

	/// Elapsed-time estimation ///
	clock_t startTime;
	clock_t currentTime;
	clock_t endTime;
	startTime = clock();
	//////////

	/// surface - Setting z position of atom to be moved in simulation. -///
	eps = 0.001;
	for (int i = 0; i < MovAtomZPosIn.size(); ++i) {
	  slab1ZPos.push_back(MovAtomZPosIn[i] / 2.0);
	  slab2ZPos.push_back(MovAtomZPosIn[i] / 2.0  +  0.5 );
	}
	///////////////

	energyProd = energyProductIn;
	compArray = composition;


        MkSupercell mkcorre(atomArray, 2, 2, 2), mkcorre_lattice(atomArrayLattice, 2, 2, 2);
        std::vector<Atom> atomArrayCorre = mkcorre.getSupercell();
        std::vector<Atom> atomArrayCorreLattice = mkcorre_lattice.getSupercell();
	double sizeSupercell = atomArraySupercell.size() / atomArray.size();
	double sizeSupercellCorre = atomArraySupercell.size() / atomArrayCorre.size();


	std::cout << "setClusterArray ... " ;
	setClusterArray(atomArrayCorre, atomArrayCorreLattice, sizeSupercell, sizeSupercellCorre, eci);
	currentTime = clock();
	std::cout << "done.   Total elapsed time:" << (currentTime - startTime) / CLOCKS_PER_SEC << " sec" << std::endl;

	setAtomArray(atomArraySupercell, atomArrayCorreLattice, nCell);
        int natom_corre = atomArrayCorreLattice.size();

	/// Split cell ///
	std::cout << "set Cell Bank ...";
	CellBank cellBank(atomArrayMC, splitCellIn);
	splitCellArrayAll = cellBank.getCellBank();
	numCell = 1;
	std::vector<std::vector<int> >::iterator its;
	for (int i = 0; i < splitCellIn.size(); ++i){
	  numCell *= splitCellIn[i];
	}
	numBaseAtom = atomArrayMC.size() / numCell;
	if (numBaseAtom != splitCellArrayAll.size()){
	  std::cout << "     Error: cannot divide cell. Please check the name: SPLITCELL." << std::endl;
	  exit(0);
	}
	its = splitCellArrayAll.begin();
	while (its != splitCellArrayAll.end()){
	  if ((*its).size() != numCell){
	    std::cout << "     Error: cannot find equivalent atoms. Please check the name: SPLITCELL." << std::endl;
	    exit(0);	    
	  }
	  ++its;
	}
	std::cout << "done." << std::endl;
	///////////////

	/// Search atoms to be moved in base cell (surf only)
	int baseAtom;
	vector3D<double> pos;
	its = splitCellArrayAll.begin();
	while (its != splitCellArrayAll.end()){
	  baseAtom = (*its)[0];
	  pos = atomArrayMC[baseAtom].getPosition();
	  for (int i = 0; i < slab1ZPos.size(); ++i ) {
	    if ( (pos.z >= slab1ZPos[i] - eps) and (pos.z <= slab1ZPos[i] + eps)) {
	      baseMovAtomArray.push_back(baseAtom);
	    } 
	  }
	  for (int i = 0; i < slab2ZPos.size(); ++i ) {
	    if ( (pos.z >= slab2ZPos[i] - eps) and (pos.z <= slab2ZPos[i] + eps)) {
	      baseMovAtomArray.push_back(baseAtom);
	    } 
	  }
	  ++its;
	}
	std::cout << "Number of atoms to be moved in base-cell: " << baseMovAtomArray.size() << std::endl;
	/////////////////////////////////////////////////////

	/// Delete information in splitCellArrayAll whose base-atom is not moved = splitCellArray (moved atom only) ////
	for ( int i = 0 ; i < splitCellArrayAll.size() ; ++i){
	  for ( int j = 0 ; j < baseMovAtomArray.size() ; ++j){
	    if (splitCellArrayAll[i][0] == baseMovAtomArray[j]){
	      splitCellArray.push_back(splitCellArrayAll[i]);
	    }
	  }
	}
	//////////////////////////////////////////////////////////////////////////

	/// Estimate composition and setting NatomArray /////////////////
	for (int i = 0; i < composition.size(); ++i){
	  NatomArray.push_back(static_cast<int>(baseMovAtomArray.size() * composition[i] + 0.001));
	}
	int nSumTmp = 0;
	for (int i = 0; i < NatomArray.size(); ++i){
	  nSumTmp += NatomArray[i];
	}
	NatomArray.push_back(baseMovAtomArray.size() - nSumTmp);
	std::cout << "Number of atoms to be moved for each elements in base-cell: ";
	for (int i = 0; i < NatomArray.size() ; ++i){
	  std::cout << NatomArray[i] << "  ";
	}
	std::cout << std::endl;
	/////////////////////////////////////////////////////////////////
	

	std::cout << "setPosition ... " ;	
	setPosition(natom_corre, positionOut, eci);
	currentTime = clock();
	std::cout << "done.   Total elapsed time:" << (currentTime - startTime) / CLOCKS_PER_SEC << " sec" << std::endl;

}

GSSURFFAST::~GSSURFFAST(){}

void GSSURFFAST::start(){
  clock_t startTime;
  clock_t endTime;
  startTime = clock();
  std::cout << "gsSearch ... " ;
  gsSearch();
  endTime = clock();
  std::cout << "done.   gsSearch Total elapsed time: " << (endTime-startTime) / CLOCKS_PER_SEC << " sec" << std::endl;
}

void GSSURFFAST::setClusterArray(std::vector<Atom>& atomArrayCorre, std::vector<Atom>& atomArrayCorreLattice,
		double& supercellSize, double & supercellSizeCorre, std::map<int,double>& eci){

        ReadInput input;
        int numSublattice = input.getNumSublattice();
        const int maxCluster = input.getMaxCluster();
        const double tol = input.getTol();
        const double truncatePair = input.getTruncatePair();
        const double truncateTriangle = input.getTruncateTriangle();
        const double truncateTetrahedron = input.getTruncateTetrahedron();
        const double truncateMultiplet = input.getTruncateMultiplet();
        std::vector<std::string> nameArray = input.getAtomName();
        spinArray = input.getSpinArray();
        int clusterOut = 1;

        int natom_index = atomArrayCorreLattice.size();
        EquivAtom equiv(atomArrayCorreLattice, natom_index);
        std::set<std::set<int> > equivAtom = equiv.getEquivAtom();
        std::vector<int> independentAtom = equiv.getIndependentAtom();
        std::vector<std::vector<int> > equivArray = equiv.getEquivArray();
        equivArray.resize(atomArrayCorreLattice.size());
	int unitCellAtom = atomArrayCorre.size()/8;


        MultiComponent multi(maxCluster, spinArray);
        std::vector<std::map<std::vector<int>, std::vector<double> > > correlationArray;
        GetCorrelation correlation(atomArrayCorre, atomArrayCorreLattice,
                independentAtom, equivArray, truncatePair,truncateTriangle,
                truncateTetrahedron, truncateMultiplet, tol, maxCluster,
                unitCellAtom, clusterOut, equivAtom, multi);
        correlationArray = correlation.getCorrelation();


	std::map<int,double>::iterator it = eci.begin();
	std::map<std::vector<int>, std::vector<double> >::iterator it2;

	Cluster tmp( 1 , 1,  0, (*it).second * supercellSize);
	clusterArray.push_back(tmp);
	++it;
	if (spinArray.size()==2){
        	while ( it != eci.end()){
			int index = (*it).first;
			if (index !=0){
				it2 = (correlationArray[index-1]).begin();
				std::vector<double> sigma = (*it2).second;
				Cluster tmp(sigma[0] * supercellSizeCorre, sigma[1] * supercellSizeCorre,\
				index, (*it).second * supercellSize);
				clusterArray.push_back(tmp);
				++it;
			}
		}
	}

}

void GSSURFFAST::setAtomArray(std::vector<Atom>& atomArraySupercell, std::vector<Atom>& atomArrayCorreLattice,
	const std::vector<int>& nCell){

	std::vector<Atom> atomArrayTmp, atomArrayTmp2;
	atomArrayTmp = atomArrayCorreLattice;

	for (int i = 0; i < atomArraySupercell.size(); ++i){
		vector3D<double> positionSupercell = atomArraySupercell[i].getPosition();
		for (int j = 0; j < atomArrayCorreLattice.size(); ++j){
			vector3D<double> position = atomArrayCorreLattice[j].getPosition();
			double x = position.x * 2 / nCell[0];
			double y = position.y * 2 / nCell[1];
			double z = position.z * 2 / nCell[2];
			if (fabs(positionSupercell.x - x) < 1e-6
			and fabs(positionSupercell.y - y) < 1e-6
			and fabs(positionSupercell.z - z) < 1e-6){
				atomArrayTmp[j] = atomArraySupercell[i];
				break;
			}
			if (j == atomArrayCorreLattice.size() - 1)
				atomArrayTmp2.push_back(atomArraySupercell[i]);
		}

	}

	for (int i = 0; i < atomArrayTmp.size(); ++i) atomArrayMC.push_back(atomArrayTmp[i]);
	for (int i = 0; i < atomArrayTmp2.size(); ++i) atomArrayMC.push_back(atomArrayTmp2[i]);

}


void GSSURFFAST::setPosition(int& natom_corre, const std::vector<std::vector<int> >& positionOut,
	std::map<int,double>& eci){ 

        EquivAtom equiv(atomArrayMC, natom_corre);
        std::set<std::set<int> > equivAtom = equiv.getEquivAtom();
        std::vector<int> independentAtom = equiv.getIndependentAtom();
        std::vector<std::vector<int> > equivArray = equiv.getEquivArray();
        equivArray.resize(natom_corre);

	std::vector<std::vector<int> > position;
	std::vector<int> pointIndex;

	int cutIndex = independentAtom.size();
	std::map<int,double>::iterator it = eci.begin();
        while ( it != eci.end()){
		int index = (*it).first;
		if ( index > cutIndex ) position.push_back(positionOut[index - cutIndex - 1]);
		++it;
		if ( index <= cutIndex and index != 0 ) pointIndex.push_back(index);
	}

	std::vector<std::set<std::set<int> > > clusterPosition;

        std::set<std::set<int> >::iterator it0;
        std::set<int>::iterator it00;
        std::set<std::set<int> > tmp0;
        std::set<int> tmp00;
	for (int i = 0; i < pointIndex.size(); ++i){
        	it0 = equivAtom.begin();
		for (int j = 1; j < pointIndex[i]; ++j) ++it0;
	        it00 = (*it0).begin();
	        while( it00 != (*it0).end()){
		        tmp00.insert(*it00);
		        tmp0.insert(tmp00);
		        tmp00.clear();
		        ++it00;
	        }
          clusterPosition.push_back(tmp0);
          tmp0.clear();
        }

        for (int i = 0; i < position.size(); ++i){
                for (int j = 0; j < equivArray[0].size(); ++j){
        		for (int k = 0; k < position[i].size(); ++k){
                                tmp00.insert(equivArray[position[i][k]][j]);
                        }
                        tmp0.insert(tmp00);
                        tmp00.clear();
                }
                clusterPosition.push_back(tmp0);
                tmp0.clear();
        }
	equivArray.clear();

	coordination.resize(atomArrayMC.size());
        for (int i = 0; i < coordination.size(); ++i){
		coordination[i].resize(clusterPosition.size());
	}

        for (int i = 0; i < clusterPosition.size(); ++i){
		it0 = clusterPosition[i].begin();
		while ( it0 != clusterPosition[i].end()){
			std::vector<int> tmpVec;
			it00 = (*it0).begin();
			while ( it00 != (*it0).end()){
				tmpVec.push_back(*it00);
				++it00;
			}

			it00 = (*it0).begin();
			while ( it00 != (*it0).end()){
				coordination[*it00][i].push_back(tmpVec);
				++it00;
			}
			++it0;
			tmpVec.clear();
		}
	}

}



void GSSURFFAST::calcNewSpinConfig(std::vector<std::pair<int, int> > spinFlipArray, long double &energyIn) {
  int spinNew, spinOld;
  int atom1, atomC1, atmp;
  
  for ( int i = 0; i < spinFlipArray.size(); ++i){
    atmp = spinFlipArray[i].first;  // atoms to be moved in base-cell (Not an index in atomArrayMC, but index in splitCellArray.)
    atom1 = splitCellArray[atmp][0]; // atoms to be moved in base-cell (Index in atomArrayMC)
    spinOld = atomArrayMC[atom1].getSpin();
    spinNew = spinFlipArray[i].second;
    std::vector<double> diffSpinArray;
    long double diffEnergy = 0;

    /// BEGIN: Loop for each cluster ////
    for (int j = 0; j < clusterArray.size() - 1; ++j){
      double diffSpinSum = 0.0;
      double spinOldCluster, spinNewCluster;
      
      /// BEGIN: Loop for each atom in each split cell ///
      for (int k = 0; k < numCell; ++k){
	atomC1 = splitCellArray[atmp][k];
	spinOldCluster = calcSpin(coordination[atomC1][j], atomC1, spinOld); 
	spinNewCluster = calcSpin(coordination[atomC1][j], atomC1, spinNew);
	diffSpinSum += spinNewCluster - spinOldCluster;
	atomArrayMC[atomC1].setSpin(spinNew); /// Put spin to new one
	
      }
      /// END: Loop for each atom in each split cell ///
      
      /// BEGIN: Loop for putting spin to original one for next cluster calc. ///
      for (int k = 0; k < numCell; ++k){
	atomC1 = splitCellArray[atmp][k];
	atomArrayMC[atomC1].setSpin(spinOld);
      }
      /// END: Loop for putting spin to original one for next cluster calc. ///      

      diffSpinArray.push_back(diffSpinSum);
      diffEnergy +=  clusterArray[j+1].getECI() * diffSpinSum / clusterArray[j+1].getNumCluster();
    }
    /// END: Loop for each cluster /////

    energyIn += diffEnergy;
    /// BEGIN: Loop for putting spin to new value ///
    for (int k = 0; k < numCell; ++k){
      atomC1 = splitCellArray[atmp][k];
      atomArrayMC[atomC1].setSpin(spinNew);
    }
    /// END: Loop for putting spin to new value ///    

    for (int m = 0; m < clusterArray.size() - 1; ++m){
      clusterArray[m+1].addSpin(diffSpinArray[m]);
    }
    diffSpinArray.clear();

  }

}


void GSSURFFAST::gsSearch(){

  long double energy(0);
  std::set<int> energyArray;
  // std::vector<long double> energyArray;
  long double energyMin = 1000000000;
  std::vector<Atom> atomArrayGS;

  for (int i = 0; i < clusterArray.size(); ++i) {
    energy += clusterArray[i].getECI() * clusterArray[i].getSpin() / clusterArray[i].getNumCluster();
  }
  std::cout << "E (for input POSCAR) = "<< energy << std::endl;


  /// Spin Bank /////////////////////////
  std::vector<std::pair<int, int> > initSpin, nextSpin;
  int Nconf;
  SpinBank   spinBank(spinArray, NatomArray);
  Nconf    = spinBank.getNconf();
  initSpin = spinBank.getInitialSpin();

  std::cout << "Number of configurations: " << Nconf + 1 << std::endl;
  ///////////////////////////////////////

  /// Calc energy for initial conf of GS search ///
  calcNewSpinConfig(initSpin, energy);
  energyArray.insert(static_cast<int>(energy * energyProd));
  if ( energy < energyMin) {
    energyMin = energy;
    atomArrayGS.clear();
    atomArrayGS = atomArrayMC;
  }
  /////////////////////////////////////////////////


  /// BEGIN: Main loop for GS search ///
  for ( int i = 0; i < Nconf; ++i){
    nextSpin = spinBank.getNextSpin();
    calcNewSpinConfig(nextSpin, energy);
    energyArray.insert(static_cast<int>(energy * energyProd));    
    if ( energy < energyMin) {
      energyMin = energy;
      atomArrayGS.clear();
      atomArrayGS = atomArrayMC;
    }
  }
  /// END: Main loop for GS search /////



  // WRITING OUTPUT FILE FOR gssurf.out	

  std::ofstream engFile, posFile;

  engFile.open("gsSurf.energy.out",std::ios::out);
  std::set<int>::iterator it0;
  it0 = energyArray.begin();
  double etmp;
  while ( it0 != energyArray.end()){
    etmp = (*it0) / energyProd;
    for ( int j = 0; j < compArray.size(); ++j){
      engFile << compArray[j] << "  ";
    }
    engFile << etmp << std::endl;
    ++it0;
  }
  engFile << "# GS-energy: " << energyMin << std::endl;
  engFile << std::endl;
  engFile.close();

  posFile.open("gsSurf.pos.out",std::ios::out);  
  for (int i = 0; i < atomArrayGS.size(); ++i){
    vector3D<double> position = atomArrayGS[i].getPosition();
    posFile << position.x << " " << position.y << " " << position.z << " ";
    posFile << atomArrayGS[i].getSpin() << std::endl;
  }
  posFile << std::endl;
  posFile.close();

}





double GSSURFFAST::calcSpin(std::vector<std::vector<int> > atomVector, int flipAtom, int flipSpin){
  
  double spinSum = 0;
  for (int i = 0; i < atomVector.size(); ++i){
    double spin = 1;
    for (int j = 0; j < atomVector[i].size(); ++j){
      if (flipAtom == atomVector[i][j]){
	spin *= flipSpin;
      }else{
	spin *= atomArrayMC[atomVector[i][j]].getSpin();
      }
    }
    spinSum += spin; 
  }
  return spinSum;
}