/* Various functions for processing bands */


/* Copyright (c) 2020, 2025 MJ Rutter
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as
 * published by the Free Software Foundation, either version 3
 * of the Licence, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, see http://www.gnu.org/licenses/
 */

#include<stdio.h>
#include<stdlib.h> /* malloc */
#include<math.h>
#include<string.h>
#include<ctype.h>

#include "c2xsf.h"

void gcoeff_write(double *psi, int *pwgrid, int nplwv, int fft[3],
		  int gamma, int goff[3],
		  struct unit_cell *c, struct contents *m, struct kpts *kpt,
		  int ikpt, int isppol, int nb, double *eval, double occ,
		  struct es *e);

void wavecar_write(double *psi, int *pwgrid, int nplwv, int fft[3],
                   int gamma, int goff[3],
                   struct unit_cell *c, struct contents *m, struct kpts *kp,
                   int ikpt, int isppol, int nb,
                   double *eval, double occ, struct es *e);

void band_store(struct grid **gp, double *dptr, double occ, double wkpt,
		int nspr, int ns, int k, int b, struct es *elect,
		struct contents *m, int fft[3]);
void band2real(double *psi, double *out, int nfft[3], double kpoint[3]);
void lgrid_sph(char *spec, struct es *e, struct unit_cell *c,
	       struct contents *m, int nfft[3]);
void lgrid_cyl(struct es *e, struct unit_cell *c,
	       struct contents *m, int nfft[3]);
void lgrid_1d(struct es *e, struct unit_cell *c,
	       struct contents *m, int nfft[3]);
void reduce2d_z(double *data, int grid[3], char axis);

void density1d(double *x, double *restrict out, int *fft, int axis){
  int i,j,k,ngx,ngy,ngz;
  double *ptr;

  ngx=fft[0];
  ngy=fft[1];
  ngz=fft[2];
  ptr=x;

  for(i=0;i<2*fft[axis];i++) out[i]=0;
  
  if (axis==0){
    for(i=0;i<ngx;i++){
      for(j=0;j<ngy;j++){
	for(k=0;k<ngz;k++){
	  out[2*i]+=(*ptr)*(*ptr)+(*(ptr+1)*(*(ptr+1)));
	  ptr+=2;
	}
      }
    }
  }
  else if (axis==1){
    for(i=0;i<ngx;i++){
      for(j=0;j<ngy;j++){
	for(k=0;k<ngz;k++){
	  out[2*j]+=(*ptr)*(*ptr)+(*(ptr+1)*(*(ptr+1)));
	  ptr+=2;
	}
      }
    }
  }
  else if (axis==2){
    for(i=0;i<ngx;i++){
      for(j=0;j<ngy;j++){
	for(k=0;k<ngz;k++){
	  out[2*k]+=(*ptr)*(*ptr)+(*(ptr+1)*(*(ptr+1)));
	  ptr+=2;
	}
      }
    }
  }
  else error_exit("Invalid axis in density1d");
}

double *band2grid(double *dptr, int fft[3], int *pwgrid, int npw, int gamma){
  double *psi_g;
  int i,offset,n0,n1,n2,nfftpts,goff[3];

  nfftpts=fft[0]*fft[1]*fft[2];
  
  for(i=0;i<3;i++) goff[i]=0;
  if (gamma){
    gamma++;
    if (gamma&1) goff[0]=1;
    if (gamma>5) goff[1]=1;
    if ((gamma-2)&2) goff[2]=1;
  }
  
  psi_g=malloc(2*nfftpts*sizeof(double));
  if (!psi_g) error_exit("Malloc error for psi");

  for(i=0;i<2*nfftpts;i++) psi_g[i]=0;
  for(i=0;i<npw;i++){
    offset=pwgrid[3*i+2]+fft[2]*(pwgrid[3*i+1]+
                                 fft[1]*pwgrid[3*i]);
    if ((offset<0)||(offset>nfftpts)){
      fprintf(stderr,"Impossible offset in band2grid off=%d i=%d\n",
              offset,i);
      exit(1);
    }
    psi_g[2*offset]=dptr[2*i];
    psi_g[2*offset+1]=dptr[2*i+1];
  }

  if (gamma>1){ /* construct psi(-k)=conjg(psi(k)) */
    if (debug>1) fprintf(stderr,"Gamma point storage type %d\n",gamma);
    for(i=0;i<npw;i++){
      if ((gamma==2)&&(pwgrid[3*i]==0)&&
          (pwgrid[3*i+1]==0)&&(pwgrid[3*i+2]==0)) continue;
      n0=fft[2]-pwgrid[3*i+2]-goff[2];
      if (n0==fft[2]) n0=0;
      n1=fft[1]-pwgrid[3*i+1]-goff[1];
      if (n1==fft[1]) n1=0;
      n2=fft[0]-pwgrid[3*i]-goff[0];
      if (n2==fft[0]) n2=0;
      offset=n0+fft[2]*(n1+fft[1]*n2);
      if ((offset<0)||(offset>nfftpts)){
        fprintf(stderr,
                "Impossible -offset in band2grid off=%d i=%d\n",
                offset,i);
        exit(1);
      }
      psi_g[2*offset]=dptr[2*i];
      psi_g[2*offset+1]=-dptr[2*i+1];
    }
  }

  return psi_g;
  
}


