/* mddisk.c */

/*
  Copyright (c) 1997 D.C. Rapaport 

  Basic interactive 2-D soft-disk MD simulation; based on article
  published in Computers in Physics (1997).
*/

# include <math.h>               /* standard C includes */
# include <stdlib.h>
# include <stdio.h>

# include <Xm/ArrowB.h>          /* Motif widget includes */
# include <Xm/DrawingA.h>
# include <Xm/Form.h>
# include <Xm/Frame.h>
# include <Xm/Label.h>
# include <Xm/PushB.h>
# include <Xm/RowColumn.h>
# include <Xm/Scale.h>
# include <Xm/Text.h>
# include <Xm/MwmUtil.h>         /* extra for toplevel shell */

/* operations for vector algebra */
typedef struct {double x, y;} VecD;
typedef struct {int x, y;} VecI;
#define VSet(v, sx, sy) v.x = sx, v.y = sy
#define VZero(v) VSet (v, 0, 0)
#define VSCopy(v2, s1, v1) v2.x = (s1) * v1.x, v2.y = (s1) * v1.y
#define VScale(v, s) VSCopy (v, s, v)
#define VAdd(v1, v2, v3) v1.x = v2.x + v3.x, v1.y = v2.y + v3.y
#define VSub(v1, v2, v3) v1.x = v2.x - v3.x, v1.y = v2.y - v3.y
#define VMul(v1, v2, v3) v1.x = v2.x * v3.x, v1.y = v2.y * v3.y
#define VDiv(v1, v2, v3) v1.x = v2.x / v3.x, v1.y = v2.y / v3.y
#define VSAdd(v1, v2, s3, v3) \
   v1.x = v2.x + (s3) * v3.x, v1.y = v2.y + (s3) * v3.y
#define VVSAdd(v1, s2, v2) VSAdd (v1, v1, s2, v2)
#define VProd(v) (v.x * v.y)
#define VLinear(p, s) p.y * s.x + p.x
#define VLenSq(v) (SQ (v.x) + SQ (v.y))

/* misc definitions */
#define SQ(x) ((x) * (x))       
#define T_STRING(s) XmRString, s, strlen (s) + 1
#define ALLOC(y, x, t) y = (t *) realloc (y, (x) * sizeof (t))


/* prototypes */
void SetupGraphics      ( int, char ** );
void SetDefaults        ( void );
void InitializeRun      ( void );
void SingleStep         ( void );
void ComputeForces      ( void );
void LeapfrogIntegrate  ( void );
void CheckPeriodicWrap  ( void );
void CalcEnergy         ( void );
void InitialState       ( void );
void SetNewTemp         ( double );
void SetNewDens         ( double );
void Redraw             ( void );
void Refresh            ( void );
Boolean MdCycle         ( void );

/* callbacks */
void CbButton ( void );
void CbScale ( void );
void CbArrow ( void );
void CbExpose ( void );


/* resources to allow user specification at runtime */
typedef struct {
  int wSize;
  Pixel pixel[2];
} AppData;
AppData appData;

#define COLOR_RES(t, n, v) {t, t, XmRPixel, sizeof (Pixel), \
   XtOffsetOf (AppData, pixel[n]), XmRString, (char *) v}
XtResource resources[] = {
  {"wSize", "wSize", XmRInt, sizeof (int),
   XtOffsetOf (AppData, wSize), XmRImmediate, (XtPointer) 650},
  COLOR_RES ("color0", 0, "black"),
  COLOR_RES ("color1", 1, "green"),
};
XrmOptionDescRec options[] = {
  {"-s", "wSize", XrmoptionSepArg, NULL},
  {"-c0", "color0", XrmoptionSepArg, NULL},
  {"-c1", "color1", XrmoptionSepArg, NULL},
};


/* global variables */
XtAppContext app;               /* application context */
Display *dpy;                   /* display struct */
GC gc;                          /* graphic context */
Pixmap pixmap;                  /* pixmap struct */
Widget wButtonS, wScaleA, wScaleD, wDraw, wText; /* globally accessed widgets */

