package lpv.simulation.propagator;

import javax.vecmath.Matrix3f;
import javax.vecmath.Point3i;
import javax.vecmath.Vector2f;
import javax.vecmath.Vector3f;
import javax.vecmath.Vector4f;

import lpv.LPV;

public class AndreasKirschPropagator implements Propagator {
	// Original Author: Andreas Kirsch
	// /home/dt07jo1/Documents/LPVPrototype/lpv.fx
	// http://blog.blackhc.net/2010/07/light-propagation-volumes/
	
	private static final float Pi = 3.1415926f;

	private static final float SH_c0 = 0.282094792f; // 1 / 2sqrt(pi)
	private static final float SH_c1 = 0.488602512f; // sqrt(3/pi) / 2
	
	// my values from maple
	private static final float SH_cosLobe_c0 = 0.886226925f; // sqrt(pi)/2
	private static final float SH_cosLobe_c1 = 1.02332671f; // sqrt(pi/3)
	
	// the old crytek paper uses different values:
	/*private static final float SH_cosLobe_c0 = 0.25f;
	private static final float SH_cosLobe_c1 = 0.5f;*/
	
	
	// orientation = [ right | up | forward ] = [ x | y | z ]
	private static final Matrix3f[] Prop_neighbourOrientations = new Matrix3f[] {
		// Z+
		new Matrix3f( 
			1, 0, 0,
			0, 1, 0,
			0, 0, 1
		),
		// Z-
		new Matrix3f( 
			-1, 0, 0,
			0, 1, 0,
			0, 0, -1
		),
		// X+
		new Matrix3f( 
			0, 0, 1,
			0, 1, 0,
			-1, 0, 0
		),
		// X-
		new Matrix3f( 
			0, 0, -1,
			0, 1, 0,
			1, 0, 0
		),
		// Y+
		new Matrix3f( 
			1, 0, 0,
			0, 0, 1,
			0, -1, 0
		),
		// Y-
		new Matrix3f( 
			1, 0, 0,
			0, 0, -1,
			0, 1, 0
		)
	};
	
	
	public AndreasKirschPropagator() {
		
	}
	
	@Override
	public void propagate(LPV in, LPV out, LPV acc, LPV geometry, Point3i pos, int iteration) {
		Vector4f contribution = Prop_gatherContributions(in, pos);
		
		// simply copy for now
		if(iteration == 0) {
			contribution.add(in.get(pos.x, pos.y, pos.z));
		}
		out.get(pos.x, pos.y, pos.z).set(contribution);
		acc.get(pos.x, pos.y, pos.z).add(contribution);
	}
	
	
	private Vector4f Prop_gatherContributions(LPV in, Point3i cellIndex) {
		
		Vector4f contribution = new Vector4f(0.0f, 0.0f, 0.0f, 0.0f);
		
		for(int neighbour = 0 ; neighbour < 6 ; neighbour++) {
			Matrix3f orientation = Prop_neighbourOrientations[neighbour];
			
			// TODO: transpose all orientation matrices and use row indexing instead? ie int3( orientation[2] )
			Vector3f mainDirection = new Vector3f(0, 0, 1);
			orientation.transform(mainDirection);
			
			Point3i neighbourIndex = new Point3i(
					(int)(cellIndex.x - mainDirection.x), 
					(int)(cellIndex.y - mainDirection.y), 
					(int)(cellIndex.z - mainDirection.z));
			
			
			Vector4f neighbourCoeffs;
			if(neighbourIndex.x >= 0 && neighbourIndex.x < in.getWidth() && 
					neighbourIndex.y >= 0 && neighbourIndex.y < in.getHeight() && 
					neighbourIndex.z >= 0 && neighbourIndex.z < in.getDepth()) {
				neighbourCoeffs = new Vector4f(in.get(neighbourIndex.x, neighbourIndex.y, neighbourIndex.z));
			} else {
				neighbourCoeffs = new Vector4f(0.0f, 0.0f, 0.0f, 0.0f);
			}
			
			
			// crytek papers says that subtended solid angle has to be divided by 4 Pi, I disagree
			/*float directFaceSubtendedSolidAngle = 0.03188428f; // 0.4006696846f / 4Pi;
			float sideFaceSubtendedSolidAngle = 0.03369559f; // 0.4234413544f / 4Pi;*/
			
			/*float directFaceSubtendedSolidAngle = 4.0f / 18.0f;
			float sideFaceSubtendedSolidAngle = 4.0f / 36.0f;*/
			
			// again the contributions modelled below are approximating integrals over intensity distributions, ie total flux
			// we then backmodel the flux as a hemisphere (area) light, so to convert from flux to intensity we divide by Pi
			// and we do it here because its easier
			// TODO: since we're approximating with low order SH coeffs, maybe use a different/more precise *normalization* constant?
			float directFaceSubtendedSolidAngle = 0.4006696846f / Pi;
			float sideFaceSubtendedSolidAngle = 0.4234413544f / Pi;
			
			float transmittanceValue = 1.0f; // no occlusion
			float reflectionValue = 0.0f; // no reflection

			float occludedDirectFaceContribution = transmittanceValue * directFaceSubtendedSolidAngle;
			
			// TODO: simplify?
			Vector4f mainDirectionCosineLobeSH = SH_evaluateCosineLobe_direct(mainDirection);
			Vector4f mainDirectionSH = SH_evaluate_direct(mainDirection);
			
			Vector4f tmp = new Vector4f(mainDirectionCosineLobeSH);
			tmp.scale(occludedDirectFaceContribution * neighbourCoeffs.dot(mainDirectionSH));
			contribution.add(tmp);
			
			for( int sideFace = 0 ; sideFace < 4 ; sideFace++ ) {
				Vector3f evalDirection = Prop_getEvalSideDirection(sideFace, orientation);
				Vector3f reprojDirection = Prop_getReprojSideDirection(sideFace, orientation);
				
				float occludedSideFaceContribution = transmittanceValue * sideFaceSubtendedSolidAngle;
				
				Vector4f reprojDirectionCosineLobeSH = SH_evaluateCosineLobe_direct(reprojDirection);
				Vector4f evalDirectionSH = SH_evaluate_direct(evalDirection);
				
				
				tmp = new Vector4f(reprojDirectionCosineLobeSH);
				tmp.scale(occludedSideFaceContribution * neighbourCoeffs.dot(evalDirectionSH));
				contribution.add(tmp);
			}
		}
		
		return contribution;
	}
	
	
	