void ldos_1d(struct es *e, double *psi, int nfft[3],
	      struct unit_cell *c, struct contents *m,
	      int ikpt, int ispin, int ib){
  int i,axis,ngpt,nfftpts;
  double *psi1d,wt;
  
  if ((debug)&&(ib==0)&&(ispin==0))
    fprintf(stderr,"processing kpoint %d\n",ikpt);
  
  axis=e->ltype&3;
  if (axis==0){
    lgrid_1d(e,c,m,nfft);
    axis=e->ltype&3;
  }
  axis-=1;

  ngpt=nfft[axis];
  psi1d=malloc(2*ngpt*sizeof(double));
  if (!psi1d) error_exit("malloc error for psi1d");
  density1d(psi,psi1d,nfft,axis);
  fft1d(psi1d,ngpt,-1);

  wt=0;
  for(i=0;i<2*ngpt;i++)
    wt+=psi1d[i]*e->lwork[i];
  
  nfftpts=nfft[0]*nfft[1]*nfft[2];
  if (debug>1) fprintf(stderr,"Band %d kpt %d ldos int=%lf wt=%lf g0=%lf\n",
		       ib,ikpt,wt/nfftpts,wt/psi1d[0],psi1d[0]);
  
  wt/=psi1d[0];
  e->lwt[ib+e->nbands*ispin+e->nbands*e->nspins*ikpt]=wt;
  free(psi1d);
}

void ldos_sph(struct es *e, double *psi, int nfft[3],
	      struct unit_cell *c, struct contents *m,
	      int ikpt, int ispin, int ib){
  int i,ffft[3],ngpts;
  double result,*rgrid;

  /*
  if ((debug)&&(ib==0)&&(ispin==0)&&(ikpt==0))
    fprintf(stderr,"ldos_sph: z=(%lf,%lf,%lf) r=%lf\n",
	    origin_frac[0],origin_frac[1],origin_frac[2],rad);
  */
  
  if ((debug)&&(ib==0)&&(ispin==0))
    fprintf(stderr,"processing kpoint %d\n",ikpt);
  
  /* Calculate density */
  ngpts=nfft[0]*nfft[1]*nfft[2];
  for(i=0;i<2*ngpts;i+=2){
    psi[i]=psi[i]*psi[i]+psi[i+1]*psi[i+1];
    psi[i+1]=0;
  }

  ffft[0]=nfft[2];
  ffft[1]=nfft[1];
  ffft[2]=nfft[0];
  fft3d(psi,ffft,-1);
  rgrid=psi;

  /* Calculate and save grid factors */

  if ((nfft[0]!=e->lwrk_grid[0])||(nfft[1]!=e->lwrk_grid[1])||
      (nfft[2]!=e->lwrk_grid[2])||(e->lwork==NULL)){
    lgrid_sph(e->ldos,e,c,m,nfft);
  }

  result=0;
  for(i=0;i<2*ngpts;i++)
    result+=rgrid[i]*e->lwork[i];

  /* If divided by ngpts, then
   * result would be the integral over a spherical volume
   * rgrid[0] would be the average value in the cell, so
   * rgrid[0]*c->vol would be the integral over the cell
   *
   * Conventional normalisation is that rgrid[0]/ngpts == 1 (approx)
   * (exact for norm conserving pseudopots)
   */
  
  result/=c->vol;
  
  if (debug>1) fprintf(stderr,"Band %d kpt %d ldos int=%lf wt=%lf g0=%lf\n",
		       ib,ikpt,result/ngpts,result/rgrid[0],rgrid[0]);
  
  e->lwt[ib+e->nbands*ispin+e->nbands*e->nspins*ikpt]=result/rgrid[0];
}