VecD *r = NULL, *rv = NULL, *ra = NULL, region;
VecI cells, initUcell, wDrawSize;
double deltaT, rCut, temperature, density, uSum, kinEnergy, potEnergy;
int *cellList = NULL, nAtom, changeTemp, changeDens, stepCount, stepDraw,
   nAtomEdge, running;


int main ( int argc, char **argv )
/*
  Purpose:

    MAIN is the main program for MDDISK.
*/
{
/* 
  Initialize graphics 
*/
  SetupGraphics ( argc, argv );
/* 
  Initialize computation.
*/
  InitializeRun ();
/*
  Specify work procedure.
*/
  XtAppAddWorkProc ( app, MdCycle, 0 );
/*
  Main processing and event loop.
*/
  XtAppMainLoop ( app );

  return 0;
}


void SetupGraphics ( int argc, char **argv )
/*
  Purpose:

    SETUPGRAPHICS sets up the graphics.
*/
{
    Widget wArrow, wForm, wFrame, wRowCol, wRowCol2, wTop;
    int j, n, scrn;
    char *buttonName[] = {"Temp: ", "Dens: "};
  
    /* initialize graphics, toplevel shell widget - no resizing allowed */
    wTop = XtVaAppInitialize (&app, "Xmd", options, XtNumber (options),
       &argc, argv, NULL, 
       XmNmwmFunctions, MWM_FUNC_ALL | MWM_FUNC_RESIZE | MWM_FUNC_MAXIMIZE,
       NULL);
    XtGetApplicationResources (wTop, &appData, resources,
       XtNumber (resources), NULL, 0);
    /* various default and initial settings */
    SetDefaults ();
    /* form widget */
    wForm = XtVaCreateManagedWidget ("", xmFormWidgetClass, wTop,
       XmNmarginWidth,            5,
       XmNmarginHeight,           5,
       NULL);
    /* vertical rowcolumn widget */
    wRowCol = XtVaCreateManagedWidget ("", xmRowColumnWidgetClass, wForm,
       XmNpacking,                XmPACK_TIGHT,
       XmNspacing,                15,
       XmNorientation,            XmVERTICAL,
       XmNleftAttachment,         XmATTACH_FORM,
       XmNtopAttachment,          XmATTACH_FORM,
       NULL);
    /* scale widget for atom count */
    wScaleA = XtVaCreateManagedWidget ("", xmScaleWidgetClass, wRowCol, 
       XmNorientation,            XmHORIZONTAL,
       XmNwidth,                  60,
       XmNshowValue,              True,  
       XmNminimum,                10,
       XmNmaximum,                60,
       XmNvalue,                  nAtomEdge,
       XtVaTypedArg,              XmNtitleString, T_STRING ("Atoms"),
       NULL);
    XtAddCallback (wScaleA, XmNvalueChangedCallback, CbScale, (XtPointer) 0);
    /* scale widget for update rate */
    wScaleD = XtVaCreateManagedWidget ("", xmScaleWidgetClass, wRowCol, 
       XmNorientation,            XmHORIZONTAL,
       XmNwidth,                  60,
       XmNshowValue,              True,  
       XmNminimum,                1,
       XmNmaximum,                100,
       XmNvalue,                  stepDraw,
       XtVaTypedArg,              XmNtitleString, T_STRING ("Update"),
       NULL);
    XtAddCallback (wScaleD, XmNvalueChangedCallback, CbScale, (XtPointer) 1);
    /* two sets of arrow pairs for changing temperature and density */
    for (n = 0; n < 2; n ++) {
        /* frame widget */
        wFrame = XtVaCreateManagedWidget ("", xmFrameWidgetClass, wRowCol,
           XmNshadowType,           XmSHADOW_ETCHED_IN,
           NULL);
        /* horizontal rowcolumn widget */
        wRowCol2 = XtVaCreateManagedWidget ("", xmRowColumnWidgetClass, wFrame,
           XmNorientation,          XmHORIZONTAL,
           XmNspacing,              5,
           NULL);
        /* label widget */
        XtVaCreateManagedWidget (buttonName[n], xmLabelWidgetClass, wRowCol2,
           NULL);
        /* arrowbutton widgets */
        for (j = 0; j < 2; j ++) {
            wArrow = XtVaCreateManagedWidget ("", xmArrowButtonWidgetClass,
               wRowCol2,
               XmNarrowDirection,     (j > 0) ? XmARROW_UP : XmARROW_DOWN,
               NULL);
            XtAddCallback (wArrow, XmNarmCallback, CbArrow,
               (XtPointer) (3 * n + 2 * j));
            XtAddCallback (wArrow, XmNdisarmCallback, CbArrow,
               (XtPointer) (3 * n + 1));
        }
    }
    /* pushbutton widget */
    wButtonS = XtVaCreateManagedWidget ("", xmPushButtonWidgetClass,
       wRowCol,
       XtVaTypedArg,              XmNlabelString, T_STRING ("Start"),
       NULL);
    XtAddCallback (wButtonS, XmNactivateCallback, CbButton, 0);
    /* frame widget */
    wFrame = XtVaCreateManagedWidget ("", xmFrameWidgetClass, wForm,
       XmNrightAttachment,        XmATTACH_FORM,
       XmNtopAttachment,          XmATTACH_FORM,
       XmNbottomAttachment,       XmATTACH_FORM,
       XmNshadowType,             XmSHADOW_IN,
       NULL);
    /* drawingarea widget */
    wDraw = XtVaCreateManagedWidget ("", xmDrawingAreaWidgetClass,
       wFrame, 
       XmNwidth,                  wDrawSize.x,
       XmNheight,                 wDrawSize.y,
       NULL);
    XtAddCallback (wDraw, XmNexposeCallback, CbExpose, 0);
    /* text widget for output */
    wText = XtVaCreateManagedWidget ("", xmTextWidgetClass, wForm,
       XmNcolumns,                12,
       XmNrows,                   4,
       XmNeditable,               False,
       XmNeditMode,               XmMULTI_LINE_EDIT,
       XmNcursorPositionVisible,  False,
       XmNhighlightThickness,     0,
       XmNshadowType,             XmSHADOW_IN,
       XmNleftAttachment,         XmATTACH_FORM,
       XmNrightAttachment,        XmATTACH_WIDGET,
       XmNrightWidget,            wFrame,
       XmNrightOffset,            5,
       XmNtopAttachment,          XmATTACH_WIDGET,
       XmNtopWidget,              wRowCol,
       XmNtopOffset,              40,
       NULL);
    /* create graphic context (gc) */
    dpy = XtDisplay (wTop);
    scrn = DefaultScreen (dpy);
    gc = XCreateGC (dpy, RootWindow (dpy, scrn), 0, NULL);
    /* create and clear pixmap for offscreen graphics */
    pixmap = XCreatePixmap (dpy, RootWindow (dpy, scrn),
       wDrawSize.x, wDrawSize.y, DefaultDepth (dpy, scrn));
    XSetForeground (dpy, gc, appData.pixel[0]);
    XFillRectangle (dpy, pixmap, gc, 0, 0, wDrawSize.x, wDrawSize.y);
    /* realize entire widget tree */
    XtRealizeWidget (wTop);
}