	// right up
	private static final Vector2f[] Prop_side = new Vector2f[] { new Vector2f(1.0f, 0.0f), new Vector2f(0.0f, 1.0f), new Vector2f(-1.0f, 0.0f), new Vector2f(0.0f, -1.0f) };

	// orientation = [ right | up | forward ] = [ x | y | z ]
	Vector3f Prop_getEvalSideDirection(int index, Matrix3f orientation) {
		float smallComponent = 0.4472135f; // 1 / sqrt(5)
		float bigComponent = 0.894427f; // 2 / sqrt(5)
		
		Vector2f side = Prop_side[index];
		// *either* x = 0 or y = 0
		Vector3f ret = new Vector3f(side.x * smallComponent, side.y * smallComponent, bigComponent);
		orientation.transform(ret);
		return ret;
	}

	Vector3f Prop_getReprojSideDirection(int index, Matrix3f orientation) {
		Vector2f side = Prop_side[index];
		Vector3f ret = new Vector3f(side.x, side.y, 0);
		orientation.transform(ret);
		return ret;
	}
	
	
	private static Vector4f SH_evaluate(Vector3f direction) {
		direction = new Vector3f(direction);
		direction.normalize();
		return new Vector4f(SH_c0, -SH_c1 * direction.y, SH_c1 * direction.z, -SH_c1 * direction.x);
	}

	// no normalization
	private static Vector4f SH_evaluate_direct(Vector3f direction) {	
		return new Vector4f(SH_c0, -SH_c1 * direction.y, SH_c1 * direction.z, -SH_c1 * direction.x);
	}

	// uses homogenous coordinate to become linear - and no normalization!
	private static Vector4f SH_evaluate(Vector4f direction) {	
		return new Vector4f(SH_c0 * direction.w, -SH_c1 * direction.y, SH_c1 * direction.z, -SH_c1 *direction.x);
	}
	
	
	private static Vector4f SH_evaluateCosineLobe(Vector3f direction) {
		direction = new Vector3f(direction);
		direction.normalize();
		
		return new Vector4f(SH_cosLobe_c0, -SH_cosLobe_c1 * direction.y, SH_cosLobe_c1 * direction.z, -SH_cosLobe_c1 * direction.x);
	}

	// no normalization
	private static Vector4f SH_evaluateCosineLobe_direct(Vector3f direction) {
		return new Vector4f(SH_cosLobe_c0, -SH_cosLobe_c1 * direction.y, SH_cosLobe_c1 * direction.z, -SH_cosLobe_c1 * direction.x);
	}

	// uses homogenous coordinate to become linear - and no normalization!
	private static Vector4f SH_evaluateCosineLobe_linear(Vector4f direction) {
		return new Vector4f(SH_cosLobe_c0 * direction.w, -SH_cosLobe_c1 * direction.y, SH_cosLobe_c1 * direction.z, -SH_cosLobe_c1 * direction.x);
	}
}