void ldos_cyl(struct es *e, double *psi, int nfft[3],
	      struct unit_cell *c, struct contents *m,
	      int ikpt, int ispin, int ib){
  int i,ffft[3],ngpts,axis;
  double result,*rgrid;

  if ((debug)&&(ib==0)&&(ispin==0))
    fprintf(stderr,"processing kpoint %d\n",ikpt);

  axis=e->ltype&3;
  if (axis==0){
    lgrid_cyl(e,c,m,nfft);
    axis=e->ltype&3;
  }
  axis+='a'-1;
  
  /* Calculate density */
  ngpts=nfft[0]*nfft[1]*nfft[2];
  for(i=0;i<2*ngpts;i+=2){
    psi[i]=psi[i]*psi[i]+psi[i+1]*psi[i+1];
    psi[i+1]=0;
  }

  reduce2d_z(psi,nfft,axis);

  ffft[0]=nfft[2];
  ffft[1]=nfft[1];
  ffft[2]=nfft[0];
  fft3d(psi,ffft,-1);
  rgrid=psi;

  result=0;
  for(i=0;i<2*nfft[0]*nfft[1];i++)
    result+=rgrid[i]*e->lwork[i];

  result/=c->vol;
  
  if (debug>1) fprintf(stderr,"Band %d kpt %d ldos int=%lf wt=%lf g0=%lf\n",
		       ib,ikpt,result/ngpts,result/rgrid[0],rgrid[0]);
  
  e->lwt[ib+e->nbands*ispin+e->nbands*e->nspins*ikpt]=result/rgrid[0];
  
}

/* Use abinit's definition of gamma, less one.
 *  gamma=0  -- not gamma
 *  gamma=1  -- gamma
 *  gamma>1  -- abinit's istwfk = gamma+1
 *
 * ikpt, ispinor, isppol and nb all start from 0, not 1
 *
 * isppol=-1 means both components of a spinor wavefunction being presented
 *   for gcoeff or wavecar output
 *
 */

