/*********************************************************************
 *                     CS433 MP3 (CHARM++)
 *                Submitted By : Sameer Kumar.
 *  ---------------------------------------------------------------
 *
 *  This program has been written in charm++. It creates an array of 
 *  chares. Each chare in this array is responsible for one cell. 
 *  Each cell first sends its coordinates to half of its neighbors. 
 *  (Entry method send_coord()). On receiving coordinates (entry method 
 *  recv_coord )from its neighbors, the chare calculates the interaction 
 *  between its cell and the received neighbor. It then sends forces 
 *  (entry method  recv forces) back to that cell. 
 *  On receiving the forces the cell updates its forces. After all 
 *  interactions have been computed, and all forces received 
 *  (indicated by MAX_COUNT) the chare calls update_cell to calculate 
 *  the kinetic energy. The potential and kinetic energies of all the 
 *  cells are added at the end of the step and printed. 
 *
 *  This program uses variable sized messages for communication and 
 *  only coordinates and forces are exchanged, not the whole atom.
 *
 *********************************************************************/

#include <iostream.h>

#include "littleMD3.h"
#include <sys/time.h>
#include <unistd.h>

#include "CommLB.h"
CkChareID mid;
CProxy_CellElement aid;
SimParams params;

const int neighbors_send[13][3] = { {-1,-1,-1}, {-1,-1, 0}, {-1,-1, 1},
				    {-1, 0,-1}, {-1, 0, 0}, {-1, 0, 1},
				    {-1, 1,-1}, {-1, 1, 0}, {-1, 1, 1},
				    { 0,-1,-1}, { 0,-1, 0}, { 0,-1, 1},
				    { 0, 0,-1} };

const int neighbors_recv[13][3] = { { 1, 1, 1}, { 1, 1, 0}, { 1, 1,-1},
				    { 1, 0, 1}, { 1, 0, 0}, { 1, 0,-1},
				    { 1,-1, 1}, { 1,-1, 0}, { 1,-1,-1},
				    { 0, 1, 1}, { 0, 1, 0}, { 0, 1,-1},
				    { 0, 0, 1} };

// The handler called after reduction is finished. 
void reduction_handler(void *proxy, int size, void *data){

    double *energy = (double *)data;
    double kin_energy = energy[1];
    double pot_energy = energy[0];

    CProxy_main mproxy(mid);
    Output * out = new Output;
    out->pot_energy = pot_energy;
    out->kin_energy = kin_energy;
    mproxy.maindone(out);    
}

int nsteps = 0;
main::main(CkArgMsg *msg)
{
    int argc = msg->argc;
    char **argv = msg->argv;

    if(argc > 1)
	nsteps = atoi(argv[1]);

    struct timeval tv;
    gettimeofday(&tv, NULL);
    
    //CkPrintf("# %d \n", tv.tv_sec);
    CkPrintf("In main\n");
    manager_init();
    //CreateRandCentLB();
    CreateCommLB();

    CkPrintf("before Init\n");
    init();        // Initialize the params struct and create the array.  
    start();       // Send messages to the array members, to send coords.
    CkPrintf("After main\n");
}

//Barrier for all array members to have finished.

void main::maindone(Output *out){
    
    static int nfinished = 0;
    pot_energy += out->pot_energy;
    kin_energy += out->kin_energy;

    delete out;
    nfinished ++;

    if (nfinished < params.cell_dim_x * params.cell_dim_y * params.cell_dim_z)
	return;

    nfinished = 0;
    
    CkPrintf("Step %d PE = %le  KE = %le Energy = %le\n", steps_completed
		 , pot_energy, kin_energy, kin_energy + pot_energy);
    
    steps_completed ++;

    pot_energy = 0.0;
    kin_energy = 0.0;

    if(steps_completed < params.steps)
	start();
    else 
	CkExit();
}


void setParamsDebug()
{
    params.steps = 10;
    if(nsteps > 0)
	params.steps = nsteps;

    params.n_atoms = 5000;
    params.min_x = 0;
    params.min_y = 0;
    params.min_z = 0;
    params.max_x = 40;
    params.max_y = 40;
    params.max_z = 20;
    params.cutoff = 15;
    params.margin = 2;
    params.step_sz = 1.;
}

void setParamsReal()
{
    params.steps = 10;
    if(nsteps > 0)
	params.steps = nsteps;
    params.n_atoms = 5000;
    params.min_x = 0;
    params.min_y = 0;
    params.min_z = 0;
    params.max_x = 300;
    params.max_y = 300;
    params.max_z = 300;
    params.cutoff = 15;
    params.margin = 2;
    params.step_sz = 1.;
}

