////////////////////////////////////////////////////////////////////////////////////////////////////////
// LANGEVIN DYNAMICS SIMULATION OF ACTIVE ROTORS
////////////////////////////////////////////////////////////////////////////////////////////////////////
// Source code accompanying the manuscript "Phase Separation in a Binary Fluid of Macroscopic Rotors"
// written 2016-2017 by Michael Engel, Christian Scholz
//
// Compilation:   g++ -O3 -std=c++11 -o LD_Rotors LD_Rotors.cpp
// Execution:     echo "seed 128 steps 200000 maxtime 1000 standalone" | ./LD_Rotors > trajectory_128.pos
////////////////////////////////////////////////////////////////////////////////////////////////////////

// including libraries
#include <cmath>
#include <fstream>
#include <iostream>
#include <iomanip>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <random>
#include <vector>

// initialization parameters
const static double MASS = 0.969e-3;
const static double INERTIA = 2.21e-8;
const static double YOUNG = 2.3e9;  // Young's modulus for dissipative force
const static double MU = 0.35;      // Poisson's ratio
const static double TAUA = 1e-5;    // Viscous relaxation time

// physics parameters
const static double TDRIVE = 8.6e-6;
const static double GAMMAT = 0.0106;
const static double GAMMAR = 8.97e-7;
const static bool   NOISE = true;
const static double DT = 7.2e-6;
const static double DR = 0.19;

// particle size
const static double DL = 15e-3;             // m
const static double DS = 3.4e-3;            // m
const static double AD = 2.5e-3;            // m
const static double VH = 2e-3;              // m

// system size and initial setup
const static double SYSDIAM = 0.47;         // m
const static double STARTANGLE = 0.185;     // rad
const static double DDELTA = 0.0015;

// particle colors
const static int GCOLOR = 0xffa000;
const static int RCOLOR = 0x9933cc;

// contolling the simulation
const static double TIMESTEP = 5e-6;
static double maxtime = 1000;               // maximum physical time in seconds
static int stepsUntilOutput = 1000;         // how frequent to output

// RDF data
const static double RMAX = 0.5;
const static int NBINS = 10000;
const static int FREQUENCY = 100;
const static int PEAKWIDTH = 30;
static int rcounter;
static std::vector<int> gR;
static std::vector<int> gAB;
const static bool   CALCULATEDOMAINSIZE = true;

// simulation control
static int seed = 1;
static bool standalone = false;

// simulation counters and data
int steps;
double confinementRadius = SYSDIAM / 2;
double etrans = 0;
double erot = 0;

struct Vector2d
{
	double x;
	double y;
	
	Vector2d()
	{
		x = 0.0;
		y = 0.0;
	}
	
	Vector2d(const double& r, const double& s)
	{
		x = r;
		y = s;
	}
	
	double& operator [](const int& k)
	{
		return ((&x)[k]);
	}
	
	Vector2d& operator =(const Vector2d& v)
	{
		x = v.x;
		y = v.y;
		return *this;
	}
	
	Vector2d& operator +=(const Vector2d& v)
	{
		x += v.x;
		y += v.y;
		return (*this);
	}
	
	Vector2d& operator -=(const Vector2d& v)
	{
		x -= v.x;
		y -= v.y;
		return (*this);
	}
	
	Vector2d& operator *=(const double& t)
	{
		x *= t;
		y *= t;
		return (*this);
	}
	
	Vector2d& operator /=(const double& t)
	{
		return (*this *= (1.0 / t));
	}
	
	double operator %=(const Vector2d& v)
	{
		return (x * v.y - y * v.x);
	}
	
	const Vector2d operator -(void) const
	{
		return (Vector2d(-x, -y));
	}
	
	const Vector2d operator +(const Vector2d& v) const
	{
		return (Vector2d(*this) += v);
	}
	
	const Vector2d operator -(const Vector2d& v) const
	{
		return (Vector2d(*this) -= v);
	}
	
	const Vector2d operator *(const double& t) const
	{
		return (Vector2d(*this) *= t);
	}
	
	const Vector2d operator /(const double& t) const
	{
		return (Vector2d(*this) /= t);
	}
	
	double operator *(const Vector2d& v) const
	{
		return (x * v.x + y * v.y);
	}
	
	double operator %(const Vector2d& v) const
	{
		return (Vector2d(*this) %= v);
	}
	
	const Vector2d proj(const Vector2d& v)
	{
		double prod = (x * v.x + y * v.y) / (x * x + y * y);
		return (Vector2d(x * prod, y * prod));
	}
	
