/*----------------------------calculate.c--------------------------
  This file contains routines for calculating taking the potential
  energy of the molecule, and the forces on each atom, which are
  calculated using the Tersoff-Brenner potential and feeding these
  to the conjugate gradient code, in order to allow the atoms to
  relax to their minimum energy configuration. This software was
  written by Stephen Jordan.
  -----------------------------------------------------------------*/

#include "calculate.h"
#include "cell.h"
#include "defs.h"
#include "diags.h"
#include "potential.h"
#include "force.h"
#include "conjgrad.h"
#include <stdio.h>
#include <math.h>
#include <malloc.h>

cell *cp;
defs *dp;
diags *diap;
int stageg;

double func(double *x) {
  double returnval;
  int i, q;
  double *xtmp;
  xtmp = (double *)malloc(3*cp->atoms*sizeof(double));
  if(xtmp == NULL) printf("Unable to allocate temporary coordinates!\n");
  for(i = 0; i < 3*cp->atoms; i++) {
    xtmp[i] = cp->x[i];
    cp->x[i] = x[i];
  }
  if(dp->relative[0] == 1 || dp->relative[1] == 1 || dp->relative[2] == 1) {
    cm(cp, diap, dp);
    for(i = 0; i < cp->atoms; i++) {
      for(q = 0; q < 3; q++) {
        if(dp->relative[q] == 1) cp->atom_array[i].x[q] -= diap->cm[q];
      }
    }
  }
  returnval = energy(cp->bond_array, cp->bonds);
  energy_tally(cp, dp, diap, stageg);
  for(i = 0; i < dp->f_num; i++) returnval += diap->obj_energy[i];
  for(i = 0; i < 3*cp->atoms; i++) cp->x[i] = xtmp[i];
  free(xtmp);
  return returnval;
}


void dfunc(double *x, double *d) {
  int i, q;
  double *xtmp;
  xtmp = (double *)malloc(3*cp->atoms*sizeof(double));
  if(xtmp == NULL) printf("Unable to allocate temporary coordinates!\n");
  for(i = 0; i < 3*cp->atoms; i++) {
    xtmp[i] = cp->x[i];
    cp->x[i] = x[i];
  }
  if(dp->relative[0] == 1 || dp->relative[1] == 1 || dp->relative[2] == 1) {
    cm(cp, diap, dp);
    for(i = 0; i < cp->atoms; i++) {
      for(q = 0; q < 3; q++) {
        if(dp->relative[q] == 1) cp->atom_array[i].x[q] -= diap->cm[q];
      }
    }
  }
  nftb(cp);  
  for(i = 0; i < cp->atoms; i++) f_ext(cp, dp, diap, i, cp->atom_array[i].f_x, stageg);
  for(i = 0; i < 3*cp->atoms; i++) d[i] = -cp->f_x[i] - cp->f_tb[i]; 
  for(i = 0; i < 3*cp->atoms; i++) cp->x[i] = xtmp[i];
  free(xtmp);
}

void calculate(cell *C, defs *d, diags *dia, int stage) {
  double final_e;
  int i;
  double f[3];
  double ff;
  cp = C;
  dp = d;
  diap = dia;
  stageg = stage;
  conjgrad(C->x, 3*C->atoms, d->threshold, &(dia->iter), &dia->tb_energy, func, dfunc);
  //bring f_tb and f_x up to date and calculate max_f
  dia->max_f = 0;
  nftb(cp);
  for(i = 0; i < cp->atoms; i++) {
    f_ext(C, d, dia, i, C->atom_array[i].f_x, stage);
    f[0] = C->atom_array[i].f_tb[0] + C->atom_array[i].f_x[0];
    f[1] = C->atom_array[i].f_tb[1] + C->atom_array[i].f_x[1];
    f[2] = C->atom_array[i].f_tb[2] + C->atom_array[i].f_x[2];
    ff = f[0]*f[0] + f[1]*f[1] + f[2]*f[2];
    if(ff > dia->max_f) dia->max_f = ff; //more efficient to sqrt outside loop
  }
  dia->max_f = sqrt(dia->max_f);
  cm(C, dia, d);
}

double mass(const cell *C, int i, const defs *d) {
  //mass of Carbon 12 in eV*S^2/A^2:
  if(C->atom_array[i].type == 1) return 1.2437e-27; 
  //mass of Hydrogen 1 in eV*S^2/A^2:
  if(C->atom_array[i].type == 2) return 1.0364e-28; 
  //otherwise
  msglog("Invalid type in mass!\n", d);
  return 1.0;
}
 

//move the cone so that the x and y components of the
//center of mass are zero
void center(cell *C, const diags *dia) {
  int i;
  for(i = 0; i < C->atoms; i++) {
    C->atom_array[i].x[0] -= dia->cm[0];
    C->atom_array[i].x[1] -= dia->cm[1];
  }
}

void cm(const cell *C, diags *dia, const defs *d) {
  int i;
  double mi;
  double total_mass;
  total_mass = 0.0;
  dia->cm[0] = 0.0;
  dia->cm[1] = 0.0;
  dia->cm[2] = 0.0;
  for(i = 0; i < C->atoms; i++) {
    mi = mass(C, i, d);
    dia->cm[0] += C->atom_array[i].x[0]*mi;
    dia->cm[1] += C->atom_array[i].x[1]*mi;
    dia->cm[2] += C->atom_array[i].x[2]*mi;
    total_mass += mi;
  }
  dia->cm[0] /= total_mass;
  dia->cm[1] /= total_mass;
  dia->cm[2] /= total_mass;
}
