/*--------------------------------cones.c---------------------------------
  This program generates the coordinates and connectivity of carbon
  nanocones. The output format is in .abc format. This software was written
  by Stephen Jordan.
  ------------------------------------------------------------------------*/

#include "cell.h"
#include "defs.h"
#include "diags.h"
#include "abcio.h"
#include "tbinit.h"
#include "calculate.h"
#include "potential.h"
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <malloc.h>

#define CC_BOND (double)1.421
#define RAD_THREE_OVR_TWO (double)0.86602540378
#define RAD_THREE (double) 1.7320508075689
#define PI (double)3.14159265
#define CH_BOND (double)1.08  

void graphene(cell *C, int krange) {  
  int k1, k2;
  int i;
  int num;
  C->atoms = 2*(2*krange+1)*(2*krange+1);
  C->atom_array = (atom *)malloc(C->atoms*sizeof(atom));
  if(C->atom_array == NULL) printf("Unable to allocate graphene atoms!\n");
  C->x = (double *)malloc(3*C->atoms*sizeof(double));
  if(C->x == NULL) printf("Unable to allocate graphene coordinates!\n");
  for(i = 0; i < C->atoms; i++) C->atom_array[i].x = C->x + 3*i;
  for(k1 = -krange; k1 <= krange; k1++) {
    for(k2 = -krange; k2 <= krange; k2++) {
      num = (k1+krange)*(2*krange+1)+k2+krange;
      C->atom_array[2*num].x[0]   = RAD_THREE_OVR_TWO*CC_BOND*(k1+k2);
      C->atom_array[2*num].x[1]   = 1.5*CC_BOND*(k1-k2) - 2.0*CC_BOND;
      C->atom_array[2*num].x[2]   = 0;
      C->atom_array[2*num].type = 1;
      C->atom_array[2*num+1].x[0] = RAD_THREE_OVR_TWO*CC_BOND*(k1+k2);
      C->atom_array[2*num+1].x[1] = 1.5*CC_BOND*(k1-k2)+CC_BOND - 2.0*CC_BOND;
      C->atom_array[2*num+1].x[2] = 0;
      C->atom_array[2*num+1].type = 1;
    }
  }
}


void cart_to_cyl(const double *cart, double *cyl) {
  cyl[0] = sqrt(cart[0]*cart[0]+cart[1]*cart[1]);
  cyl[1] = atan2(cart[0], cart[1]);
  cyl[2] = cart[2];
}

void cyl_to_cart(const double *cyl, double *cart) {
  cart[0] = cyl[0]*cos(cyl[1]);
  cart[1] = cyl[0]*sin(cyl[1]);
  cart[2] = cyl[2];
}


void curl3(cell *C) {
  double tcut, fcut;
  double cart[3], cyl[3];
  int i, j;
  for(i = 0; i < C->atoms; i++) {
    //load the graphene
    cart[0] = C->atom_array[i].x[0];
    cart[1] = C->atom_array[i].x[1];
    cart[2] = C->atom_array[i].x[2];
    //rotate it
    cart_to_cyl(cart, cyl);
    cyl[1] += PI/6.0;
    if(cyl[1] > PI) cyl[1] -= 2*PI;
    //make it into a cone
    if(cyl[0] > CC_BOND*RAD_THREE) {
      tcut = PI/6.0-asin(CC_BOND*RAD_THREE/(2.0*cyl[0]));
      fcut = PI/(PI-2*tcut);
      //remove three wedges
      if((cyl[1] < 2*tcut && cyl[1] > 0) ||
         (cyl[1] < -(2.0/3.0)*PI + 2*tcut && cyl[1] > -(2.0/3.0)*PI) ||
         (cyl[1] < (2.0/3.0)*PI + 2*tcut && cyl[1] > (2.0/3.0)*PI))
        C->atom_array[i].type = 0;
      for(j = 0; j < 3; j++) {
        //stretch the remaining stuff to fill a gap
        if(cyl[1] > 2*tcut && cyl[1] < (2.0/3.0)*PI) {
          cyl[1] = (cyl[1] - 2*tcut)/(1-3.0*tcut/PI) + 2*tcut;
          if(cyl[0] > CC_BOND*4) {
            cyl[0] -= (cyl[1] - 2*tcut)*0.5;
            cyl[1] += (cyl[1] - 2*tcut)*0.01;
          }
        }
        //rotate
        cyl[1] += (2.0/3.0)*PI;
        if(cyl[1] > PI) cyl[1] -= 2*PI;
      }
    }  
    if(cyl[0] > CC_BOND*4) {
      cyl[2] = 0.5*sqrt(3.0)*cyl[0];
      cyl[0] *= 0.5;  
    }
    else {
      cyl[0] *= 0.8 - 0.03*cyl[0];
      cyl[2] = 1.7 + 0.1*cyl[0] + 0.2*cyl[0]*cyl[0];
      cyl[0] *= 0.9;
    }
    cyl_to_cart(cyl, cart);
    C->atom_array[i].x[0] = cart[0];
    C->atom_array[i].x[1] = cart[1];
    C->atom_array[i].x[2] = cart[2];
  }
}

