package extractors.hog;

import Jama.Matrix;
import io.*;

import java.io.*;
import java.util.*;
import java.util.concurrent.*;

import util.Model;
import fragmenter.*;
import geometric.Triangle;

/**
 * This class generates distance fields from a set of meshes utilizing a hierarchical AABB-structure. It
 * accelerates computation by processing several meshes in parallel.
 * 
 */
public class DistanceFieldGenerator {

	// resolution of the distance fields
	private int resx, resy, resz;

	// static fields needed for parallelization
	private static DistanceFieldGenerator dfg;
	private static long start;
	private static Queue<File> meshes;

	/**
	 * Constructs a {@link DistanceFieldGenerator} with the specified resolution.
	 * 
	 * @param resx
	 * @param resy
	 * @param resz
	 */
	public DistanceFieldGenerator(int resx, int resy, int resz) {
		this.resx = resx;
		this.resy = resy;
		this.resz = resz;
	}

	/**
	 * Generates a distance field as 3D double array from the given mesh with a AABB-structure.
	 * 
	 * @param mesh
	 * @return
	 */
	public double[][][] generateDistanceField(Model mesh) {
		// build the octree
		OctNode octree = buildOctree(mesh);

		double[][][] field = new double[resx][resy][resz];
		double dx = 2.0 / resx;
		double dy = 2.0 / resy;
		double dz = 2.0 / resz;

		// the point representing the center of a voxel
		Matrix point = new Matrix(3, 1);

		for (int x = 0; x < resx; x++) {
			point.set(0, 0, (x + .5) * dx - 1);
			for (int y = 0; y < resy; y++) {
				point.set(1, 0, (y + .5) * dy - 1);
				for (int z = 0; z < resz; z++) {
					point.set(2, 0, (z + .5) * dz - 1);
					// get the distance to the mesh
					field[x][y][z] = getDistance(octree, point);
				}
			}
		}

		return field;
	}

	/**
	 * Generates a distance field as 3D double array from the given mesh with a brute force approach. This is
	 * mainly for test purposes.
	 * 
	 * @param mesh
	 * @return
	 */
	public double[][][] generateBruteForce(Model mesh) {
		double[][][] field = new double[resx][resy][resz];
		double dx = 2.0 / resx;
		double dy = 2.0 / resy;
		double dz = 2.0 / resz;

		Matrix point = new Matrix(3, 1);

		int cnt = 0;
		for (int x = 0; x < resx; x++) {
			point.set(0, 0, (x + .5) * dx - 1);
			for (int y = 0; y < resy; y++) {
				point.set(1, 0, (y + .5) * dy - 1);
				for (int z = 0; z < resz; z++) {
					cnt++;
					point.set(2, 0, (z + .5) * dz - 1);
					field[x][y][z] = distanceBruteForce(mesh, point);
				}
			}
		}

		return field;
	}

	/**
	 * Computes the distance of the given point to the mesh in a brute force manner and returns it. This is
	 * mainly for test purposes.
	 * 
	 * @param mesh
	 * @param point
	 * @return
	 */
	private double distanceBruteForce(Model mesh, Matrix point) {
		double distance = Double.MAX_VALUE;
		for (Triangle triangle : mesh.getTriangles()) {
			double tmp = triangle.distanceToPoint(point);
			if (tmp < distance) {
				distance = tmp;
			}
		}

		return distance;
	}