void main::init(){

    // Set some startup parameters
    // These values can be customized
    
    setParamsDebug();
    //setParamsReal();
    
    pot_energy = 0;
    kin_energy = 0;
    
    const double pdx = params.max_x - params.min_x;
    const double pdy = params.max_y - params.min_y;
    const double pdz = params.max_z - params.min_z;
    
    params.cell_dim_x = 1+ (int) (pdx / (params.cutoff + params.margin));
    params.cell_dim_y = 1+ (int) (pdy / (params.cutoff + params.margin));
    params.cell_dim_z = 1+ (int) (pdz / (params.cutoff + params.margin));
    
    // Handle to main.
    mid = thishandle;

    // Allocate the array of chares.

    aid = CProxy_CellElement::ckNew( params.cell_dim_x * params.cell_dim_y 
    		       * params.cell_dim_z );
    
    //    CkPrintf("NELEMENTS = %d\n", params.cell_dim_x * params.cell_dim_y 
    // * params.cell_dim_z ); 
    
    // CProxy_CellElement cell_proxy(aid);    

    steps_completed = 0;
    /*
    CkRegisterArrayReductionHandler(cell_proxy, reduction_handler, 
				    (void *)&cell_proxy);
    */
}

void main::start(){
    
    CProxy_CellElement cell_proxy(aid);

    for(int pos = 0; pos < params.cell_dim_x * params.cell_dim_y * 
	    params.cell_dim_z; pos++)
	cell_proxy[pos].send_coord();
}

// Create contsructor.

CellElement::CellElement() {
    
    recv_count = 0;
    
    pot_energy = 0.0;
    kin_energy = 0.0;
    first_step = true;

    int xi,yi,zi;
    
    zi = thisIndex % params.cell_dim_z;
    int div = thisIndex / params.cell_dim_z;
    
    xi = div % params.cell_dim_x;
    yi = div / params.cell_dim_x; 
    
    init_cell(&cell,xi,yi,zi);
    
    atom_ptr = cell->atoms;

    MAX_COUNT = 0;           // Total number of neighborr of the cell.

    for(int nbor=0; nbor < 13; nbor++) {

	const int xn = xi + neighbors_send[nbor][0];
	const int yn = yi + neighbors_send[nbor][1];
	const int zn = zi + neighbors_send[nbor][2];
	
	if (xn < 0 || xn >= params.cell_dim_x)
	    continue;
	
	if (yn < 0 || yn >= params.cell_dim_y)
	    continue;
	
	if (zn < 0 || zn >= params.cell_dim_z)
	    continue;
	MAX_COUNT++;
    }

    for(int nbor=0; nbor < 13; nbor++) {

	const int xn = xi + neighbors_recv[nbor][0];
	const int yn = yi + neighbors_recv[nbor][1];
	const int zn = zi + neighbors_recv[nbor][2];
	
	if (xn < 0 || xn >= params.cell_dim_x)
	    continue;
	
	if (yn < 0 || yn >= params.cell_dim_y)
	    continue;
	
	if (zn < 0 || zn >= params.cell_dim_z)
	    continue;
	MAX_COUNT ++;
    }

    usesAtSync=CmiTrue;
}

CellElement::CellElement(CkMigrateMessage *){
}

void CellElement::pup(PUP::er &p) {
    ArrayElement1D::pup(p);//Pack superclass
    
    p(recv_count);
    p(MAX_COUNT);
    p(pot_energy);
    p(kin_energy);
    p(first_step);

    if (p.isUnpacking()) {
	cell = new Cell;
    }

    p((char*) cell, sizeof(Cell));
    if (p.isUnpacking()) {
	atom_ptr = new Atom[cell->n_atoms];
    }
    p((char*) atom_ptr, sizeof(Atom) * cell->n_atoms);
    
    cell->atoms = atom_ptr;
}

// Send coordinates to neighbors.

void CellElement::send_coord(){

    int nbor,xi,yi,zi;
    Msg * mesg;

    //    CkPrintf("In send coord\n");
    for(nbor=0; nbor < 13; nbor++) {
	xi = cell->my_x;
	yi = cell->my_y;
	zi = cell->my_z;

	const int xn = xi + neighbors_send[nbor][0];
	const int yn = yi + neighbors_send[nbor][1];
	const int zn = zi + neighbors_send[nbor][2];
	
	if (xn < 0 || xn >= params.cell_dim_x)
	    continue;
	
	if (yn < 0 || yn >= params.cell_dim_y)
	    continue;
	
	if (zn < 0 || zn >= params.cell_dim_z)
	    continue;

	int size;
	size = cell->n_atoms;	
	mesg = new(&size,0) Msg;
	
	mesg->x = xi;
	mesg->y = yi;
	mesg->z = zi;
	mesg->n_elem = cell->n_atoms;

	for(int count = 0; count < cell->n_atoms; count++){
	    mesg->data[count].m = cell->atoms[count].m;
	    mesg->data[count].q = cell->atoms[count].q;
	    mesg->data[count].val_x = cell->atoms[count].x;
	    mesg->data[count].val_y = cell->atoms[count].y;
	    mesg->data[count].val_z = cell->atoms[count].z;
	}
	
	CProxy_CellElement cell_proxy(aid);
	int pos = (yn * params.cell_dim_x + xn) * params.cell_dim_z + zn; 
	cell_proxy[pos].recv_coord(mesg);
    }
}