	const double norm() const
	{
		return sqrt(x * x + y * y);
	}
};

const Vector2d operator *(const double& t, const Vector2d& v)
{
	return (v * t);
}

// define 2d Vector type
typedef Vector2d Vector;

// cell list of objects is implemented as an array of linked list
class LinkedListArray
{
public:
	LinkedListArray()
	{
	}
	
	// initialize
	// 'arraySize': number of array entries
	// 'size': total number of objects sorted into the array
	LinkedListArray(int arraySize, int size)
	{
		cellStart = std::vector<int>(arraySize, -1);
		cellIndex = std::vector<int>(size, -1);
		cellNext  = std::vector<int>(size, -1);
		cellPrev  = std::vector<int>(size, -1);
	}
	
	// (re)assign an object into one of the cells
	// 'obj': object that is (re)assigned
	// 'index': index of cell it is (re)assigned to
	void set(int obj, int index)
	{
		// Step 1: check if cell changed
		int oldIndex = cellIndex[obj];
		if (oldIndex == index) return;
		cellIndex[obj] = index;
		
		// Step 2: remove from cell
		if (oldIndex >= 0)
		{
			int prev = cellPrev[obj];
			int next = cellNext[obj];
			// update forward link
			if (next >= 0) cellPrev[next] = prev;
			// update backward link to another obj
			if (prev >= 0) cellNext[prev] = next;
			// update backward link to a start
			else cellStart[oldIndex] = next;
		}
		
		// Step 3: fill into cell
		int cell = cellStart[index];
		// update forward link from start
		cellStart[index] = obj;
		// update backward link
		if (cell >= 0) cellPrev[cell] = obj;
		// create links from obj
		cellNext[obj] = cell;
		cellPrev[obj] = -1;
	}
	
	// returns the first object when iterating over all objects in a cell
	// 'index': cell index to iterate over
	int iterator(int index)
	{
		return cellStart[index];
	}
	
	// returns next object in the iteration
	// 'obj': current object
	int next(int obj)
	{
		return cellNext[obj];
	}
	
private:
	// start of the list for each cell
	std::vector<int> cellStart;
	// cell of each objevt
	std::vector<int> cellIndex;
	// forward link
	std::vector<int> cellNext;
	// backward link
	std::vector<int> cellPrev;
};

// cell list data
double boxLength;
double cellLength;
int cellNumber;
LinkedListArray cellList;

// declare random number generator
std::default_random_engine generator(seed);
std::normal_distribution<double> distribution(0.0, 1.0);

// shape data
struct Sphere
{
    Sphere(Vector x0, double r0)
    {
        x = x0;
        r = r0;
    }
    // relative position
    Vector x;
    // radius
    double r;
};

// NB: it is assumed that the center of mass is at (0,0)
struct Shape: std::vector<Sphere>
{
	Shape(char t0, int color0, double mass0, double inertia0)
	{
		t = t0;
		color = color0;
		mass = mass0;
		inertia = inertia0;
	}
	// parameters
	char t;
	int color;
	double mass, inertia;
    double area()
    {
        double area = 0.0;
        for (int i = 0; i < size(); i++)
        {
            area += at(i).r * at(i).r;
        }
        return area;
    }
    double circumradius()
    {
        double rmax = 0.0;
        for (int i = 0; i < size(); i++)
        {
            double r = at(i).r + at(i).x.norm();
            if (r > rmax) rmax = r;
        }
        return rmax;
    }
};
std::vector<Shape> shapes;

// particle data
struct Particle
{
    Particle(int shape0)
    {
        shape = shape0;
        spheres.resize(shapes[shape0].size());
    }
    // position, velocity, force, previous position
    Vector x, v, f;
    // angle, angular velocity (omega), torque, previous angle
    double a, o, t;
    // geometry
    int shape;
    std::vector<Vector> spheres;
};
std::vector<Particle> particles;

// calculates packing density
double getDensity()
{
    double area = 0.0;
    for (int i = 0; i < particles.size(); i++)
    {
        area += shapes[particles[i].shape].area();
    }
    return area / (confinementRadius * confinementRadius);
}

// returns the index in the cell list
int getCell(Vector& x)
{
    int cx = (int)((x.x + boxLength * 0.5) / cellLength);
    int cy = (int)((x.y + boxLength * 0.5) / cellLength);
    return cy * cellNumber + cx;
}

