package edu.rice.hj.example.applications.barneshut;

import edu.rice.hj.api.SuspendableException;

import java.util.Random;

import static edu.rice.hj.Module1.*;

/**
 * BarnessHut class of the Barnes Hut application.
 *
 * @author <a href="http://shams.web.rice.edu/">Shams Imam</a> (shams@rice.edu)
 */
public final class BarnessHut {
    private static int nbodies; // number of bodies in system
    private static int ntimesteps; // number of time steps to run
    private static double dtime; // length of one time step
    private static double eps; // potential softening parameter
    private static double tol; // tolerance for stopping recursion, should be less than 0.57 for 3D case to bound error

    private static double dthf, epssq, itolsq;
    private static OctTreeLeafNodeData body[]; // the n bodies
    private static double diameter, centerx, centery, centerz;
    private static int curr;

    static final long[] nodesTraversed_100K = {71891025, 147078212};
    static final double[] posX_100K = {0.4989361610840987, 0.5114754484309701};
    static final double[] posY_100K = {0.5003786980143968, 0.5128978196831631};
    static final double[] posZ_100K = {0.5005633448563184, 0.5130668272583695};

    static final long[] nodesTraversed_50K = {31731122, 64845694};
    static final double[] posX_50K = {0.49890703243720386, 0.5114613898364733};
    static final double[] posY_50K = {0.5015303882283104, 0.5140502713460989};
    static final double[] posZ_50K = {0.5025002040538423, 0.5149475631876476};

    static final long[] nodesTraversed_10K = {4586829, 9345701};
    static final double[] posX_10K = {0.4995124926425938, 0.5120021242194736};
    static final double[] posY_10K = {0.4971952727192604, 0.5097297518450723};
    static final double[] posZ_10K = {0.5038525463784792, 0.5163728449901297};

    private static void RandomInput(final int num, final boolean print) {
        double vx, vy, vz;
        final Random rand = new Random(0);
        nbodies = num;
        ntimesteps = 1;
        dtime = 0.025;
        eps = 0.05;
        tol = 0.5;
        dthf = 0.5 * dtime;
        epssq = eps * eps;
        itolsq = 1.0 / (tol * tol);
        if (print) {
            System.err.println("configuration: " + nbodies + " bodies, " + ntimesteps + " time steps");
            System.err.println();
        }
        body = new OctTreeLeafNodeData[nbodies];
        for (int i = 0; i < nbodies; i++) {
            body[i] = new OctTreeLeafNodeData();
        }
        for (int i = 0; i < nbodies; i++) {
            body[i].mass = 1.0E-04;
            body[i].posx = rand.nextDouble();
            body[i].posy = rand.nextDouble();
            body[i].posz = rand.nextDouble();
            vx = rand.nextDouble();
            vy = rand.nextDouble();
            vz = rand.nextDouble();
            body[i].setVelocity(vx, vy, vz);
        }
    }

    /*
     * Computes a bounding box around all the bodies.
     */
    private static void ComputeCenterAndDiameter() {
        double minx, miny, minz;
        double maxx, maxy, maxz;
        double posx, posy, posz;
        minx = miny = minz = Double.MAX_VALUE;
        maxx = maxy = maxz = Double.MIN_VALUE;
        for (int i = 0; i < nbodies; i++) {
            posx = body[i].posx;
            posy = body[i].posy;
            posz = body[i].posz;
            if (minx > posx) {
                minx = posx;
            }
            if (miny > posy) {
                miny = posy;
            }
            if (minz > posz) {
                minz = posz;
            }
            if (maxx < posx) {
                maxx = posx;
            }
            if (maxy < posy) {
                maxy = posy;
            }
            if (maxz < posz) {
                maxz = posz;
            }
        }
        diameter = maxx - minx;
        if (diameter < (maxy - miny)) {
            diameter = (maxy - miny);
        }
        if (diameter < (maxz - minz)) {
            diameter = (maxz - minz);
        }
        centerx = (maxx + minx) * 0.5;
        centery = (maxy + miny) * 0.5;
        centerz = (maxz + minz) * 0.5;
    }