Cell * construct_cell(Msg * mesg){
    Cell * cell = new Cell;

    cell->my_x = mesg->x;
    cell->my_y = mesg->y;
    cell->my_z = mesg->z;

    cell->n_atoms = mesg->n_elem;

    cell->atoms = new Atom[mesg->n_elem];
    for(int count = 0; count < cell->n_atoms; count++){
	cell->atoms[count].m = mesg->data[count].m; 
	cell->atoms[count].q = mesg->data[count].q;
	cell->atoms[count].x = mesg->data[count].val_x;
	cell->atoms[count].y = mesg->data[count].val_y;
	cell->atoms[count].z = mesg->data[count].val_z;

	cell->atoms[count].fx = 0.0;
	cell->atoms[count].fy = 0.0;
	cell->atoms[count].fz = 0.0;	
    }
    return cell;
}

void free_cell(Cell * cell){
    delete cell->atoms;
    delete cell;
}

// Receive coordinates from neighbors and compute the interaction.
void CellElement::recv_coord(Msg *mesg){
    
    //    CkPrintf("In recv coord\n");
    Cell *nbr = construct_cell(mesg);
    
    pot_energy += cell_neighbor(cell,nbr);
 
    for(int count = 0; count < nbr->n_atoms; count++){
	mesg->data[count].val_x = nbr->atoms[count].fx;
	mesg->data[count].val_y = nbr->atoms[count].fy;
	mesg->data[count].val_z = nbr->atoms[count].fz;
    }
    
    CProxy_CellElement cell_proxy(aid);

    int pos = (nbr->my_y * params.cell_dim_x + nbr->my_x) * 
	params.cell_dim_z + nbr->my_z; 

    free_cell(nbr);

    cell_proxy[pos].recv_forces(mesg);
    recv_count++;

    if(recv_count == MAX_COUNT){
	//	CkPrintf("itr done\n");
	recv_count = 0;
	pot_energy += cell_self(cell);
	kin_energy = update_cell(cell, first_step);
	first_step = false;
	/*
	double *energy = new double[2];
	
	energy[0] = pot_energy;
	energy[1] = kin_energy;

	//contribute(2*sizeof(double),(void *)energy, CkReduction_sum_double);
	
	CProxy_main mproxy(mid);
	Output * out = new Output;
	out->pot_energy = pot_energy;
	out->kin_energy = kin_energy;    
    
	mproxy.maindone(out);    

	pot_energy = 0.0;
	kin_energy = 0.0;
	*/
	AtSync();
    }
}

// Receive forces from neighbor and increment internal atom's forces.
void CellElement::recv_forces(Msg *msg){
    
    //    CkPrintf("recv forces\n");
    int n_elem = msg->n_elem;
    
    for(int count = 0; count < n_elem; count++){
	cell->atoms[count].fx += msg->data[count].val_x;
	cell->atoms[count].fy += msg->data[count].val_y;
	cell->atoms[count].fz += msg->data[count].val_z;
    }

    recv_count ++;
    delete msg;

    if(recv_count == MAX_COUNT){
	//	CkPrintf("itr done\n");
	recv_count = 0;
	pot_energy += cell_self(cell);
	kin_energy = update_cell(cell, first_step);
	first_step = false;
	/*
	double *energy = new double[2];
	
	energy[0] = pot_energy;
	energy[1] = kin_energy;

       	//contribute(2*sizeof(double),(void *)energy,CkReduction_sum_double);
	
	CProxy_main mproxy(mid);
	Output * out = new Output;
	out->pot_energy = pot_energy;
	out->kin_energy = kin_energy;    
    
	mproxy.maindone(out);    
	
	pot_energy = 0.0;
	kin_energy = 0.0;
	*/
	AtSync();
    }
}

void CellElement::ResumeFromSync(void){
    //    CkPrintf("Finished Load Balancing\n");
    //        CProxy_CellElement cell_proxy(aid);
    /* 
    if((CkMyPe() == 0) && (thisIndex == 10))
	CkPrintf("# About to resume %5.3lf\n", CmiWallTimer());
    */
    CProxy_main mproxy(mid);
    Output * out = new Output;
    out->pot_energy = pot_energy;
    out->kin_energy = kin_energy;    
    
    mproxy.maindone(out);    
    
    pot_energy = 0.0;
    kin_energy = 0.0;
}

// Creation routines for variable size messages.
void* Msg::alloc(int mnum, size_t size, int *sizes, int priobits){
    
    int total_size = size + sizes[0] * sizeof(Data);

    Msg * mesg = (Msg *)CkAllocMsg(mnum, total_size, priobits);

    mesg->data = (Data*)((char*)mesg + sizeof(Msg));
    
    return (void *)mesg;
}

void* Msg::pack(Msg *mesg){
    return (void *) mesg;
}

Msg* Msg::unpack(void *buf){
    Msg *mesg = (Msg*)buf;

    mesg->data = (Data*)((char*)mesg + sizeof(Msg));
    return mesg;
}

#include "littleMD.def.h"

