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

static int cells_equal(struct unit_cell *c, struct unit_cell *c2){
  double abc[6],abc2[6];
  int i,ok;

  if ((!c->basis)||(!c2->basis)) return 1;

  basis2abc(c->basis,abc);
  basis2abc(c2->basis,abc2);

  ok=1;
  for(i=0;i<6;i++)
    if (!aeq(abc[i],abc2[i])) ok=0;

  if ((debug>1)&&(ok==0)){
    fprintf(stderr,"Cells %f %f %f %f %f %f\n"
	    "  and %f %f %f %f %f %f compare different\n",
	    abc[0],abc[1],abc[2],abc[3],abc[4],abc[5],
	    abc2[0],abc2[1],abc2[2],abc2[3],abc2[4],abc2[5]);
  }
  
  return ok;
  
}

void data_combine(int op, struct unit_cell *c,
		  struct contents *m, struct kpts *k, struct symmetry *s,
		  struct grid *g, struct es *elect, struct time_series *ts,
		  int *i_grid,struct unit_cell *c2,
		  struct contents *m2, struct kpts *k2, struct symmetry *s2,
		  struct grid *g2, struct es *elect2, struct time_series *ts2){

  struct grid *gptr,*gptr2,ng;
  int i,j,ic,ngpts,ok;
  char *name2;

  name2=(char*)dict_get(m2->dict,"in_file");
  if (!name2) name2="(unknown)";

  if (op==C2X_MERGE_PATH){
    if (k2->path_nkpt){
      k->path_nkpt=k2->path_nkpt;
      k->path_kpts=k2->path_kpts;
      elect->path_eval=elect2->path_eval;
      elect->path_occ=elect2->path_occ;
      elect->path_nbands=elect2->path_nbands;
      if (debug) fprintf(stderr,"Taking path data from %s\n",name2);
    }
    else if (k2->bs_n){
      k->path_nkpt=k2->bs_n;
      k->path_kpts=k2->bs_kpts;
      elect->path_eval=elect2->eval;
      elect->path_occ=elect2->occ;
      elect->path_nbands=elect2->nbands;
      if (debug) fprintf(stderr,"Treating bs path data from %s as path\n",
			 name2);
    }
    else if (k2->n){
      k->path_nkpt=k2->n;
      k->path_kpts=k2->kpts;
      elect->path_eval=elect2->eval;
      elect->path_occ=elect2->occ;
      elect->path_nbands=elect2->nbands;
      if (debug) fprintf(stderr,"Treating kpt data from %s as path\n",
			 name2);
    }
    if (k2->path_n){
      k->path_n=k2->path_n;
      k->path=k2->path;
      k->path_spacing=k2->path_spacing;
      k->break_n=k2->break_n;
      k->ibreaks=k2->ibreaks;
      k->breaks=k2->breaks;
    }
  }
  if (op==C2X_MERGE){
    if ((!c->basis)&&(c2->basis)){
      *c=*c2;
      if (debug) fprintf(stderr,"Taking cell from %s\n",name2);
    }
    if ((m->n)&&(m2->n)&&(m->n!=m2->n)){
      fprintf(stderr,"Warning: both cells contain atoms, but have a different "
	      "number of atoms.\nResult is likely to be nonsense.");
    }
    else if ((!m->n)&&(m2->n)){
      *m=*m2;
      if (debug) fprintf(stderr,"Taking atoms from %s\n",name2);
    }
    else if (m2->forces){
      ok=0;
      for(i=0;i<m->n;i++){
	if ((vmod2(m->atoms[i].force)==0)&&(vmod2(m2->atoms[i].force)!=0)&&
	    (atom_in_list(m->atoms+i,m2->atoms+i,1,c->basis)==0)){
	  for(j=0;j<3;j++)
	    m->atoms[i].force[j]=m2->atoms[i].force[j];
	  ok=1;
	  fprintf(stderr,"Copied force for atom %d from %s\n",i,name2);
	}
      }
      if (ok) m->forces=1;
    }
    
    if (((!k->n)&&(!k->bs_n)&&(!k->path_nkpt)&&(k->path_n))&&
	((k2->n)||(!k2->bs_n)||(!k2->path_nkpt)||(k2->path_n))){
      *k=*k2;
      *elect=*elect2;
      if (debug) fprintf(stderr,"Taking all kpoints and es from %s\n",name2);
    }
    else{
      if ((!k->n)&&(k2->n)){
	k->n=k2->n;
	k->kpts=k2->kpts;
	elect->nspins=elect2->nspins;
	if (!elect->e_fermi) elect->e_fermi=elect2->e_fermi;
	elect->nbands=elect2->nbands;
	elect->occ=elect2->occ;
	elect->eval=elect2->eval;
	if (debug) fprintf(stderr,"Taking SC kpoints and es from %s\n",name2);
      }
      if ((!k->path_nkpt)&&(k2->path_nkpt)){
	k->path_nkpt=k2->path_nkpt;
	k->path_kpts=k2->path_kpts;
	elect->nbands=elect2->nbands;
	elect->path_occ=elect2->path_occ;
	elect->path_eval=elect2->path_eval;
	if (debug) fprintf(stderr,"Taking path kpts and es from %s\n",name2);
      }
    }
    if ((!s->n)&&(s2->n)){
      *s=*s2;
      if (debug) fprintf(stderr,"Taking symmetry from %s\n",name2);
    }
    
    if (g2->data){
      if (!cells_equal(c,c2)){
	fprintf(stderr,"Unable to merge data from different cells\n");
	return;
      }
      gptr=g;
      while(gptr->next) gptr=gptr->next;
      while(g2->data){
	gptr=grid_new(gptr);
	gptr->data=g2->data;
	for(i=0;i<3;i++) gptr->size[i]=g2->size[i];
	gptr->comps=g2->comps;
	gptr->name=g2->name;
	gptr->origin_abs=g2->origin_abs;
	if (debug)
	  fprintf(stderr,"Merging %s from %s\n",
		  (g2->name)?g2->name:"unknown data",name2);
	g2=g2->next;
      }
    }
    return;
  }

  gptr=g;
  gptr2=g2;
  while((gptr->data)&&(gptr2->data)){
    if (i_grid){
      if ((i_grid[0]==0)&&(i_grid[1]==0)&&(i_grid[2]==0))
	for(i=0;i<3;i++)
	  i_grid[i]=max(gptr->size[i],gptr2->size[i]);
      ok=1;
      for(i=0;i<3;i++)
	if (i_grid[i]!=gptr->size[i]) ok=0;
      if (!ok){
	fprintf(stderr,"Interpolating %s from %s\n",gptr->name,
		(char*)dict_get(m->dict,"in_file"));
	for(i=0;i<3;i++) ng.size[i]=i_grid[i];

	interpolate3d(gptr,&ng);

	free(gptr->data);
	gptr->data=ng.data;
	for(i=0;i<3;i++) gptr->size[i]=ng.size[i];
      }
      
      ok=1;
      for(i=0;i<3;i++)
	if (i_grid[i]!=gptr2->size[i]) ok=0;
      if (!ok){
	fprintf(stderr,"Interpolating %s from %s\n",gptr2->name,
		(char*)dict_get(m2->dict,"in_file"));
	for(i=0;i<3;i++) ng.size[i]=i_grid[i];

	interpolate3d(gptr2,&ng);

	free(gptr2->data);
	gptr2->data=ng.data;
	for(i=0;i<3;i++) gptr2->size[i]=ng.size[i];
      }
    }
    if ((gptr->size[0]!=gptr2->size[0])||
	(gptr->size[1]!=gptr2->size[1])||
	(gptr->size[2]!=gptr2->size[2])){
      fprintf(stderr,"Unable to work with different grid sizes,"
	      " have %dx%dx%d and %dx%dx%d\n",
	      gptr->size[0],gptr->size[1],gptr->size[2],
	      gptr2->size[0],gptr2->size[1],gptr2->size[2]);
      exit(1);
    }
    if (!cells_equal(c,c2))
      error_exit("Unable to combine data from different cells");

    ngpts=gptr->size[0]*gptr->size[1]*gptr->size[2]*gptr->comps;
    if (op==C2X_ADD){
      if (gptr->comps!=gptr2->comps)
	error_exit("Unable to add sets with different numbers of components");
      if (debug)
	fprintf(stderr,"Adding %s from %s to %s\n",
		(gptr2->name)?gptr2->name:"unknown data",name2,
		(gptr->name)?gptr->name:"unknown data");
      for(i=0;i<ngpts;i++)
	gptr->data[i]+=gptr2->data[i];
      free(gptr2->data);
      gptr2->data=NULL;
      gptr2=gptr2->next;
    }
    else if (op==C2X_DIFF){
      if (gptr->comps!=gptr2->comps)
	error_exit("Unable to subtract sets with different "
		   "numbers of components");
      if (debug)
	fprintf(stderr,"Subtracting %s from %s from %s\n",
		(gptr2->name)?gptr2->name:"unknown data",name2,
		(gptr->name)?gptr->name:"unknown data");
      for(i=0;i<ngpts;i++)
	gptr->data[i]-=gptr2->data[i];
      free(gptr2->data);
      gptr2->data=NULL;
      gptr2=gptr2->next;
    }
    else if (op==C2X_MASK){
      if (debug)
	fprintf(stderr,"Masking data %s with mask %s\n",name2,
		(gptr2->name)?gptr2->name:"unknown data");

      if (gptr->comps==gptr2->comps)
	for(i=0;i<ngpts;i++)
	  gptr->data[i]*=gptr2->data[i];
      else if(gptr2->comps==1){
	ngpts=gptr->size[0]*gptr->size[1]*gptr->size[2];
	for(ic=0;ic<gptr->comps;ic++)
	  for(i=0;i<ngpts;i++)
	    gptr->data[i+ic*ngpts]*=gptr2->data[i];
      }
    }
    gptr=gptr->next;
  }

  if (op==C2X_MASK){
    free(gptr2->data);
    gptr2->data=NULL;
  }

}