    /*
     * Recursively inserts a body into the octree.
     */
    private static void Insert(final NodeBarnessHut root, final OctTreeLeafNodeData b,
                               final double r) {
        double x = 0.0, y = 0.0, z = 0.0;
        final OctTreeNodeData n = root.getData();
        int i = 0;
        if (n.posx < b.posx) {
            i = 1;
            x = r;
        }
        if (n.posy < b.posy) {
            i += 2;
            y = r;
        }
        if (n.posz < b.posz) {
            i += 4;
            z = r;
        }
        final NodeBarnessHut child = root.getChild(i);
        if (child == null) {
            final NodeBarnessHut newnode = new NodeBarnessHut(b);
            root.setChild(i, newnode);
        } else {
            final double rh = 0.5 * r;
            final OctTreeNodeData ch = child.getData();
            if (!(ch.isLeaf())) {
                Insert(child, b, rh);
            } else {
                final NodeBarnessHut newnode = new NodeBarnessHut(new OctTreeNodeData(n.posx - rh + x, n.posy - rh + y, n.posz
                        - rh + z));
                Insert(newnode, b, rh);
                Insert(newnode, (OctTreeLeafNodeData) ch, rh);
                root.setChild(i, newnode);
            }
        }
    }

    /*
     * Traverses the tree bottom up to compute the total mass and the center of
     * mass of all the bodies in the subtree rooted in each internal octree node.
     */
    private static void ComputeCenterOfMass(final NodeBarnessHut root) {
        double m, px = 0.0, py = 0.0, pz = 0.0;
        final OctTreeNodeData n = root.getData();
        int j = 0;
        n.mass = 0.0;
        for (int i = 0; i < 8; i++) {
            final NodeBarnessHut child = root.getChild(i);
            if (child != null) {
                if (i != j) {
                    root.setChild(i, null);
                    root.setChild(j, child);
                }
                j++;
                final OctTreeNodeData ch = child.getData();
                if (ch.isLeaf()) {
                    body[curr++] = (OctTreeLeafNodeData) ch;
                } else {
                    ComputeCenterOfMass(child);
                }
                m = ch.mass;
                n.mass += m;
                px += ch.posx * m;
                py += ch.posy * m;
                pz += ch.posz * m;
            }
        }
        m = 1.0 / n.mass;
        n.posx = px * m;
        n.posy = py * m;
        n.posz = pz * m;
    }

    /*
     * Calculates the force acting on a body
     */
    private static void ComputeForce(final OctTreeLeafNodeData nd, final NodeBarnessHut root, final double size, final double itolsq, final int step, final double dthf, final double epssq) {
        final double ax;
        final double ay;
        final double az;
        ax = nd.accx;
        ay = nd.accy;
        az = nd.accz;
        nd.accx = 0.0;
        nd.accy = 0.0;
        nd.accz = 0.0;
        RecurseForce(nd, root, size * size * itolsq, epssq);
        if (step > 0) {
            nd.velx += (nd.accx - ax) * dthf;
            nd.vely += (nd.accy - ay) * dthf;
            nd.velz += (nd.accz - az) * dthf;
        }
    }

    /*
     * Helper method to compute the force acting on a body (recursively walks the
     * tree)
     */

    /**
     * <p>RecurseForce.</p>
     *
     * @param nd    a {@link edu.rice.hj.example.applications.barneshut.OctTreeLeafNodeData} object.
     * @param nn    a {@link edu.rice.hj.example.applications.barneshut.NodeBarnessHut} object.
     * @param dsq   a double.
     * @param epssq a double.
     */
    public static void RecurseForce(final OctTreeLeafNodeData nd, final NodeBarnessHut nn, final double dsq, final double epssq) {
        if (nn == null) {
            return;
        }
        final double drx;
        final double dry;
        final double drz;
        double drsq;
        final double nphi;
        final double scale;
        final double idr;
        final OctTreeNodeData n = nn.getData();
        drx = n.posx - nd.posx;
        dry = n.posy - nd.posy;
        drz = n.posz - nd.posz;
        drsq = drx * drx + dry * dry + drz * drz;
        nd.nodesTraversed.incrementAndGet();
        if (drsq < dsq) {
            if (!(n.isLeaf())) { // n is a cell
                RecurseForce(nd, nn.child0, dsq * 0.25, epssq);
                RecurseForce(nd, nn.child1, dsq * 0.25, epssq);
                RecurseForce(nd, nn.child2, dsq * 0.25, epssq);
                RecurseForce(nd, nn.child3, dsq * 0.25, epssq);
                RecurseForce(nd, nn.child4, dsq * 0.25, epssq);
                RecurseForce(nd, nn.child5, dsq * 0.25, epssq);
                RecurseForce(nd, nn.child6, dsq * 0.25, epssq);
                RecurseForce(nd, nn.child7, dsq * 0.25, epssq);

            } else { // n is a body
                if (n != nd) {
                    drsq += epssq;
                    idr = 1 / Math.sqrt(drsq);
                    nphi = n.mass * idr;
                    scale = nphi * idr * idr;
                    nd.lock.lock();
                    nd.accx += drx * scale;
                    nd.accy += dry * scale;
                    nd.accz += drz * scale;
                    nd.lock.unlock();
                }
            }
        } else { // node is far enough away, don't recurse any deeper
            drsq += epssq;
            idr = 1 / Math.sqrt(drsq);
            nphi = n.mass * idr;
            scale = nphi * idr * idr;
            nd.lock.lock();
            nd.accx += drx * scale;
            nd.accy += dry * scale;
            nd.accz += drz * scale;
            nd.lock.unlock();
        }
    }