void CbArrow (Widget w, XtPointer clientData, XtPointer callData)
/*
  Purpose:

    CBARROW is the arrowbutton widgets callback function.
*/
{
    int id;
  
    id = (int) clientData;
    /* check which button and whether press/release */
    if (id < 3) changeTemp = id - 1;
    else changeDens = id - 4;
}

void CbScale (Widget w, XtPointer clientData, XtPointer callData)
/*
  Purpose:

    CBSCALE is the scale widget value-changed callback function.
*/
{
    int id, v;
  
    id = (int) clientData;
    v = ((XmScaleCallbackStruct *) callData)->value;
    /* check which scale */
    if (id == 0) {
        nAtomEdge = v;
        InitializeRun ();
    } else stepDraw = v;
}


void CbButton (Widget w, XtPointer clientData, XtPointer callData)
/*
  Purpose:

    CBBUTTON is the pushbutton widget callback function.
*/
{
    running = ! running;
    /* change button label */
    XtVaSetValues (wButtonS, XtVaTypedArg, XmNlabelString,
       T_STRING (running ? "Stop" : "Start"), NULL);
    /* allow/prevent changes to atom number */
    XtSetSensitive (wScaleA, ! running);
}


void CbExpose (Widget w, XtPointer clientData,
   XmDrawingAreaCallbackStruct *cbs)