	/**
	 * Computes the octree from a mesh and returns it.
	 * 
	 * @param mesh
	 * @return
	 */
	private OctNode buildOctree(Model mesh) {
		// compute bounds
		double dx = 2 / (resx - 4);
		double dy = 2 / (resy - 4);
		double dz = 2 / (resz - 4);
		double[][] llbArr = { { -1 - 2 * dx }, { -1 - 2 * dy }, { -1 - 2 * dz } };
		double[][] trfArr = { { 1 + 2 * dx }, { 1 + dy }, { 1 + dz } };

		// set root
		OctNode root = new OctNode(null, new Matrix(llbArr), new Matrix(trfArr), mesh);
		Queue<OctNode> queue = new LinkedList<OctNode>();
		queue.add(root);

		while (!queue.isEmpty()) {
			OctNode current = queue.poll();
			if (!current.isEmpty() && current.getMesh().getTriangles().length > 2 && relativeVolumeSize(current) < 0.95) {
				// split the mesh
				Fragmenter f = new BBOctFragmenter(current.getCenter());

				f.doFragmentation(current.getMesh());
				for (int i = 0; i < f.getFragments().length; i++) {
					// create the child
					OctNode child = new OctNode(current,
							getChildLlb(current.getLowerLeftBack(), current.getCenter(), i), getChildTrf(current
									.getTopRightFront(), current.getCenter(), i), f.getFragments()[i]);

					// adjust bounds if necessary
					child.recalculateBounds();

					// add child to tree and to queue
					current.setChild(i, child);
					queue.add(child);
				}
				// clear the mesh
				if (!current.clearMesh()) {
					throw new RuntimeException("Tried to clear mesh of leaf.");
				}
			}
		}
		return root;
	}

	/**
	 * Returns the relative volume size of the given node to its parent.
	 * 
	 * @param current
	 * @return
	 */
	private double relativeVolumeSize(OctNode current) {
		if (current.isRoot()) {
			return 0;
		}
		return current.getVolume() / current.getParent().getVolume();
	}

	/**
	 * Return the left, lower, back bound of a child depending on its index i.
	 * 
	 * @param parentLlb
	 * @param center
	 * @param i
	 * @return
	 */
	private Matrix getChildLlb(Matrix parentLlb, Matrix center, int i) {
		Matrix llb;
		switch (i) {
			case 0:
				return parentLlb.copy();
			case 1:
				llb = parentLlb.copy();
				llb.set(2, 0, center.get(2, 0));
				return llb;
			case 2:
				llb = parentLlb.copy();
				llb.set(1, 0, center.get(1, 0));
				return llb;
			case 3:
				llb = center.copy();
				llb.set(0, 0, parentLlb.get(0, 0));
				return llb;
			case 4:
				llb = parentLlb.copy();
				llb.set(0, 0, center.get(0, 0));
				return llb;
			case 5:
				llb = center.copy();
				llb.set(1, 0, parentLlb.get(1, 0));
				return llb;
			case 6:
				llb = center.copy();
				llb.set(2, 0, parentLlb.get(2, 0));
				return llb;
			case 7:
				return center.copy();
			default:
				throw new IllegalArgumentException("i must be between 0 and 7 (inclusive), but was " + i);
		}
	}

	/**
	 * Return the top, right, front bound of a child depending on its index i.
	 * 
	 * @param parentLlb
	 * @param center
	 * @param i
	 * @return
	 */
	private Matrix getChildTrf(Matrix parentTrf, Matrix center, int i) {
		Matrix trf;
		switch (i) {
			case 0:
				return center.copy();
			case 1:
				trf = center.copy();
				trf.set(2, 0, parentTrf.get(2, 0));
				return trf;
			case 2:
				trf = center.copy();
				trf.set(1, 0, parentTrf.get(1, 0));
				return trf;
			case 3:
				trf = parentTrf.copy();
				trf.set(0, 0, center.get(0, 0));
				return trf;
			case 4:
				trf = center.copy();
				trf.set(0, 0, parentTrf.get(0, 0));
				return trf;
			case 5:
				trf = parentTrf.copy();
				trf.set(1, 0, center.get(1, 0));
				return trf;
			case 6:
				trf = parentTrf.copy();
				trf.set(2, 0, center.get(2, 0));
				return trf;
			case 7:
				return parentTrf.copy();
			default:
				throw new IllegalArgumentException("i must be between 0 and 7 (inclusive), but was " + i);
		}
	}

