#include <stdio.h>
#include <unistd.h>
#include <math.h>
#include "charm++.h"
#include "LBDatabase.h"

#include "CentralLB.h"
#include "RandCentLB.h"
#include "CommLB.h"
#include "manager.h"

#define NEW_UPDATE_CELL

struct Atom {
    double m;     // Mass in AMU (= 1.66E-27 kg)
    double q;     // Charge in units of elemental charge ( = 1.602E-19C)
    double x, y, z;     // Position in Angstroms ( = 1E-10 m)
    double vx, vy, vz;    // Velocity in angstroms/fs ( = 1E5m/s)
    double vhx, vhy, vhz;  // Half-step velocity, used for Verlet integration
    double fx, fy, fz;    // Force in newtons
};

struct Cell {
    int n_atoms;
    int my_x, my_y, my_z;  // index of the cell
    double min_x, min_y, min_z;
    double max_x, max_y, max_z;   // probably unnecessary?
    Atom* atoms;   // array to be allocated ;
};

struct SimParams {
    int steps;
    int n_atoms;
    double min_x, min_y, min_z;
    double max_x, max_y, max_z;  // Angstroms
    int cell_dim_x, cell_dim_y, cell_dim_z;
    double cutoff;  // Angstroms
    double margin;  // Angstroms
    double step_sz; // femtoseconds (=1E-15s)
};

// A few global variables 
extern SimParams params;

// Functions in littleMD.C

void calc_forces(void);
void update_coordinates(void);

// Functions in utilities.C
void init_cell(Cell** cell, int xi, int yi, int zi);
double cell_self(Cell* this_cell);
double cell_neighbor(Cell* cell1, Cell* cell2);
double calc_force(Atom *atom1 , Atom* atom2);

double update_cell(Cell* this_cell, int first_step);

// structure to hold coordinates/forces of atoms.
struct Data{
    double m,q,val_x,val_y,val_z;
};

#include "littleMD.decl.h"

// Message for cells to send output to main.
class Output : public CMessage_Output {
 public:
    double pot_energy,kin_energy;
};

// Var size message to send cooredinates and forces.
class Msg : public CMessage_Msg {
 public:
    int x,y,z;
    int n_elem;
    Data * data;
    
    static void *alloc(int mnum, size_t size, int *sizes, int priobits);
    static void *pack(Msg *msg);
    static Msg *unpack(void *buf);  
};

class main : public Chare{
    double pot_energy;
    double kin_energy;
    int steps_completed;   // Number of steps completed.
    int n_recvd;           // Number of outputs recvd.
    
 public:
    main(CkArgMsg *);
    void init();
    void start();
    void maindone(Output *);
};

// The cell array element.
class CellElement : public ArrayElement1D {
    Cell * cell;
    Atom * atom_ptr;
    int recv_count, MAX_COUNT;
    double pot_energy,kin_energy;
    bool first_step;

 public:
    CellElement(void);
    CellElement(CkMigrateMessage *);

    //Function to send and recv, coordinates and forces.
    void send_coord(); 
    void recv_coord(Msg *);
    void recv_forces(Msg *);
    void pup(PUP::er &p);
    void ResumeFromSync(void); 
};



