#include <Python.h>
#include <numpy/arrayobject.h>
#include "structmember.h"
#include <mpi.h>

#include "allvars.h"
#include "proto.h"
#include "ptreelib.h"



/******************************************************************************

TREE OBJECT

*******************************************************************************/




static void
Tree_dealloc(Tree* self)
{
        
    /* here, we have to deallocate all what we have allocated */
    
    free(self->Exportflag);
    
    ngb_treefree(self);
    
    force_treefree(self);
    
    domain_deallocate(self);
    
    //free(self->Nodes_base);
    //free(self->Nextnode);
    //free(self->Father);
    //free(self->Ngblist);
    
    
    free(self->CommBuffer);
    free(self->P);
    free(self->SphP);
    
    self->ob_type->tp_free((PyObject*)self);
    
}

static PyObject *
Tree_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
{
    Tree *self;

    self = (Tree *)type->tp_alloc(type, 0);
    if (self != NULL) 
      {
	
        /* not clear what I have to put here */
	
	
      }

    return (PyObject *)self;
}

static int
Tree_init(Tree *self, PyObject *args, PyObject *kwds)
{

    import_array();

    int i;
    size_t bytes;

    PyArrayObject *ntype,*pos,*vel,*mass,*num,*tpe;

    static char *kwlist[] = {"npart", "pos","vel","mass","num","tpe", NULL};

    if (! PyArg_ParseTupleAndKeywords(args, kwds, "|OOOOOO",kwlist,&ntype,&pos,&vel,&mass,&num,&tpe))
        return -1; 

    /* MPI init */

    MPI_Comm_rank(MPI_COMM_WORLD, &self->ThisTask);
    MPI_Comm_size(MPI_COMM_WORLD, &self->NTask);
        
    for(self->PTask = 0; self->NTask > (1 << self->PTask); self->PTask++);


    /* parameters */				/* should be done differently */
    self->All.PartAllocFactor 		= 2.0;
    self->All.TreeAllocFactor 		= 2.0;
    self->All.ErrTolTheta     		= 0.7;
    self->All.DesNumNgb       		= 33;
    self->All.MaxNumNgbDeviation 	= 3;
    self->All.MinGasHsmlFractional 	= 0.25;
    self->All.BufferSize        	= 30.;
    self->All.OutputInfo		= 0;		/* output info */
    self->All.PeanoHilbertOrder		= 0;		/* peano hilbert order */
    self->All.ComovingIntegrationOn	= 0;
    self->All.PeriodicBoundariesOn      = 0;
    self->All.TypeOfOpeningCriterion	= 0;
    self->All.ErrTolForceAcc		= 0.005;
    
    /* physical ctes */
    self->All.OmegaLambda		= 0;
    self->All.Omega0			= 0;
    self->All.Hubble 			= 0;
    self->All.G 			= 1;
    
    /* count number of particles */

    self->NtypeLocal[0] = *(int*) (ntype->data + 0*(ntype->strides[0]));
    self->NtypeLocal[1] = *(int*) (ntype->data + 1*(ntype->strides[0]));
    self->NtypeLocal[2] = *(int*) (ntype->data + 2*(ntype->strides[0]));
    self->NtypeLocal[3] = *(int*) (ntype->data + 3*(ntype->strides[0]));
    self->NtypeLocal[4] = *(int*) (ntype->data + 4*(ntype->strides[0]));
    self->NtypeLocal[5] = *(int*) (ntype->data + 5*(ntype->strides[0]));
    
    self->NumPart = 0;
    self->N_gas   = self->NtypeLocal[0];
    for (i = 0; i < 6; i++) 
      self->NumPart += self->NtypeLocal[i];
         
    MPI_Allreduce(&self->NumPart, &self->All.TotNumPart, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
    MPI_Allreduce(&self->N_gas, &self->All.TotN_gas, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
    
    for (i = 0; i < 6; i++) 
      self->All.SofteningTable[i] = 0.1;			/* a changer !!!! */

    for (i = 0; i < 6; i++) 
      self->All.ForceSoftening[i] = 0.1;			/* a changer !!!! */
    
    
    /* update parameters */    

    self->All.MaxPart	 = self->All.PartAllocFactor * (self->All.TotNumPart / self->NTask);   
    self->All.MaxPartSph = self->All.PartAllocFactor * (self->All.TotN_gas / self->NTask);
    self->All.MinGasHsml = self->All.MinGasHsmlFractional * self->All.ForceSoftening[0];
    

    self->All.BunchSizeDomain =
      (self->All.BufferSize * 1024 * 1024) / (sizeof(struct particle_data) + sizeof(struct sph_particle_data) +
   					sizeof(peanokey));

    if(self->All.BunchSizeDomain & 1)
      self->All.BunchSizeDomain -= 1;	  /* make sure that All.BunchSizeDomain is even 
   				             --> 8-byte alignment of DomainKeyBuf for 64bit processors */


    self->All.BunchSizeForce =
     (self->All.BufferSize * 1024 * 1024) / (sizeof(struct gravdata_index) + 2 * sizeof(struct gravdata_in));

    if(self->All.BunchSizeForce & 1)
      self->All.BunchSizeForce -= 1;	/* make sure that All.BunchSizeForce is an even number 
				   --> 8-byte alignment for 64bit processors */






    self->first_flag = 0;

    /*********************/
    /* some allocation   */
    /*********************/


    if(!(self->CommBuffer = malloc(bytes = self->All.BufferSize * 1024 * 1024)))
      {
    	printf("failed to allocate memory for `CommBuffer' (%g MB).\n", bytes / (1024.0 * 1024.0));
    	endrun(2);
      }


    self->Exportflag = malloc(self->NTask * sizeof(char)); 
    self->GravDataIndexTable = (struct gravdata_index *) self->CommBuffer;
    self->GravDataIn = (struct gravdata_in *) (self->GravDataIndexTable + self->All.BunchSizeForce);
    self->GravDataGet = self->GravDataIn + self->All.BunchSizeForce;
    self->GravDataOut = self->GravDataIn;
    self->GravDataResult = self->GravDataGet;
       
    
    /*********************/
    /* create P          */
    /*********************/
    
    if(!(self->P = malloc(bytes = self->All.MaxPart * sizeof(struct particle_data))))
      {
        printf("failed to allocate memory for `P' (%g MB).\n", bytes / (1024.0 * 1024.0));
        endrun(self,1);
      }    

    if(!(self->SphP = malloc(bytes = self->All.MaxPartSph * sizeof(struct sph_particle_data))))
      {
        printf("failed to allocate memory for `SphP' (%g MB) %d.\n", bytes / (1024.0 * 1024.0), sizeof(struct sph_particle_data));
        endrun(1);
      }
    
    /*********************/
    /* init P            */
    /*********************/

    for (i = 0; i < pos->dimensions[0]; i++) 
      {
        self->P[i].Pos[0] = *(float *) (pos->data + i*(pos->strides[0]) + 0*pos->strides[1]);
        self->P[i].Pos[1] = *(float *) (pos->data + i*(pos->strides[0]) + 1*pos->strides[1]);
        self->P[i].Pos[2] = *(float *) (pos->data + i*(pos->strides[0]) + 2*pos->strides[1]);
        self->P[i].Vel[0] = *(float *) (vel->data + i*(vel->strides[0]) + 0*vel->strides[1]);
        self->P[i].Vel[1] = *(float *) (vel->data + i*(vel->strides[0]) + 1*vel->strides[1]);
        self->P[i].Vel[2] = *(float *) (vel->data + i*(vel->strides[0]) + 2*vel->strides[1]);	
        self->P[i].Mass   = *(float *) (mass->data + i*(mass->strides[0]));
        self->P[i].ID   = *(unsigned int *) (num->data + i*(num->strides[0]));
	self->P[i].Type   = *(int *)        (tpe->data + i*(tpe->strides[0]));										/* this should be changed... */ 
      }

    /***************************************
     * domain decomposition construction   * 
    /***************************************/
    
    domain_allocate(self);
    domain_Decomposition(self);


    return 0;
}



static PyObject *
Tree_info(Tree* self)
{

    //static PyObject *format = NULL;
    //PyObject *args, *result;

   
    printf("(%d) NumPart = %d\n",self->ThisTask,self->NumPart);
    printf("(%d) N_gas   = %d\n",self->ThisTask,self->N_gas);
    
    printf("(%d) DomainLen	= %g\n",self->ThisTask,self->DomainLen);    
    printf("(%d) DomainCenter x = %g\n",self->ThisTask,self->DomainCenter[0]);
    printf("(%d) DomainCenter y = %g\n",self->ThisTask,self->DomainCenter[1]);
    printf("(%d) DomainCenter z = %g\n",self->ThisTask,self->DomainCenter[2]);
    printf("(%d) DomainCorner x = %g\n",self->ThisTask,self->DomainCorner[0]);
    printf("(%d) DomainCorner y = %g\n",self->ThisTask,self->DomainCorner[1]);
    printf("(%d) DomainCorner z = %g\n",self->ThisTask,self->DomainCorner[2]);
    printf("(%d) NTopnodes = %d\n",self->ThisTask,self->NTopnodes);
    printf("(%d) NTopleaves = %d\n",self->ThisTask,self->NTopleaves);

	
    return Py_BuildValue("i",1);
}


static PyObject *
Tree_GetExchanges(Tree* self)
{

    PyArrayObject *a_num,*a_procs;
    int   ld[1];
    int   i;

    /* create a NumPy object */	  
    ld[0]= self->NSend;
    a_num  = (PyArrayObject *) PyArray_FromDims(1,ld,PyArray_INT);
    a_procs= (PyArrayObject *) PyArray_FromDims(1,ld,PyArray_INT);     
    
    for (i = 0; i < a_num->dimensions[0]; i++) 
      {
	*(int *)(a_num->data   + i*(a_num->strides[0]))    = self->DomainIdProc[i].ID;
	*(int *)(a_procs->data + i*(a_procs->strides[0]))  = self->DomainIdProc[i].Proc;
      }    	
	
	
    return Py_BuildValue("OO",a_num,a_procs);
}


static PyObject *
Tree_SetParticles(Tree* self, PyObject *args)
{

    int i;

    PyArrayObject *ntype,*pos,*vel,*mass,*num,*tpe;


    if (! PyArg_ParseTuple(args, "OOOOOO",&ntype,&pos,&vel,&mass,&num,&tpe))
        return Py_BuildValue("i",-1);
        
    /* count number of particles */

    self->NtypeLocal[0] = *(int*) (ntype->data + 0*(ntype->strides[0]));
    self->NtypeLocal[1] = *(int*) (ntype->data + 1*(ntype->strides[0]));
    self->NtypeLocal[2] = *(int*) (ntype->data + 2*(ntype->strides[0]));
    self->NtypeLocal[3] = *(int*) (ntype->data + 3*(ntype->strides[0]));
    self->NtypeLocal[4] = *(int*) (ntype->data + 4*(ntype->strides[0]));
    self->NtypeLocal[5] = *(int*) (ntype->data + 5*(ntype->strides[0]));
    
    self->NumPart = 0;
    self->N_gas   = self->NtypeLocal[0];
    for (i = 0; i < 6; i++) 
      self->NumPart += self->NtypeLocal[i];
      
	 
    MPI_Allreduce(&self->NumPart, &self->All.TotNumPart, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
    MPI_Allreduce(&self->N_gas, &self->All.TotN_gas, 1, MPI_INT, MPI_SUM, MPI_COMM_WORLD);
    

    /*********************/
    /* init P		 */
    /*********************/

    for (i = 0; i < pos->dimensions[0]; i++) 
      {
	self->P[i].Pos[0] = *(float *) (pos->data + i*(pos->strides[0]) + 0*pos->strides[1]);
	self->P[i].Pos[1] = *(float *) (pos->data + i*(pos->strides[0]) + 1*pos->strides[1]);
	self->P[i].Pos[2] = *(float *) (pos->data + i*(pos->strides[0]) + 2*pos->strides[1]);
	self->P[i].Vel[0] = *(float *) (vel->data + i*(vel->strides[0]) + 0*vel->strides[1]);
	self->P[i].Vel[1] = *(float *) (vel->data + i*(vel->strides[0]) + 1*vel->strides[1]);
	self->P[i].Vel[2] = *(float *) (vel->data + i*(vel->strides[0]) + 2*vel->strides[1]); 
	self->P[i].Mass   = *(float *) (mass->data + i*(mass->strides[0]));
	self->P[i].ID	= *(unsigned int *) (num->data + i*(num->strides[0]));
      self->P[i].Type	= *(int *)	  (tpe->data + i*(tpe->strides[0]));									      /* this should be changed... */ 
      }

    return Py_BuildValue("i",1);
}


static PyObject *
Tree_BuildTree(Tree* self)
{

    /************************
    /* tree construction    *
    /************************/    
    force_treeallocate(self,self->All.TreeAllocFactor * self->All.MaxPart, self->All.MaxPart);
    force_treebuild(self,self->NumPart);
    
    /************************
    /* ngb                  *
    /************************/
    ngb_treeallocate(self,self->NumPart);
	
    return Py_BuildValue("i",1);
}

static PyObject *
Tree_AllPotential(Tree* self)
{
  compute_potential(self);
  printf("\n %g %g %g %g\n",self->P[0].Pos[0],self->P[0].Pos[1],self->P[0].Pos[2],self->P[0].Potential);
  return Py_BuildValue("i",1);
}


static PyObject *
Tree_AllAcceleration(Tree* self)
{

  /* gravitational acceleration */
  gravity_tree(self);

  return Py_BuildValue("i",1);
}

static PyObject *
Tree_Potential(Tree* self, PyObject *args)
{


    PyArrayObject *pos;
    float eps;

    if (! PyArg_ParseTuple(args, "Of",&pos,&eps))
        return PyString_FromString("error");
        
    PyArrayObject *pot;
    int i;
    int   ld[1];
    int input_dimension;
    size_t bytes;
    
    input_dimension =pos->nd;
    
    if (input_dimension != 2)
      PyErr_SetString(PyExc_ValueError,"dimension of first argument must be 2");

    if (pos->descr->type_num != PyArray_FLOAT)
      PyErr_SetString(PyExc_ValueError,"argument 1 must be of type Float32");

    
    /* create a NumPy object */	  
    ld[0]=pos->dimensions[0];
    pot = (PyArrayObject *) PyArray_FromDims(1,ld,PyArray_FLOAT);  
    

    self->NumPartQ = pos->dimensions[0];
    self->ForceSofteningQ = eps;
  
    if(!(self->Q = malloc(bytes = self->NumPartQ * sizeof(struct particle_data))))
      {
    	printf("failed to allocate memory for `Q' (%g MB).\n", bytes / (1024.0 * 1024.0));
    	endrun(self,1);
      }    
  
  
    for (i = 0; i < pos->dimensions[0]; i++) 
      {
        self->Q[i].Pos[0] = *(float *) (pos->data + i*(pos->strides[0]) + 0*pos->strides[1]);
        self->Q[i].Pos[1] = *(float *) (pos->data + i*(pos->strides[0]) + 1*pos->strides[1]);
        self->Q[i].Pos[2] = *(float *) (pos->data + i*(pos->strides[0]) + 2*pos->strides[1]);
	self->Q[i].Type = 0;	
	self->Q[i].Potential = 0;
      }

    
    compute_potential_sub(self);
        

    for (i = 0; i < pos->dimensions[0]; i++) 
      {
	*(float *)(pot->data + i*(pot->strides[0]))  = self->Q[i].Potential;
      }  

    
    free(self->Q);
    
    return PyArray_Return(pot);
}


static PyObject *
Tree_Acceleration(Tree* self, PyObject *args)
{


    PyArrayObject *pos;
    float eps;

    if (! PyArg_ParseTuple(args, "Of",&pos,&eps))
        return PyString_FromString("error");
        
    PyArrayObject *acc;
    int i;
    int   ld[1];
    int input_dimension;
    size_t bytes;
    
    input_dimension =pos->nd;
    
    if (input_dimension != 2)
      PyErr_SetString(PyExc_ValueError,"dimension of first argument must be 2");

    if (pos->descr->type_num != PyArray_FLOAT)
      PyErr_SetString(PyExc_ValueError,"argument 1 must be of type Float32");

    
    /* create a NumPy object */	  
    ld[0]=pos->dimensions[0];
    ld[1]=pos->dimensions[1];
    acc = (PyArrayObject *) PyArray_FromDims(pos->nd,ld,pos->descr->type_num); 
    acc = (PyArrayObject *) PyArray_SimpleNew(pos->nd,pos->dimensions,pos->descr->type_num); 
    

    self->NumPartQ = pos->dimensions[0];
    self->ForceSofteningQ = eps;
  
    if(!(self->Q = malloc(bytes = self->NumPartQ * sizeof(struct particle_data))))
      {
    	printf("failed to allocate memory for `Q' (%g MB).\n", bytes / (1024.0 * 1024.0));
    	endrun(self,1);
      }    
  
  
    for (i = 0; i < pos->dimensions[0]; i++) 
      {
        self->Q[i].Pos[0] = *(float *) (pos->data + i*(pos->strides[0]) + 0*pos->strides[1]);
        self->Q[i].Pos[1] = *(float *) (pos->data + i*(pos->strides[0]) + 1*pos->strides[1]);
        self->Q[i].Pos[2] = *(float *) (pos->data + i*(pos->strides[0]) + 2*pos->strides[1]);
	self->Q[i].Type = 0;
	self->Q[i].GravAccel[0] = 0;
	self->Q[i].GravAccel[1] = 0;
	self->Q[i].GravAccel[2] = 0;	
      }

    
    gravity_tree_sub(self);
        

    for (i = 0; i < pos->dimensions[0]; i++) 
      {
	*(float *)(acc->data + i*(acc->strides[0]) + 0*acc->strides[1]) = self->Q[i].GravAccel[0];
        *(float *)(acc->data + i*(acc->strides[0]) + 1*acc->strides[1]) = self->Q[i].GravAccel[1];
        *(float *)(acc->data + i*(acc->strides[0]) + 2*acc->strides[1]) = self->Q[i].GravAccel[2];
      }  

    
    free(self->Q);
    
    return PyArray_Return(acc);
}


static PyMemberDef Tree_members[] = {

    //{"first", T_OBJECT_EX, offsetof(Tree, first), 0,
    // "first name"},
    //{"list", T_OBJECT_EX, offsetof(Tree, list), 0,
    // "list of"},     
    //{"number", T_INT, offsetof(Tree, number), 0,
    // "Tree number"},
     
    {NULL}  /* Sentinel */
};



static PyMethodDef Tree_methods[] = {

    {"info", (PyCFunction)Tree_info, METH_NOARGS,
     "Return some info"
    },

    {"GetExchanges", (PyCFunction)Tree_GetExchanges, METH_NOARGS,
    "This function returns the list of particles that have been exchanged and the corresponding processor."
    },    

    {"SetParticles", (PyCFunction)Tree_SetParticles, METH_VARARGS,
    "Set values of particles"
    },    

    {"BuildTree", (PyCFunction)Tree_BuildTree, METH_NOARGS,
     "Build the tree"
    },    

    {"Potential", (PyCFunction)Tree_Potential, METH_VARARGS,
     "Computes the potential at a given position using the tree"
    },    

    {"Acceleration", (PyCFunction)Tree_Acceleration, METH_VARARGS,
     "Computes the acceleration at a given position using the tree"
    },    

    {"AllPotential", (PyCFunction)Tree_AllPotential, METH_VARARGS,
     "Computes the potential for each particle"
    },    

    {"AllAcceleration", (PyCFunction)Tree_AllAcceleration, METH_VARARGS,
     "Computes the acceleration for each particle"
    },    
                    
    {NULL}  /* Sentinel */
};


static PyTypeObject TreeType = {
    PyObject_HEAD_INIT(NULL)
    0,                         /*ob_size*/
    "tree.Tree",               /*tp_name*/
    sizeof(Tree),              /*tp_basicsize*/
    0,                         /*tp_itemsize*/
    (destructor)Tree_dealloc,  /*tp_dealloc*/
    0,                         /*tp_print*/
    0,                         /*tp_getattr*/
    0,                         /*tp_setattr*/
    0,                         /*tp_compare*/
    0,                         /*tp_repr*/
    0,                         /*tp_as_number*/
    0,                         /*tp_as_sequence*/
    0,                         /*tp_as_mapping*/
    0,                         /*tp_hash */
    0,                         /*tp_call*/
    0,                         /*tp_str*/
    0,                         /*tp_getattro*/
    0,                         /*tp_setattro*/
    0,                         /*tp_as_buffer*/
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /*tp_flags*/
    "Tree objects",            /* tp_doc */
    0,		               /* tp_traverse */
    0,		               /* tp_clear */
    0,		               /* tp_richcompare */
    0,		               /* tp_weaklistoffset */
    0,		               /* tp_iter */
    0,		               /* tp_iternext */
    Tree_methods,              /* tp_methods */
    Tree_members,              /* tp_members */
    0,                         /* tp_getset */
    0,                         /* tp_base */
    0,                         /* tp_dict */
    0,                         /* tp_descr_get */
    0,                         /* tp_descr_set */
    0,                         /* tp_dictoffset */
    (initproc)Tree_init,       /* tp_init */
    0,                         /* tp_alloc */
    Tree_new,                  /* tp_new */
};

static PyMethodDef module_methods[] = {
    {NULL}  /* Sentinel */
};

#ifndef PyMODINIT_FUNC	/* declarations for DLL import/export */
#define PyMODINIT_FUNC void
#endif
PyMODINIT_FUNC
initptreelib(void) 
{
    PyObject* m;

    if (PyType_Ready(&TreeType) < 0)
        return;

    m = Py_InitModule3("ptreelib", module_methods,
                       "Example module that creates an extension type.");

    if (m == NULL)
      return;

    Py_INCREF(&TreeType);
    PyModule_AddObject(m, "Tree", (PyObject *)&TreeType);
}