void band_process(double *dptr, int fft[3], int *pwgrid, int npw, int gamma,
		  struct unit_cell *c,
		  struct grid **gp, struct es *elect, struct kpts *kp,
		  struct contents *m, int ikpt, int ispinor, int isppol,
		  int nb, int *i_grid){
  int i,off;
  int nfftpts,ffft[3],nfft[3],goff[3];
  double scale,occ,*psi,*kpoint,eval[2];

  off=elect->nbands*elect->nbspins*ikpt+isppol*elect->nbands;
  if (elect->occ)
    occ=elect->occ[off+nb];
  else
    occ=1;
  
  for(i=0;i<3;i++) nfft[i]=fft[i];
  nfftpts=fft[0]*fft[1]*fft[2];

  if (nfftpts==0) error_exit("No grid size in band_process.");
  
  for(i=0;i<3;i++) goff[i]=0;
  if (gamma){
    gamma++;
    if (gamma&1) goff[0]=1;
    if (gamma>5) goff[1]=1;
    if ((gamma-2)&2) goff[2]=1;
    gamma--;
  }
  
  scale=1.0;
  kpoint=kp->kpts[ikpt].frac;

  if (flags&GCOEFF){
    if (elect->eval)
      eval[0]=elect->eval[off+nb];
    else
      eval[0]=1;
    eval[1]=0;
    if (dict_get(m->dict,"wavecar_output"))
      wavecar_write(dptr,pwgrid,npw,fft,gamma,goff,c,m,kp,ikpt,isppol,nb+1,
                    eval,occ,elect);
    else
      gcoeff_write(dptr,pwgrid,npw,fft,gamma,goff,c,m,kp,ikpt,isppol,nb+1,
                    eval,occ,elect);
    return;
  }

  if (isppol==-1){
    fprintf(stderr,"Unable to process spinor wavefunction\n");
    return;
  }
  
  psi=band2grid(dptr,fft,pwgrid,npw,gamma);

  if ((aeq(kpoint[0],0)||aeq(fabs(kpoint[0]),0.5))&&
      (aeq(kpoint[1],0)||aeq(fabs(kpoint[1]),0.5))&&
      (aeq(kpoint[2],0)||aeq(fabs(kpoint[2]),0.5))&&
      (flags&BANDPARITY)) inv_parity(psi,fft,nb+1,kpoint);

  /* Was the parity all we were requested to report? */
  if (!(flags&BANDS)){
    free(psi);
    return;
  }
  
  /* Padding */
            
  if (i_grid){
    for(i=0;i<3;i++) nfft[i]=i_grid[i];
    if(debug>1)
      fprintf(stderr,"Padding wavefunction onto %dx%dx%d grid\n",
	      nfft[0],nfft[1],nfft[2]);
    if ((fft[0]==nfft[0])&&(fft[1]==nfft[1])&&(fft[2]==nfft[2])){
      if (debug>1)
	fprintf(stderr,"Skipping null padding operation\n");
    }
    else{
      pad_recip(psi,fft,&dptr,nfft);
      nfftpts=nfft[0]*nfft[1]*nfft[2];
      free(psi);
      psi=dptr;
    }
  }

  ffft[0]=nfft[2];
  ffft[1]=nfft[1];
  ffft[2]=nfft[0];
  fft3d(psi,ffft,1);
  
  if (elect->ldos){
    if (!elect->lwt){
      elect->lwt=malloc(elect->nbands*elect->nspins*kp->n*sizeof(double));
      if (!elect->lwt) error_exit("malloc error for lwt");
      fprintf(stderr,"Created ldos weights array\n");
    }
    if (elect->ltype==0){
      /* Sphere and cylinder specifications contain a colon */
      if (strchr(elect->ldos,':')==NULL)
	elect->ltype=LDOS_SLAB;
      /* Sphere specifications never contain a minus sign */
      else if (strchr(elect->ldos,'-')==NULL)
	elect->ltype=LDOS_SPH;
      else
	elect->ltype=LDOS_CYL;
    }

    if (elect->ltype==LDOS_SPH){
      ldos_sph(elect,psi,nfft,c,m,ikpt,isppol,nb);
      free(psi);
      return;
    }
    if (elect->ltype&LDOS_CYL){
      ldos_cyl(elect,psi,nfft,c,m,ikpt,isppol,nb);
      free(psi);
      return;
    }
    ldos_1d(elect,psi,nfft,c,m,ikpt,isppol,nb);
    free(psi);
    return;
  }

  dptr=malloc(nfftpts*sizeof(double));
  if(!dptr) error_exit("Malloc error for grid data");
  band2real(psi,dptr,nfft,kp->kpts[ikpt].frac);
  free(psi);

  /* Do we need to rescale? */
  if (((flags&RAW)==0)&&((flags&BANDPHASE)==0)){ /* Yes */
    if (flags&BANDDEN) scale=1/c->vol;
    else scale=1/sqrt(c->vol);
    if (debug>2) fprintf(stderr,"Scaling wavefun by %f\n",scale);
    for(i=0;i<nfftpts;i++) dptr[i]*=scale;
  }

  band_store(gp,dptr,occ,kp->kpts[ikpt].wt,
	     ispinor,isppol,ikpt+1,nb+1,elect,m,nfft);

}

  
void band2real(double *psi, double *out, int nfft[3], double kpoint[3]){
  double phase_r,phase_r2,phase_i,phase_i2,phi,dtmp;
  double min,max,sum;
  int i,ii,jj,kk,ind,nfft_pts;

  nfft_pts=nfft[0]*nfft[1]*nfft[2];
  
  if ((!(flags&BANDDEN))&&
      ((kpoint[0]!=0)||(kpoint[1]!=0)||(kpoint[2]!=0))){ /* want psi,
							    but not at gamma! */
    if (debug)
      fprintf(stderr,"unwinding psi for non-gamma k-point...\n");
    for(ii=0;ii<nfft[0];ii++){
      for(jj=0;jj<nfft[1];jj++){
	for(kk=0;kk<nfft[2];kk++){
	  phi=2*M_PI*((ii*kpoint[0])/nfft[0]+
		      (jj*kpoint[1])/nfft[1]+
		      (kk*kpoint[2])/nfft[2]);
	  phase_r=cos(phi);
	  phase_i=sin(phi);
	  ind=2*(kk+nfft[2]*(jj+ii*nfft[1]));
	  dtmp=psi[ind];
	  psi[ind]=phase_r*psi[ind]-phase_i*psi[ind+1];
	  psi[ind+1]=phase_r*psi[ind+1]+phase_i*dtmp;
	}
      }
    }
  }
  phase_r=phase_i=phase_r2=phase_i2=0;
  for(i=0;i<nfft_pts;i++){
    if (psi[2*i]>0){
      phase_r+=psi[2*i];
      phase_i-=psi[2*i+1];
    }else{
      phase_r2-=psi[2*i];
      phase_i2+=psi[2*i+1];
    }
  } 
  phase_r+=phase_r2;
  phase_i+=phase_i2;
  dtmp=sqrt(phase_r*phase_r+phase_i*phase_i);
  phase_r/=dtmp;
  phase_i/=dtmp;
  ii=0;
  max=-1e300;
  min=1e300;
  sum=0;
  for (i=0;i<nfft_pts;i++){
    if (flags&BANDPHASE){
      out[i]=atan2(psi[2*i+1],psi[2*i]);
    }
    else if (flags&BANDREAL){
      out[i]=psi[2*i];
    }
    else if (flags&BANDIMAG){
      out[i]=psi[2*i+1];
    }
    else
      if (flags&BANDDEN)
	out[i]=psi[2*i]*psi[2*i]+psi[2*i+1]*psi[2*i+1];
      else{
	out[i]=psi[2*i]*phase_r-psi[2*i+1]*phase_i;
	dtmp=psi[2*i]*phase_i+psi[2*i+1]*phase_r;
	if((fabs(dtmp)>.05))ii++;
      }
    sum+=out[i];
    if(out[i]<min) min=out[i];
    if(out[i]>max) max=out[i];
  }
  if (debug>2) fprintf(stderr,"Min=%g Max=%g Sum=%g\n",min,max,sum);
  if ((debug>1)&&(ii>0)) fprintf(stderr,"Warning: %d components with "
			      " imaginary part >0.05\n",ii);

}