// initialize cell list
void initCellList()
{
	// search for maximum feature size of any particle
	cellLength = 0.0;
	for (int i = 0; i < shapes.size(); i++)
	{
		double r = shapes[i].circumradius() * 2.0;
		if (r > cellLength) cellLength = r;
	}
	
	// initialize cell list parameters
	boxLength  = (confinementRadius + cellLength) * 2.0;
	cellNumber = (int)floor(boxLength / cellLength);
	cellLength = boxLength / cellNumber;
	cellList   = LinkedListArray(cellNumber * cellNumber, particles.size());
	
	// sort atoms into cell list
	for (int i = 0; i < particles.size(); i++)
	{
		cellList.set(i, getCell(particles[i].x));
	}
}

// output frame to file
void outputData(std::ostream& out)
{
	double domainSize = 0.0;
	// the following part is a bit slow and not optimized
	if (CALCULATEDOMAINSIZE)
	{
		// convolute rdf data with a Gaussian to smoothen the data
		std::vector<double> gAB2(NBINS, 0.0);
		std::vector<double> gR2 (NBINS, 0.0);
		std::vector<double> norm(NBINS, 0.0);
		for (int p1 = 1; p1 < NBINS; p1++)
		{
			for (int p2 = 1; p2 < NBINS; p2++)
			{
				double gauss = exp(-(p1 - p2) * (p1 - p2) / (2.0 * PEAKWIDTH * PEAKWIDTH));
				norm[p1] += gauss;
				gAB2[p1] += gauss * gAB[p2];
				gR2[p1]  += gauss * gR[p2];
			}
		}
		// normalize distribution
		for (int p = 0; p <= NBINS; p++) if (norm[p] != 0.0)
		{
			gAB2[p] /= norm[p];
			gR2[p] /= norm[p];
		}
		
		int gg = 1;
		// while below zero...
		while (gAB2[gg] <= 0.0 && gg < NBINS) gg++;
		// while above zero...
		while (gAB2[gg] >= 0.0 && gg < NBINS) gg++;
		double dg = (gAB2[gg - 1] != gAB2[gg] ? gAB2[gg - 1] / (gAB2[gg - 1] - gAB2[gg]) : 0.0);
		domainSize = RMAX * sqrt((gg + dg) / NBINS);
	}
	
	// output header
	out << std::setprecision(6);
    out << "#[data] Steps DomainSize SeedNo fps Time Radius Density Etrans Erot Ekin\n";
    out << steps << " ";
	out << domainSize << " ";
	out << seed << " ";
	out << 1.0 / (stepsUntilOutput * TIMESTEP) << " ";
	out << steps * TIMESTEP << " ";
	out << confinementRadius << " ";
	out << getDensity() << " ";
	out << etrans << " ";
	out << erot << " ";
	out << (etrans + erot) << "\n";
    out << "#[done]\n";

    // output general data
	out << "dimension 3\n";
    out << "sphere " << (confinementRadius * 2.0) << " ffffff 0 0 " << confinementRadius << "\n";
	out << "dimension 2\n";
	out << "box " << boxLength << " " << boxLength <<  "\n";
    for (int i = 0; i < shapes.size(); i++)
    {
		Shape& shape = shapes[i];

		out << "def " << shape.t << " \"sphereU " << shape.size();
        for (int j = 0; j < shape.size(); j++)
        {
            Sphere& s = shape[j];
			out << " " << (s.r * 2.0) << " " << s.x.x << " " << s.x.y;
        }
		out << " " << std::hex << shape.color << std::dec << "\"\n";
    }

    // output particle data
    for (int i = 0; i < particles.size(); i++)
    {
        Particle& p = particles[i];
		Shape& shape = shapes[p.shape];

		out << shape.t;
		out << " " << p.x.x << " " << p.x.y << " " << cos(p.a * 0.5) << " " << sin(p.a * 0.5) << "\n";
    }

	// output footer
	out << "eof\n";
	std::flush(out);
}

// linear dashpot interaction
bool FDISC(double r, double sigma, double reff, Vector& rrel, Vector& vrel, double& force)
{
	double xi = sigma - r;
	if (xi < 0) return false;
	double xilog = (1.0 + MU) + log(xi / (4.0 * reff));
	double prop = M_PI * YOUNG * VH / xilog;

	force = prop * xi; //remember: value of log is negative, VH is thickness of disc
	if (TAUA != 0)
	{
		force += - TAUA * prop * (1.0 - 1.0 / xilog) * (vrel * rrel / rrel.norm());
	}
    return true;
}

