EP2-Team/src/Octree.java

import codedraw.CodeDraw;

import java.awt.*;

/**
 * + -> >= 0
 * - -> < 0
 * 0 -> x+ y+ z+
 * 1 -> x- y+ z+
 * 2 -> x+ y- z+
 * 3 -> x- y- z+
 * 4 -> x+ y+ z-
 * 5 -> x- y+ z-
 * 6 -> x+ y- z-
 * 7 -> x- y- z-
 */
public class Octree {
    private OctreeItem root;
    private final Vector center;
    private final double size;

    public Octree(double size) {
        this(new Vector(), size);
    }

    public Octree(Vector center, double size) {
        this.center = center;
        this.size = size;
    }

    public void add(Body body) {
        if (root == null) {
            root = new OctreeLeaf(center, size, body);
        } else {
            root = root.add(body);
        }
    }

    public void applyForces(Body[] bodies, double t) {
        if (root == null) return;
        root.preCalc();
        Vector[] forces = new Vector[bodies.length];
        for (int i = 0; i < bodies.length; i++) {
            forces[i] = root.getForcesOnBody(bodies[i], t);
        }
        for (int i = 0; i < bodies.length; i++) {
            bodies[i].move(forces[i]);
        }
    }

    public void draw(CodeDraw cd) {
        if (root == null) return;
        root.draw(cd);
    }
}

abstract class OctreeItem {
    protected final Vector center;
    protected final double size;

    protected Body pseudoBody;

    protected OctreeItem(Vector center, double size) {
        this.center = center;
        this.size = size;
    }

    abstract protected OctreeNode add(Body body);

    abstract protected void preCalc();

    abstract protected Vector getForcesOnBody(Body body, double t);

    abstract protected void draw(CodeDraw cd);
}

class OctreeNode extends OctreeItem {
    private final OctreeItem[] children;
    private final Vector p1;
    private final Vector p2 ;

    protected OctreeNode(Vector center, double size) {
        super(center, size);
        children = new OctreeItem[8];
        p1 = center.plus(new Vector(size / 2));
        p2 = center.minus(new Vector(size / 2));
    }

    @Override
    protected OctreeNode add(Body body) {
        int num = getOctantNum(body.massCenter());
        if (num < 0) return this;
        OctreeItem oct = children[num];
        if (oct == null) {
            children[num] = new OctreeLeaf(getOctantCenter(num), size / 2, body);
        } else {
            children[num] = oct.add(body);
        }
        return this;
    }

    @Override
    protected void preCalc() {
        double mass = 0;
        for (OctreeItem oct : children) {
            if (oct == null) continue;
            oct.preCalc();
            mass += oct.pseudoBody.mass();
        }
        Vector massCenter = new Vector();
        for (OctreeItem oct : children) {
            if (oct == null) continue;
            massCenter.add(oct.pseudoBody.massCenter().times(oct.pseudoBody.mass() / mass));
        }
        this.pseudoBody = new Body(mass, massCenter);
    }

    @Override
    protected Vector getForcesOnBody(Body body, double t) {
        double r = pseudoBody.massCenter().distanceTo(body.massCenter());
        if (r == 0) {
            return new Vector();
        } else if (size / r < t) {
            return body.gravitationalForce(pseudoBody);
        }
        Vector force = new Vector();
        for (OctreeItem child : children) {
            if (child == null) continue;
            force.add(child.getForcesOnBody(body, t));
        }
        return force;
    }

    @Override
    protected void draw(CodeDraw cd) {
        for (OctreeItem child : children) {
            if (child == null) continue;
            child.draw(cd);
        }
        /*
        Vector p1 = center.minus(new Vector(size / 2));
        Vector p2 = new Vector(size).minus(new Vector(Simulation.SECTION_SIZE / 2));
        cd.setColor(Color.CYAN);
        cd.drawRectangle(p1.getScreenX(cd), p1.getScreenY(cd), p2.getScreenX(cd), p2.getScreenY(cd));
        */
    }

    private int getOctantNum(Vector bodyPos) {
        if (!isInOctant(bodyPos)) return -1;
        Vector pos = bodyPos.minus(this.center);
        return ((pos.getX() < 0) ? 1 : 0) | ((pos.getY() < 0) ? 2 : 0) | ((pos.getZ() < 0) ? 4 : 0);
    }

    private boolean isInOctant(Vector pos) {
        return  pos.getX() <= p1.getX() && pos.getY() <= p1.getY() && pos.getZ() <= p1.getZ() &&
                pos.getX() >= p2.getX() && pos.getY() >= p2.getY() && pos.getZ() >= p2.getZ();
    }

    private Vector getOctantCenter(int octNum) {
        return this.center.plus(new Vector(
                ((octNum & 1) != 0) ? -1 : 1,
                ((octNum & 2) != 0) ? -1 : 1,
                ((octNum & 4) != 0) ? -1 : 1)
                .times(this.size / 4));
    }
}

class OctreeLeaf extends OctreeItem {
    private final Body body;

    public OctreeLeaf(Vector center, double size, Body body) {
        super(center, size);
        this.body = body;
    }

    @Override
    protected OctreeNode add(Body body) {
        return new OctreeNode(this.center, this.size).add(this.body).add(body);
    }

    @Override
    protected void preCalc() {
        this.pseudoBody = new Body(this.body);
    }

    @Override
    protected Vector getForcesOnBody(Body body, double t) {
        return body.gravitationalForce(pseudoBody);
    }

    @Override
    protected void draw(CodeDraw cd) {
        Vector p1 = center.minus(new Vector(size / 2));
        Vector p2 = new Vector(size).minus(new Vector(Simulation.SECTION_SIZE / 2));
        cd.setColor(Color.GREEN);
        cd.drawRectangle(p1.getScreenX(cd), p1.getScreenY(cd), p2.getScreenX(cd), p2.getScreenY(cd));
    }
}