#define CBUFF 100
void band_store(struct grid **gp, double *dptr, double occ, double wkpt,
		int nspr, int ns, int k, int b, struct es *elect,
		struct contents *m, int fft[3]){
  double w;
  int i,nfft_pts;
  char cbuff[CBUFF+1];
  struct grid *g;

  g=*gp;
  nfft_pts=fft[0]*fft[1]*fft[2];
  w=1;

  /* Do we need to weight? */
  if ((flags&OCC_WEIGHT)||(flags&K_WEIGHT)){
    if (flags&OCC_WEIGHT) w=occ;
    if (flags&K_WEIGHT) w*=wkpt;
    /* If we want densities, and we do not have spins, each
     * band is doubly occupied */
    if ((elect->nspins==1)&&(elect->nspinors==1)&&
	(flags&BANDDEN)&&(flags&OCC_WEIGHT))
      w*=2;
    if ((w!=1)&&(!(flags&BANDDEN))) w=sqrt(w);
    if (debug)
      fprintf(stderr,"Using weight %f for ns=%d k=%d band=%d\n",
	      w,ns,k,b);
    if (debug>1)
      fprintf(stderr,"  kpt weight %f occupancy %f\n",wkpt,occ);
    
    if (w!=1)
      for(i=0;i<nfft_pts;i++) dptr[i]*=w;
  }

  if (debug) {
    fprintf(stderr,"Processing band %d kpoint %d",b,k);
    if (elect->nspinors==2)
      fprintf(stderr," spin %d\n",nspr);
    else if (elect->nspins==2)
      fprintf(stderr," spin %d\n",ns);
    else
      fprintf(stderr,"\n");
  }
  
  if (!(flags&ACCUMULATE)){
    g->data=dptr;
    for(i=0;i<3;i++) g->size[i]=fft[i];
    g->name=malloc(40);
    if (!g->name) error_exit("malloc error for name");
    if (elect->nspinors==2)
      sprintf(g->name,"band_vs%d_k%d_b%d",nspr,k,b);
    else if (elect->nspins==2)
      sprintf(g->name,"band_s%d_k%d_b%d",ns,k,b);
    else
      sprintf(g->name,"band_k%d_b%d",k,b);
    g=grid_new(g);
    if ((flags&OCC_WEIGHT)||(flags&K_WEIGHT)){
      snprintf(cbuff,CBUFF,
	       "Weight %f used for spin %d kpt %d band %d",
	       w,(elect->nspinors==2)?nspr:ns,k,b);
      if (m) add_cmt(m->comment,cbuff);
    }
  }else{  /* Are accumulating */
    if (!g->data){  /* This is the first set */
      g->data=dptr;
      for(i=0;i<3;i++) g->size[i]=fft[i];
      g->name=malloc(40);
      if (!g->name) error_exit("malloc error for name");
      sprintf(g->name,"bands"); /* Don't move to a new grid */
    }else{
      for(i=0;i<nfft_pts;i++) g->data[i]+=dptr[i];
      free(dptr);
    }
  }
  *gp=g;
}