// run simulation for one step
void runSimulation()
{
    etrans = 0.0;
    erot = 0.0;
	
	double FT = sqrt(2.0 * DT * GAMMAT * GAMMAT / TIMESTEP);
	double FR = sqrt(2.0 * DR * GAMMAR * GAMMAR / TIMESTEP);
    // add friction and random forces
    for (int i = 0; i < particles.size(); i++)
    {
        Particle& p = particles[i];
		int color = shapes[p.shape].color;
        
        if (color == RCOLOR || color == GCOLOR)
        {
			// add driving torque
			if (color == RCOLOR) p.t =  TDRIVE;
			if (color == GCOLOR) p.t = -TDRIVE;
			
			// random force
            if (NOISE)
            {
				// random force
                p.f.x += FT * distribution(generator);
                p.f.y += FT * distribution(generator);
				p.t   += FR * distribution(generator);
            }
			// frictional force
			p.f += -GAMMAT * p.v;
			p.t += -GAMMAR * p.o;
        }
		
    }

    // calculate sphere positions
    for (int i = 0; i < particles.size(); i++)
    {
        Particle& p = particles[i];
        Shape& shape = shapes[p.shape];

        double c = cos(p.a);
        double s = sin(p.a);
        for (int l = 0; l < shape.size(); l++)
        {
            double x = shape[l].x.x;
            double y = shape[l].x.y;
            p.spheres[l] = Vector(c * x - s * y,
								  s * x + c * y);
        }
    }

    // confinement forces
    for (int i = 0; i < particles.size(); i++)
    {
        Particle& p = particles[i];
        Shape& shape = shapes[p.shape];

        for (int l = 0; l < shape.size(); l++)
        {
            Vector& s = p.spheres[l];

            // force center
            Vector dr = p.x + s;
            double r = dr.norm();
            double R = shape[l].r;
            double delta = confinementRadius - r;

            // relative velocity
            Vector dv = p.v;    // translation
            dv.x -= p.o * s.y;  // + rotation
            dv.y += p.o * s.x;  // + rotation

            // force
            double f0;
            if (FDISC(-2.0 * (R+r), -2.0 * confinementRadius, R, dr, dv, f0) == false) continue;
            
            // update
			Vector f = dr * (2.0 * f0 / r);
            p.f += f;
            p.t += s % f;
        }

        // calculate kinetic energy
        etrans += 0.5 * (shape.mass * p.v * p.v);
        erot += 0.5 * (shape.inertia * p.o * p.o);
    }


    // interaction forces
    for (int i1 = 0; i1 < particles.size(); i1++)
    {
        Particle& p1 = particles[i1];
        Shape& shape1 = shapes[p1.shape];

        // init cell list access
        int cc = getCell(p1.x) - cellNumber - 1;

        // loop over relevant cells
        for (int j = 3; j > 0; j--, cc += cellNumber - 3)
        {
            for (int k = 3; k > 0; k--, cc++)
            {
                // now loop over all particles in this cell
                int i2 = cellList.iterator(cc);
                while (i2 >= 0)
                {
                    // avoid double counting
                    if (i1 < i2)
                    {
                        Particle& p2 = particles[i2];
                        Shape& shape2 = shapes[p2.shape];

                        Vector dx = p2.x - p1.x;
                        if (dx * dx < cellLength * cellLength)
                        {
                            for (int l1 = 0; l1 < shape1.size(); l1++)
                            {
                                double r1 = shape1[l1].r;
                                Vector& s1 = p1.spheres[l1];

                                for (int l2 = 0; l2 < shape2.size(); l2++)
                                {
                                    double r2 = shape2[l2].r;
                                    Vector& s2 = p2.spheres[l2];

                                    // separating vector
                                    Vector dr = dx + s2 - s1;
                                    double r = dr.norm();
                                    
                                    // relative velocity
                                    Vector dv = p2.v - p1.v;    // translation
                                    dv.x -= p2.o * s2.y - p1.o * s1.y;  // + rotation
                                    dv.y += p2.o * s2.x - p1.o * s1.x;  // + rotation

                                    // force
                                    double f0;
                                    if (FDISC(r, r1 + r2, 2.0 * (r1 * r2) / (r1 + r2), dr, dv, f0) == false) continue;
									Vector f = dr * (f0 / r);

                                    // update forces
                                    p1.f += f;
                                    p2.f -= f;
                                    p1.t += s1 % f;
                                    p2.t -= s2 % f;
                                }
                            }
                        }
                    }
                    i2 = cellList.next(i2);
                }
            }
        }
    }

    // update
    for (int i = 0; i < particles.size(); i++)
    {
        Particle& p = particles[i];
		Shape& shape = shapes[p.shape];

        // Euler method
		p.v += (TIMESTEP / shape.mass   ) * p.f;
        p.o += (TIMESTEP / shape.inertia) * p.t;
        p.x += TIMESTEP * p.v;
        p.a += TIMESTEP * p.o;
        p.f  = Vector();
        p.t  = 0.0;

        // update cell list
        cellList.set(i, getCell(p.x));
    }
}

