#include "fix_subsys_min.h"

#include "stdio.h"
#include "string.h"
#include "atom.h"
#include "force.h"
#include "update.h"
#include "respa.h"
#include "error.h"
#include "memory.h"
#include "modify.h"
#include <math.h>
#include "comm.h"
#include "atom_vec.h"
#include "min.h"
#include "output.h"
#include "group.h"
#include "thermo.h"
#include "compute_pe.h"

#include <set>
#include <vector>

#include "fix_drude.h"
#include "fix_setforce.h"

using namespace LAMMPS_NS;
using namespace FixConst;

enum{MAXITER,MAXEVAL,ETOL,FTOL,DOWNHILL,ZEROALPHA,ZEROFORCE,ZEROQUAD};

FixSubsysMin::FixSubsysMin(LAMMPS *lmp, int narg, char **arg) : Fix(lmp,narg,arg)
{
  fix_drude = NULL;

  if (narg < 6) error->all(FLERR,"Illegal fix Subsys/Min command");

  update->etol = force->numeric(FLERR,arg[3]);
  update->ftol = force->numeric(FLERR,arg[4]);
  mymaxiter = force->inumeric(FLERR,arg[5]);
  update->max_eval = force->inumeric(FLERR, arg[6]);

  printf("Subsys/Min: energy tolerance: %14.8f, force tol %14.8f, maxiter: %i, max_eval: %i\n", update->etol, update->ftol, mymaxiter, update->max_eval);

  if (update->etol < 0.0 || update->ftol < 0.0)
    error->all(FLERR,"Illegal fix Subsys/Min");

}

void FixSubsysMin::init()
{
    
  first_time = 1;

  //FU| create group from all - group->names[igroup] that shall be frozen during submin step
  char **grparg = new char*[4];

  printf("IGROUP is %i\n", igroup);
  grparg[0] = (char *) "FRZGRP";
  grparg[1] = (char *) "subtract";
  grparg[2] = (char *) "all";
  grparg[3] = group->names[igroup];

  group->assign(4,grparg);

  delete [] grparg;

  //FU| create fix that freezes all the dofs not of interested to us
  char **fixarg = new char*[6];
  fixarg[0] = (char *) "SETFORCE";
  fixarg[1] = (char *) "FRZGRP";
  fixarg[2] = (char *) "setforce";
  fixarg[3] = (char *) "0.";
  fixarg[4] = (char *) "0.";
  fixarg[5] = (char *) "0.";

  //FU| add fix only to min_post_force, do not want in the integration loop (need to be friends with modify class to be able to do this here)
  modify->add_fix(6,fixarg);
  delete [] fixarg;
  fix_setforce = (FixSetForce *) modify->fix[modify->nfix-1];

  modify->list_init(MIN_POST_FORCE,modify->n_min_post_force,modify->list_min_post_force);

}

void FixSubsysMin::pre_force(int vflag)
{

    printf("STARTING MINIMIZATION\n");

  update->whichflag = 2;

  //FU| This needs to be moved into, e.g., setup function
  if (first_time) {

    update->minimize->init();
    //FU| if not included always stopping because of non-downhill search direction
    update->minimize->setup(); // i.e., energy/forces not calculated without this

    //FU| in case of FIRE minimization we also need to save an intermittent copy of the velocities (because FIRE uses atom->v
    on_fire = 0;
    if ( strncmp(update->minimize_style, "fire", 2) == 0)
      on_fire = 1;

    if (on_fire) {
      allocate();
      save_velocities();
    }

    first_time = 0;

  }
  else {
  //FU| initial force calculation, as in update->minimize->setup, but without the additional stuff
  //FU| if not done minimization stopping due to non-downhill search direction
  //FU| and needed to get actual forces
    update->minimize->energy_force(1);
    update->minimize->ecurrent = update->minimize->pe_compute->compute_scalar();
    update->minimize->einitial = update->minimize->ecurrent;
    update->minimize->fnorm2_init = sqrt(update->minimize->fnorm_sqr());
    update->minimize->fnorminf_init = update->minimize->fnorm_inf();

    if ( on_fire ) save_velocities();
  }

  //FU| resetting minimization counters
  update->minimize->niter = 0;
  update->minimize->neval = 0;

  //FU| keep for reset
  int tsold = update->ntimestep;
  int toutp = output->next;

  //FU| avoid getting output every minimization step
  output->next = 0;

  //FU| perform the actual minimization
  stop_condition = update->minimize->iterate(mymaxiter);
  stopstr = update->minimize->stopstrings(stop_condition);

  //FU| reset timestep to what it was before minimization
  update->ntimestep = tsold;

  //FU| reset stuff that got changed in update->minimize->energy_force
  update->minimize->ev_set(update->ntimestep);
  output->next = toutp;

  //FU| the 
  if ( on_fire ) restore_velocities();

  printf("NUMBER OF ITERATIONS: %i\n", update->minimize->niter);
  printf("NUMBER OF FORCE EVALUATIONS: %i\n", update->minimize->neval);
  printf("STOPPING BECAUSE OF %s\n", stopstr);

  //FU| not sure if that's really needed
  update->whichflag = 1;

  printf("STOPPING MINIMIZATION\n");

  //FU| that doesn't seem to help in any way? how are velocities synchronized?

  // zero_subsys_vel();
  // comm->forward_comm();
}

void FixSubsysMin::zero_subsys_vel()
{
  double **v = atom->v;
  int nlocal = atom->nlocal;
  int *mask = atom->mask;

  for ( int i = 0; i<nlocal; i++ )
      if ( mask[i] & groupbit )
          for ( int j; j<3; j++ )
              v[i][j] = 0.;
}

int FixSubsysMin::setmask()
{
  int mask = 0;
  mask |= PRE_FORCE;
  return mask;
}

FixSubsysMin::~FixSubsysMin()
{
  if ( on_fire ) 
    memory->destroy(save_v);
}

void FixSubsysMin::save_velocities()
{

  //FU| check in communications what we need to do to be sure that we copy the right things
  //FU| maybe we can just set/reset a pointer for FIRE minimization, so that we don't overwrite/use the original atom->v data, but an empty array

  double **v = atom->v;

  copy_velocities(save_v, v);
}

void FixSubsysMin::restore_velocities()
{
    double **v = atom->v;
    copy_velocities(v, save_v);
    //FU| needed? again, how are velocities synchronized?
    comm->forward_comm();
}

void FixSubsysMin::copy_velocities(double **out, double **in)
{
  int nlocal = atom->nlocal;

  //FU| taken from min_fire.cpp setup_style

  for (int i = 0; i < nlocal; i++)
      for (int j = 0; j<3; j++ )
          out[i][j] = in[i][j];
}

void FixSubsysMin::allocate()
{
  int natoms = atom->natoms;
  memory->create(save_v, natoms, 3, "fix_subsys_min:save_v");
}