/* Find maximum |g|^2 in a pwgrid of gvectors */
double g2max(double recip[3][3], int *pwgrid, int nplwv, int fft[3],
             double *kpt){
  int i,j;
  double g2,g2m,gv[3],x[3];

  g2m=0;
  for(i=0;i<nplwv;i++){
    for(j=0;j<3;j++){
      x[j]=pwgrid[3*i+j]+kpt[j];
      if (x[j]>fft[j]/2) x[j]-=fft[j];
    }
    for(j=0;j<3;j++)
      gv[j]=x[0]*recip[0][j]+x[1]*recip[1][j]+x[2]*recip[2][j];
    g2=vmod2(gv);
    g2m=max(g2m,g2);
  }

  return g2m;

}

void lgrid_sph(char *spec, struct es *e, struct unit_cell *c,
	       struct contents *m, int nfft[3]){
  int i,j,k,ii,jj,ngx,ngy,ngz,ngpts,atno,norigin;
  double *origin_frac,gv[3],gva[3],radius,modk,gs,theta,dot,cos_dot,sin_dot;
  char *cptr,sym[5];

  /* Do we have atom* syntax? */
  atno=0;
  origin_frac=NULL;
  cptr=strchr(spec,'*');
  if ((cptr)&&(cptr-spec<4)&&(isalpha(*spec))){
    strncpy(sym,spec,cptr-spec);
    sym[cptr-spec]=0;
    atno=atsym2no(sym);
    norigin=0;
    for(i=0;i<m->n;i++){
      if (m->atoms[i].atno==atno){
	origin_frac=realloc(origin_frac,(norigin+1)*3*sizeof(double));
	if (!origin_frac) error_exit("realloc error in lgrid_sph");
	for(j=0;j<3;j++)
	  origin_frac[3*norigin+j]=m->atoms[i].frac[j];
	norigin++;
      }
    }
    cptr++;
  }
  else{
    origin_frac=malloc(3*sizeof(double));
    if (!origin_frac) error_exit("malloc error in lgrid_sph");
    cptr=spec;
    lscan(&cptr,m,origin_frac);
    norigin=1;
  }
  if (*cptr!=':') error_exit("Failed to find colon in ldos_spec");
  cptr++;
  sscanf(cptr,"%lf%n",&radius,&i);
  cptr+=i;
  if (*cptr=='B'){
    radius*=BOHR;
    cptr++;
  }
  if ((*cptr)&&(*cptr!='+')){
    fprintf(stderr,"Warning: ignoring unexpected characters in ldos: %s\n",
	    cptr);
  }
  if (debug)
    for(i=0;i<norigin;i++)
      fprintf(stderr,"ldos_sph: z=(%lf,%lf,%lf) r=%lf\n",
	      origin_frac[3*i],origin_frac[3*i+1],origin_frac[3*i+2],radius);
  
  ngx=nfft[0];
  ngy=nfft[1];
  ngz=nfft[2];
  ngpts=ngx*ngy*ngz;

  if ((ngx!=e->lwrk_grid[0])||(ngy!=e->lwrk_grid[1])||(ngz!=e->lwrk_grid[2])||
      (e->lwork==NULL)){
    e->lwork=realloc(e->lwork,2*ngpts*sizeof(double));
    if (!e->lwork) error_exit("malloc error in ldos_sph");
    for(i=0;i<2*ngpts;i++) e->lwork[i]=0;
    if (debug) fprintf(stderr,"ldos work array allocated\n");
    e->lwrk_grid[0]=ngx;
    e->lwrk_grid[1]=ngy;
    e->lwrk_grid[2]=ngz;
  }
  
  e->lwork[0]+=4*M_PI*norigin*radius*radius*radius/3;
  e->lwork[1]=0;
  for(i=0;i<ngx;i++){
    gv[0]=i;
    if (gv[0]>ngx/2) gv[0]=gv[0]-ngx;
    for(j=0;j<ngy;j++){
      gv[1]=j;
      if (gv[1]>ngy/2) gv[1]=gv[1]-ngy;
      for(k=0;k<ngz;k++){
	if ((i==0)&&(j==0)&&(k==0)) continue;
	gv[2]=k;
	if (gv[2]>ngz/2) gv[2]=gv[2]-ngz;
	  
	for(ii=0;ii<3;ii++){
	  gva[ii]=0;
	  for(jj=0;jj<3;jj++)
	    gva[ii]+=gv[jj]*c->recip[jj][ii];
	}

	cos_dot=sin_dot=0;
	for(jj=0;jj<norigin;jj++){
	  dot=0;
	  for(ii=0;ii<3;ii++)
	    dot+=gv[ii]*origin_frac[3*jj+ii];
	  cos_dot+=cos(2*M_PI*dot);
	  sin_dot+=sin(2*M_PI*dot);
	}

	modk=2*M_PI*sqrt(gva[0]*gva[0]+gva[1]*gva[1]+gva[2]*gva[2]);
	theta=radius*modk;

	gs=4*M_PI*(sin(theta)/modk-radius*cos(theta))/(modk*modk);
	e->lwork[2*k+2*j*ngz+2*i*ngz*ngy]+=cos_dot*gs;
	e->lwork[2*k+2*j*ngz+2*i*ngz*ngy+1]+=-sin_dot*gs;
      }
    }
  }
  
  free(origin_frac);
  if (*cptr=='+')
    lgrid_sph(cptr+1,e,c,m,nfft);
  
}