void initPositionsGrid(std::vector<Particle> &particles)
{
	steps = 0;
	particles = std::vector<Particle>();

	bool alternate = true;
	int NX = 20;
	int NY = 21;
	for (int i = 0; i < NX; i++)
	{
		for (int j = 0; j < NY; j++)
		{
			// position new particle without overlap
			Particle p = Particle(alternate ? 1 : 0);
			alternate = !alternate;
			
			p.x.x = (i - 0.5 * (NX - 1)) * (DL + DDELTA / 2);
			p.x.y = (j - 0.5 * (NY - 1)) * (DL + DDELTA / 2);
			p.a = M_PI * STARTANGLE;
			
			if (sqrt(p.x * p.x) >= confinementRadius - (DL + DDELTA / 2)) continue;
			particles.push_back(p);
		}
	}

	initCellList();
}

int main(int argc, char** argv)
{
    // define rotor shape
    Shape shape('A', GCOLOR, MASS, INERTIA);
	shape.push_back(Sphere(Vector( 0,  0), DL/2));
    shape.push_back(Sphere(Vector( (DL+DS)/2,  0), DS/2));
    shape.push_back(Sphere(Vector( 0,  (DL+DS)/2), DS/2));
    shape.push_back(Sphere(Vector((-DL-DS)/2,  0), DS/2));
    shape.push_back(Sphere(Vector( 0, (-DL-DS)/2), DS/2));
    shapes.push_back(shape);

	shape.t = 'B';
	shape.color = RCOLOR;
	shapes.push_back(shape);

	// position particles on grid
    initPositionsGrid(particles);
	
	while (true)
	{
		char cmd[256];
		if (steps * TIMESTEP >= maxtime) standalone = false;
		if (standalone) strcpy(cmd, "next"); else std::cin >> cmd;
		
		// exit simulation
		if (strcasecmp(cmd, "quit") == 0) break;
		
		// run simulation up to next output
		if (strcasecmp(cmd, "next") == 0)
		{
			rcounter = 0;
			gR  = std::vector<int>(NBINS, 0);
			gAB = std::vector<int>(NBINS, 0);
			
			for (int i = 0; i < stepsUntilOutput; i++)
			{
				// perform simulation
				runSimulation();
				steps++;

				// update RDF
				if (steps % FREQUENCY == 0) for (int ig = 0; ig < particles.size(); ig++)
				{
					rcounter++;
					Vector& x1 = particles[ig].x;
					char t1 = shapes[particles[ig].shape].t;
					for (int jg = 0; jg < ig; jg++)
					{
						Vector& x2 = particles[jg].x;
						char t2 = shapes[particles[jg].shape].t;
						double dx = x2.x - x1.x;
						double dy = x2.y - x1.y;
						double d2 = dx * dx + dy * dy;
						if (d2 <= (RMAX * RMAX))
						{
							int bin = floor(d2 * (NBINS / (RMAX * RMAX)));
							gR [bin] ++;
							gAB[bin] += (t1 == t2 ? +1 : -1);
						}
					}
				}
			}
			
			outputData(std::cout);
		}

		// number of MD steps between outputs
		if (strcasecmp(cmd, "steps") == 0)
		{
			std::cin >> stepsUntilOutput;
		}
		
		// number of MD steps between outputs
		if (strcasecmp(cmd, "maxtime") == 0)
		{
			std::cin >> maxtime;
		}
		
		// set random number generator seed
		if (strcasecmp(cmd, "seed") == 0)
		{
			std::cin >> seed;
			generator.seed(seed);

			// slightly randomize the system
			for (int i = 0; i < particles.size(); i++)
			{
				Particle& p = particles[i];
				p.x.x += 1e-9 * distribution(generator);
				p.x.y += 1e-9 * distribution(generator);
			}
		}

		// set random number generator seed
		if (strcasecmp(cmd, "standalone") == 0)
		{
			standalone = true;
		}
	}
}