	/**
	 * Computes the distance of a point to a mesh organized in an octree.
	 * 
	 * @param octree
	 * @param point
	 * @return
	 */
	private double getDistance(OctNode octree, Matrix point) {
		double distance = Double.MAX_VALUE;
		Queue<OctNode> queue = new LinkedList<OctNode>();
		// start at root
		queue.add(octree);
		while (!queue.isEmpty()) {
			OctNode current = queue.poll();
			if (current.distanceToPoint(point) < distance) {
				// if the distance to the current bounding box is smaller than current smallest distance
				if (current.isLeaf()) {
					// if node is leaf, iterate over triangles and save minimum
					for (Triangle triangle : current.getMesh().getTriangles()) {
						double tmp = triangle.distanceToPoint(point);
						if (tmp < distance) {
							distance = tmp;
						}
					}
				} else {
					// else add non-empty children to the queue
					for (OctNode octNode : current.getChildren()) {
						if (!octNode.isEmpty()) queue.add(octNode);
					}
				}
			}
		}
		return distance;
	}

	/**
	 * Saves the given distance field as floats in a binary file.
	 * 
	 * @param field
	 * @param inPath
	 * @throws IOException
	 */
	public void saveFieldAsFloats(double[][][] field, String inPath) throws IOException {
		String finalPath = getOutfilePath(inPath);
		DataOutputStream out = new DataOutputStream(new FileOutputStream(finalPath));
		for (int x = 0; x < resx; x++) {
			for (int y = 0; y < resy; y++) {
				for (int z = 0; z < resz; z++) {
					out.writeFloat((float) field[x][y][z]);
				}
			}
		}
		out.flush();
		out.close();
	}

	/**
	 * Saves the given distance field as doubles in a binary file.
	 * 
	 * @param field
	 * @param inPath
	 * @throws IOException
	 */
	public void saveField(double[][][] field, String inPath) throws IOException {
		String finalPath = getOutfilePath(inPath);
		DataOutputStream out = new DataOutputStream(new FileOutputStream(finalPath));
		for (int x = 0; x < resx; x++) {
			for (int y = 0; y < resy; y++) {
				for (int z = 0; z < resz; z++) {
					out.writeDouble(field[x][y][z]);
				}
			}
		}
		out.flush();
		out.close();
	}

	/**
	 * Saves the given distance field as shorts in a binary file.
	 * 
	 * @param field
	 * @param inPath
	 * @throws IOException
	 */
	public void saveFieldAsShorts(double[][][] field, String inPath) throws IOException {
		double maxDistance = Math.sqrt(12);
		String finalPath = getOutfilePath(inPath);
		DataOutputStream out = new DataOutputStream(new FileOutputStream(finalPath));
		for (int x = 0; x < resx; x++) {
			for (int y = 0; y < resy; y++) {
				for (int z = 0; z < resz; z++) {
					double val = field[x][y][z] / maxDistance;
					short shVal = (short) (val * Short.MAX_VALUE);
					out.writeShort(shVal);
				}
			}
		}
		out.flush();
		out.close();

		// generate description file for rendering program (test purposes)
		String txtPath = finalPath.substring(0, finalPath.lastIndexOf(".")).concat(".txt");
		PrintWriter pw = new PrintWriter(txtPath);
		pw.println(resx);
		pw.println(resy);
		pw.println(resz);
		pw.println(1);
		pw.println(1);
		pw.println(1);
		pw.println(16);
		pw.flush();
		pw.close();
	}

	/**
	 * Reads a distance field from a binary file (doubles).
	 * 
	 * @param file
	 * @return
	 * @throws IOException
	 */
	public double[][][] readFromDoubles(String file) throws IOException {
		double[][][] field = new double[resx][resy][resz];
		DataInputStream in = new DataInputStream(new FileInputStream(file));
		for (int x = 0; x < resx; x++) {
			for (int y = 0; y < resy; y++) {
				for (int z = 0; z < resz; z++) {
					field[x][y][z] = in.readDouble();
				}
			}
		}
		in.close();
		return field;
	}

