#include<iostream>
#include<fstream>
#include<cstdlib>
#include<iomanip>
#include<unistd.h>
#include<cstdio>
#include<string.h>
#include<sstream>
#include<math.h>
#include<complex>
#include<stdlib.h>     /* srand, rand */
#include<time.h>
#include<iomanip>
#include<ctime>

using namespace std;

void transformDtoC(double* array, double a, double z){
//transform direct coordinates to cartesian
// [x] =[a	a*cos(gamma)] [u]
// [y] =[0	a*sin(gamma)] [v]
  double* dummy;	dummy = new double[3];
  dummy[0] = a*array[0] + a*cos(120*3.14159265/180)*array[1];
  dummy[1] = 0*array[0] + a*sin(120*3.14159265/180)*array[1];
  dummy[2] = array[2]*z;
  for( int i=0; i<3; i++) array[i] = dummy[i];
}
void transformCtoD(double* array, double a, double z){
//transform cartesian to direct
// [u] =[1/a	-cos(gamma)/a*sin(gamma)][x]
// [v] =[0	1/(a*sin(gamma) 	][y]
  double* dummy;	dummy = new double[3];
  dummy[0] = ((1./a)*array[0]) -(cos(120*3.14159265/180)/(a*sin(120*3.14159265/180)))*array[1];
  dummy[1] = 0*array[0] + (1./(a*sin(120*3.14159265/180)))*array[1];
  dummy[2] = array[2]/z;
  for( int i=0; i<3; i++) array[i] = dummy[i];
}
void transform(double* pos_pol, double* pos_cart){
  
  //transform molecular [X,Y,Z,r,theta,phi] to [X1,Y1,Z1,X2,Y2,Z2]
  double sin_t = sin(pos_pol[4]*3.14159265/180);
  double cos_t = cos(pos_pol[4]*3.14159265/180);
  double sin_p = sin(pos_pol[5]*3.14159265/180);
  double cos_p = cos(pos_pol[5]*3.14159265/180);

  double fac1 = 0.5;//MASS_B/(MASS_A+MASS_B);
  double fac2 = 0.5;//MASS_A/(MASS_A+MASS_B);

//  double fac1 = 0.97235024; double fac2 = 0.027649686;

  //poscart 1 - 3 is for H1-ATOM
  //poscart 4 - 6 is for H2-ATOM
  pos_cart[0] = pos_pol[0] + fac1*pos_pol[3]   * sin_t*cos_p;
  pos_cart[3] = pos_pol[0] - fac2*pos_pol[3]   * sin_t*cos_p;
  pos_cart[1] = pos_pol[1] + fac1*pos_pol[3]   * sin_t*sin_p;
  pos_cart[4] = pos_pol[1] - fac2*pos_pol[3]   * sin_t*sin_p;
  pos_cart[2] = pos_pol[2] + fac1*pos_pol[3]   * cos_t;
  pos_cart[5] = pos_pol[2] - fac2*pos_pol[3]   * cos_t;

  //if(pos_cart[2] < -3.627185) pos_cart[2] = -3.627185;
  //if(pos_cart[5] < -3.627185) pos_cart[5] = -3.627185; 

  //cout << "CART ";
  //for(int i = 0; i < 6; i++) cout << pos_cart[i] << " ";
  //cout << "    \n\n";

};

