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

  Class for surface canonical MC simulation (fix the bulk layer atomic configuration)

        Modified by k-fleak

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

#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"


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

class CMCSURF{
  ///  ///
protected:
  //////////
	std::vector<Cluster> clusterArray;
	std::vector<Atom> atomArrayMC;
	vector4 coordination;
  ///  //.
  std::vector<int> tempMC;
  int step1MC;
  int step2MC;

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

	std::ofstream cmcFile;
//	std::ofstream cmcEnergy;
//	std::ofstream cmcEnergyAve;
	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 monteCarlo(const std::vector<int>& temp, int nStepAnn, int nStepEqv);
	virtual double calcSpin(std::vector<std::vector<int> > atomVector, int flipAtom, int flipSpin, int anotherAtom);
	//////////

public: 

	CMCSURF(std::vector<Atom>& atomArray, std::vector<Atom>& atomArrayLattice, 
		std::vector<Atom>& atomArraySupercell, const std::vector<int>& temp, const int& nStepAnn, 
		const int& nStepEqv, const std::vector<int>& nCell, std::map<int, double>& eci,
 		const std::vector<std::vector<int> >& positionOut, std::vector<double> MovAtomZPosIn);
	~CMCSURF();
  ///  ///
  virtual void start();
  //////////

};

CMCSURF::CMCSURF(std::vector<Atom>& atomArray, std::vector<Atom>& atomArrayLattice, 
	std::vector<Atom>& atomArraySupercell, const std::vector<int>& temp, const int& nStepAnn, 
	const int& nStepEqv, const std::vector<int>& nCell, std::map<int, double>& eci,
 	const std::vector<std::vector<int> >& positionOut, std::vector<double> MovAtomZPosIn){

	///  ///
	clock_t startTime;
	clock_t currentTime;
	clock_t endTime;
	startTime = clock();
	//////////

	/// surface - Setting z position of atom to be fixed in MC. - ///
	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 );
	}
	///////////////

        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();

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

	///  ///
	step1MC = nStepAnn * atomArrayMC.size();
	step2MC = nStepEqv * atomArrayMC.size();
	for (int j = 0; j < temp.size(); ++j) tempMC.push_back(temp[j]);
	//////////

}

CMCSURF::~CMCSURF(){}

void CMCSURF::start(){
  clock_t startTime;
  clock_t endTime;
  startTime = clock();
  std::cout << "monteCarlo ... " ;
  monteCarlo(tempMC, step1MC, step2MC);
  endTime = clock();
  std::cout << "done.   monteCarlo elapsed time: " << (endTime-startTime) / CLOCKS_PER_SEC << " sec" << std::endl;
}