    /*
     * Advances a body's velocity and position by one time step
     */
    private static void Advance(final double dthf, final double dtime) {
        double dvelx, dvely, dvelz;
        double velhx, velhy, velhz;

        for (int i = 0; i < nbodies; i++) {
            final OctTreeLeafNodeData nd = body[i];
            body[i] = nd;
            dvelx = nd.accx * dthf;
            dvely = nd.accy * dthf;
            dvelz = nd.accz * dthf;
            velhx = nd.velx + dvelx;
            velhy = nd.vely + dvely;
            velhz = nd.velz + dvelz;
            nd.posx += velhx * dtime;
            nd.posy += velhy * dtime;
            nd.posz += velhz * dtime;
            nd.velx = velhx + dvelx;
            nd.vely = velhy + dvely;
            nd.velz = velhz + dvelz;
        }
    }

    /**
     * BarnessHut method to launch the binary search tree code.
     *
     * @param args an array of {@link String} objects.
     */
    public static void main(final String[] args) {

        System.out.println("USAGE: java BarnessHut <numBodies>");
        final int size = args.length > 0 ? Integer.parseInt(args[0]) : 100000;

        final long[] harnessTime = {0};
        launchHabaneroApp(() -> {

            System.out.println("HJ Threads=" + numWorkerThreads());
            System.out.println("Input: size = " + size);
            RandomInput(size, true);

            OctTreeNodeData res = null;

            int step = 0;

            int inner = 5;
            int outter = 3;
            if (args.length > 1) {
                inner = Integer.parseInt(args[1]);
            }
            if (args.length > 2) {
                outter = Integer.parseInt(args[2]);
            }

            final long start = System.nanoTime();
            boolean harness = false;
            for (int ii = 0; ii < outter; ii++) {
                if (ii + 1 == outter) {
                    System.gc();
                    harness = true;
                }
                for (int jj = 0; jj < inner; jj++) {
                    System.out.println("========================== ITERATION (" + ii + "." + jj + ") ==================================");
                    final long startTime = System.nanoTime();
                    final int s = step;
                    ComputeCenterAndDiameter();
                    final NodeBarnessHut root = new NodeBarnessHut(new OctTreeNodeData(centerx, centery, centerz)); // create the
                    final double radius = diameter * 0.5;
                    for (int i = 0; i < nbodies; i++) {
                        Insert(root, body[i], radius); // grow the tree by inserting each body
                    }
                    curr = 0;
                    // summarize subtree info in each internal node (plus restructure tree and sort bodies for performance reasons)
                    ComputeCenterOfMass(root);

                    ParallelFor(root, step);

                    Advance(dthf, dtime); // advance the position and velocity of each body

                    // print center of mass for this timestep
                    res = root.getData();

                    long totalNodesTraversed = 0;
                    for (int i = 0; i < body.length; i++) {
                        totalNodesTraversed += body[i].nodesTraversed.get();
                    }
                    System.out.println("totalNodesTraversed: " + totalNodesTraversed);
                    String status = "PASSED";
                    if (nbodies == 100000 && s < 2) {
                        // check expected values
                        if (totalNodesTraversed != nodesTraversed_100K[s]) {
                            System.out.println("FAILED: Exptected nodes = " + nodesTraversed_100K[s] + " But obtained = " + totalNodesTraversed);
                            status = "FAILED";
                        }
                        if (res.posx != posX_100K[s]) {
                            System.out.println("FAILED: Exptected posX = " + posX_100K[s] + " But obtained = " + res.posx);
                            status = "FAILED";
                        }

                        if (res.posy != posY_100K[s]) {
                            System.out.println("FAILED: Exptected posY = " + posY_100K[s] + " But obtained = " + res.posy);
                            status = "FAILED";
                        }

                        if (res.posz != posZ_100K[s]) {
                            System.out.println("FAILED: Exptected posZ = " + posZ_100K[s] + " But obtained = " + res.posz);
                            status = "FAILED";
                        }
                    } else if (nbodies == 50000 && s < 2) {
                        // check expected values
                        if (totalNodesTraversed != nodesTraversed_50K[s]) {
                            System.out.println("FAILED: Exptected nodes = " + nodesTraversed_50K[s] + " But obtained = " + totalNodesTraversed);
                            status = "FAILED";
                        }
                        if (res.posx != posX_50K[s]) {
                            System.out.println("FAILED: Exptected posX = " + posX_50K[s] + " But obtained = " + res.posx);
                            status = "FAILED";
                        }

                        if (res.posy != posY_50K[s]) {
                            System.out.println("FAILED: Exptected posY = " + posY_50K[s] + " But obtained = " + res.posy);
                            status = "FAILED";
                        }

                        if (res.posz != posZ_50K[s]) {
                            System.out.println("FAILED: Exptected posZ = " + posZ_50K[s] + " But obtained = " + res.posz);
                            status = "FAILED";
                        }
                    } else if (nbodies == 10000 && s < 2) {
                        // check expected values
                        if (totalNodesTraversed != nodesTraversed_10K[s]) {
                            System.out.println("FAILED: Exptected nodes = " + nodesTraversed_10K[s] + " But obtained = " + totalNodesTraversed);
                            status = "FAILED";
                        }
                        if (res.posx != posX_10K[s]) {
                            System.out.println("FAILED: Exptected posX = " + posX_10K[s] + " But obtained = " + res.posx);
                            status = "FAILED";
                        }

                        if (res.posy != posY_10K[s]) {
                            System.out.println("FAILED: Exptected posY = " + posY_10K[s] + " But obtained = " + res.posy);
                            status = "FAILED";
                        }

                        if (res.posz != posZ_10K[s]) {
                            System.out.println("FAILED: Exptected posZ = " + posZ_10K[s] + " But obtained = " + res.posz);
                            status = "FAILED";
                        }
                    }
                    final long time = System.nanoTime() - startTime;
                    final double secs = ((double) time) / 1e9;
                    System.out.println("BarnessHut: (Timestep " + step + ") " + status + ". Time: " + secs + " secs");
                    step++;
                    if (harness) {
                        harnessTime[0] += time;
                    }
                }
            }

        });

        System.out.println("Harness ended...");
        System.out.println("============================ MMTk Statistics Totals ============================");
        System.out.println("time.mu");
        System.out.println(harnessTime[0]);
        System.out.println("Total time: " + harnessTime[0] + " ms");
        System.out.println("------------------------------ End MMTk Statistics -----------------------------");
        System.out.println("===== TEST PASSED in 0 msec =====");
    }

    /**
     * <p>ComputeForceAll.</p>
     *
     * @param i    a int.
     * @param root a {@link edu.rice.hj.example.applications.barneshut.NodeBarnessHut} object.
     * @param s    a int.
     */
    public static void ComputeForceAll(final int i, final NodeBarnessHut root, final int s) {
        ComputeForce(body[i], root, diameter, itolsq, s, dthf, epssq);
    }

    // a substitute for parallel for annotations on for statements, which are not supported in Java 6

    /**
     * <p>ParallelFor.</p>
     *
     * @param root a {@link edu.rice.hj.example.applications.barneshut.NodeBarnessHut} object.
     * @param s    a int.
     */
    public static void ParallelFor(final NodeBarnessHut root, final int s) throws SuspendableException {
        forallChunked(0, body.length - 1, (i) -> {
            ComputeForceAll(i, root, s);
        });
    }
}