/*
  Purpose:

    CBEXPOSE is the drawing area widget expose callback function,
*/
{
    /* ensure last expose event in sequence */
    if (((XExposeEvent *) cbs->event)->count == 0) Refresh ();
}


void SetDefaults ( void )
/*
  Purpose:

    SETDEFAULTS sets fixed and default values.
*/
{
    wDrawSize.x = appData.wSize;
    if (wDrawSize.x < 500) wDrawSize.x = 500;
    wDrawSize.y = wDrawSize.x;
    nAtomEdge = 20;
    stepDraw = 10;
    deltaT = 0.004;
    rCut = pow (2., 1./6.);
}


void InitializeRun ( void )
/*
  Purpose:

    INITIALIZERUN initializes the computation.
*/
{
    temperature = 1.;
    density = 0.8;
    VSet (initUcell, nAtomEdge, nAtomEdge / sqrt (3.));
    region.x = initUcell.x / sqrt (density * sqrt (3.) / 2.);
    region.y = region.x;
    nAtom = 2 * VProd (initUcell);
    density = nAtom / VProd (region);
    VSCopy (cells, 1. / rCut, region);
    /* memory (re)allocation */
    ALLOC (r, nAtom, VecD);
    ALLOC (rv, nAtom, VecD);
    ALLOC (ra, nAtom, VecD);
    ALLOC (cellList, nAtom + VProd (cells), int);
    /* initial state */
    InitialState ();
    ComputeForces ();
    CalcEnergy ();
    running = 0;
    changeTemp = 0;
    changeDens = 0;
    Redraw ();
}


Boolean MdCycle ( void )
/*
  Purpose:

    MDCYCLE is the work procedure.
*/
{
    if (running) SingleStep ();
    return (False);
}


void SingleStep ( void )
/*
  Purpose:

    SINGLESTEP takes a single MD step.
*/
{
    ++ stepCount;
    ComputeForces ();
    LeapfrogIntegrate ();
    CheckPeriodicWrap ();
    CalcEnergy ();
    if (stepCount % stepDraw == 0) {
        /* only do value changes when redrawing */
        if (changeTemp != 0) SetNewTemp (1. + 0.01 * changeTemp);
        if (changeDens != 0) SetNewDens (1. - 0.01 * changeDens);
        Redraw ();
    }
}


/* convenience macros */
#define DO_ATOM for (n = 0; n < nAtom; n ++)
#define DO_CELL(j, m) for (j = cellList[m]; j >= 0; j = cellList[j])
#define WRAP_C(t)     \
   if (m2v.t >= cells.t) m2v.t = 0,  shift.t = region.t;    \
   else if (m2v.t < 0) m2v.t = cells.t - 1,  shift.t = - region.t

