/* Convert a scalar data set to its vector gradient.
 * 
 * Copyright (c) 2023 MJ Rutter
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "c2xsf.h"

/* Doing this by FFTs is surely better, so this is here for testing
 * purposes.
 */

void grad_real(struct grid *g, struct unit_cell *c){
  int i,j,k,ic,size,ii_left[3],ii_right[3];
  double *data,*dptr1,*dptr2,*dptr3,scale[3],abc[6];
  char *name;
  
  if (debug) fprintf(stderr,"grad_real called\n");
  if ((!g)||(!g->data)) error_exit("No data passed to grad");
  if (g->comps!=1) error_exit("Can take grad of scalars only");

  if ((!c)||(!c->basis)) error_exit("need basis for grad");

  if ((fabs(c->basis[0][1])>tol*tol)||(fabs(c->basis[0][2])>tol*tol)||
      (fabs(c->basis[1][0])>tol*tol)||(fabs(c->basis[1][2])>tol*tol)||
      (fabs(c->basis[2][0])>tol*tol)||(fabs(c->basis[2][1])>tol*tol))
    error_exit("grad_real requires that a & x, b & y and c & z are parallel");

  
  size=g->size[0]*g->size[1]*g->size[2];
  data=malloc(3*size*sizeof(double));
  if (!data) error_exit("Malloc error for new grid in grad");

  basis2abc(c->basis,abc);
  for(i=0;i<3;i++)
    scale[i]=0.5*g->size[i]/abc[i];
  
  for(ic=0;ic<3;ic++){
    for(i=0;i<g->size[0];i++){
      for(j=0;j<g->size[1];j++){
	for(k=0;k<g->size[2];k++){
	  dptr1=data+((i*g->size[1])+j)*g->size[2]+k;
	  ii_left[0]=ii_right[0]=i;
	  ii_left[1]=ii_right[1]=j;
	  ii_left[2]=ii_right[2]=k;
	  ii_left[ic]=(ii_left[ic]+g->size[ic]-1)%g->size[ic];
	  ii_right[ic]=(ii_right[ic]+1)%g->size[ic];
	  dptr2=g->data+((ii_left[0]*g->size[1])+ii_left[1])*g->size[2]+
	    ii_left[2];
	  dptr3=g->data+((ii_right[0]*g->size[1])+ii_right[1])*g->size[2]+
	    ii_right[2];
	  *(dptr1+ic*size)=scale[ic]*(*dptr3-*dptr2);
	}
      }
    }
  }

  free(g->data);
  g->data=data;
  g->comps=3;
  if (g->name){
    name=malloc(strlen(g->name)+7);
    if (!name) error_exit("malloc error for new grid name");
    strcpy(name,"grad(");
    strcat(name,g->name);
    strcat(name,")");
    /* Can't free(g->name), as it is not always malloc'ed */
    g->name=name;
  }
}

