/* ----------------------------------------------------------------------
   LAMMPS - Large-scale Atomic/Molecular Massively Parallel Simulator
   http://lammps.sandia.gov, Sandia National Laboratories
   Steve Plimpton, sjplimp@sandia.gov

   Copyright (2003) Sandia Corporation.  Under the terms of Contract
   DE-AC04-94AL85000 with Sandia Corporation, the U.S. Government retains
   certain rights in this software.  This software is distributed under 
   the GNU General Public License.

   See the README file in the top-level LAMMPS directory.
------------------------------------------------------------------------- */

#include "comm.h"
#include "math.h"
#include "string.h"
#include "stdio.h"
#include "stdlib.h"
#include "compute_tc.h"
#include "atom.h"
#include "input.h"
#include "variable.h"
#include "update.h"
#include "force.h"
#include "domain.h"
#include "group.h"
#include "modify.h"
#include "memory.h"
#include "error.h"

using namespace LAMMPS_NS;

#define INVOKED_SCALAR 1
#define INVOKED_VECTOR 2

/* ---------------------------------------------------------------------- */

ComputeTC::ComputeTC(LAMMPS *lmp, int narg, char **arg) :
  Compute(lmp, narg, arg)
{
  if (narg != 11 && narg != 12) error->all("Illegal compute tc command");

  scalar_flag = 1;
  extscalar = 0; 

  int n;
  if (strncmp(arg[3],"c_",2) == 0) {
    n = strlen(arg[3]);
    id_temp = new char[n];
    strcpy(id_temp,&arg[3][2]);
  } else {
    error->all("Could not find c_ for compute temp ID in compute tc command");
  }  

  if (strncmp(arg[4],"c_",2) == 0) {
    n = strlen(arg[4]);
    id_flux = new char[n];
    strcpy(id_flux,&arg[4][2]);
  } else {
    error->all("Could not find c_ for compute heat/flux ID in compute tc command");
  }  

  if (strncmp(arg[5],"v_",2) == 0) {
    n = strlen(arg[5]);
    id_variable1 = new char[n];
    strcpy(id_variable1,&arg[5][2]);
  } else {
    error->all("Could not find v_ for variable1 ID in compute tc command");
  }

  if (strncmp(arg[6],"v_",2) == 0) {
    n = strlen(arg[6]);
    id_variable2 = new char[n];
    strcpy(id_variable2,&arg[6][2]);
  } else {
    error->all("Could not find v_ for variable2 ID in compute tc command");
  }

  n = strlen(arg[7]) + 1;
  direction = new char[n];
  strcpy(direction,arg[7]);
  if (strcmp(direction,"x") != 0 && strcmp(direction,"y") != 0 && strcmp(direction,"z") != 0 && strcmp(direction,"iso") != 0) error->all("Wrong direction value for compute tc command");

  n = strlen(arg[8]) + 1;
  portion = new char[n];
  strcpy(portion,arg[8]);
  if (strcmp(portion,"first") != 0 && strcmp(portion,"second") != 0) error->all("Wrong heat flux portion value for compute tc command");

  mbig = atoi(arg[9]);
  if (mbig <= 0) error->all("Wrong mbig value for compute tc command"); 

  tc_ntimestep_start = atoi(arg[10]);
  if (tc_ntimestep_start < 0) error->all("Wrong tc_ntimestep_start value for compute tc command");

  if(narg == 12) tc_ntimestep_frequency = atoi(arg[11]);
  else tc_ntimestep_frequency = 100000;

  int itemp = modify->find_compute(id_temp);
  if (itemp < 0)
    error->all("Could not find compute temp ID for compute tc command");
  if (strcmp(modify->compute[itemp]->style,"temp") != 0)
    error->all("Compute tc compute ID does not compute temp");

  int iflux = modify->find_compute(id_flux);
  if (iflux < 0)
    error->all("Could not find compute heat/flux ID for compute tc command");
  if (strcmp(modify->compute[iflux]->style,"heat/flux") != 0)
    error->all("Compute tc compute ID does not compute heat/flux");

  int ivariable1 = input->variable->find(id_variable1);
  if (ivariable1 < 0) error->all("Could not find variable1 ID for compute tc command");
  if (input->variable->equalstyle(ivariable1) == 0)
    error->all("Compute tc variable1 is not equal-style variable");

  int ivariable2 = input->variable->find(id_variable2);
  if (ivariable2 < 0) error->all("Could not find variable2 ID for compute tc command");
  if (input->variable->equalstyle(ivariable2) == 0)
    error->all("Compute tc variable2 is not equal-style variable");

  ac = NULL;
  acc = NULL;
  tcc = NULL;
  jflux = NULL;

  thermal_conductivity = 0.0;
}

