#include "fix_orient_crystal.h"

#include "fix.h"
#include "atom.h"
#include "domain.h"
#include "region.h"
#include "input.h"
#include "variable.h"
#include "neighbor.h"
#include "force.h"
#include "error.h"
#include "memory.h"
#include "comm.h"
#include "math_const.h"

#include "update.h"
#include <cmath>
#include <cstring>
using namespace LAMMPS_NS;
using namespace FixConst;
using namespace MathConst;
//using namespace MathConst;

FixOrientCrystal::FixOrientCrystal(LAMMPS *lmp, int narg, char **arg) :
  Fix(lmp, narg, arg), iregion(nullptr), idregion(nullptr)
{
  debug_flag=0;
  scalar_flag = 1;         
  extscalar = 1;  // 0 =  intensive, 1 = extensive
  ntypes = 0;
  params = nullptr;

  int iarg = 3;
  iregion = domain->get_region_by_id(arg[iarg]);
  if (!iregion) error->all(FLERR,"Region ID {} for fix orient/crystal does not exist", arg[iarg]);
  idregion = utils::strdup(arg[iarg]);
  iarg++;
  if(debug_flag){
    fprintf(screen, "FIX_ORIENT_CRYSTAL: initialized\n");
    fprintf(screen, "  region: %s\n", idregion);
    }
// 解析多个 type/k/r0 块
  while (iarg < narg) {
      if (strcmp(arg[iarg], "type") != 0) 
          error->all(FLERR, "Expected 'type' keyword, current = {}", arg[iarg]);
      iarg++;
      
      // 扩展数组
      ntypes++;
      params = (OrientParams*) memory->srealloc(params, ntypes*sizeof(OrientParams), "orient:params");
      auto& p = params[ntypes-1];
      
      // btype
      p.btype = utils::inumeric(FLERR, arg[iarg++], false, lmp);
      
      // k
      p.k = utils::numeric(FLERR, arg[iarg++], false, lmp);
      
      // theta0 (ADD)
      p.theta0 = utils::numeric(FLERR, arg[iarg++], false, lmp);
      
      // 预计算 c0, s0
      double theta0_rad = p.theta0 * DEG2RAD;
      p.c0 = cos(theta0_rad);
      p.s0 = sin(theta0_rad);
      
      // r0x r0y r0z
      p.r0[0] = utils::numeric(FLERR, arg[iarg++], false, lmp);
      p.r0[1] = utils::numeric(FLERR, arg[iarg++], false, lmp);
      p.r0[2] = utils::numeric(FLERR, arg[iarg++], false, lmp);
      
      // normalize r0
      double r0mag = sqrt(p.r0[0]*p.r0[0] + p.r0[1]*p.r0[1] + p.r0[2]*p.r0[2]);
      if (r0mag == 0) error->all(FLERR,"Zero crystal direction vector");
      p.r0[0] /= r0mag; p.r0[1] /= r0mag; p.r0[2] /= r0mag;
      if(debug_flag){
        fprintf(screen, 
          "DEBUG: secton: type=%d, preset angle %f, "
          "oreintation: (%f, %f, %f), e=%f,c=%f, flag=%d\n",
          p.btype, p.theta0, p.r0[0], p.r0[1], p.r0[2], p.c0,p.s0,iarg);
      }

      
  }
}

FixOrientCrystal::~FixOrientCrystal()
{
  delete[] idregion;  // only delete the name string
  // don't delete iregion - region is owned by domain
}

int FixOrientCrystal::setmask()
{
  int mask = 0;
  mask |= POST_FORCE;
  return mask;
}

void FixOrientCrystal::init()
{
  // Re-get region pointer (in case region was redefined since construction)
  iregion = domain->get_region_by_id(idregion);
  if (!iregion) error->all(FLERR,"Region ID {} for fix orient/crystal does not exist", idregion);

  // check bond type valid
  for (int it = 0; it < ntypes; it++) {
    if (params[it].btype < 1 || params[it].btype > atom->nbondtypes)
      error->all(FLERR,"Invalid bond type {} in fix orient/crystal", params[it].btype);
    // if(debug_flag)  
    //   fprintf(screen, "FIX_ORIENT_CRYSTAL: init done, region %s found\n", idregion);
  }
}

void FixOrientCrystal::setup(int vflag)
{
  post_force(vflag);
}

