/* 2D molecular dynamics simulation of a Lennard-Jones fluid

   Copyright 2010-2014 Cheng Zhang

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 2 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   A copy of the GNU General Public License can be found in
   <http://www.gnu.org/licenses/>.
*/
import java.awt.*;
import java.awt.event.*;
import java.awt.geom.*;
import java.awt.image.*;
import javax.swing.*;
import static java.lang.Math.*;
import java.text.*;
import java.util.Random;



public class LJ2MDApp extends JApplet implements ActionListener
{
  double rho_init = 0.8;
  LJ2MD md = new LJ2MD(rho_init);

  int delay = 100;
  int speed = 10; // md steps per frame
  Timer timer;

  XYCanvas canvas;
  JPanel spnl, cpnl;
  JTextField tNum   = new JTextField("    " + md.N);
  JTextField tTemp  = new JTextField("    " + md.kT);
  JTextField tRho   = new JTextField("    " + rho_init);
  JTextField tSpeed   = new JTextField("   " + speed);
  JButton    bReset = new JButton("Reset");
  JToggleButton bStart = new JToggleButton("Start", false);
  JCheckBox  bTstat = new JCheckBox("Thermostat", true);
  JLabel     lStatus = new JLabel("Status");
  JTextField tAvK   = new JTextField("   0");
  JTextField tAvU   = new JTextField("   0");
  JTextField tAvp    = new JTextField("   0");
  JButton    bRetime = new JButton("Reset time");

  /** Applet initialization */
  public void init() {
    Container box = getContentPane();
    box.setLayout(new BorderLayout());

    cpnl = new JPanel(); // create a panel for controls
    cpnl.setLayout(new GridLayout(18, 2));
    box.add(cpnl, BorderLayout.EAST);

    // add controls
    cpnl.add(bStart);
    bStart.addActionListener(this);

    cpnl.add(bReset);
    bReset.addActionListener(this);

    cpnl.add(new JLabel(" N:"));
    tNum.addActionListener(this);
    cpnl.add(tNum);

    cpnl.add(new JLabel(" Temperature:"));
    tTemp.addActionListener(this);
    cpnl.add(tTemp);

    cpnl.add(new JLabel(" Density:"));
    tRho.addActionListener(this);
    cpnl.add(tRho);

    cpnl.add(new JLabel(" Steps/frame:"));
    tSpeed.addActionListener(this);
    cpnl.add(tSpeed);

    cpnl.add(bTstat);
    bTstat.addActionListener(this);

    cpnl.add(new JLabel(" < K/N >:"));
    tAvK.setEditable(false);
    cpnl.add(tAvK);

    cpnl.add(new JLabel(" < U/N >:"));
    tAvU.setEditable(false);
    cpnl.add(tAvU);

    cpnl.add(new JLabel(" Av pressure:"));
    tAvp.setEditable(false);
    cpnl.add(tAvp);

    cpnl.add(bRetime);
    bRetime.addActionListener(this);

    spnl = new JPanel(); // status panel
    box.add(spnl, BorderLayout.SOUTH);
    lStatus.setFont(new Font("Courier", 0, 12));
    spnl.add(lStatus);

    canvas = new XYCanvas();
    canvas.setmd(md);
    canvas.setBackground(Color.WHITE);
    box.add(canvas, BorderLayout.CENTER);

    timer = new Timer(delay, this);
    timer.start();
    timer.stop();
  }

  public void update(Graphics g) {
    paint(g);
  }

  DecimalFormat df = new DecimalFormat("0.000");
  //int frame = 0;
  public void paint(Graphics g) {
    //System.out.println("frame: " + (frame++));
    lStatus.setText("t = " + df.format(md.dt*md.step) + ", "
         + "N = " + md.N + ", "
         + "E/N = " + df.format(md.E/md.N) + ", "
         + "U/N = " + df.format(md.U/md.N) + ", "
         + "K/N = " + df.format(md.K/md.N) + ", "
         + "p = " + df.format(md.p));
    tAvK.setText(" " + df.format(md.avK.getAve()/md.N));
    tAvU.setText(" " + df.format(md.avU.getAve()/md.N));
    tAvp.setText(" " + df.format(md.avp.getAve()));
    canvas.repaint();
    spnl.repaint();
    cpnl.repaint();
  }

