/**
 * Using the Storm 3D library for lightning.
 */

use ui;
use ui:s3d;
use core:geometry;
use graphics;

/**
 * Vertex type for lights. Contains both the screen and the world coordinates so that we can do
 * lightning in world space.
 */
vertex LightIntermediate {
	// Screen space position.
	Vector4 pos : position, perspective;

	// World position. Used for lightning.
	Vector world : perspective;

	// Normal direction transformed to world space. For lightning.
	Vector normal : perspective;

	// Vertex color. Used for ambient and diffuse lightning.
	Color color : perspective;

	// Texture coordinates, if we want to apply textures.
	Point tex : perspective;
}


/**
 * Pipeline for lights.
 */
pipeline LightPipeline {
	// Here, the source vertex type is ui:s3d:Vertex to make it possible to use Model.
	data (Vertex vertex) -> LightIntermediate -> Output;

	// Transform from model space to world space. We have one version for positions and one for normals.
	// We could compute the normal transform matrix in the shader, but it is cheaper to compute it
	// on the CPU once before running the shaders.
	Transform worldPosition;
	Transform worldNormal;

	// Set the world matrices easily.
	assign world(Transform tfm) {
		worldPosition = tfm;
		worldNormal = tfm.forNormal();
	}

	// Transform from world space to camera space.
	Transform camera;

	// Camera position in world space.
	Vector cameraPos;

	// Position of the light in world space.
	Vector lightPos;

	// Ambient light strength.
	Float ambientStrength;

	// Specular strength and exponent.
	Float specularStrength;
	Float specularExponent;

	// Specular color.
	Color specularColor;
}


// Vertex shader - transforms from model space to camera space.
vertex shader for LightPipeline: LightIntermediate lightVertexShader(Vertex point) {
	Vector4 worldPos = Vector4(point.pos) * worldPosition;
	Vector4 screenPos = worldPos * camera;

	// Transform the normal. Note that we need to normalize it at some point. We do it later,
	// however, since the interpolation will mess up the normalization anyway.
	Vector normal = point.normal * worldNormal;

	LightIntermediate(screenPos, Vector(worldPos), normal, point.color, point.tex);
}

// Pixel shader - performs lightning.
pixel shader for LightPipeline: Output lightPixelShader(LightIntermediate point) {
	// Compute the direction vector to the light source and the normal. Note that we need to
	// normalize the normal at this point.
	Vector cameraDir = (cameraPos - point.world).normalized;
	Vector lightDir = (lightPos - point.world).normalized;
	Vector normal = point.normal.normalized;

	// Ambient light:
	Color light = point.color * ambientStrength;

	// Diffuse light:
	light += point.color * max(normal * lightDir, 0.0);

	// Specular light:
	Vector reflectDir = lightDir.reflect(normal);
	Float specular = pow(max(cameraDir * reflectDir, 0.0), specularExponent);
	light += specularColor * specularStrength * specular;

	Output(light);
}


/**
 * The painter.
 */
class LightPainter extends SamplePainter3D {
	// The pipeline itself.
	LightPipeline pipeline;

	// Model to render.
	Model model;

	// Transform to account for the camera position.
	Transform cameraPos;

	// Create.
	init() {
		init {
			// Create the pipeline and attach shaders.
			pipeline = LightPipeline:create(lightVertexShader, lightPixelShader);

			// Create the default camera transform.
			cameraPos = rotateX(-35 deg) * translate(Vector(0, 0, 3));

			// Load a teapot to render.
			model = resUrl("teapot.obj").loadOBJ.withColor(red);
		}

		// Set ambient light strength.
		pipeline.ambientStrength = 0.1;

		// Set specular light properties.
		pipeline.specularStrength = 0.8;
		pipeline.specularExponent = 32.0;
		pipeline.specularColor = white;
	}

	// Render.
	Bool render(Size size, Graphics3D g) : override {
		Float time = animationTime(20 s);

		// Clear the old buffer.
		g.clearDepth();
		g.fill(black);

		// Set the world transforms.
		pipeline.world = rotateY(360 deg * time * 4) * scale(0.5);

		// Set the camera transform.
		pipeline.camera = cameraPos * perspectiveProject(100 deg, size.aspect, 0.1, 100);

		// Let the light rotate above the model in a big circle, but slower.
		pipeline.lightPos = Vector(0, 10, -10) * rotateY(360 deg * -time);

		// Compute the camera position, for specular highlights.
		pipeline.cameraPos = Vector(0, 0, 0) * cameraPos.inverted;

		// Now we can render the model.
		pipeline.render(g, model);

		true;
	}
}


void light() {
	runSample("3D Lightning", LightPainter());
}