void curl2(cell *C) {
  double tcut, fcut;
  double cart[3], cyl[3];
  int i;
  int immune;
  for(i = 0; i < C->atoms; i++) {
    immune = 0;
    //load the graphene
    cart[0] = C->atom_array[i].x[0];
    cart[1] = C->atom_array[i].x[1];
    cart[2] = C->atom_array[i].x[2];
    //rotate it
    cart_to_cyl(cart, cyl);
    cyl[1] += PI/6.0;
    if(cyl[1] > PI) cyl[1] -= 2*PI;
    //make it into a cone
    if(cyl[0] > CC_BOND*RAD_THREE) {
      tcut = PI/6.0-asin(CC_BOND*RAD_THREE/(2.0*cyl[0]));
      fcut = PI/(PI-2*tcut);
      if((cyl[1] < 2*tcut && cyl[1] > 0) || cyl[1] < -PI + 2*tcut) C->atom_array[i].type = 0;
      if(cyl[1] < 0) {
        cyl[1] *= fcut;
        cyl[0] -= 0.8*cyl[1]/PI;  //fudge factor
        cyl[1] += cyl[1]/PI*0.01; //fudge factor
      }
      else {
        cyl[1] -= (PI-cyl[1])*2*tcut/(PI-2*tcut);
        cyl[0] += 0.8*(PI-cyl[1])/PI;  //fudge factor 
        cyl[1] -= (PI-cyl[1])/PI*0.01; //fudge factor
      }
      cyl[0] *= 0.6215;
      cyl[2] = 1.118*cyl[0]-0.5;//sqrt(5.0/9.0)*cyl[0];
    }
    else {
      cyl[2] += 1.0; //fudge factor
      cyl[0] *= 0.8;
      immune = 1;
    }
    cyl_to_cart(cyl, cart);
    if(immune = 0) {
      C->atom_array[i].x[0] = cart[0]*1.4; //fudge factor
      C->atom_array[i].x[1] = cart[1]*1.4; //fudge factor
      C->atom_array[i].x[2] = cart[2]*1.4; //fudge factor
    }
    else {
      C->atom_array[i].x[0] = cart[0];
      C->atom_array[i].x[1] = cart[1];
      C->atom_array[i].x[2] = cart[2];
    }
  }
}

void curl(cell *C) {
  double cart[3], cyl[3];
  int i;
  for( i = 0; i < C->atoms; i++) {
    //load the graphene
    cart[0] = C->atom_array[i].x[0];
    cart[1] = C->atom_array[i].x[1];
    cart[2] = C->atom_array[i].x[2];
    //rotate it according to "section"
    cart_to_cyl(cart, cyl);
    cyl[1] += PI/6.0;
    if(cyl[1] > PI) cyl[1] -= 2.0*PI;
    //make it into a cone
    if(cyl[1] > 5.0*PI/6.0 || cyl[1] < -5.0*PI/6.0) C->atom_array[i].type = 0;
    cyl[1] *= 1.2;
    cyl[2] = sqrt(11.0/36.0)*cyl[0];
    cyl[0] *= 5.0/6.0;
    //store the result
    cyl_to_cart(cyl, cart);
    C->atom_array[i].x[0] = cart[0];
    C->atom_array[i].x[1] = cart[1];
    C->atom_array[i].x[2] = cart[2];
  }
}

