#include <constant.h>
#include <matplotlibcpp.h>
#include <util/argparse.h>
#include <utils.h>

#include <Eigen/Dense>
#include <cmath>
#include <ctime>
#include <fstream>
#include <iostream>
#include <random>
#include <vector>

using namespace std;
using Eigen::Vector2d;
namespace plt = matplotlibcpp;

#define PI GSL_CONST_NUM_PI
#define kb GSL_CONST_MKSA_BOLTZMANN

#define sigma 3.405e-10
#define epsilon 1.65e-21
#define mass (6.63382e-26 / 39.95)
#define tem (epsilon / kb)
#define vel sqrt(epsilon / mass)
#define tim (sigma * sqrt(mass / epsilon))
#define M 39.95

bool debug;
int N, W, Epoch;
double L, T, dt, eps;

class Molecule {
public:
    Molecule(){};
    Molecule(double m, Vector2d r0, Vector2d v0) {
        r = r0;
        v = v0;
        cnt = 0;
        length = 0;
        lastCrash = 0;
        this->m = m;
    };
    ~Molecule(){};

public:
    int cnt, lastCrash;
    double m, length;
    Vector2d a, aPre, r, v;

public:
    void warp() {
        r[0] -= L * floor(r[0] / L);
        r[1] -= L * floor(r[1] / L);
    }
    void crash(int index) {
        cnt += 1;
        length += v.norm() * (index - lastCrash) * dt;
        lastCrash = index;
    }
    void stretch(double factor) { v *= factor; }
};

Vector2d separateDistance(Vector2d ra, Vector2d rb) {
    double dx = ra[0] - rb[0];
    double dy = ra[1] - rb[1];
    if (dx > L / 2)
        dx -= L;
    else if (dx < -L / 2)
        dx += L;
    if (dy > L / 2)
        dy -= L;
    else if (dy < -L / 2)
        dy += L;
    return Vector2d(dx, dy);
}

double lennardJones(Vector2d r) { return 4 * (pow(1. / r.norm(), 12) - pow(1. / r.norm(), 6)); }

Vector2d lennardJonesFource(Vector2d r) {
    double fm = 48 * (pow(1. / r.norm(), 14) - 0.5 * pow(1. / r.norm(), 8));
    return fm * r;
}

double getKinetic(vector<Molecule> mol) {
    double Ek = 0;
    for (int i = 0; i < N; i++) Ek += 0.5 * mol[i].m * mol[i].v.squaredNorm();
    return Ek;
}

double getPotential(vector<Molecule> mol) {
    double Vr = 0;
    for (int i = 0; i < N; i++)
        for (int j = i + 1; j < N; j++) Vr += lennardJones(separateDistance(mol[i].r, mol[j].r));
    return Vr;
}

double getEnergy(vector<Molecule> mol) {
    double Ek = getKinetic(mol);
    double Vr = getPotential(mol);
    return Ek + Vr;
}

double getTemperature(vector<Molecule> mol) {
    double Ek = getKinetic(mol);
    return Ek / (N - 1);
}

double getPressure(vector<Molecule> mol) {
    double Te = getTemperature(mol);
    double p = N * Te;

    Vector2d r, f;
    for (int i = 1; i < N; i++)
        for (int j = 0; j < i; j++) {
            r = separateDistance(mol[i].r, mol[j].r);
            f = lennardJonesFource(r);
            p += r.dot(f) / 2;
        }
    return p / pow(L, 2);
}