	/**
	 * Reads the distance field of a model given as offFile from the respective binary file instead of
	 * calculating it. It is assumed, that the conversion already took place.
	 * 
	 * @param offFile
	 * @return
	 * @throws IOException
	 */
	public double[][][] readFromConvertedModel(File offFile) throws IOException {
		String path = getOutfilePath(offFile.getPath());
		return readFromDoubles(path);
	}

	/**
	 * Utility method for creating filenames of the distance fields depending on the path of the offFile.
	 * 
	 * @param inPath
	 * @return
	 */
	private String getOutfilePath(String inPath) {
		String path = inPath.substring(0, inPath.lastIndexOf("."));
		return path.concat("_" + resx + "x" + resy + "x" + resz + ".raw");
	}

	/**
	 * main method for distance field computation of a benchmark.
	 * 
	 * @param args
	 */
	public static void main(String[] args) {
		// parse arguments
		String root = args[0];
		int res = Integer.parseInt(args[1]);
		int threads = Integer.parseInt(args[2]);

		dfg = new DistanceFieldGenerator(res, res, res);
		System.out.println("Converting " + root + "with res " + res + "...");
		start = System.currentTimeMillis();
		meshes = new LinkedList<File>();

		// collect all .off-files
		process(new File(root));

		System.out.println("Starting processing ...");
		// process all .off-files
		ExecutorService es = Executors.newFixedThreadPool(threads);
		while (!meshes.isEmpty()) {
			es.execute(dfg.new DoOff(meshes.poll()));
		}

		try {
			es.shutdown();
			if (!es.awaitTermination(7, TimeUnit.DAYS)) {
				es.shutdownNow();
			}
		} catch (InterruptedException e) {
			es.shutdownNow();
			Thread.currentThread().interrupt();
		}

		System.out.println("Done in " + prettyprint(System.currentTimeMillis() - start) + ".");
	}

	/**
	 * Collects .off-files recursively
	 * 
	 * @param file
	 */
	private static void process(File file) {
		File[] dirList = file.listFiles(new DirFilter());
		File[] offList = file.listFiles(new OFFFilenameFilter());
		for (File dir : dirList) {
			process(dir);
		}

		for (File off : offList) {
			meshes.add(off);
		}
	}

	/**
	 * This class represents a thread that calculates the distance field for a .off-file.
	 * 
	 * 
	 */
	private class DoOff implements Runnable {

		// the file to calculate the distance filed from
		File off;

		public DoOff(File file) {
			off = file;
		}

		@Override
		public void run() {
			try {
				System.out.println("reading the model ... ");
				// read the model
				Model m = ModelIO.loadModel(off);

				System.out.println("normalize the model ... ");
				// normalize the model
				m.normalize();
				if (!new File(dfg.getOutfilePath(off.getPath())).exists()) {
					// if distance field has not been calculated - calculate
					double[][][] field = dfg.generateDistanceField(m);
					dfg.saveField(field, off.getPath());
					System.out.println("Finished " + off.getName());
				}
			} catch (CorruptOFFFileException ex) {
				// this program can only handle triangles
				System.out.println(off.getPath() + ": Illegal Face. Ignored.");
			} catch (IOException e) {
				e.printStackTrace();
			}
		}

	}

	/**
	 * @return the resx
	 */
	public int getResx() {
		return resx;
	}

	/**
	 * @return the resy
	 */
	public int getResy() {
		return resy;
	}

	/**
	 * @return the resz
	 */
	public int getResz() {
		return resz;
	}

	/**
	 * Return a time length in a pretty format.
	 * 
	 * @param etl
	 * @return
	 */
	public static String prettyprint(double etl) {
		int hours = (int) etl / 3600000;
		int min = (int) (etl % 3600000) / 60000;
		int sec = (int) (etl % 60000) / 1000;

		return hours + "h " + min + "m " + sec + "s";
	}
}