void CMCSURF::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();
        std::vector<int> 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 CMCSURF::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 CMCSURF::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 CMCSURF::monteCarlo(const std::vector<int>& temp, int nStepAnn, int nStepEqv){
//////////

	const double k_B = 0.86171 * pow (10,-4);
	unsigned long init = time(NULL) + getpid();
	init_genrand(init);


	///  ///
	std::vector<double> correlationArray;
	correlationArray.resize(clusterArray.size());
	//////////

	
	long double energy(0), energyAve(0), energySquare(0);
	int natom = atomArrayMC.size();

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

	std::cout << "E (initial) = "<< energy << std::endl;

	// BEGIN MAIN LOOP

	for (int iterTemp = 0; iterTemp < temp.size() + 1; ++iterTemp){
		int tempTrial, nStep;
		if (iterTemp == temp.size()){
			tempTrial = temp[temp.size()-1];
			nStep = nStepEqv;
		}
		else {
			tempTrial = temp[iterTemp];
			nStep = nStepAnn;
		}
	

		std::cout << "TEMPERATURE = "<< tempTrial << " " << "NSTEP = "<< nStep << std::endl;
		int stepSum = 1;

		while (stepSum <= nStep){

			int atom1, atom2;
			int spin1, spin2;
			vector3D<double> pos1, pos2;
			int flag1, flag2;
			atom1 = genrand_int32() % natom;
			atom2 = genrand_int32() % natom;
			spin1 = atomArrayMC[atom1].getSpin();
			spin2 = atomArrayMC[atom2].getSpin();
			/// surf ///
			flag1 = 0;
			flag2 = 0;
			pos1  = atomArrayMC[atom1].getPosition();
			pos2  = atomArrayMC[atom2].getPosition();
			if (spin1 != spin2) {
			  for (int i = 0; i < slab1ZPos.size(); ++i ) {
			    if ( (pos1.z >= slab1ZPos[i] - eps) and (pos1.z <= slab1ZPos[i] + eps)) {
			      flag1 = 1;
			    } 
			    if ( (pos2.z >= slab1ZPos[i] - eps) and (pos2.z <= slab1ZPos[i] + eps) ) {
			      flag2 = 1;
			    }
			  }
			  for (int i = 0; i < slab2ZPos.size(); ++i ) {
			    if ( (pos1.z >= slab2ZPos[i] - eps) and (pos1.z <= slab2ZPos[i] + eps)) {
			      flag1 = -1;
			    }
			    if ( (pos2.z >= slab2ZPos[i] - eps) and (pos2.z <= slab2ZPos[i] + eps) ) {
			      flag2 = -1;
			    }
			  }
			}
			////////////

			if ( (flag1 * flag2) == 1 )  { 
				//if (stepSum%100000 ==0) std::cout << "STEP"<< stepSum << " : ";

				double spinOld, spinNew, diffSpin;
				std::vector<double> diffSpinArray;
				long double diffEnergy = 0;

				///  ///
				for (int i = 0; i < clusterArray.size() - 1; ++i){
				//////////
				
					///  ///
					spinOld = calcSpin(coordination[atom1][i], atom1, spin1, atom2) \
						+ calcSpin(coordination[atom2][i], atom2, spin2, atom1); 

					spinNew = calcSpin(coordination[atom1][i], atom1, spin2, atom2) \
						+ calcSpin(coordination[atom2][i], atom2, spin1, atom1);
					//////////
				
					diffSpin = spinNew - spinOld;
				
					diffSpinArray.push_back(diffSpin);

					///  ///
					diffEnergy +=  clusterArray[i+1].getECI() * diffSpin / clusterArray[i+1].getNumCluster();
					//////////

				}

				double tmpProb = -diffEnergy / k_B / tempTrial;
				double probability = exp ( tmpProb );
				double cut = genrand_real2();


				if (probability > cut ) {
					energy += diffEnergy;
					atomArrayMC[atom1].setSpin(spin2);
					atomArrayMC[atom2].setSpin(spin1);

					///  ///
					for (int i = 0; i < clusterArray.size() - 1; ++i) clusterArray[i+1].addSpin(diffSpinArray[i]);
					//////////
				}


	
				/*********************************************************************************/
				// check the correlation functions 
				/*********************************************************************************/
	
			//	for (int i = 0; i < clusterArray.size(); ++i){
			//		double spin_tmp = clusterArray[i].getSpin();
			//		double clusterNum_tmp = clusterArray[i].getNumCluster();
			//		if (fabs(spin_tmp) > clusterNum_tmp ){	
			//			std::cerr << "Error: Correlation Functions are over one!\n";
			//			std::cout << "Error: " << clusterArray[i].getSerialIndex() << " Correlation Functions are over one! " << std::endl;
			//	                exit(8);
     			//		}
			//	}	
				/*********************************************************************************/
				// Average 
				/*********************************************************************************/

				if (iterTemp == temp.size()){
					for (int i = 1; i < clusterArray.size(); ++i){
						double spin = clusterArray[i].getSpin();
						double numCluster = clusterArray[i].getNumCluster();

						///  ///
						correlationArray[i] += spin / numCluster / nStepEqv;
						//////////

					}
					//	energyAve += energy;	
					//      energySquare += energy * energy;	
				}
	
				diffSpinArray.clear();
				++stepSum;

			}
	
		}
		std::cout << "E = "<< energy << " " << std::endl;
	}

	// END OF MAIN LOOP	
	std::cout << " End of loop " << std::endl;

	// WRITING OUTPUT FILE FOR cmc.out	

	cmcFile.open("cmc.out",std::ios::out);
//
//	energyAve /= NstepEq;
//	energySquare /= NstepEq;
//	
	cmcFile << "------------------------------------------" << std::endl;
	cmcFile << "	Averages of correlation functions     " << std::endl;
	cmcFile << "------------------------------------------" << std::endl;
	energyAve = 0;
	correlationArray[0] = 1;
	for (int i = 0; i < clusterArray.size(); ++i){
		int index = clusterArray[i].getIndex();
		double numCluster = clusterArray[i].getNumCluster();
		//double correlation = correlationArray[i] / numCluster / NstepEq;
		double correlation = correlationArray[i];
		energyAve += clusterArray[i].getECI() * correlation;

		cmcFile << index << " : " << correlation << std::endl;
	}



	/*	std::cout << "------------------------------------------" << std::endl;
	std::cout << "	Averages of correlation functions     " << std::endl;
	std::cout << "------------------------------------------" << std::endl;

	correlationArray[0] = 1;
	for (int i = 0; i < correlationArray.size(); ++i){
		std::cout << correlationArray[i] << std::endl;
		}*/


//
	cmcFile << "------------------------------------------" << std::endl;
	cmcFile << "	       Averages of energy             " << std::endl;
	cmcFile << "------------------------------------------" << std::endl;
//	
	cmcFile.precision(10);
//	cmcFile << "Number of steps of average = " << stepAverage << std::endl;
	cmcFile << "Average of energy = "<< energyAve << " eV" << std::endl;
//	cmcFile << "Average of square of energy = "<< energySquare << " eV^2" << std::endl;
//	//cmcFile << "Boltzmann constant = "<< k_B << std::endl;
//	cmcFile << "Temperature = "<< tempEnd <nn< std::endl;
//	long double squareTemp = k_B * tempEnd * tempEnd;
//	long double squareEnergyAve = energyAve * energyAve;
//	long double diffEnergyAve = energySquare - squareEnergyAve;
//	cmcFile << "Specific heat at constant volume  = "<< diffEnergyAve / squareTemp  << std::endl;
//
	cmcFile << std::endl;
//
	cmcFile << "------------------------------------------" << std::endl;
	cmcFile << "	            Position                  " << std::endl;
	cmcFile << "------------------------------------------" << std::endl;
//
	for (int i = 0; i < atomArrayMC.size(); ++i){
		vector3D<double> position = atomArrayMC[i].getPosition();
		cmcFile << position.x << " " << position.y << " " << position.z << " ";
		cmcFile << atomArrayMC[i].getSpin() << std::endl;

	}
	cmcFile << std::endl;
	cmcFile.close();
}

double CMCSURF::calcSpin(std::vector<std::vector<int> > atomVector, int flipAtom, int flipSpin, int anotherAtom){

	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 if(anotherAtom == atomVector[i][j]){
				spin = 0.0;
			}else{
				spin *= atomArrayMC[atomVector[i][j]].getSpin();
			}
			//////////
		}
		spinSum += spin; 
	}

	return spinSum;
}