void lgrid_cyl(struct es *e, struct unit_cell *c,
	       struct contents *m, int nfft[3]){
  int i,j,ii,jj,ngx,ngy,ngpts;
  double origin_frac[3],gv[3],gva[3],modk,dot,cos_dot,sin_dot;
  double x,z,radius1,radius2,recip[2][3],org[2];
  char *cptr,axis;

  axis='c';
  cptr=e->ldos;
  lscan(&cptr,m,origin_frac);

  if (*cptr!=':') error_exit("Failed to find colon in ldos_spec");
  cptr++;

  if (*cptr!='r') error_exit("Failed to find r in lgrid_cyl");
  cptr++;

  if(!sscanf(cptr,"%lf%n",&radius1,&i))
    error_exit("Error scanning for first radius in lgrid_cyl");
  cptr+=i;
  if (*cptr=='B'){
    radius1*=BOHR;
    cptr++;
  }

  if (*cptr!='-') error_exit("Failed to find - in lgrid_cyl");
  cptr++;
  if(!sscanf(cptr,"%lf%n",&radius2,&i))
    error_exit("Error scanning for second radius in lgrid_cyl");
  cptr+=i;
  if (*cptr=='B'){
    radius2*=BOHR;
    cptr++;
  }

  if ((*cptr>='a')&&(*cptr<='c')){
    axis=*cptr;
    cptr++;
  }

  if (!is_orthog(axis-'a',c->basis))
    fprintf(stderr,"WARNING: cell axis %c not orthogonal to other axes\n",
	    axis);
  
  if (*cptr){
    fprintf(stderr,"Warning: ignoring unexpected characters in ldos: %s\n",
	    cptr);
  }

  if (debug)
    fprintf(stderr,"ldos_cyl: z=(%lf,%lf,%lf) r1=%lf r2=%lf\n",
	    origin_frac[0],origin_frac[1],origin_frac[2],radius1,radius2);

  z=0;
  if (axis=='c'){
    ngx=nfft[0];
    ngy=nfft[1];
    for(i=0;i<3;i++) recip[0][i]=c->recip[0][i];
    for(i=0;i<3;i++) recip[1][i]=c->recip[1][i];
    for(i=0;i<3;i++) z+=c->basis[2][i]*c->basis[2][i];
    org[0]=origin_frac[0];
    org[1]=origin_frac[1];
  }
  else if (axis=='b'){
    ngx=nfft[0];
    ngy=nfft[2];
    for(i=0;i<3;i++) recip[0][i]=c->recip[0][i];
    for(i=0;i<3;i++) recip[1][i]=c->recip[2][i];
    for(i=0;i<3;i++) z+=c->basis[1][i]*c->basis[1][i];
    org[0]=origin_frac[0];
    org[1]=origin_frac[2];
  }
  else if (axis=='a'){
    ngx=nfft[1];
    ngy=nfft[2];
    for(i=0;i<3;i++) recip[0][i]=c->recip[1][i];
    for(i=0;i<3;i++) recip[1][i]=c->recip[2][i];
    for(i=0;i<3;i++) z+=c->basis[0][i]*c->basis[0][i];
    org[0]=origin_frac[1];
    org[1]=origin_frac[2];
  }
   else
    error_exit("Impossible axis");

  z=sqrt(z);
  e->ltype=LDOS_CYL+1+(axis-'a');
  
  ngpts=ngx*ngy;
  e->lwork=realloc(e->lwork,2*ngpts*sizeof(double));
  if (!e->lwork) error_exit("malloc error in lgrid_cyl");
  for(i=0;i<2*ngpts;i++) e->lwork[i]=0;
  if (debug) fprintf(stderr,"ldos work array allocated\n");
  e->lwrk_grid[0]=ngx;
  e->lwrk_grid[1]=ngy;
  e->lwrk_grid[2]=1;

  e->lwork[0]=M_PI*(radius2*radius2-radius1*radius1)*z;
  e->lwork[1]=0;

  for(i=0;i<ngx;i++){
    gv[0]=i;
    if (gv[0]>ngx/2) gv[0]=gv[0]-ngx;
    for(j=0;j<ngy;j++){
      if ((i==0)&&(j==0)) continue;
      gv[1]=j;
      if (gv[1]>ngy/2) gv[1]=gv[1]-ngy;
      
      for(ii=0;ii<3;ii++){
	gva[ii]=0;
	for(jj=0;jj<2;jj++)
	  gva[ii]+=gv[jj]*recip[jj][ii];
      }

      dot=gv[0]*org[0]+gv[1]*org[1];
      cos_dot=cos(2*M_PI*dot);
      sin_dot=sin(2*M_PI*dot);

      modk=2*M_PI*sqrt(gva[0]*gva[0]+gva[1]*gva[1]+gva[2]*gva[2]);

      x=2*M_PI*(radius2*j1(radius2*modk)-radius1*j1(radius1*modk))*z/modk;

      e->lwork[2*(i*ngy+j)]=x*cos_dot;
      e->lwork[2*(i*ngy+j)+1]=-x*sin_dot;
      
    }
  }

}