int main(int argc, char* argv[]){
//input a3D uc layers Dvac fix layerpos[layer] R position Z theta phi fix
//how to use the program
//cout << "input: a3D [A] -- Supercell [int] -- layers [int] -- vacuum distance [A](int) \n";
//cout << "BtH Barrier - E[eV]=0.636  Z[A]=1.172  r[A]=1.03 -->DOI:10.1021/acs/jpca.5b06008  \n\n";
//------------------------------------------------------------------------------
if( argc == 1){
  cout << "ERROR: not enough arguments\n";
  cout << "Create POSCAR with H2 \n"
       << "input: a3D [A](double) -- Supercell [int] -- layers [int] -- vacuum distance [A](int)  -- fix -- layerpositions [] [] []... -- Atom/Molecule -- u -- v -- t -- p --Z --r\n";
  return 0;
}

double 	a3D = atof( argv[1] ); 	//bluk constant
int 	SCd = atoi( argv[2] ); 	//supercell dimension i.e. 1x1 2x2 etc
int 	lay = atoi( argv[3] ); 	//number of layers
int 	vac = atoi( argv[4] );  //vacuum distance in Angstrom
//char fixslab= argv[5];
//rough interlayer distance Cu atoms in direct coordinates
//double CuddA= 2.1; 
double CuddA= a3D / sqrt(3);  //bulk Z distance between atoms
double Cudd = CuddA / ( (lay * CuddA) + vac);	 	

//#Cu(111) primitive unit cell - Direct coordinates 
double Cu1u = 0.;	double Cu1v = 0.;	double Cu1z = 0*Cudd;
double Cu2u = 1./3.;	double Cu2v = 2./3.;	double Cu2z = 1*Cudd;
double Cu3u = 2./3.; 	double Cu3v = 1./3.;	double Cu3z = 2*Cudd;
double Cu4u = 0.;	double Cu4v = 0.;	double Cu4z = 3*Cudd;

//store in array; fourth line not neccesary --> dumb programming
//fcc (111) only has A B C 
double* CuCell;	CuCell = new double[4*3];
CuCell[0] = Cu1u;	CuCell[1] = Cu1v;	CuCell[2] = Cu1z;
CuCell[3] = Cu2u;	CuCell[4] = Cu2v;	CuCell[5] = Cu2z;
CuCell[6] = Cu3u;	CuCell[7] = Cu3v;	CuCell[8] = Cu3z;
CuCell[9] = Cu4u;	CuCell[10]= Cu4v;	CuCell[11]= Cu4z;

//create slab to store all x,y,z coordinates needed for poscar file
double* slab;	slab = new double[lay*SCd*SCd*3]; 

//array to hold polar and cartesian coordinates
double* pos_pol;	pos_pol = new double[6];
double* pos_pol2;	pos_pol2= new double[6];
double* pos_cart;	pos_cart = new double[6];


//barrier position Mark wijzenbroek
double* H1;	H1 = new double[3];
double* H2;	H2 = new double[3];
//barrier position christina
double* Hcom;     Hcom = new double[3];
double* H2c;     H2c = new double[3];
//gasphase
double* H1g;     H1g = new double[3];
double* H2g;     H2g = new double[3];

//save to POSCAR
ofstream out;   out.open("POSCAR");

//use count to change position in slab array; do not sum i,j,k
int i,j,k,count=0,layer,barrier=0;
for(i=0; i<lay; i++){ 		//number of layers
  for(j=0; j<SCd; j++){		//supercell repetition in u
    for(k=0; k<SCd; k++){	//supercell repetition in v

//make sure that more then 4 layers is not a problem
if	( i%3 == 0 ) layer=0;
else if ( i%3 == 1 ) layer=1;
else if ( i%3 == 2 ) layer=2;
else cout << "ERROR";
slab[count]  	= (CuCell[0 + layer*3] + j) / SCd;
count++;
slab[count]	= (CuCell[1 + layer*3] + k) / SCd;
count++;
//cout << "argv[5]= " << argv[5] << "\n";
//if the keyword fix is used, then the provided layer positions are used
if( !strcmp(argv[5],"fix") ){
  slab[count]	= atof( argv[6+i] ) / ( (lay * CuddA) + vac);
  //cout << "fix ifstatement \n";
}
else{
  //otherwise the bulk interlayer distance is used a3D / sqrt(3)
  slab[count]  	= i*Cudd;//CuCell[2 + i*3];
}
count++;

//calculate barrier position
//barrier
if( i == (lay -1) && barrier==0 ){
  barrier = 1;
/*
Slab[count-3] is now the X position in direct coordinates of the 1st TOP atom
of the top most layer. 
1) move position relative to TOP --> e.g. TOP = 0,0 in u & v
2) Select Atom or Molecule
3) If molecule then transform COM position from direct to cartesian
   		- take Z r theta and phi --> transform to cartesian.
		- trans from to two direct coordinate positions for both H's 
4) If atom -> take Z from commandline argument. 
5) print POSCAR. 

Commandline arguments:
argv[1]		= a3D
argv[2]		= Unit Cell size
argv[3]		= number of layers
argv[4]		= Dvac
argv[5]		= fix keyword to take commandline layerpositions.
argv[6+i]	= z layering of slab
argv[6+i+1]	= Atom or Molecule swith
argv[6+i+2]	= u position of H or H2 com
argv[6+i+3]	= v position of H or H2 com
argv[6+i+4]	= theta H2
argv[6+i+5]	= phi H2
argv[6+i+6]	= Z position of H or H2 com
argv[6+i+7]	= R

*/

//First Store TOP (u,v)=(0,0) in H1
  Hcom[0] = slab[count-3];
  Hcom[1] = slab[count-2];
  Hcom[2] = slab[count-1]; 
//Hcom is now the origin against which we can shift using u,v,Z,r,theta,phi

  //only reason for bigger supercell is convergence --> u and v always within one primitive cell --> scale using SCd
  Hcom[0] += (atof( argv[6+i+2] ) / SCd ); // u / supercell dimension
  Hcom[1] += (atof( argv[6+i+3] ) / SCd ); // v / supercell dimension
  Hcom[2] += (atof(argv[6+i+6] ) / ( (lay * CuddA) + vac)); //Z scaled to correct direct coordinates
//Hcom is now the COM coordinate of H2 or the coordinate of H

  //Only do molecule calculation when specified to prevent atof from causing segfault
  if( !strcmp(argv[6+i+1],"Molecule") ){
    H1[0] = Hcom[0];
    H1[1] = Hcom[1];
    H1[2] = Hcom[2];
    //transfrom center of mass to cartesian
    transformDtoC( H1, (a3D/sqrt(2))*SCd, (lay*CuddA + vac) );

    pos_pol[0] = H1[0];//x
    pos_pol[1] = H1[1];//y
    pos_pol[2] = H1[2];//z
    pos_pol[3] = atof( argv[6+i+7] ); // req
    pos_pol[4] = atof( argv[6+i+4] ); //theta
    pos_pol[5] = atof( argv[6+i+5] ); //phi    
    //pos_pol now has molecular coordinates

    transform( pos_pol, pos_cart);
 
    H1[0] = pos_cart[0];
    H1[1] = pos_cart[1];
    H1[2] = pos_cart[2];
    H2[0] = pos_cart[3];
    H2[1] = pos_cart[4];
    H2[2] = pos_cart[5];
     
    transformCtoD( H1, (a3D/sqrt(2))*SCd, (lay*CuddA + vac) );
    transformCtoD( H2, (a3D/sqrt(2))*SCd, (lay*CuddA + vac) ); 
    //now we have two atom positions in direct coordinates 
    
    //output poscar title if molecule:
    cout << "u= "<< argv[6+i+2] << " v= " << argv[6+i+3] << " Z= " << argv[6+i+6] << " r= " << argv[6+i+7] << " t= " << argv[6+i+4] << " p= " << argv[6+i+5] << "\n";
  } // IF Molecule 

  //output poscar title if ATOM
  if( !strcmp(argv[6+i+1],"Atom") ){ 
    cout << "u= "<< argv[6+i+2] << " v= " << argv[6+i+3] << " Z= " << argv[6+i+6] << "\n";
  }
}//if BARRIER
//cout << "i="<<i<<endl;
    }//i
  }//j
}//k
//------------------------------------------------------------------------------
//Print Direct and Cartesian POSCAR files
//------------------------------------------------------------------------------
//------------------------------------------------------------------------------
//fix cout precision en length
std::cout.precision(9);
std::cout.setf( std::ios::fixed, std:: ios::floatfield );

//unit cell vectors in cartesian - same format as poscar file -- unit cell vectors are fractions of the "bulk lattice constant of the supercell"
//a3D is the bulk lattice constant  
double zero = 0.; //to get correct formating
cout << a3D*SCd << "\n";
cout << setw(12) << a3D / sqrt(2) / a3D						<< "    " << zero 						<< "    " 				<< zero << " \n";
cout << setw(12) <<	-((a3D / sqrt(2))*cos(60 * 3.14159265 / 180))/a3D 	<< "    "  << ((a3D / sqrt(2))*sin(60 * 3.14159265 / 180))/a3D 	<< "    " 				<< zero << " \n";
cout << setw(12) << zero							<< "    " << zero						<< "    " << (lay*CuddA + vac) / (a3D*SCd)<< " \n";
if( !strcmp(argv[6+lay],"Atom") ){ 	
  cout << "  Cu H \n";
  cout << "   " << lay*SCd*SCd << " 1 \n";
}
else{
  cout << "  Cu H \n";
  cout << "   " << lay*SCd*SCd << " 2 \n";
}
cout << "Selective Dynamics \n";
cout << "direct \n";

//print direct coordinates of (super)Cell
for(i=0;i<lay*SCd*SCd;i++){
/*original
  if( slab[i*3+2] == 0 ){	//botom layer is always fixed
    cout << slab[i*3] << "    " << slab[i*3+1] << "    " << slab[i*3+2] << "  F  F  F \n";
  }//if
*/
  if( !strcmp(argv[5],"fix") ){
    cout << slab[i*3] << "    " << slab[i*3+1] << "    " << slab[i*3+2] << "  F  F  F \n";
  }
  else{
    if( slab[i*3+2] == 0 || slab[i*3+2] == Cudd ){
      if( slab[i*3+2] == Cudd && lay < 5){
        cout << slab[i*3] << "    " << slab[i*3+1] << "    " << slab[i*3+2] << "  F  F  T \n";
      }//if
      else{
        cout << slab[i*3] << "    " << slab[i*3+1] << "    " << slab[i*3+2] << "  F  F  F \n";
      }//else
    }//if
    else{				//of the other layers only Z coordinate is free to move
      cout << slab[i*3] << "    " << slab[i*3+1] << "    " << slab[i*3+2] << "  F  F  T \n";
    }//else
  }
}//for

//print H position of H2 positions
if( !strcmp(argv[6+lay],"Atom") ){
  cout << Hcom[0]	<< "    " << Hcom[1] << "    " << Hcom[2] << "  F  F  F \n";
}
if( !strcmp(argv[6+lay],"Molecule") ){
  cout << H1[0]	<< "    " << H1[1] << "    " << H1[2] << "  F  F  F \n"; 
  cout << H2[0]	<< "    " << H2[1] << "    " << H2[2] << "  F  F  F \n";
}
return 0;
}
    