void trim(cell *C, int krange) {
  int i;
  double rsquared;
  for(i = 0; i < C->atoms; i++) {
    rsquared = C->atom_array[i].x[0]*C->atom_array[i].x[0] +
               C->atom_array[i].x[1]*C->atom_array[i].x[1];
    if(rsquared > CC_BOND*CC_BOND*(double)krange*(double)krange)
      C->atom_array[i].type = 0;
  }
}

void trim2(cell *C, int krange) {
  int i;
  for(i = 0; i < C->atoms; i++)
    if(C->atom_array[i].x[2] > 1.5*krange)
      C->atom_array[i].type = 0;
}

void trim3(cell *C, int krange) {
  int i;
  for(i = 0; i < C->atoms; i++)
    if(C->atom_array[i].x[2] > 1.8*krange)
      C->atom_array[i].type = 0;
}


void cleanup(cell *C) {
  atom *new_array;
  double *new_x;
  int i, j, k;
  int kept = 0;
  for(i = 0; i < C->atoms; i++) if(C->atom_array[i].type != 0) kept++;
  new_array = (atom *)malloc(kept*sizeof(atom));
  if(new_array == NULL) printf("Unable to allocate cleaned-up atom array!\n");
  new_x = (double *)malloc(3*kept*sizeof(double));
  if(new_x == NULL) printf("Unable to allocate cleaned-up coordinates!\n");
  for(i = 0; i < kept; i++) new_array[i].x = new_x + 3*i;
  j = 0;
  for(i = 0; i < C->atoms; i++) {
    if(C->atom_array[i].type != 0) {
      new_array[j].x[0] = C->atom_array[i].x[0];
      new_array[j].x[1] = C->atom_array[i].x[1];
      new_array[j].x[2] = C->atom_array[i].x[2];
      new_array[j].type = C->atom_array[i].type;
      new_array[j].neighbs = C->atom_array[i].neighbs;
      for(k = 0; k < new_array[j].neighbs; k++)
        new_array[j].neighb_i[k] = C->atom_array[i].neighb_i[k];
      j++;
    }
  }
  free(C->atom_array);
  free(C->x);
  C->atom_array = new_array;
  C->x = new_x;
  C->atoms = kept;
}


void scale(cell *C, double factor) {
  int i;
  for(i = 0; i < C->atoms; i++) {
    C->atom_array[i].x[0] *= factor;
    C->atom_array[i].x[1] *= factor;
    C->atom_array[i].x[2] *= factor;
  }
}      
  

int two_carbon(const cell *C) {
  int i;
  int returnval = 0;
  for(i = 0; i < C->atoms; i++)
    if(C->atom_array[i].type == 1 && C->atom_array[i].neighbs == 2) returnval++;
  return returnval;
}


void hydrogenate(cell *C) {
  int atom_num;
  int new_atom;
  int q;
  double dx[3];
  double mag_delta;
  double scale_factor;
  for(atom_num = 0; atom_num < C->atoms; atom_num++) {
    if(C->atom_array[atom_num].type == 1) {
      if(C->atom_array[atom_num].neighbs == 1) {
        C->atom_array[atom_num].type = 2; 
        for(q = 0; q < 3; q++) {
          dx[q] = (CH_BOND/CC_BOND)*(C->atom_array[atom_num].x[q] - 
            C->atom_array[C->atom_array[atom_num].neighb_i[0]].x[q]);
          C->atom_array[atom_num].x[q] =
            C->atom_array[C->atom_array[atom_num].neighb_i[0]].x[q] + dx[q];
        }
      }
      if(C->atom_array[atom_num].neighbs == 2) {
        new_atom = 0;
        for(q = 0; q < 3; q++) {
          dx[q] = 2*C->atom_array[atom_num].x[q] -
            C->atom_array[C->atom_array[atom_num].neighb_i[0]].x[q] -
            C->atom_array[C->atom_array[atom_num].neighb_i[1]].x[q];
        } 
        mag_delta = sqrt(dx[0]*dx[0] + dx[1]*dx[1] + dx[2]*dx[2]);
        scale_factor = CH_BOND/mag_delta;
        while(C->atom_array[new_atom].type != 0) new_atom++;
        C->atom_array[new_atom].type = 2;
        for(q = 0; q < 3; q++)
          C->atom_array[new_atom].x[q] = C->atom_array[atom_num].x[q] + scale_factor*dx[q];
        C->atom_array[new_atom].neighbs = 1;
        C->atom_array[new_atom].neighb_i[0] = &C->atom_array[atom_num] - C->atom_array;
        C->atom_array[atom_num].neighb_i[2] = &C->atom_array[new_atom] - C->atom_array;
        C->atom_array[atom_num].neighbs = 3; 
      }
    }
  }
}