void grad(struct grid *g, struct unit_cell *c){
  int i,j,k,jj,ix,iy,iz,ic,nn[3],ii[3],ngx,ngy,ngz,ioff,size;
  double *data,*rgrid1,*rgrid2,*ptr,*ptr2,abc[6],gcomp;
  char *name;
  
  if (debug) fprintf(stderr,"grad called\n");
  if ((!g)||(!g->data)) error_exit("No data passed to grad");
  if (g->comps!=1) error_exit("Can take grad of scalars only");

  if ((!c)||(!c->basis)) error_exit("need basis for grad");
  
  size=g->size[0]*g->size[1]*g->size[2];
  data=malloc(3*size*sizeof(double));
  if (!data) error_exit("Malloc error for new grid in grad");

  /* Copy to complex grid for FFT */

  rgrid1=malloc(2*size*sizeof(double));
  if (!rgrid1) error_exit("Malloc error for FFT grid in grad");
  rgrid2=malloc(2*size*sizeof(double));
  if (!rgrid2) error_exit("Malloc error for second FFT grid in grad");

  ptr=g->data;
  ptr2=rgrid1;
  for(i=0;i<size;i++){
    *(ptr2++)=*(ptr++);
    *(ptr2++)=0;
  }

  nn[0]=g->size[2];
  nn[1]=g->size[1];
  nn[2]=g->size[0];
  
  fft3d(rgrid1,nn,-1);

  ngx=g->size[0];
  ngy=g->size[1];
  ngz=g->size[2];

  basis2abc(c->basis,abc);

  for(ic=0;ic<3;ic++){
    for(i=0;i<g->size[0];i++){
      ix=i;
      if (ix>ngx/2) ix=ix-ngx;
      ii[0]=ix;
      for(j=0;j<g->size[1];j++){
	iy=j;
	if (iy>ngy/2) iy=iy-ngy;
	ii[1]=iy;
	for(k=0;k<g->size[2];k++){
	  iz=k;
	  if (iz>ngz/2) iz=iz-ngz;
	  ii[2]=iz;
	  ioff=i*ngz*ngy+j*ngz+k;
	  gcomp=0;
	  for(jj=0;jj<3;jj++)
	    gcomp+=ii[jj]*c->recip[jj][ic];
	  gcomp*=2*M_PI;
	  rgrid2[2*ioff]=-rgrid1[2*ioff+1]*gcomp/size;
	  rgrid2[2*ioff+1]=rgrid1[2*ioff]*gcomp/size;
	}
      }
    }

    fft3d(rgrid2,nn,1);

    ptr=data+ic*size;
    for(i=0;i<size;i++){
      *(ptr++)=rgrid2[2*i];
    }
  }
  
  free(rgrid2);
  free(rgrid1);

  free(g->data);
  g->data=data;
  g->comps=3;
  if (g->name){
    name=malloc(strlen(g->name)+7);
    if (!name) error_exit("malloc error for new grid name");
    strcpy(name,"grad(");
    strcat(name,g->name);
    strcat(name,")");
    /* Can't free(g->name), as it is not always malloc'ed */
    g->name=name;
  }
  
}

/* reduce multi-component dataset to single component */
void component(struct grid *g, int c){
  int size;
  char *name;
  
  if ((c<1)||(c>g->comps)){
    fprintf(stderr,"Cannot extract component %d from set with %d components\n",
	    c,g->comps);
    exit(1);
  }

  if (debug)
    fprintf(stderr,"Extracting component %d from set with %d components\n",
	    c,g->comps);
  
  if (g->comps==1) return;
  
  size=g->size[0]*g->size[1]*g->size[2];
  
  if (c>1)
    memcpy(g->data,g->data+(c-1)*size,size*sizeof(double));
  
  g->data=realloc(g->data,size*sizeof(double));
  g->comps=1;

  if (g->name){
    name=malloc(strlen(g->name)+8);
    if (!name) error_exit("malloc error for new grid name");
    snprintf(name,strlen(g->name)+8,"%s_comp=%d",g->name,c);
    /* Can't free(g->name), as it is not always malloc'ed */
    g->name=name;
  }
  
}

/* reduce (multi-component) dataset to its modulus */

void gmod(struct grid *g){
  int i,c,size;
  double x,xmin;
  char *name;
  
  size=g->size[0]*g->size[1]*g->size[2];

  if (debug) fprintf(stderr,
		     "Calculating modulus of grid with %d component%s\n",
		     g->comps,(g->comps==1)?"":"s");
  
  if (g->comps==1){
    xmin=0;
    for(i=0;i<size;i++){
      xmin=min(xmin,g->data[i]);
      g->data[i]=fabs(g->data[i]);
    }
    if (xmin==0) return; /* Don't rename data if unchanged */
  }
  else{
    for(i=0;i<size;i++){
      x=0;
      for(c=0;c<g->comps;c++)
	x+=g->data[i+c*size]*g->data[i+c*size];
      g->data[i]=sqrt(x);
    }
    g->data=realloc(g->data,size*sizeof(double));
    g->comps=1;
  }

  if (g->name){
    name=malloc(strlen(g->name)+3);
    if (!name) error_exit("malloc error for new grid name");
    snprintf(name,strlen(g->name)+3,"|%s|",g->name);
    /* Can't free(g->name), as it is not always malloc'ed */
    g->name=name;
  }
  
}

/* Doing this by FFTs is surely better, so this is here for testing
 * purposes.
 */

