/* Read a QE charge-density.dat file */

/* Copyright (c) 2023 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/
 */ 


/* N.B. QE stores densities in reciprocal space */

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

void fft3d(double *c, int *ngptar, int dir);

void qe_rho_read(FILE* infile, struct unit_cell *c, struct contents *m,
                 struct kpts *k, struct symmetry *s, struct grid *g,
		 struct es *elect, int *i_grid){
  int j,ns,gamma_only,nspin;
  unsigned int i,tmp,ngm_g;
  int *mill_g;
  double *rho_g,recip[3][3],*data,*data2,(*old_basis)[3];
  int ngx_min,ngx_max,ngy_min,ngy_max,ngz_min,ngz_max;
  int ngx,ngy,ngz,offset,fft[3];
  int nx,ny,nz;
  
  if (debug) fprintf(stderr,"QE binary reader called\n");

  data=data2=rho_g=NULL;
  old_basis=c->basis;
  
  fread(&tmp,4,1,infile);
  if (tmp!=12) error_exit("Unexpected start to QE binary file");

  fread(&gamma_only,4,1,infile);
  fread(&ngm_g,4,1,infile);
  fread(&nspin,4,1,infile);

  if ((nspin==1)&&(flags&SPINDEN))
    error_exit("Spin density requested, but nspin=1 in density file");

  if (debug) fprintf(stderr,"QE density file with %d spins\n",nspin);

  elect->nspins=nspin;
  
  fread(&tmp,4,1,infile);
  if (tmp!=12) error_exit("Unexpected data in QE binary file");

  fread(&tmp,4,1,infile);
  if (tmp!=72) error_exit("Unexpected data in QE binary file");

  fread(recip,72,1,infile);

  fread(&tmp,4,1,infile);
  if (tmp!=72) error_exit("Unexpected data in QE binary file");

  mill_g=malloc(3*ngm_g*sizeof(int));
  if (!mill_g) error_exit("Malloc error for mill_g in qe_bin_read()");

  fread(&tmp,4,1,infile);
  if (tmp!=3*ngm_g*4) error_exit("Unexpected data in QE binary file");
  fread(mill_g,3*ngm_g*4,1,infile);
  fread(&tmp,4,1,infile);
  if (tmp!=3*ngm_g*4) error_exit("Unexpected data in QE binary file");

  if (debug>1) fprintf(stderr,"Header and remapping grid read\n");
  
  /* We have reciprocal axes, in odd units. First convert the units */

  for(i=0;i<3;i++)
    for(j=0;j<3;j++)
      recip[i][j]/=2*M_PI*BOHR;

  /* Given that a real2rec routine and a rec2real routine would be identical,
     abuse existing real2rec... */
  
  c->basis=malloc(72);
  if (!c->basis) error_exit("Malloc error for basis in qe_bin_read()");
  
  for(i=0;i<3;i++)
    for(j=0;j<3;j++)
      c->basis[i][j]=recip[i][j];

  real2rec(c);

  /* Now we need to exchange real and recip, and correct the cell volume */
  
  for(i=0;i<3;i++)
    for(j=0;j<3;j++)
      c->basis[i][j]=c->recip[i][j];
  
  real2rec(c);

  if (old_basis){
    if (debug>1) fprintf(stderr,"Checking bases\n");
    tmp=0;
    for(i=0;i<3;i++)
      for(j=0;j<3;j++)
	if (!aeq(c->basis[i][j],old_basis[i][j])) tmp=1;
    if (tmp){
      fprintf(stderr,"Warning: basis in charge-density.dat differs from "
	      "previous basis\n");
      fprintf(stderr,"Previous basis:\n");
      print_basis(old_basis);
      fprintf(stderr,"New basis:\n");
      print_basis(c->basis);
    }
    free(old_basis);
    old_basis=NULL;
  }
  
  /* Next work out our FFT grid size */

  ngx_min=ngx_max=ngy_min=ngy_max=ngz_min=ngz_max=0;

  for(i=0;i<ngm_g;i++){
    ngx_min=min(ngx_min,mill_g[3*i]);
    ngx_max=max(ngx_max,mill_g[3*i]);
    ngy_min=min(ngy_min,mill_g[3*i+1]);
    ngy_max=max(ngy_max,mill_g[3*i+1]);
    ngz_min=min(ngz_min,mill_g[3*i+2]);
    ngz_max=max(ngz_max,mill_g[3*i+2]);
  }

  if (debug) fprintf(stderr,"grid extent in file %d:%d,%d:%d,%d:%d\n",
	  ngx_min,ngx_max,ngy_min,ngy_max,ngz_min,ngz_max);  

  ngx=ngx_max-ngx_min+1;
  to235(&ngx);
  ngy=ngy_max-ngy_min+1;
  to235(&ngy);
  ngz=ngz_max-ngz_min+1;
  to235(&ngz);
  
  if (i_grid){
    if (i_grid[0]!=0) ngx=i_grid[0];
    if (i_grid[1]!=0) ngy=i_grid[1];
    if (i_grid[2]!=0) ngy=i_grid[2];
  }
  else{
    if (debug) fprintf(stderr,"FFT grid used %dx%dx%d\n",ngx,ngy,ngz);
  }
  ngx_max=ngx/2;
  ngx_min=-(ngx-1)/2;
  ngy_max=ngy/2;
  ngy_min=-(ngy-1)/2;
  ngz_max=ngz/2;
  ngz_min=-(ngz-1)/2;


  rho_g=malloc(ngm_g*2*sizeof(double));  /* rho_g is complex */
  if (!rho_g) error_exit("Malloc error for rho_g in qe_bin_read()");

  data=malloc(ngx*ngy*ngz*2*sizeof(double));
  if (!data) error_exit("Error in malloc for data in qe_bin_read");

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

  for(ns=0;ns<nspin;ns++){
    if ((ns>0)&&((flags&SPINDEN)==0)) break;
    fread(&tmp,4,1,infile);
    if (tmp!=ngm_g*2*sizeof(double))
      error_exit("Unexpected size for density data in QE binary file");
    if (((ns==0)&&(flags&CHDEN))||((ns>0)&&(flags&SPINDEN))){
      if (fread(rho_g,ngm_g*16,1,infile)!=1){
	fprintf(stderr,"End of file reading ");
	if (ns)
	  fprintf(stderr,"spin component %d\n",ns);
	else
	  fprintf(stderr,"charge density\n");
	exit(1);
      }
      fread(&tmp,4,1,infile);
      if (debug>1){
	if (ns==0)
	  fprintf(stderr,"Read charge density\n");
	else
	  fprintf(stderr,"Read spin component %d\n",ns);
      }
    }
    else{
      fseek(infile,ngm_g*16+4,SEEK_CUR);
      if (debug>1){
	if (ns==0)
	  fprintf(stderr,"Skipped charge density\n");
	else
	  fprintf(stderr,"Skipped spin component %d\n",ns);
      }
      continue;
    }
    
    for(i=0;i<ngx*ngy*ngz*2;i++) data[i]=0.0;

    for(i=0;i<ngm_g;i++){
      nx=mill_g[3*i];
      if (nx>ngx_max) continue;
      if (nx<ngx_min) continue;
      if (nx<0) nx+=ngx;
      ny=mill_g[3*i+1];
      if (ny>ngy_max) continue;
      if (ny<ngy_min) continue;
      if (ny<0) ny+=ngy;
      nz=mill_g[3*i+2];
      if (nz>ngz_max) continue;
      if (nz<ngz_min) continue;
      if (nz<0) nz+=ngz;
      offset=nz+ngz*(ny+ngy*nx);
      data[2*offset]=rho_g[2*i];
      data[2*offset+1]=rho_g[2*i+1];
    }

    fft3d(data,fft,1);

    if (ns==0){
      data2=malloc(ngx*ngy*ngz*sizeof(double));
      g->data=data2;
    }
    else if (ns==1){
      data2=malloc(ngx*ngy*ngz*(nspin-1)*sizeof(double));
      g->data=data2;
    }
    else
      data2+=ngx*ngy*ngz;

    if (!data2) error_exit("Error in malloc for data2 in qe_bin_read");

    for(i=0;i<ngx*ngy*ngz;i++)
      data2[i]=data[2*i]/(BOHR*BOHR*BOHR);

    if (ns<=1){
      g->size[0]=ngx;
      g->size[1]=ngy;
      g->size[2]=ngz;
      if (ns==0)
	g->name="Charge";
      else{
	if (nspin==2)
	  g->name="Spin";
	else{
	  g->comps=nspin-1;
	  g->name="Vector spin";
	}
      }
    }
    if ((ns==0)&&(flags&SPINDEN)){
      g->next=malloc(sizeof(struct grid));
      if (!g->next) error_exit("malloc failure for struct grid");
      init_grid(g->next);
      g=g->next;
    }
  }
  
  free(data);
  free(rho_g);
  free(mill_g);
}