/* ---------------------------------------------------------------------- */

ComputeTC::~ComputeTC()
{ 
  delete [] id_temp;
  delete [] id_flux;
  delete [] id_variable1;
  delete [] id_variable2;
  delete [] direction;
  delete [] portion;
  memory->destroy_1d_double_array(ac, 1);
  memory->destroy_1d_double_array(acc, 1);
  memory->destroy_1d_double_array(tcc, 1);
  memory->destroy_2d_double_array(jflux,1);
}

/* ---------------------------------------------------------------------- */

void ComputeTC::init()
{
  // error checks

  int itemp = modify->find_compute(id_temp);
  if (itemp < 0)
    error->all("Could not find compute temp ID for compute tc command");
  c_temp = modify->compute[itemp];

  int iflux = modify->find_compute(id_flux);
  if (iflux < 0)
    error->all("Could not find compute heat/flux ID for compute tc command");
  c_flux = modify->compute[iflux];

  int ivariable1 = input->variable->find(id_variable1);
  if (ivariable1 < 0) error->all("Could not find variable1 ID for compute tc command");

  int ivariable2 = input->variable->find(id_variable2);
  if (ivariable2 < 0) error->all("Could not find variable2 ID for compute tc command");

}

/* ---------------------------------------------------------------------- */

double ComputeTC::compute_scalar()
{
  invoked_scalar = update->ntimestep;

  int ntimestep = update->ntimestep;
  int nsteps = update->nsteps;
  if (ntimestep == 0) {
    memory->destroy_1d_double_array(ac, 1);
    memory->destroy_1d_double_array(acc, 1);
    memory->destroy_1d_double_array(tcc, 1);
    memory->destroy_2d_double_array(jflux,1);
    jflux = memory->create_2d_double_array(6,1,mbig+1,"compute tc:jflux");
    ac    = memory->create_1d_double_array(1,mbig,"compute tc:ac");
    acc   = memory->create_1d_double_array(1,mbig,"compute tc:acc");
    tcc   = memory->create_1d_double_array(1,mbig,"compute tc:tcc");
    for (int m = 1; m <= mbig; m++) {
      ac[m]  = 0.0;
      acc[m] = 0.0;
      tcc[m] = 0.0;
    }
  }

  // invoke 2 computes if they haven't been already

  if (!(c_temp->invoked_flag & INVOKED_SCALAR)) {
    c_temp->compute_scalar();
    c_temp->invoked_flag |= INVOKED_SCALAR;
  }
  if (!(c_flux->invoked_flag & INVOKED_VECTOR)) {
    c_flux->compute_vector();
    c_flux->invoked_flag |= INVOKED_VECTOR;
  }

  double temp = c_temp->scalar;
  double *jt  = c_flux->vector;

  factor1 = input->variable->compute_equal(0);
  factor2 = input->variable->compute_equal(1);

  if (strcmp(portion,"first") == 0) {
    if (strcmp(direction,"x") == 0) {
      if (ntimestep >= 1 && ntimestep <= mbig+1) {
        jflux[0][ntimestep] = jt[0];
        jflux[1][ntimestep] = jt[1];
        jflux[2][ntimestep] = jt[2];
        for (int m = 1; m <= ntimestep-1; m++) {    
          ac[m] += jflux[0][ntimestep-m]*jflux[0][ntimestep];
        }
      } else if (ntimestep > mbig+1) {
        for (int k = 1; k <= mbig; k++) {
          jflux[0][k] = jflux[0][k+1];
          jflux[1][k] = jflux[1][k+1];
          jflux[2][k] = jflux[2][k+1];
        }
        jflux[0][mbig+1] = jt[0];
        jflux[1][mbig+1] = jt[1];
        jflux[2][mbig+1] = jt[2];
        for (int m = 1; m <= mbig; m++) {
          ac[m] += jflux[0][mbig-m+1]*jflux[0][mbig+1];
        }
      }
    } else if (strcmp(direction,"y") == 0) {
      if (ntimestep >= 1 && ntimestep <= mbig+1) {
        jflux[0][ntimestep] = jt[0];
        jflux[1][ntimestep] = jt[1];
        jflux[2][ntimestep] = jt[2];
        for (int m = 1; m <= ntimestep-1; m++) {    
          ac[m] += jflux[1][ntimestep-m]*jflux[1][ntimestep];
        }
      } else if (ntimestep > mbig+1) {
        for (int k = 1; k <= mbig; k++) {
          jflux[0][k] = jflux[0][k+1];
          jflux[1][k] = jflux[1][k+1];
          jflux[2][k] = jflux[2][k+1];
        }
        jflux[0][mbig+1] = jt[0];
        jflux[1][mbig+1] = jt[1];
        jflux[2][mbig+1] = jt[2];
        for (int m = 1; m <= mbig; m++) {
          ac[m] += jflux[1][mbig-m+1]*jflux[1][mbig+1];
        }
      }
    } else if (strcmp(direction,"z") == 0) {
      if (ntimestep >= 1 && ntimestep <= mbig+1) {
        jflux[0][ntimestep] = jt[0];
        jflux[1][ntimestep] = jt[1];
        jflux[2][ntimestep] = jt[2];
        for (int m = 1; m <= ntimestep-1; m++) {    
          ac[m] += jflux[2][ntimestep-m]*jflux[2][ntimestep];
        }
      } else if (ntimestep > mbig+1) {
        for (int k = 1; k <= mbig; k++) {
          jflux[0][k] = jflux[0][k+1];
          jflux[1][k] = jflux[1][k+1];
          jflux[2][k] = jflux[2][k+1];
        }
        jflux[0][mbig+1] = jt[0];
        jflux[1][mbig+1] = jt[1];
        jflux[2][mbig+1] = jt[2];
        for (int m = 1; m <= mbig; m++) {
          ac[m] += jflux[2][mbig-m+1]*jflux[2][mbig+1];
        }
      }
    } else if (strcmp(direction,"iso") == 0) {
      if (ntimestep >= 1 && ntimestep <= mbig+1) {
        jflux[0][ntimestep] = jt[0];
        jflux[1][ntimestep] = jt[1];
        jflux[2][ntimestep] = jt[2];
        for (int m = 1; m <= ntimestep-1; m++) {    
          ac[m] += jflux[0][ntimestep-m]*jflux[0][ntimestep] + jflux[1][ntimestep-m]*jflux[1][ntimestep] + jflux[2][ntimestep-m]*jflux[2][ntimestep];
        }
      } else if (ntimestep > mbig+1) {
        for (int k = 1; k <= mbig; k++) {
          jflux[0][k] = jflux[0][k+1];
          jflux[1][k] = jflux[1][k+1];
          jflux[2][k] = jflux[2][k+1];
        }
        jflux[0][mbig+1] = jt[0];
        jflux[1][mbig+1] = jt[1];
        jflux[2][mbig+1] = jt[2];
        for (int m = 1; m <= mbig; m++) {
          ac[m] += jflux[0][mbig-m+1]*jflux[0][mbig+1] + jflux[1][mbig-m+1]*jflux[1][mbig+1] + jflux[2][mbig-m+1]*jflux[2][mbig+1];
        }
      }
    }
  } else if (strcmp(portion,"second") == 0) {
    if (strcmp(direction,"x") == 0) {
      if (ntimestep >= 1 && ntimestep <= mbig+1) {
        jflux[3][ntimestep] = jt[3];
        jflux[4][ntimestep] = jt[4];
        jflux[5][ntimestep] = jt[5];
        for (int m = 1; m <= ntimestep-1; m++) {    
          ac[m] += jflux[3][ntimestep-m]*jflux[3][ntimestep];
        }
      } else if (ntimestep > mbig+1) {
        for (int k = 1; k <= mbig; k++) {
          jflux[3][k] = jflux[3][k+1];
          jflux[4][k] = jflux[4][k+1];
          jflux[5][k] = jflux[5][k+1];
        }
        jflux[3][mbig+1] = jt[3];
        jflux[4][mbig+1] = jt[4];
        jflux[5][mbig+1] = jt[5];
        for (int m = 1; m <= mbig; m++) {
          ac[m] += jflux[3][mbig-m+1]*jflux[3][mbig+1];
        }
      }
    } else if (strcmp(direction,"y") == 0) {
      if (ntimestep >= 1 && ntimestep <= mbig+1) {
        jflux[3][ntimestep] = jt[3];
        jflux[4][ntimestep] = jt[4];
        jflux[5][ntimestep] = jt[5];
        for (int m = 1; m <= ntimestep-1; m++) {    
          ac[m] += jflux[4][ntimestep-m]*jflux[4][ntimestep];
        }
      } else if (ntimestep > mbig+1) {
        for (int k = 1; k <= mbig; k++) {
          jflux[3][k] = jflux[3][k+1];
          jflux[4][k] = jflux[4][k+1];
          jflux[5][k] = jflux[5][k+1];
        }
        jflux[3][mbig+1] = jt[3];
        jflux[4][mbig+1] = jt[4];
        jflux[5][mbig+1] = jt[5];
        for (int m = 1; m <= mbig; m++) {
          ac[m] += jflux[4][mbig-m+1]*jflux[4][mbig+1];
        }
      }
    } else if (strcmp(direction,"z") == 0) {
      if (ntimestep >= 1 && ntimestep <= mbig+1) {
        jflux[3][ntimestep] = jt[3];
        jflux[4][ntimestep] = jt[4];
        jflux[5][ntimestep] = jt[5];
        for (int m = 1; m <= ntimestep-1; m++) {    
          ac[m] += jflux[5][ntimestep-m]*jflux[5][ntimestep];
        }
      } else if (ntimestep > mbig+1) {
        for (int k = 1; k <= mbig; k++) {
          jflux[3][k] = jflux[3][k+1];
          jflux[4][k] = jflux[4][k+1];
          jflux[5][k] = jflux[5][k+1];
        }
        jflux[3][mbig+1] = jt[3];
        jflux[4][mbig+1] = jt[4];
        jflux[5][mbig+1] = jt[5];
        for (int m = 1; m <= mbig; m++) {
          ac[m] += jflux[5][mbig-m+1]*jflux[5][mbig+1];
        }
      }
    } else if (strcmp(direction,"iso") == 0) {
      if (ntimestep >= 1 && ntimestep <= mbig+1) {
        jflux[3][ntimestep] = jt[3];
        jflux[4][ntimestep] = jt[4];
        jflux[5][ntimestep] = jt[5];
        for (int m = 1; m <= ntimestep-1; m++) {    
          ac[m] += jflux[3][ntimestep-m]*jflux[3][ntimestep] + jflux[4][ntimestep-m]*jflux[4][ntimestep] + jflux[5][ntimestep-m]*jflux[5][ntimestep];
        }
      } else if (ntimestep > mbig+1) {
        for (int k = 1; k <= mbig; k++) {
          jflux[3][k] = jflux[3][k+1];
          jflux[4][k] = jflux[4][k+1];
          jflux[5][k] = jflux[5][k+1];
        }
        jflux[3][mbig+1] = jt[3];
        jflux[4][mbig+1] = jt[4];
        jflux[5][mbig+1] = jt[5];
        for (int m = 1; m <= mbig; m++) {
          ac[m] += jflux[3][mbig-m+1]*jflux[3][mbig+1] + jflux[4][mbig-m+1]*jflux[4][mbig+1] + jflux[5][mbig-m+1]*jflux[5][mbig+1];
        }
      }
    }
  }

  for (int m = 1; m <= mbig; m++) {
    acc[m] = ac[m];
  }

  double volume = domain->xprd * domain->yprd * domain->zprd;
  if(strcmp(direction,"iso") == 0) {
    factor = update->dt/3.0/force->boltz/volume/temp/temp;
  } else {
    factor = update->dt/force->boltz/volume/temp/temp;
  }

  if (ntimestep >= 2 && ntimestep <= mbig+1) {
    for (int m = 1; m <= ntimestep-1; m++) {
      acc[m] /= (ntimestep - m);
      if(m==1) tcc[m] = acc[m] * factor;
      if(m>1) tcc[m] = tcc[m-1] + acc[m] * factor;
    }
    thermal_conductivity = tcc[ntimestep-1];
  } else if (ntimestep > mbig+1) {
    for (int m = 1; m <= mbig; m++) {
      acc[m] /= (ntimestep - m);
      if(m==1) tcc[m] = acc[m] * factor;
      if(m>1) tcc[m] = tcc[m-1] + acc[m] * factor;
    }
    thermal_conductivity = tcc[mbig];
  }
  if (comm->me == 0) {
    if(ntimestep>=tc_ntimestep_start && ntimestep%tc_ntimestep_frequency==0 || ntimestep==nsteps) {
      FILE *fp;
      char *str1 = new char[10];
      char *space = new char[2];
      sprintf(str1, "%s%d%s","ac",ntimestep,".dat");    
      sprintf(space, "  "); 
      int n = strlen(str1) + 1;
      char *str = new char[n];    
      strcpy(str,str1);
      fp = fopen(str,"w");
      if (fp == NULL) error->one("Cannot open ac.dat");
      for (int m = 1; m <= mbig; m++) { 
        fprintf(fp,"%d%s%20.10f\n",m,space,acc[m]*factor1);
      }
      fclose(fp);

      sprintf(str1, "%s%d%s","tc",ntimestep,".dat"); 
      strcpy(str,str1);
      fp = fopen(str,"w");
      if (fp == NULL) error->one("Cannot open tc.dat");
      for (int m = 1; m <= mbig; m++) { 
        fprintf(fp,"%d%s%20.10f\n",m,space,tcc[m]*factor2);
      }
      fclose(fp);
    }
  }
  scalar = thermal_conductivity*factor2;
  return scalar;
}