void grad2_real(struct grid *g, struct unit_cell *c){
  int i,j,k,ic,size,ii_left[3],ii_right[3];
  double *data,*dptr0,*dptr1,*dptr2,*dptr3,scale[3],abc[6];
  char *name;
  
  if (debug) fprintf(stderr,"grad2_real called\n");
  if ((!g)||(!g->data)) error_exit("No data passed to grad2");
  if (g->comps!=1) error_exit("Can take grad2 of scalars only");

  if ((!c)||(!c->basis)) error_exit("need basis for grad2");

  if ((fabs(c->basis[0][1])>tol*tol)||(fabs(c->basis[0][2])>tol*tol)||
      (fabs(c->basis[1][0])>tol*tol)||(fabs(c->basis[1][2])>tol*tol)||
      (fabs(c->basis[2][0])>tol*tol)||(fabs(c->basis[2][1])>tol*tol))
    error_exit("grad_real requires that a & x, b & y and c & z are parallel");

  
  size=g->size[0]*g->size[1]*g->size[2];
  data=malloc(size*sizeof(double));
  if (!data) error_exit("Malloc error for new grid in grad");
  for(i=0;i<size;i++) data[i]=0;
  
  basis2abc(c->basis,abc);
  for(i=0;i<3;i++)
    scale[i]=g->size[i]*g->size[i]/(abc[i]*abc[i]);
  
  for(ic=0;ic<3;ic++){
    for(i=0;i<g->size[0];i++){
      for(j=0;j<g->size[1];j++){
	for(k=0;k<g->size[2];k++){
	  dptr0=data+((i*g->size[1])+j)*g->size[2]+k;
	  ii_left[0]=ii_right[0]=i;
	  ii_left[1]=ii_right[1]=j;
	  ii_left[2]=ii_right[2]=k;
	  ii_left[ic]=(ii_left[ic]+g->size[ic]-1)%g->size[ic];
	  ii_right[ic]=(ii_right[ic]+1)%g->size[ic];
	  dptr1=g->data+((i*g->size[1])+j)*g->size[2]+k;
	  dptr2=g->data+((ii_left[0]*g->size[1])+ii_left[1])*g->size[2]+
	    ii_left[2];
	  dptr3=g->data+((ii_right[0]*g->size[1])+ii_right[1])*g->size[2]+
	    ii_right[2];
	  *dptr0+=scale[ic]*(*dptr3-2*(*dptr1)+*dptr2);
	}
      }
    }
  }

  free(g->data);
  g->data=data;
  if (g->name){
    name=malloc(strlen(g->name)+8);
    if (!name) error_exit("malloc error for new grid name");
    strcpy(name,"grad2(");
    strcat(name,g->name);
    strcat(name,")");
    /* Can't free(g->name), as it is not always malloc'ed */
    g->name=name;
  }
}

void grad2(struct grid *g, struct unit_cell *c){
  int i,j,k,jj,ix,iy,iz,nn[3],ii[3],ngx,ngy,ngz,ioff,size;
  double *rgrid1,*ptr,*ptr2,scale,abc[6],gvec[3],gsq;
  char *name;
  
  if (debug) fprintf(stderr,"grad2 called\n");
  if ((!g)||(!g->data)) error_exit("No data passed to grad2");
  if (g->comps!=1) error_exit("Can take grad2 of scalars only");

  if ((!c)||(!c->basis)) error_exit("need basis for grad2");
  
  size=g->size[0]*g->size[1]*g->size[2];

  /* Copy to complex grid for FFT */

  rgrid1=malloc(2*size*sizeof(double));
  if (!rgrid1) error_exit("Malloc error for FFT grid in grad");

  ptr=g->data;
  ptr2=rgrid1;
  for(i=0;i<size;i++){
    *(ptr2++)=*(ptr++);
    *(ptr2++)=0;
  }

  nn[0]=g->size[2];
  nn[1]=g->size[1];
  nn[2]=g->size[0];
  
  fft3d(rgrid1,nn,-1);

  ngx=g->size[0];
  ngy=g->size[1];
  ngz=g->size[2];

  basis2abc(c->basis,abc);

  for(i=0;i<g->size[0];i++){
    ix=i;
    if (ix>ngx/2) ix=ix-ngx;
    ii[0]=ix;
    for(j=0;j<g->size[1];j++){
      iy=j;
      if (iy>ngy/2) iy=iy-ngy;
      ii[1]=iy;
      for(k=0;k<g->size[2];k++){
	iz=k;
	if (iz>ngz/2) iz=iz-ngz;
	ii[2]=iz;
	ioff=i*ngz*ngy+j*ngz+k;
	for(jj=0;jj<3;jj++)
	  gvec[jj]=ii[0]*c->recip[0][jj]+ii[1]*c->recip[1][jj]+
	    ii[2]*c->recip[2][jj];
        gsq=4*M_PI*M_PI*(gvec[0]*gvec[0]+gvec[1]*gvec[1]+gvec[2]*gvec[2]);

	rgrid1[2*ioff]*=-gsq;
	rgrid1[2*ioff+1]*=-gsq;
      }
    }
  }

  fft3d(rgrid1,nn,1);

  scale=1./(ngx*ngy*ngz);
  ptr=g->data;
  for(i=0;i<size;i++){
    *(ptr++)=rgrid1[2*i]*scale;
  }
  
  free(rgrid1);

  if (g->name){
    name=malloc(strlen(g->name)+8);
    if (!name) error_exit("malloc error for new grid name");
    strcpy(name,"grad2(");
    strcat(name,g->name);
    strcat(name,")");
    /* Can't free(g->name), as it is not always malloc'ed */
    g->name=name;
  }
  
}