int main(int argc, char *argv[]) {
  FILE *fp;
  int krange;
  int digits;
  int pentagons;
  int i;
  int newatoms;
  double tb_energy;
  cell C;
  defs d;
  diags dia;
  if(argc != 5) {
    printf("Usage: cones <size> <digits> <pentagons> filename.abc\n");
    return EXIT_FAILURE;
  }
  fp = fopen(argv[4],"w");
  if(fp == NULL) {
    printf("Unable to create %s\n", argv[4]);
    return EXIT_FAILURE;
  }
  sscanf(argv[1], "%i", &krange);
  sscanf(argv[2], "%i", &digits);
  sscanf(argv[3], "%i", &pentagons);
  if(pentagons != 1 && pentagons != 2 && pentagons != 3) {
    printf("Only one, two, and three pentagon cones are currently supported.\n");
    return EXIT_FAILURE;
  }
  //generate a graphene sheet
  graphene(&C, krange);
  //curl it into a cone
  if(pentagons == 1) {
    curl(&C);
    trim(&C, krange);
  }
  if(pentagons == 2) {
    curl2(&C);
    trim2(&C, krange);
    scale(&C, 1.1);
  }
  if(pentagons == 3) {
    curl3(&C);
    trim3(&C, krange);
    scale(&C, 1.3);
  }
  //remove the excess atoms
  cleanup(&C);
  //find the neighbors
  find_neighb_i(&C);
  //cap the dangling bonds with hydrogen
  newatoms = two_carbon(&C);
  C.atom_array = (atom *)realloc(C.atom_array, (C.atoms + newatoms)*sizeof(atom));
  if(C.atom_array == NULL) printf("Unable to reallocate atom array!\n");
  C.x = (double *)realloc(C.x, 3*(C.atoms + newatoms)*sizeof(double));
  if(C.x == NULL) printf("Unable to reallocate coordinates!\n");
  for(i = 0; i < C.atoms + newatoms; i++) C.atom_array[i].x = C.x + 3*i; 
  for(i = C.atoms; i < C.atoms + newatoms; i++) C.atom_array[i].type = 0;
  C.atoms += newatoms;
  hydrogenate(&C);
  saveabc(fp, &C, "COMM 6 created by cones\n", digits, 0);
  fclose(fp);
  //stop here if relaxed structure is not desired
  return EXIT_SUCCESS;
  //destroy_cell(&C); //assumes that f_x and f_tb are allocated
  free(C.atom_array);
  free(C.x);
  C.atoms = 0;
  C.bonds = 0;
  C.allocated = 0;
  //relax the structure
  fp = fopen(argv[4], "r");
  loadabc(fp, &C, 0);
  fclose(fp);
  fp = fopen(argv[4], "w");
  tbinit(&C, 0);  
  printf("atoms = %i\tbonds = %i\n", C.atoms, C.bonds);
  d.start   = 0;
  d.stop    = 0;
  d.offset  = 0; 
  d.verbose = 0;
  d.digits  = digits;
  strcpy(d.logfile, "conemaker.log");
  d.automatic = 0;
  d.threshold = 1e-6;
  d.f_num = 0;
  d.tip = 0;
  calculate(&C, &d, &dia, 0);
  find_neighb_i(&C); //update bonding
  tb_energy = energy(C.bond_array, C.bonds);
  printf("iterations = %i\tenergy = %f\n", dia.iter, tb_energy);
  //output to disk
  saveabc(fp, &C, "COMM 6 Created by cones\n", digits, 0);
  fclose(fp);
  destroy_cell(&C);
  return EXIT_SUCCESS;
}