void FixOrientCrystal::post_force(int vflag)
{
  double **x = atom->x;
  double **f = atom->f;
  int **bondlist = neighbor->bondlist;
  int nbondlist = neighbor->nbondlist;
  int nlocal = atom->nlocal;

  int nmatch_type = 0;
  int nmatch_region = 0;
  int ndebug = 0;
  const int max_debug = 100;
  total_energy = 0.0;
  for (int n = 0; n < nbondlist; n++) {
    int i = bondlist[n][0];
    int j = bondlist[n][1];
    int type = bondlist[n][2];
    
    // Step 1: check bond type
    int iparam = -1;
    for (int it = 0; it < ntypes; it++) {
        if (params[it].btype == type) {
            iparam = it;
            break;
        }
    }
    if (iparam == -1) continue;

    nmatch_type++;
    auto& p = params[iparam];
    double k = p.k;
    double* r0 = p.r0;
    double c0_stored = p.c0;
    double s0 = p.s0;
    // if (i >= nlocal && j >= nlocal) {
    //   nghost_ghost++;
    //   continue;
    // }
    // Step 2: get bond vector with minimum image
    double dx = x[i][0] - x[j][0];
    double dy = x[i][1] - x[j][1];
    double dz = x[i][2] - x[j][2];
    domain->minimum_image(dx, dy, dz);
    
    // Step 3: bond midpoint (using j as base)
    double xmid[3];
    xmid[0] = x[j][0] + 0.5 * dx;
    xmid[1] = x[j][1] + 0.5 * dy;
    xmid[2] = x[j][2] + 0.5 * dz;
    // Step 4: remap midpoint to main box
    domain->remap(xmid);
    
    if (debug_flag && ndebug < max_debug) {
    
      fprintf(screen, 
        "DEBUG: timestep =" BIGINT_FORMAT "bond %d: atoms %d-%d , "
        "dx(%f, %f, %f),i =(%f, %f, %f),   j=(%f, %f, %f)\n",
        update->ntimestep,n, atom->tag[i], atom->tag[j], 
        dx, dy, dz, x[i][0],x[i][1],x[i][2],x[j][0],x[j][1],x[j][2]);
      ndebug++;
    }

    
    // Step 5: check region (use iregion->match like fix_deposit)
    if (!iregion->match(xmid[0], xmid[1], xmid[2])) continue;
    nmatch_region++;
    double rsq = dx*dx + dy*dy + dz*dz;
    double r = sqrt(rsq);
    if (r < 1e-10) continue; //should issue an warning of such short length bond?
    double rinv = 1.0 / r;
    
    double rx = dx * rinv;
    double ry = dy * rinv;
    double rz = dz * rinv;
    
    //calculate cos(theta) as c
    double c = r0[0]*rx + r0[1]*ry + r0[2]*rz;
    
    //inital energy and force
    double e_bond = 0;
    double fx = 0;
    double fy = 0;
    double fz = 0;

    if (p.theta0 > 1e-2) //in this case the theta0 is deviated from reference crystal direction
    {
          //calculate sin theta as s
        double s = 1.0- c*c;
        if (s < 0) s = 0;
        s = sqrt(s);

      //symtrical of reference of cos(theta) and cos(theta0), they should be in same sgin to minimize the energy
      //which means the reerence angle of theta0 and pi-theta0 is the same, beacus of the forward/backward direction of vector
      //for bond and crystal oreintation is meaningless.
      double c0 = fabs(c0_stored);
      if ( c < 0) {
        c0 = -1*fabs(c0_stored);
      }
      //caldulate cos(theta-theta0) as c_c0
      double c_c0 = c*c0+s*s0;
      //E_bond based on c_c0 as k*(1-c_c0^2)
      e_bond = k * (1.0 - c_c0*c_c0);
      
      // seperately trate the sin(theta) close to 0 and normal situation, use 0.001 as thershold?
      if (s>0.001) {
        //calculate dcos(theta)/dx,dy,dz
        double dc_dx = rinv*(r0[0]-c*rx);
        double dc_dy = rinv*(r0[1]-c*ry);
        double dc_dz = rinv*(r0[2]-c*rz);
        //calculate dsin/dx,dy,dz
        //which is d(1-cos^2(x))^0.5/dx or y or z
        //which is 0.5*(1-cos^2(x))^-0.5*-2*cos(theta)*dcos(theta)/dx
        //which is -dcos/dx*cos(theta)/sin(theta)
        double c_s= -1*c/s;
        double ds_dx = c_s*dc_dx;
        double ds_dy = c_s*dc_dy;
        double ds_dz = c_s*dc_dz;
        
        //calculate f=-de_bond/dx,dy,dz
        //which is k*(2*c_c0)*c_c0/dx,dy,dz
        double force_prefix = k*2*c_c0;
        //which is k*(2*c_c0)*(dc/dx*c0+ds_dx*s0)
        fx = force_prefix * (dc_dx*c0 + ds_dx * s0);
        fy = force_prefix * (dc_dy*c0 + ds_dy * s0);
        fz = force_prefix * (dc_dz*c0 + ds_dz * s0);
      }
      //s close to zero, which means is highly close to refrencing crystal oreintation direction
      //abort force implication for numerical stablility as dsin_dx based on 1/sin(theta)
      //else{
        //do nothing, fx,y,yz is already 0 e_bond is already calculate

      //}



    }
    //here we have theta0 is not deviated from reference crystal direction
    //energy backward to E=k*(1-cos^2(theta))
    else{
      e_bond = k * (1.0 - c*c);
      
      
      double fp = 2.0 * k * c * rinv;
      fx = fp * (r0[0] - c*rx);
      fy = fp * (r0[1] - c*ry);
      fz = fp * (r0[2] - c*rz);

    // if(debug_flag){
    //     fprintf(screen, 
    //       "DEBUG: theta0=0 branch "
          
    //     );
    //   }
    }


    
    total_energy += e_bond;
    f[i][0] += fx;
    f[i][1] += fy;
    f[i][2] += fz;

    
    f[j][0] -= fx;
    f[j][1] -= fy;
    f[j][2] -= fz;
    

    
    //Debug output
    if (debug_flag && ndebug < max_debug) {
      double intersection = (rx*fx+ry*fy+rz*fz)/sqrt(fx*fx+fy*fy+fz*fz);
      fprintf(screen, 
        "DEBUG: timestep =" BIGINT_FORMAT "bond %d: atoms %d-%d , "
        "bond_dirct =(%f, %f, %f),   f=(%f, %f, %f) ,f-b=%f\n",
        update->ntimestep,n, atom->tag[i], atom->tag[j], 
        rx,ry,rz,fx,fy,fz,intersection);
      ndebug++;
    }
  }
  
  // if (debug_flag) {
  //   fprintf(screen, 
  //     "FIX_ORIENT_CRYSTAL: total bonds %d, "
  //     "type match %d, region match %d\n",
  //     nbondlist,  nmatch_type, nmatch_region);
  // }

}

double FixOrientCrystal::compute_scalar()
{
  double total;
  MPI_Allreduce(&total_energy, &total, 1, MPI_DOUBLE, MPI_SUM, world);
  return total;
}