  /** handle timer and control events */
  public void actionPerformed(ActionEvent e) {
    Object src = e.getSource();
    if (src == timer) {
      for (int i = 0; i < speed; i++) // integrate a few steps
        md.vv();
      repaint();
      return;
    }

    if (src == bTstat) md.thermostat = !md.thermostat;

    if (src == tTemp) {
      double kT = Double.parseDouble(tTemp.getText().trim());
      if (kT < 1e-8) { kT = 1e-8; tTemp.setText("  " + kT); }
      md.kT = kT;
      md.clearData();
    }

    if (src == tRho) {
      double rho = Double.parseDouble(tRho.getText().trim());
      if (rho < 1e-3) {rho = 1e-3; tRho.setText("   " + rho); }
      if (rho > 1.2) {rho = 1.2; tRho.setText("   " + rho); }
      md.setDensity(rho);
    }

    if (src == tSpeed) {
      speed = Integer.parseInt(tSpeed.getText().trim());
      if (speed < 1) { speed = 1; tSpeed.setText("   " +speed); }
    }

    if (src == bRetime) md.clearData();

    if (src == bStart) {
      boolean on = bStart.isSelected();
      if (on) {
        timer.restart();
        bStart.setText("Pause");
      } else {
        timer.stop();
        bStart.setText("Resume");
      }
    }

    if (src == tNum) {
      int n = Integer.parseInt(tNum.getText().trim());
      if (n < 2) { n = 2; tNum.setText(" " + n); }
      md.N = n;
      md.init(md.rho);
    }

    if (src == bReset) {
      if (timer.isRunning()) timer.stop();
      bStart.setSelected(false);
      bStart.setText("Start");
      md.init(md.rho);
    }
    repaint();
  }
  public void start() { if (bStart.isSelected()) timer.restart(); }
  public void stop() { if (timer.isRunning()) timer.stop(); }
}




class LJ2MD {
  static final int D = 2;
  public int N = 200; // number of particles
  private int DOF = N*D - D*(D+1)/2;
  double rho;

  double dt = 0.002; // time step for integrating Newton's equation
  double L = 10.0; // side of the box
  double m = 1.0; // mass of each particle
  double kT = 1.0; // temperature times Boltzmann constant
  double r[][]; // position (0, 1)
  double v[][], f[][]; // velocity and force
  double K, U, E; // kinetic, potential, and total energy
  double Vir, p; // virial and pressure
  boolean thermostat = true; // use thermostat
  double Utail; // potential energy tail correction
  double ptail; // pressure tail correction
  int step; // simulation steps
  Ave avU = new Ave(), avK = new Ave(), avp = new Ave();

  LJ2MD(double den) {
    init(den);
  }

  /** initialize system */
  void init(double den) {
    int i, j, id;

    DOF = N*D - D*(D+1)/2;
    step = 0;
    r = new double[N][D];
    v = new double[N][D];
    f = new double[N][D];
    for (i = 0; i < N; i++) v[i][0] = v[i][1] = 0;

    int N1 = (int) (pow(2*N, 1./D)-1e-6) + 1;
    double a = L/N1;
    for (id = 0, i = 0; i < N1 && id < N; i++) // place particles on a square lattice
      for (j = 0; j < N1 && id < N; j++) {
        if ((i + j) % 2 != 0) continue;
        r[id][0] = a * (i + .5)/L;
        r[id][1] = a * (j + .5)/L;
        id++;
      }
    setDensity(den);
    rmcom();
  }

  void setDensity(double den) {
    L = pow(N/(rho = den), 1.0/D);

    /* compute shifts and tail corrections */
    double irc = 1./(L*.5), irc3 = irc*irc*irc, irc6 = irc3*irc3;
    Utail = PI*rho*N*(0.4*irc6 - 1)*irc3*irc;
    ptail = PI*rho*rho*(1.6*irc6 - 2)*irc3*irc;

    clearData();
  }