void ComputeForces ( void )
/*
  Purpose:

    COMPUTEFORCES computes force and potential energy.

  Discussion:

    The routine uses the cell method, and allows for periodic boundaries.
*/
{
    VecD dr, invWid, rs, shift;
    VecI cc, m1v, m2v, iof[] = {0,0, 1,0, 1,1, 0,1, -1,1};
    double fcVal, rr, rrCut, rri, rri3;
    int c, j1, j2, m1, m1x, m1y, m2, n, offset;
  
    rrCut = SQ (rCut);
    for (n = nAtom; n < nAtom + VProd (cells); n ++) cellList[n] = -1;
    VDiv (invWid, cells, region);
    /* assign atoms to cells */
    DO_ATOM {
        VSAdd (rs, r[n], 0.5, region);
        VMul (cc, rs, invWid);
        c = VLinear (cc, cells) + nAtom;
        cellList[n] = cellList[c];
        cellList[c] = n;
    }
    /* zero accelerations */
    DO_ATOM VZero (ra[n]);
    uSum = 0.;
    /* loop over all cells */
    for (m1y = 0; m1y < cells.y; m1y ++) {
        for (m1x = 0; m1x < cells.x; m1x ++) {
            VSet (m1v, m1x, m1y);
            m1 = VLinear (m1v, cells) + nAtom;
            /* loop over offsets between adjacent cells */
            for (offset = 0; offset < 5; offset ++) {
                VAdd (m2v, m1v, iof[offset]);
                /* allow for periodic wraparound */
                VZero (shift);
                WRAP_C (x);
                WRAP_C (y);
                m2 = VLinear (m2v, cells) + nAtom;
                /* loops over atoms in each cell */
                DO_CELL (j1, m1) {
                    DO_CELL (j2, m2) {
                        /* ensure each atom pair only examined once */
                        if (m1 != m2 || j2 < j1) {
                            VSub (dr, r[j1], r[j2]);
                            VSub (dr, dr, shift);
                            rr = VLenSq (dr);
                            /* if atoms within interaction range */
                            if (rr < rrCut) {
                                rri = 1. / rr;
                                rri3 = rri * rri * rri;
                                fcVal = 48. * rri3 * (rri3 - 0.5) * rri;
                                /* update accelerations and energy sum */
                                VVSAdd (ra[j1], fcVal, dr);
                                VVSAdd (ra[j2], - fcVal, dr);
                                uSum += 4. * rri3 * (rri3 - 1.) + 1.;
                            }
                        }
                    }
                }
            }
        }
    }
}

void LeapfrogIntegrate ( void )
/*
  Purpose:

    LEAPFROGINTEGRATE integrates over a single time step.
*/
{
    int n;
  
    DO_ATOM {
        VVSAdd (rv[n], deltaT, ra[n]);
        VVSAdd (r[n], deltaT, rv[n]);
    }
}


/* convenience macro */
#define WRAP_R(t) if (r[n].t >= 0.5 * region.t) r[n].t -= region.t;  \
   else if (r[n].t < - 0.5 * region.t) r[n].t += region.t

void CheckPeriodicWrap ( void )
/*
  Purpose:

    CHECKPERIODICWRAP handles periodic boundaries.
*/
{
    int n;
  
    DO_ATOM {
        WRAP_R (x);
        WRAP_R (y);
    }
}


void CalcEnergy ( void )
/*
  Purpose:

    CALCENERGY computes the energy.
*/
{
    VecD v;
    double vvSum;
    int n;
  
    vvSum = 0.;
    DO_ATOM {
        /* allow for half timestep */
        VSAdd (v, rv[n], -0.5 * deltaT, ra[n]);
        vvSum += VLenSq (v);
    }
    kinEnergy = 0.5 * vvSum / nAtom;
    potEnergy = uSum / nAtom;
}


void InitialState ( void )
/*
  Purpose:

    INITIALSTATE initializes the atom coordinates and velocities.
*/
{
    VecD p, gap, vSum;
    double ang, vMag;
    int n, nx, ny;
  
    stepCount = 0;
    VDiv (gap, region, initUcell);
    n = 0;
    for (ny = 0; ny < initUcell.y; ny ++) {
        for (nx = 0; nx < initUcell.x; nx ++) {
            /* coordinates on lattice site */
            VSet (p, nx + 0.25, ny + 0.25);
            VMul (p, p, gap);
            VVSAdd (p, -0.5, region);
            r[n ++] = p;
            VVSAdd (p, 0.5, gap);
            r[n ++] = p;
        }
    }
    vMag = sqrt (2. * temperature);
    VZero (vSum);
    DO_ATOM {
        /* velocity has fixed magnitude and random direction */
        ang = 2. * M_PI * drand48 ();
        VSet (rv[n], cos (ang), sin (ang));
        VScale (rv[n], vMag);
        VAdd (vSum, vSum, rv[n]);
    }
    /* ensure zero center of mass velocity */
    DO_ATOM VVSAdd (rv[n], - 1. / nAtom, vSum);
}