void lgrid_1d(struct es *e, struct unit_cell *c,
	       struct contents *m, int fft[3]){
  double z1,z2,scale,g,theta1,theta2,f_r,f_i;
  int i,ngpt,axis;
  char caxis;
  
  if (sscanf(e->ldos,"%lf-%lf%c",&z1,&z2,&caxis)!=3)
    error_exit("ldos_1d spec parse error");
  if ((caxis>='A')&&(caxis<='C')){
    scale=0;
    for(i=0;i<3;i++) scale+=c->basis[caxis-'A'][i]*c->basis[caxis-'A'][i];
    scale=1.0/sqrt(scale);
    z1*=scale;
    z2*=scale;
    caxis+=32;
  }
  axis=caxis-'a';
  if ((axis<0)||(axis>2)) error_exit("invalid axis in ldos");
  if (!is_orthog(axis,c->basis))
    fprintf(stderr,"WARNING: cell axis %c not orthogonal to other axes\n",
	    caxis);
  e->ltype=LDOS_SLAB+axis+1;

  if (debug)
    fprintf(stderr,"ldos_1d: axis=%c, fractional extent %lf to %lf\n",
	    axis+'a',z1,z2);
  
  ngpt=fft[axis];
  e->lwork=realloc(e->lwork,2*ngpt*sizeof(double));
  if (!e->lwork) error_exit("malloc error in lgrid_1d");
  for(i=0;i<2*ngpt;i++) e->lwork[i]=0;
  if (debug) fprintf(stderr,"ldos work array allocated\n");

  e->lwork[0]=z2-z1;
  e->lwork[1]=0;

  for(i=1;i<ngpt;i++){
    g=i;
    if (g>ngpt/2) g=g-ngpt;
    g*=2*M_PI;
    theta1=z1*g;
    theta2=z2*g;
    f_r=cos(theta2)-cos(theta1);
    f_i=sin(theta2)-sin(theta1);
    e->lwork[2*i]=f_i/g;
    e->lwork[2*i+1]=f_r/g;
  }
  
}