  void clearData() {
    step = 0;
    avU.clear();
    avK.clear();
    avp.clear();
  }

  /** integrate Newton's equation by one step, velocity Verlet */
  void vv() {
    for (int id = 0; id < N; id++) // velocity-verlet, first half
      for (int d = 0; d < D; d++) {
        v[id][d] += f[id][d]/m * dt * .5;
        r[id][d] += v[id][d] * dt/L;
      }

    force(); // compute force

    for (int id = 0; id < N; id++) // velocity-verlet, second half
      for (int d = 0; d < D; d++)
        v[id][d] += f[id][d]/m * dt * .5;

    rmcom();
    K = vrescale(0.02); // compute kinetic energy and adjust velocity
    E = K + U;
    step++;
    avU.add(U);
    avK.add(K);
    avp.add(p);
  }

  /** remove center of mass motion */
  void rmcom() {
    int j, id;

    for (j = 0; j < D; j++) { // remove the center of mass motion
      double vc = 0.0;
      for (id = 0; id < N; id++)
        vc += v[id][j];
      vc /= N;
      for (id = 0; id < N; id++)
        v[id][j] -= vc;
      //System.out.println("vc " + vc);
    }
  }

  /** calculate potential energy and force
   *  half-box cutoff, minimal distance image */
  void force() {
    int i, j, d;
    double invr2, invr6, r2, fscal;
    double xij[] = new double[D];
    double fij[] = new double[D];

    // clear the force and energy
    U = 0.0;
    Vir = 0.0;
    for (i = 0; i < N; i++)
      for (d = 0; d < D; d++)
        f[i][d] = 0.0;

    // loop over pairs of particles and compute energy
    for (i = 0; i < N - 1; i++) {
      for (j = i + 1; j < N; j++) {
        r2 = 0.;
        for (d = 0; d < D; d++) {
          xij[d] = mimage(r[i][d] - r[j][d]);
          r2 += xij[d] * xij[d];
        }
        if (r2 > .25*L*L) continue;
        invr2 = 1.0/r2;
        invr6 = invr2 * invr2 * invr2;
        fscal = invr2 * invr6 * (48.0 * invr6 - 24.0);
        for (d = 0; d < D; d++) {
          fij[d] = xij[d] * fscal;
          f[i][d] += fij[d];
          f[j][d] -= fij[d];
        }
        U += 4.0 * invr6 * (invr6 - 1.0);
        Vir += 48.0 * invr6 * (invr6 - 0.5);
      }
    }
    U += Utail;
    p = rho*kT + Vir/(2*L*L) + ptail;
  }

  Random rng = new Random();

  /** velocity rescaling thermostat */
  double vrescale(double dt) {
    double ek1 = getEKin();
    if (!thermostat) return ek1;
    double ekav = .5f*kT*DOF;
    double amp = 2*sqrt(ek1*ekav*dt/DOF);
    double ek2 = ek1 + (ekav - ek1)*dt + amp*rng.nextGaussian();
    if (ek2 < 0) ek2 = ek1;
    double s = sqrt(ek2/ek1);
    for (int i = 0; i < N; i++)
      for (int d = 0; d < D; d++)
        v[i][d] *= s;
    return ek2;
  }

  /** return kinetic energy */
  double getEKin() {
    double ekin = 0.0;
    for (int id = 0; id < N; id++)
      for (int d = 0; d < D; d++)
        ekin += v[id][d] * v[id][d];
    return .5 * m * ekin;
  }

  /** image with the minimal distance */
  double mimage(double x) {
    x = x + 1000.;
    return (x - (int) (x + 0.5)) * L;
  }

  /** wrap coordinates back into box */
  double wrap(double x) {
    x = x + 1000.0;
    return (x - (int) x)*L;
  }
}




