#include <stddef.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <string.h>

#include "../src/param.h"
#include "../src/tdef.h"
#include "../src/common.c"

#include "../src/libutility/utility.h"

int idsort(void *p1,void *p2)
{
   long g1,g2;	
   
   g1 = *(long *)p1;
   g2 = *(long *)p2;
   
   if ( g1 < g2 )
      return(-1);
   else if ( g1 > g2 )
      return(1);
   else
      return(0);
}

int main()
{
   long          i, BUFSIZE, nhalos, nkeep, *ikeep, ipart, jpart, ihalo, jhalo, ncheck, irem;
   long          *ikeep_halo, nkeep_halo;
   int           j, ld;
   double        dw[7];
   float         fw[7];
	double        *xpart, *ypart, *zpart, *vxpart, *vypart, *vzpart, *mpart;
   long          *id, nid, *idrem, nrem, *idrem2;
   double        nvkeep, nvold;
   char          infile[MAXSTRING],outfile[MAXSTRING],stereofile[MAXSTRING],halofile[MAXSTRING],dummyline[MAXSTRING];
   FILE          *fpin, *fpout, *fphalo, *fpstereo;
   int           machine_sizeof_long;
	boolean       remove_particles;
   double        r,g,b;
   int           ISEED=-123456;
   int           SWAPBYTES;
   
   printf("===============================================================\n");
   printf("  remove all but a couple of specfied halos from AMIGA binary\n");
   printf("===============================================================\n");
   
   printf("Please give name of input file:            ");
   scanf("%s",infile);
   printf("%s\n",infile);
   printf("Please give name of AHF/BDM_particles:     ");
   scanf("%s",halofile);
   printf("%s\n",halofile);
   printf("Please give number of halos to keep:       ");
   scanf("%ld",&nkeep_halo);
   printf("%ld\n",nkeep_halo);
   ikeep_halo = (long *) calloc(nkeep_halo, sizeof(long));
   for(j=0; j<nkeep_halo; j++)
     {
      printf("Please give id of %5ld.halo to keep:      ", j);
      scanf("%ld",&ikeep_halo[j]);
      printf("%ld\n",ikeep_halo[j]);
     }
   printf("Please give name for output file:          ");
   scanf("%s",outfile);
   printf("%s\n\n",outfile);
   
   strcpy(stereofile,outfile);
   strcat(stereofile,"-halos");
   
   /*========================================
      * open files
      *========================================*/
   if((fpin=fopen(infile,"rb"))==NULL)
     {
      printf("I cannot open %s\n", infile);
      exit(1);
     }
   
   if((fphalo=fopen(halofile,"rb"))==NULL)
     {
      printf("I cannot open %s\n", halofile);
      exit(1);
     }
   
   if((fpout=fopen(outfile,"w"))==NULL)
     {
      printf("I cannot open %s\n", outfile);
      exit(1);
     }
   
   if((fpstereo=fopen(stereofile,"w"))==NULL)
     {
      printf("I cannot open %s\n", stereofile);
      exit(1);
     }
   
   /* read AMIGA header */
   read_amiga_header(fpin, &io, &SWAPBYTES);
   
   
  /* print some information */
  printf("some information about input file:\n");
  printf("==================================\n");
  printf("%s\n\n",io.header.header);
  printf("multi-mass       = %d\n",io.header.multi_mass);
  printf("double precision = %d\n",io.header.double_precision);
  printf("no_part          = %ld\n",io.header.no_part);
  printf("no_vpart         = %g\n",io.header.no_vpart);
  printf("boxsize          = %g\n",io.header.boxsize);
  printf("Omega0           = %g\n",io.header.omega0);
  printf("Lambda0          = %g\n",io.header.lambda0);
  printf("z_initial        = %g (a=%lf)\n",
         1.0/io.header.a_initial-1.0,io.header.a_initial);
  printf("z_current        = %g (a=%lf)\n",
         1.0/io.header.a_current-1.0,io.header.a_current);
  printf("no_timestep      = %d\n",io.header.no_timestep);
   
  /*========================================
   * read in AMIGA binary file
   *========================================*/
  xpart  = (double *) calloc(io.header.no_part, sizeof(double));
  ypart  = (double *) calloc(io.header.no_part, sizeof(double));
  zpart  = (double *) calloc(io.header.no_part, sizeof(double));
  vxpart = (double *) calloc(io.header.no_part, sizeof(double));
  vypart = (double *) calloc(io.header.no_part, sizeof(double));
  vzpart = (double *) calloc(io.header.no_part, sizeof(double));

  if(io.header.multi_mass == 1)
    {
     BUFSIZE = 7;
     mpart   = (double *) calloc(io.header.no_part, sizeof(double));
    }

  else
    {
     BUFSIZE = 6;
     mpart   = NULL;
    }

  if(io.header.double_precision)
    {        
     for(i=0;i<io.header.no_part;i++)
       {
        for(j=0; j<BUFSIZE; j++)
           ReadDouble(fpin,&(dw[j]),SWAPBYTES);
        
        xpart[i]  = dw[0];
        ypart[i]  = dw[2];
        zpart[i]  = dw[4];
        vxpart[i] = dw[1];
        vypart[i] = dw[3];
        vzpart[i] = dw[5];
        if(io.header.multi_mass)
           mpart[i] = dw[6];
       }
    }
  else /* double precision */
    {
     for(i=0;i<io.header.no_part;i++)
       {
        for(j=0; j<BUFSIZE; j++)
           ReadFloat(fpin,&(fw[j]),SWAPBYTES);
        
        xpart[i]  = fw[0];
        ypart[i]  = fw[2];
        zpart[i]  = fw[4];
        vxpart[i] = fw[1];
        vypart[i] = fw[3];
        vzpart[i] = fw[5];
        if(io.header.multi_mass)
           mpart[i] = fw[6];
       }
    }
     
  fclose(fpin);

  /*========================================
  * read in AHF/BDM_particles file
  *========================================*/
  fprintf(stderr,"\n o reading particle id's ... \n");

  idrem  = NULL;
  nrem   = 0;
  ihalo  = 0;
  while(1)
    {
       /* read number of particle id's */
       fscanf(fphalo,"%ld",&nid);
     
       /* have we reached EOF yet? */
       if(feof(fphalo))
          break;
       
       /* allocate memory for id-array */
       id = (long *) calloc(nid, sizeof(long));
       
       /* read all particle id's */
       for(i=0; i<nid; i++)
          fscanf(fphalo,"%ld",&id[i]);
       
       /* check if particle should be removed */
       remove_particles = TRUE;
       for(jhalo=0; jhalo<nkeep_halo; jhalo++)
         {
          if(ihalo == ikeep_halo[jhalo])
            {
             remove_particles = FALSE;
             break;
            }
         }
       
       /* if not host or satellite remember particles for removal later on */
       if(remove_particles == TRUE)
         {
          /* dump particles into STEREO2 file for cross-checking */
          r = 1.0;
          g = 0.0;
          b = 0.0;
          fprintf(stderr,"    o removing halo %5ld (npart = %12ld): %4.3g %4.3g %4.3g\n",ihalo,nid,r,g,b);
          for(i=0; i<nid; i++)
             fprintf(fpstereo,"p %g %g %g %g %g %g\n",xpart[id[i]],ypart[id[i]],zpart[id[i]],r,g,b);

          /* remember particle id's to remove them later from AMIGA binary */
          nrem += nid;
          idrem = (long *) realloc(idrem, nrem*sizeof(long));
          for(i=nrem-nid; i<nrem; i++)
             idrem[i] = id[i-(nrem-nid)];
         }
       else
         {
          /* dump particles into STEREO2 file for cross-checking */
          r = 0.0;
          g = 0.0;
          b = 1.0;
          fprintf(stderr,"    o keeping  halo %5ld (npart = %12ld): %4.3g %4.3g %4.3g\n",ihalo,nid,r,g,b);
          for(i=0; i<nid; i++)
             fprintf(fpstereo,"p %g %g %g %g %g %g\n",xpart[id[i]],ypart[id[i]],zpart[id[i]],r,g,b);
         }
       
       ihalo++;
       free(id);
    }
   fclose(fphalo);
   fclose(fpstereo);
   fprintf(stderr," o ... reading particle id's done\n");

#ifdef DOUBLE_CHECK_IDS
   /* double check id's to be removed */
   {
     FILE *fptest;
     fptest = fopen("test.id","w");
     for(i=0; i<nrem; i++)
       {
        fprintf(fptest,"%ld\n",idrem[i]);
       }
     fclose(fptest);
   }
#endif


   /*=================================================
    * qsort idrem for quicker removing
    *=================================================*/
   fprintf(stderr," o doing qsort (%ld) ... ",nrem);
   qsort(idrem, nrem, sizeof(long), idsort);
   fprintf(stderr,"done\n");

#ifdef DOUBLE_CHECK_IDS2
   /* double check id's to be removed */
   {
      FILE *fptest;
      fptest = fopen("test.id2","w");
      for(i=0; i<nrem; i++)
        {
         fprintf(fptest,"%ld\n",idrem[i]);
        }
      fclose(fptest);
   }
#endif

   /*=================================================
    * make sure to have every id only once
    *=================================================*/
   irem      = 1;
   idrem2    = calloc(irem, sizeof(long));
   idrem2[0] = idrem[0];
   for(i=1; i<nrem; i++)
     {
      if(idrem[i] == idrem[i-1])
        {
         /* do not keep this id */
        }
      else
        {
         /* keep this id */
         irem++;
         
         idrem2         = realloc(idrem2, irem*sizeof(long));
         idrem2[irem-1] = idrem[i];
        }
     }
   
   /* store everything in "old" arrays again... */
   free(idrem);
   nrem  = irem;
   idrem = idrem2;
   
#ifdef DOUBLE_CHECK_IDS3
   /* double check id's to be removed */
   {
      FILE *fptest;
      fptest = fopen("test.id3","w");
      for(i=0; i<nrem; i++)
        {
         fprintf(fptest,"%ld\n",idrem[i]);
        }
      fclose(fptest);
   }
#endif
   
   /*=================================================
    * calculate new total mass in simulation box
    *=================================================*/
    nvkeep = 0.0;
    nvold  = 0.0;
    irem   = 0;
    for(ipart=0; ipart<io.header.no_part; ipart++)
      {
       if(io.header.multi_mass)
          nvold += mpart[ipart];
       else
          nvold += 1.0;
       
       if(ipart == idrem[irem])
         {
          irem++;
         }
       else
         {
          if(io.header.multi_mass)
             nvkeep += mpart[ipart];
          else
             nvkeep += 1.0;
         }
      }

   fprintf(stderr," o number of particles to be removed: %ld (%ld)\n",nrem,irem);
   fprintf(stderr," o new number of           particles: %ld (%ld)\n",io.header.no_part-nrem,io.header.no_part);
   fprintf(stderr," o new number of (virtual) particles: %g (%g)\n",nvkeep,nvold);
   fprintf(stderr," o new particle mass:                 %g (%g)\n",
           io.header.omega0 * rhoc0 * pow3(io.header.boxsize) / nvkeep,
           io.header.omega0 * rhoc0 * pow3(io.header.boxsize) / nvold);

   io.header.no_part -= nrem;
   io.header.pmass    = io.header.omega0 * rhoc0 * pow3(io.header.boxsize) / nvkeep;
   io.header.no_vpart = nvkeep;


   /*=================================================
    * write new AMIGA file
    *=================================================*/
   /* write a sizeof(long) to file (also used for BYTESWAP testing) */
   machine_sizeof_long = sizeof(long);
   fwrite(&machine_sizeof_long, sizeof(int), 1, fpout);

   fprintf(stderr," o writing sizeof(io.header) = %d to file\n",sizeof(io.header));
   if(fwrite(&(io.header), sizeof(io.header), 1, fpout) != 1)
     {
      perror("\n\ncould not write IO header\n");
      exit(1);
     }

   /* write output file */
   fprintf(stderr," o writing output file ...");
   irem = 0;
   if(io.header.double_precision)
     {
      for(ipart=0; ipart<io.header.no_part+nrem; ipart++)
        {
         if(ipart == idrem[irem])
           {
            irem++;
            goto next_dpart;
           }
         
         dw[0] = xpart[ipart];
         dw[2] = ypart[ipart];
         dw[4] = zpart[ipart];
         dw[1] = vxpart[ipart];
         dw[3] = vypart[ipart];
         dw[5] = vzpart[ipart];
         if(io.header.multi_mass)
            dw[6] = mpart[ipart];
         
         fwrite((void*)dw,sizeof(double),BUFSIZE,fpout);		 
   next_dpart:;
        }
     }
   else
     {
        for(ipart=0; ipart<io.header.no_part+nrem; ipart++)
          {
           if(ipart == idrem[irem])
             {
              irem++;
              goto next_fpart;
             }
           
           fw[0] = xpart[ipart];
           fw[2] = ypart[ipart];
           fw[4] = zpart[ipart];
           fw[1] = vxpart[ipart];
           fw[3] = vypart[ipart];
           fw[5] = vzpart[ipart];
           if(io.header.multi_mass)
              fw[6] = mpart[ipart];
           
           fwrite((void*)fw,sizeof(float),BUFSIZE,fpout);		 
   next_fpart:;
          }
     }

   fclose(fpout);
   fprintf(stderr,"done\n");

}

