#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();
  data->BoxSize = 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;
  GadgetHeader h;

  if (id >= 0) {
    int k = snprintf(0, 0, "%s.%d", fname, id)+1;
    char *out_fname = new char[k];
    snprintf(out_fname, k, "%s.%d", fname, id);
     
    f = new UnformattedRead(out_fname);
    if (f == 0)
      return 0;

    delete out_fname;

  } else {

    f = new UnformattedRead(fname);
    if (f == 0)
      return 0;

  }

  data = new SimuData;
  if (data == 0) {
    delete f;
    return 0;
  }

  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();
  data->time = 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();
  data->BoxSize = h.BoxSize = f->readReal64();
  h.Omega0 = f->readReal64();
  h.OmegaLambda = f->readReal64();
  h.HubbleParam = f->readReal64();
  f->endCheckpoint(true);

  long NumPart = 0, NumPartTotal = 0;
  for(int k=0; k<5; k++)
    {
      NumPart += h.npart[k];
      NumPartTotal += (id < 0) ? h.npart[k] : h.npartTotal[k];
    }
  data->NumPart = NumPart;
  data->TotalNumPart = NumPartTotal;

  if (loadflags & NEED_POSITION) {
      
    for (int i = 0; i < 3; i++) {
      data->Pos[i] = new float[data->NumPart];
      if (data->Pos[i] == 0) {
	delete data;
	return 0;
      }
    }

    f->beginCheckpoint();
    for(int k = 0, p = 0; k < 5; k++) {
      for(int n = 0; n < h.npart[k]; n++) {
	data->Pos[0][p] = f->readReal32();
	data->Pos[1][p] = f->readReal32();
	data->Pos[2][p] = f->readReal32();
	p++;
      }
    }
    f->endCheckpoint();
    
  } else {
    // Skip positions
    f->skip(NumPart * 3 * sizeof(float) + 2*4);
  }
  
  if (loadflags & NEED_VELOCITY) {
    for (int i = 0; i < 3; i++)
      {
	data->Vel[i] = new float[data->NumPart];
	if (data->Vel[i] == 0) 
	  {
	    delete data;
	    return 0;
	  }
      }

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

  // Skip ids
  if (loadflags & NEED_GADGET_ID) {
    f->beginCheckpoint();
    data->Id = new int[data->NumPart];
    if (data->Id == 0)
      {
	delete data;
	return 0;
      }

    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]*4);
  }

  delete f;

  return data;
}