int main(int argc, char const *argv[]) {
    auto args = util::argparser("Argon MD Simulate");
    args.set_program_name("ComSim")
        .add_help_option()
        .add_argument<double>("length", "length of box in sigma")
        .add_argument<int>("number", "number of particles")
        .add_argument<double>("temperature", "temperature in Kelvin")
        .add_option("-D", "--debug", "debug mode")
        .add_option<int>("-E", "--epoch", "maximum number of rounds simulated,\ndefault is 10000", 10000)
        .add_option<int>("-W", "--width", "width of window to get average,\ndefault is 1000", 1000)
        .add_option<double>("", "--dt", "delta t to simualte,\ndefault is 0.01", 0.01)
        .add_option<double>("", "--eps", "error t0 end simualte,\ndefault is 0.01", 0.01)
        .parse(argc, argv);

    L = args.get_argument_double("length");
    N = args.get_argument_int("number");
    T = args.get_argument_double("temperature") / tem;
    debug = args.get_option_bool("-D");
    Epoch = args.get_option_int("-E");
    W = min(args.get_option_int("-W"), Epoch);
    dt = args.get_option_double("--dt");
    eps = args.get_option_double("--eps");

    int cnt = 0;
    double fac, eqT, eqE, eqP;
    vector<Molecule> argon(N);
    Vector2d Vc = Vector2d(0, 0);
    vector<double> x(N), y(N);
    vector<double> Ti, Te, MTe, En, MEn, Pr, MPr;

    default_random_engine seed(time(NULL));
    uniform_real_distribution<double> u01(0, 1);
    normal_distribution<double> gauss(0, sqrt(T * tem * kb / (M * mass)));

    int NM = sqrt(N);
    double H = (L - 1.) / (NM - 1);
    for (int i = 0; i < N; i++) {
        double vx = gauss(seed) / vel;
        double vy = gauss(seed) / vel;

        argon[i] = Molecule(M, Vector2d((i / NM) * H + 0.5, (i % NM) * H + 0.5), Vector2d(vx, vy));
    }

    for (int i = 0; i < N; i++) Vc += argon[i].v;
    Vc /= N;
    for (int i = 0; i < N; i++) argon[i].v -= Vc;
    eqT = getTemperature(argon);

    do {
        fac = sqrt(T / eqT);
        for (int i = 0; i < N; i++) {
            argon[i].stretch(fac);
            argon[i].a = Vector2d(0, 0);
            for (int j = i + 1; j < N; j++) {
                Vector2d f = lennardJonesFource(separateDistance(argon[i].r, argon[j].r));
                argon[i].a += f;
                argon[j].a -= f;
            }
            argon[i].a /= argon[i].m;
        }

        Ti.push_back(Epoch * cnt * dt);
        Te.push_back(getTemperature(argon));
        En.push_back(getEnergy(argon));
        Pr.push_back(getPressure(argon));
        MTe.push_back(Te[Epoch * cnt]);
        MEn.push_back(En[Epoch * cnt]);
        MPr.push_back(Pr[Epoch * cnt]);

        for (int i = 1; i < Epoch; i++) {
            for (int j = 0; j < N; j++) {
                argon[j].r += argon[j].v * dt + argon[j].a * pow(dt, 2) / 2;
                argon[j].aPre = argon[j].a;
                argon[j].a = Vector2d(0, 0);
            }
            for (int j = 0; j < N; j++) {
                for (int k = j + 1; k < N; k++) {
                    Vector2d f = lennardJonesFource(separateDistance(argon[j].r, argon[k].r));
                    argon[j].a += f;
                    argon[k].a -= f;
                }
                argon[j].a /= argon[j].m;
                argon[j].v += (argon[j].aPre + argon[j].a) * dt / 2;
                argon[j].warp();
            }

            int k = i + Epoch * cnt;
            Ti.push_back(k * dt);
            Te.push_back(getTemperature(argon));
            En.push_back(getEnergy(argon));
            Pr.push_back(getPressure(argon));

            if (i >= W) {
                MTe.push_back((MTe[k - 1] * W + Te[k] - Te[k - W]) / W);
                MEn.push_back((MEn[k - 1] * W + En[k] - En[k - W]) / W);
                MPr.push_back((MPr[k - 1] * W + Pr[k] - Pr[k - W]) / W);
            } else {
                MTe.push_back((MTe[k - 1] * i + Te[k]) / (i + 1));
                MEn.push_back((MEn[k - 1] * i + En[k]) / (i + 1));
                MPr.push_back((MPr[k - 1] * i + Pr[k]) / (i + 1));
            }
        }

        cnt += 1;
        eqT = average(cutArrs(MTe, Epoch * cnt - W, Epoch * cnt - 1), W);
        eqE = average(cutArrs(MEn, Epoch * cnt - W, Epoch * cnt - 1), W);
        eqP = average(cutArrs(MPr, Epoch * cnt - W, Epoch * cnt - 1), W);
        if (debug) cout << eqT * tem << ", " << abs(eqT - T) / T << ", " << eqE << ", " << eqP << endl;

        if (abs(eqT - T) / T < eps) break;
    } while (1);
    cout << eqT * tem << ", " << abs(eqT - T) / T << ", " << eqE << ", " << eqP << endl;

    plt::figure_size(900, 600);
    plt::subplot(3, 1, 1);
    plt::named_plot("Instant", Ti, Te);
    plt::named_plot("Average", Ti, MTe);
    plt::legend();
    plt::title("Equilibrium Temperature : " + to_string(eqT));
    plt::xlabel("Time / ($\\sigma\\sqrt{m/\\epsilon}$)");
    plt::ylabel("Temperature / ($\\epsilon/k_b$)");

    plt::subplot(3, 1, 2);
    plt::named_plot("Instant", Ti, En);
    plt::named_plot("Average", Ti, MEn);
    plt::legend();
    plt::title("Equilibrium Energy : " + to_string(eqE));
    plt::xlabel("Time / ($\\sigma\\sqrt{m/\\epsilon}$)");
    plt::ylabel("Energy / ($\\epsilon$)");

    plt::subplot(3, 1, 3);
    plt::named_plot("Instant", Ti, Pr);
    plt::named_plot("Average", Ti, MPr);
    plt::legend();
    plt::title("Equilibrium Pressure : " + to_string(eqP));
    plt::xlabel("Time / ($\\sigma\\sqrt{m/\\epsilon}$)");
    plt::ylabel("Pressure / ($\\epsilon/\\sigma^2$)");

    plt::suptitle("Argon N = " + to_string(N));
    plt::tight_layout();
    plt::save("result.png");
    plt::show();

    return 0;
}

/*
    Unit of quantities:
    Length : σ (Lennard-Jones parameters)
    Energy : ε (Lennard-Jones parameters)
    Masses : m (mass of the atom)
    Velocity : sqrt(ε / m)
    Time : σ * sqrt(m / ε)
    Temperature : ε / k (k is Boltzmann's constant)
*/

/*
    2D Maxwell Distribution
    normal_distribution<double> gauss(0, sqrt(T * kb / m));
    vector<double> x(100000);
    for (int i = 0; i < 100000; i++) x[i] = sqrt(pow(gauss(seed), 2) + pow(gauss(seed), 2));
    plt::hist(x, 100);
    plt::show();
*/