/** panel for drawing particals  */
class XYCanvas extends JPanel {
  Image img; // buffered image
  Graphics gi; // context associated with the image
  Dimension szi; // size of the image
  DecimalFormat df = new DecimalFormat("##0.000");
  LJ2MD md;

  private Image imgCircle;

  public XYCanvas() { super(); }
  void setmd(LJ2MD m) { md = m; }

  protected void paintComponent(Graphics g) {
    super.paintComponent(g);

    Dimension sz = getSize();
    //System.out.println(sz.toString());
    // create a off-screen context gi,
    if (gi == null || sz.width != szi.width || sz.height != szi.height) {
      szi = sz;
      // create an image
      img = createImage(sz.width, sz.height);
      gi  = img.getGraphics();
    }

    // draw on gi instead of directly on g
    gi.setColor(Color.BLACK);
    gi.fillRect(0, 0, sz.width, sz.height);
    for (int i = 0; i < md.N; i++)
      drawCircle(gi, md.wrap(md.r[i][0]), md.wrap(md.r[i][1]));

    // put image to screen
    g.drawImage(img, 0, 0, null);
  }

  private final static int R = 40; // 2R x 2R is the image size
  private final static int hx = 15;  // (hx, hy) is the offset of the spot light from the center
  private final static int hy = 15;
  private final static float spotlightMag = .5f; // spotlight brightness, 1.f for full
  private static byte [] data;
  private static int maxr;  // maximal intensity

  static {
    data = new byte[R*2*R*2];
    int mr = 0;
    for (int Y = 2 * R; --Y >= 0;) {
      int x0 = (int) (Math.sqrt(R * R - (Y - R) * (Y - R)) + 0.5);
      int p = Y * (R * 2) + R - x0;
      for (int X = -x0; X < x0; X++) {
        int x = X + hx;
        int y = Y - R + hy;
        int r = (int) (Math.sqrt(x * x + y * y) + 0.5);
        if (r > mr)
           mr = r;
        data[p++] = r <= 0 ? 1 : (byte) r;
      }
    }
    maxr = mr;
  }

  private final int blend(int fg, int bg, float fgfactor) {
    return (int) (bg + (fg - bg) * fgfactor);
  }

  /** prepare circle image */
  private void mkCircle() {
    final int Rl = 0, Gl = 0, Bl = 255;

    byte red[] = new byte[256];
    byte green[] = new byte[256];
    byte blue[] = new byte[256];
    for (int i = maxr; i >= 1; --i) {
      float d = (float) i / maxr;
      d = 1 - (1-d) * spotlightMag;
      red[i] = (byte) blend(Rl, 255, d);
      green[i] = (byte) blend(Gl, 255, d);
      blue[i] = (byte) blend(Bl, 255, d);
    }
    // 256 color model
    IndexColorModel model = new IndexColorModel(8, maxr + 1, red, green, blue, 0);
    imgCircle = createImage(new MemoryImageSource(R*2, R*2, model, data, 0, R*2));
  }

  /** draw one particle */
  void drawCircle(Graphics g, double x, double y)
  {
    double margin = 20.0;
    double w = szi.width - margin*2;

    if (imgCircle == null) mkCircle();
    double dx = w/md.L;
    x *= dx;
    y *= dx;

    double radius = .5*1.122*w/md.L;
    int size = (int)(radius*2);
    g.drawImage(imgCircle, (int)(x + margin - size/2), (int)(y + margin - size/2), size, size, this);
  }
}





/** Average and standard deviation */
class Ave {
  double cnt, xsm, x2sm;

  public void clear() {
    cnt = xsm = x2sm = 0;
  }

  public void add(double x) {
    cnt += 1;
    xsm += x;
    x2sm += x*x;
  }

  public double getAve() {
    return (cnt > 0) ? xsm/cnt : 0;
  }

  public double getVar() {
    if (cnt <= 1) return 0;
    double av = xsm/cnt;
    return x2sm/cnt - av*av;
  }

  public double getStd() {
    return Math.sqrt( getVar() );
  }
}