void SetNewTemp ( double f )
/*
  Purpose:

    SETNEWTEMP changes the temperature.
*/
{
    double vvSum;
    int n;
  
    /* check in range */
    if (f > 1. && temperature > 5. || f < 1. && temperature < 0.02) return;
    vvSum = 0.;
    DO_ATOM {
        /* scale velocity */
        VScale (rv[n], f);
        vvSum += VLenSq (rv[n]);
    }
    temperature = vvSum / (2. * nAtom);
}


void SetNewDens ( double f )
/*
  Purpose:

    SETNEWDENS changes the density.
*/
{
    int n;
  
    /* check in range */
    if (f > 1. && density < 0.05 || f < 1. && density > 1.25) return;
    if (f < 1.) {
        /* scale coordinates */
        DO_ATOM VScale (r[n], f);
    }
    density /= SQ (f);
    VScale (region, f);
    VSCopy (cells, 1. / rCut, region);
    /* new memory allocation for cell list */
    ALLOC (cellList, nAtom + VProd (cells), int);
}


/* convenience macro */
#define FILL_DISK(x, y) XFillArc (dpy, pixmap, gc, x, y, diam, diam, \
   0, 360 * 64)

void Redraw ( void )
/*
  Purpose:

    Redraw redraws the entire window contents.
*/
{
    double scale;
    int diam, dx, dy, inxHi, inxLo, inyHi, inyLo, n, x, y;
    char buff[60];
  
    /* clear pixmap */
    XSetForeground (dpy, gc, appData.pixel[0]);
    XFillRectangle (dpy, pixmap, gc, 0, 0, wDrawSize.x, wDrawSize.y);
    XSetForeground (dpy, gc, appData.pixel[1]);
    scale = wDrawSize.x / region.x;
    diam = 0.9 * scale;
    DO_ATOM {
        /* determine coordinates in pixmap */
        x = scale * (0.5 * region.x + r[n].x) - diam / 2;
        y = scale * (0.5 * region.y - r[n].y) - diam / 2;
        /* draw disk */
        FILL_DISK (x, y);
        /* draw periodic replicas (if any) */
        inxLo = (x >= 0);
        inxHi = (x <= wDrawSize.x - diam);
        inyLo = (y >= 0);
        inyHi = (y <= wDrawSize.y - diam);
        if (! (inxLo && inxHi && inyLo && inyHi)) {
            dx = (inxLo ? -1 : +1) * wDrawSize.x;
            dy = (inyLo ? -1 : +1) * wDrawSize.y;
            if (inyLo && inyHi) FILL_DISK (x + dx, y);
            else if (inxLo && inxHi) FILL_DISK (x, y + dy);
            else {
                FILL_DISK (x + dx, y);
                FILL_DISK (x, y + dy);
                FILL_DISK (x + dx, y + dy);
            }
        }
    }
    Refresh ();
    /* write to text widget */
    sprintf (buff, "Step: %d \nE_tot: %.3f \nE_kin: %.3f \nDens: %.2f",
       stepCount, kinEnergy + potEnergy, kinEnergy, density);
    XmTextSetString (wText, buff);
    /* flush X queue */
    XSync (dpy, False);
}

void Refresh ( void )
/*
  Purpose:

    REFRESH refreshes the display, copying the pixmap to the drawing window.
*/
{
  XCopyArea (dpy, pixmap, XtWindow (wDraw), gc, 0, 0,
     wDrawSize.x, wDrawSize.y, 0, 0);
}

/* end */