////////////////////////////////////////////////////////
//
// Program to generate a "twistable" DNA molecule.
// Either generates a linear random walk, or a circular DNA
//

// key to atom types
#define ADNA 1             // dna bead type
#define APROT 2            // protein atom type

// key to angle types
#define ANGB 1         // bending angle
#define ANGT 2         // torsion angle
#define ANGE 3         // special torsion agle for end of DNA

// key to bond types
#define BDNADNA 1         // dna-dna

#define PI 3.14159265358979

//#include "generate_DNA.h"

#include <iostream>
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <vector>

using namespace std;

struct bond {
  int a,b,
    type;
bond(const int &a,const int &b,const int &type) : a(a), b(b), type(type) {}; 
};

struct angle {
  int a,b,c,
    type;
angle(int a, int b, int c, int type) : a(a), b(b), c(c), type(type) {};
};

struct quaternion {
  double q0,q1,q2,q3;
};


struct atom {
  double x,y,z,density;
  double q[4];
  int id,
    type, 
    mol,
    ellipse_flag;
atom(const double &x,const double &y,const double &z) : x(x), y(y), z(z) {};
};



int main() {

  vector<atom> atoms;
  vector<bond> bonds;
  vector<angle> angles;

  int N,                    // length of DNA
    n_prot,                 // number of proteins
    seed,                   // seed for randoms
    patch_flag,
    circ_flag,
    prot_flag,
    Nfiles;                 // number of files to make

  double x,y,z,
    box[3],                 // size of simulation box
    hbox[3],                // half the size of the simulation box
    d,                      // separation of DNA-patch
    density;                // of atoms

  char fn[50];
  ofstream ouf;

  double dx,dy,dz,
    theta,
    phi,
    radius;


  // get parameters
  cout<<"length of DNA : "<<endl; cin>>N;

 pflag:
  cout<<"Make 'twistable' DNA? Enter 1 for yes, 0 for no.  "<<endl; 
  cin>>patch_flag;
  if ( patch_flag==1 ) {cout<<"DNA cores are ellipsoids and are atom type "<<ADNA<<"."<<endl;}
  else if ( patch_flag==0 ) {;}
  else  {cout<<"Enter 1 or 0 only."<<endl;goto pflag;}

 clfag:
  cout<<"Linear or circular DNA? Enter 0 for linear, 1 for circular"<<endl; 
  cin>>circ_flag;
  if ( circ_flag==1 ) {
    cout<<"Generating circular DNA."<<endl;
    radius = 0.5*N/PI;
  } else if ( circ_flag==0 ) {;}
  else  {cout<<"Enter 1 or 0 only."<<endl;goto pflag;}

  cout<<"Number of proteins"<<endl;cin>>n_prot;
  if (n_prot>0) prot_flag=1;

  cout<<"size of box (three values) : "<<endl; cin>>box[0]>>box[1]>>box[2];
  hbox[0]=box[0]*0.5;
  hbox[1]=box[1]*0.5;
  hbox[2]=box[2]*0.5;

  if ( circ_flag==1 && radius>hbox[0] && radius>hbox[1] && radius>hbox[2]    ) {
    cout<<"Box is not large enough to fit circular DNA"<<endl;
    exit(0);
  }

  cout<<"seed for randoms"<<endl; cin>>seed;
  srand(seed);
  cout<<"number of files to generate"<<endl; cin>>Nfiles;  

  density=(1-0.1*patch_flag)*6.0/PI;


  for (int filenum=1;filenum<=Nfiles;filenum++) {

    // file name
    if (Nfiles==1) {
      sprintf(fn,"lammps.input");
    } else {
      sprintf(fn,"lammps.input_%d",filenum);
    }

    //***********************************************
    // Do DNA

    if ( circ_flag == 0 )  {
      // Linear DNA

      // first bead
      atoms.push_back( atom(0.0,0.0,0.0) );
      atoms.back().type=ADNA;
      atoms.back().id=1;
      atoms.back().mol=1;
      atoms.back().ellipse_flag=1;
      atoms.back().q[0]=1; atoms.back().q[1]=0; atoms.back().q[2]=0; atoms.back().q[3]=0;
      atoms.back().density=density;

      // middle beads
      for (int i=1;i<N;i++) {

	// position
	do {
	  theta=double(rand())/double(RAND_MAX)*PI;
	  phi=double(rand())/double(RAND_MAX)*2.0*PI;
	  dx=atoms.back().x+sin(theta)*cos(phi);
	  dy=atoms.back().y+sin(theta)*sin(phi);
	  dz=atoms.back().z+cos(theta);
	} while ( abs(dx)>hbox[0]||abs(dy)>hbox[1]||abs(dz)>hbox[2] ); // reject if outside box

	atoms.push_back( atom(dx,dy,dz) );
	atoms.back().id = atoms[atoms.size()-2].id+1;
	atoms.back().type=ADNA;
	atoms.back().mol=atoms.back().mol;
	atoms.back().ellipse_flag=1;
	atoms.back().density=density;   
	atoms.back().q[0]=1; atoms.back().q[1]=0; atoms.back().q[2]=0; atoms.back().q[3]=0;

	// add bond
	bonds.push_back( bond(atoms[atoms.size()-2].id,atoms.back().id,BDNADNA) );

	// add twist angle
	angles.push_back( angle(atoms[atoms.size()-2].id,atoms.back().id,atoms.back().id,ANGT) );

	if (i>1) { // one onto the third bead, start adding bending anlges
	  angles.push_back( angle(atoms[atoms.size()-3].id,atoms[atoms.size()-2].id,atoms.back().id,ANGB) );
	}
 
      }

      // add a final angle for the end of the linear DNA chain
      angles.push_back( angle(atoms[atoms.size()-2].id,atoms.back().id,atoms.back().id,ANGE) ); // the third id here does nothing 

    } else if ( circ_flag == 1 ) {
      // Circular DNA
      radius = 0.5*N/PI;
      phi = 2*PI / N;
      theta = 0;

      // first bead
      atoms.push_back( atom(radius,0.0,0.0) );
      atoms.back().type=ADNA;
      atoms.back().id=1;
      atoms.back().mol=1;
      atoms.back().ellipse_flag=1;
      atoms.back().q[0]=1; atoms.back().q[1]=0; atoms.back().q[2]=0; atoms.back().q[3]=0;
      atoms.back().density=density;

      // middle beads
      for (int i=1;i<N;i++) {

	theta += phi;
	dx = radius*cos(theta);
	dy = radius*sin(theta);
	dz = 0.0;

	atoms.push_back( atom(dx,dy,dz) );
	atoms.back().id = atoms[atoms.size()-2].id+1;
	atoms.back().type=ADNA;
	atoms.back().mol=atoms.back().mol;
	atoms.back().ellipse_flag=1;
	atoms.back().density=density;   
	atoms.back().q[0]=1; atoms.back().q[1]=0; atoms.back().q[2]=0; atoms.back().q[3]=0;

	// add bond
	bonds.push_back( bond(atoms[atoms.size()-2].id,atoms.back().id,BDNADNA) );

	// add twist angle
	angles.push_back( angle(atoms[atoms.size()-2].id,atoms.back().id,atoms.back().id,ANGT) );

	if (i>1) { // one onto the third bead, start adding bending anlges
	  angles.push_back( angle(atoms[atoms.size()-3].id,atoms[atoms.size()-2].id,atoms.back().id,ANGB) );
	}
 
      }

      // add bond to connect end with start
      bonds.push_back( bond(atoms.back().id,atoms[0].id,BDNADNA) );

      // add two bend angles to connect end with start
      angles.push_back( angle(atoms[atoms.size()-2].id,atoms.back().id,atoms[0].id,ANGB) );
      angles.push_back( angle(atoms.back().id,atoms[0].id,atoms[1].id,ANGB) );

      // add twist angle to connect end with start
      angles.push_back( angle(atoms.back().id,atoms[0].id,atoms[0].id,ANGE) ); // the third id here does nothing 



    } else {
      cout<<"Something went wrong"<<endl;
      exit(0);
    }
   

    // Now add proteins
    if (prot_flag) {
      int  id=atoms.back().id,
	mol=atoms.back().mol;
      for (int i=0;i<n_prot;i++) {
	x=box[0]*double(rand())/double(RAND_MAX)-hbox[0];
	y=box[1]*double(rand())/double(RAND_MAX)-hbox[1];
	z=box[2]*double(rand())/double(RAND_MAX)-hbox[2];
	// make the new atom
	id++;
	mol++;
	atoms.push_back( atom(x,y,z) );
	atoms.back().ellipse_flag=0;
	atoms.back().id=id;
	atoms.back().type=APROT;
	atoms.back().mol=mol;
	atoms.back().density=1.0;
      }
    }


    // output 

    ouf.open(fn);
    ouf<<" LAMMPS data file"<<endl;
    ouf<<endl;
    ouf<<" "<<atoms.size()<<" atoms"<<endl;
    ouf<<" "<<N<<" ellipsoids"<<endl;
    ouf<<" "<<bonds.size()<<" bonds"<<endl;
    ouf<<" "<<angles.size()<<" angles"<<endl;
    ouf<<endl;
    ouf<<" "<<1+prot_flag<<" atom types"<<endl;
    ouf<<" 1 bond types"<<endl;
    ouf<<" 3 angle types"<<endl;
    ouf<<endl;
    ouf<<" "<<-hbox[0]<<" "<<hbox[0]<<" xlo xhi"<<endl;
    ouf<<" "<<-hbox[1]<<" "<<hbox[1]<<" ylo yhi"<<endl;
    ouf<<" "<<-hbox[2]<<" "<<hbox[2]<<" zlo zhi"<<endl;
    ouf<<endl;
    ouf<<endl<<" Masses"<<endl<<endl;
    ouf<<ADNA<<" "<<1<<endl; // mass is ignored for ellipsoids
    if (prot_flag) {
      ouf<<APROT<<" "<<"1"<<endl;
    } 
    ouf<<endl<<" Atoms"<<endl<<endl;
    for (int i=0;i<atoms.size();i++) {
      ouf<<" "<<atoms[i].id<<" "<<atoms[i].type<<" "<<atoms[i].x<<" "<<atoms[i].y<<" "<<atoms[i].z<<" "<<atoms[i].mol<<" "<<atoms[i].ellipse_flag<<" "<<atoms[i].density<<endl;
    }
    ouf<<endl<<" Ellipsoids"<<endl<<endl;
    for (int i=0;i<atoms.size();i++) {
      if (atoms[i].ellipse_flag==1) {
	ouf<<" "<<atoms[i].id<<" 1 1 1 "<<atoms[i].q[0]<<" "<<atoms[i].q[1]<<" "<<atoms[i].q[2]<<" "<<atoms[i].q[3]<<endl; 
      }
    }
    ouf<<endl<<" Velocities"<<endl<<endl;
    for (int i=0;i<atoms.size();i++) {
      ouf<<" "<<atoms[i].id<<" 0 0 0 0 0 0"<<endl;
    }
    ouf<<endl<<" Bonds"<<endl<<endl;
    for (int i=0;i<bonds.size();i++) {
      ouf<<" "<<i+1<<" "<<bonds[i].type<<" "<<bonds[i].a<<" "<<bonds[i].b<<endl;
    }
    ouf<<endl<<" Angles"<<endl<<endl;
    for (int i=0;i<angles.size();i++) {
      if (angles[i].c==angles[i].b) { angles[i].c = angles[i].b+1; }
      if (circ_flag==0) {
	if (angles[i].c>atoms.size()) { angles[i].c = angles[i].a-1; }
      } else {
	if (angles[i].c>atoms.size()) { angles[i].c -= atoms.size(); }
      }
      ouf<<" "<<i+1<<" "<<angles[i].type<<" "<<angles[i].a<<" "<<angles[i].b<<" "<<angles[i].c<<endl;
    }
    ouf.close();

    // clean up
    atoms.clear(); vector<atom>().swap(atoms);
    bonds.clear(); vector<bond>().swap(bonds);
    angles.clear(); vector<angle>().swap(angles);

  }

}