#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include "load_data.hpp"
#include "loadGadget.hpp"
#include "fortran.hpp"

using namespace CosmoTool;

PurePositionData *CosmoTool::loadGadgetPosition(const char *fname)
{
  PurePositionData *data;
  int p, n;
  UnformattedRead f(fname);
  GadgetHeader h;
  
  data = new PurePositionData;
  f.beginCheckpoint();
  for (int i = 0; i < 6; i++)
    h.npart[i] = f.readInt32();
  for (int i = 0; i < 6; i++)    
    h.mass[i] = f.readReal64();
  h.time = f.readReal64();
  h.redshift = f.readReal64();
  h.flag_sfr = f.readInt32();
  h.flag_feedback = f.readInt32();
  for (int i = 0; i < 6; i++)
    h.npartTotal[i] = f.readInt32();
  h.flag_cooling = f.readInt32();
  h.num_files = f.readInt32();
  h.BoxSize = f.readReal64();
  h.Omega0 = f.readReal64();
  h.OmegaLambda = f.readReal64();
  h.HubbleParam = f.readReal64();
  f.endCheckpoint(true);
  
  data->NumPart = 0;
  for(int k=0; k<5; k++)
    data->NumPart += h.npart[k];
  
  data->pos = new FCoordinates[data->NumPart];
  
  f.beginCheckpoint();
  for(int k = 0, p = 0; k < 5; k++) {
    for(int n = 0; n < h.npart[k]; n++) {
      data->pos[p][0] = f.readReal32();
      data->pos[p][1] = f.readReal32();
      data->pos[p][2] = f.readReal32();
      p++;
    }
  }
  f.endCheckpoint();
  
  // Skip velocities
  f.skip((long)data->NumPart*3+2*4);
  // Skip ids

  return data;
}




SimuData *CosmoTool::loadGadgetMulti(const char *fname, int id, int loadflags)
{
  SimuData *data;
  int p, n;
  UnformattedRead f(fname);
  GadgetHeader h;

  data = new SimuData;

  f.beginCheckpoint();
  for (int i = 0; i < 6; i++)
    h.npart[i] = f.readInt32();
  for (int i = 0; i < 6; i++)    
    h.mass[i] = f.readReal64();
  h.time = f.readReal64();
  h.redshift = f.readReal64();
  h.flag_sfr = f.readInt32();
  h.flag_feedback = f.readInt32();
  for (int i = 0; i < 6; i++)
    h.npartTotal[i] = f.readInt32();
  h.flag_cooling = f.readInt32();
  h.num_files = f.readInt32();
  h.BoxSize = f.readReal64();
  h.Omega0 = f.readReal64();
  h.OmegaLambda = f.readReal64();
  h.HubbleParam = f.readReal64();
  f.endCheckpoint(true);

  long NumPart = 0;
  for(int k=0; k<5; k++)
    NumPart += h.npart[k];
  data->NumPart = NumPart;

  if (loadflags & NEED_POSITION) {
      
    for (int i = 0; i < 3; i++)
      data->Pos[i] = new float[data->NumPart];
    
    f.beginCheckpoint();
    for(int k = 0, p = 0; k < 5; k++) {
      for(int n = 0; n < h.npart[k]; n++) {
	data->Pos[p][0] = f.readReal32();
	data->Pos[p][1] = f.readReal32();
	data->Pos[p][2] = f.readReal32();
	p++;
      }
    }
    f.endCheckpoint();
    
  } else {
    // Skip positions
    f.skip(NumPart * 3 + 2*4);
  }
  
  if (loadflags & NEED_VELOCITY) {
    for (int i = 0; i < 3; i++)
      data->Vel[i] = new float[data->NumPart];
    
    f.beginCheckpoint();
    for(int k = 0, p = 0; k < 5; k++) {
      for(int n = 0; n < h.npart[k]; n++) {
	data->Vel[p][0] = f.readReal32();
	data->Vel[p][1] = f.readReal32();
	data->Vel[p][2] = f.readReal32();
	p++;
      }
    }
    f.endCheckpoint();
    
    // TODO: FIX THE UNITS OF THESE FUNKY VELOCITIES !!!
  } else {
    // Skip velocities
    f.skip(NumPart*3+2*4);
  }

  // Skip ids
  if (loadflags & NEED_GADGET_ID) {
    f.beginCheckpoint();
    for(int k = 0, p = 0; k < 6; k++)
      {
	for(int n = 0; n < h.npart[k]; n++)
	  {
	    data->Id[p] = f.readInt32();
	    p++;
	  }
      }
    f.endCheckpoint();
  } else {
    f.skip(2*4);
    for (int k = 0; k < 6; k++)
      f.skip(h.npart[k]);
  }

  return data;
}