void vdiv(struct grid *g, struct unit_cell *c){
  int i,j,k,jj,ix,iy,iz,ic,nn[3],ii[3],ngx,ngy,ngz,ioff,size;
  double *rgrid,*rgrid1,*ptr,*ptr2,scale,gvec[3],gsc;
  char *name;
  
  if (debug) fprintf(stderr,"div called\n");
  if ((!g)||(!g->data)) error_exit("No data passed to div");
  if (g->comps!=3) error_exit("Can take div of 3D vectors only");
  
  if ((!c)||(!c->basis)) error_exit("need basis for div");
  
  size=g->size[0]*g->size[1]*g->size[2];

  rgrid=malloc(2*size*sizeof(double));
  if (!rgrid) error_exit("Malloc error for FFT grid in grad");
  for(i=0;i<2*size;i++) rgrid[i]=0;
  rgrid1=malloc(2*size*sizeof(double));
  if (!rgrid1) error_exit("Malloc error for FFT grid in grad");

  nn[0]=g->size[2];
  nn[1]=g->size[1];
  nn[2]=g->size[0];
  
  for(ic=0;ic<3;ic++){
    ptr=g->data+ic*size;
    ptr2=rgrid1;
    for(i=0;i<size;i++){
      *(ptr2++)=*(ptr++);
      *(ptr2++)=0;
    }
    fft3d(rgrid1,nn,-1);

    ngx=g->size[0];
    ngy=g->size[1];
    ngz=g->size[2];

    for(i=0;i<g->size[0];i++){
      ix=i;
      if (ix>ngx/2) ix=ix-ngx;
      ii[0]=ix;
      for(j=0;j<g->size[1];j++){
	iy=j;
	if (iy>ngy/2) iy=iy-ngy;
	ii[1]=iy;
	for(k=0;k<g->size[2];k++){
	  iz=k;
	  if (iz>ngz/2) iz=iz-ngz;
	  ii[2]=iz;
	  ioff=i*ngz*ngy+j*ngz+k;
	  jj=ic;
	  gvec[jj]=ii[0]*c->recip[0][jj]+ii[1]*c->recip[1][jj]+
	    ii[2]*c->recip[2][jj];
	  gsc=2*M_PI*gvec[ic];
	  rgrid[2*ioff]+=-gsc*rgrid1[2*ioff+1];
	  rgrid[2*ioff+1]+=gsc*rgrid1[2*ioff];
	}
      }
    }
  }

  free(rgrid1);
  fft3d(rgrid,nn,1);

  scale=1./(ngx*ngy*ngz);
  ptr=g->data;
  for(i=0;i<size;i++){
    *(ptr++)=rgrid[2*i]*scale;
  }
  free(rgrid);

  g->data=realloc(g->data,size*sizeof(double));
  g->comps=1;

  if (g->name){
    name=malloc(strlen(g->name)+6);
    if (!name) error_exit("malloc error for new grid name");
    strcpy(name,"div(");
    strcat(name,g->name);
    strcat(name,")");
    /* Can't free(g->name), as it is not always malloc'ed */
    g->name=name;
  }

  
}
