/* Copyright (c) 2015 saharan
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy of this
 * software and associated documentation  * files (the "Software"), to deal in the Software
 * without restriction, including without limitation the rights to use, copy,  * modify, merge,
 * publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to
 * whom the Software is furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in all copies or
 * substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
 * PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR
 * ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
 * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
package {
	import com.adobe.utils.*;
	import flash.display.*;
	import flash.display3D.*;
	import flash.display3D.textures.*;
	import flash.events.*;
	import flash.net.*;
	import flash.geom.*;
	import flash.ui.*;
	import flash.utils.*;
	import net.hires.debug.*;
	
	/**
	 * GPU Fluid
	 * @author saharan
	 */
	[SWF(width = "512", height = "512", frameRate = "60")]
	public class GPUFluid extends Sprite {
		private var pmouseX:Number; // 直前のマウス座標
		private var pmouseY:Number;
		private var pressed:Boolean;
		private var s3d:Stage3D;
		private var c3d:Context3D;
		
		// color texture: r, g, b, UNUSED
		private var color1:Texture; // dst
		private var color2:Texture; // src
		
		// velocity texture: vx, vy, divergence, pressure
		private var velocity1:Texture; // dst
		private var velocity2:Texture; // src
		
		private var pressure:Texture;
		private var advectColorPG:Program3D;
		private var advectVelocityPG:Program3D;
		private var calcDivergencePG:Program3D;
		private var calcPressurePG:Program3D;
		private var applyPressurePG:Program3D;
		private var displayColorPG:Program3D;
		private var displayVelocityPG:Program3D;
		
		private var vtxB:VertexBuffer3D;
		private var idxB:IndexBuffer3D;
		
		private var fragmentConsts:Vector.<Number>;
		
		private const texSize:int = 512;
		
		private var displayType:int = 0;
		
		public function GPUFluid() {
			if (stage) init();
			else addEventListener(Event.ADDED_TO_STAGE, init);
		}
		
		private function init(e:Event = null): void {
			removeEventListener(Event.ADDED_TO_STAGE, init);
			stage.align = StageAlign.TOP_LEFT;
			stage.scaleMode = StageScaleMode.NO_SCALE;
			stage.addEventListener(MouseEvent.MOUSE_DOWN, function(e:Event):void { pressed = true; });
			stage.addEventListener(MouseEvent.MOUSE_UP, function(e:Event):void { pressed = false; });
			stage.addEventListener(KeyboardEvent.KEY_DOWN, function(e:KeyboardEvent):void {
				switch (e.keyCode) {
				case Keyboard.NUMBER_1:
					displayType = 0;
					break;
				case Keyboard.NUMBER_2:
					displayType = 1;
					break;
				}
			});
			s3d = stage.stage3Ds[0];
			s3d.addEventListener(Event.CONTEXT3D_CREATE, start);
			s3d.requestContext3D(Context3DRenderMode.AUTO, Context3DProfile.STANDARD_CONSTRAINED);
			var stats:Stats = new Stats();
			stats.mouseEnabled = false;
			stats.alpha = 0.8;
			addChild(stats);
		}
		
		private function start(e:Event):void {
			c3d = s3d.context3D;
			c3d.enableErrorChecking = true;
			c3d.configureBackBuffer(texSize, texSize, 0, false);
			c3d.setCulling(Context3DTriangleFace.NONE);
			//
			color1 = c3d.createTexture(texSize, texSize, Context3DTextureFormat.RGBA_HALF_FLOAT, true);
			color2 = c3d.createTexture(texSize, texSize, Context3DTextureFormat.RGBA_HALF_FLOAT, true);
			velocity1 = c3d.createTexture(texSize, texSize, Context3DTextureFormat.RGBA_HALF_FLOAT, true);
			velocity2 = c3d.createTexture(texSize, texSize, Context3DTextureFormat.RGBA_HALF_FLOAT, true);
			reset();
			//
			idxB = c3d.createIndexBuffer(6);
			idxB.uploadFromVector(Vector.<uint>([
				0, 1, 2,
				1, 3, 2,
			]), 0, 6);
			vtxB = c3d.createVertexBuffer(4, 5);
			vtxB.uploadFromVector(Vector.<Number>([
				-1, 1, 0, 0, 0,
				1, 1, 0, 1, 0,
				-1, -1, 0, 0, 1,
				1, -1, 0, 1, 1,
			]), 0, 4);
			c3d.setVertexBufferAt(0, vtxB, 0, Context3DVertexBufferFormat.FLOAT_3); // pos
			c3d.setVertexBufferAt(1, vtxB, 3, Context3DVertexBufferFormat.FLOAT_2); // uv
			
			//
			var agal:AGALMiniAssembler = new AGALMiniAssembler();
			var vtx:String =
				"mov op, va0 \n" +
				"mov v0, va1 \n"
			;
			//
			displayColorPG = createPG(agal, vtx, displayColorShader);
			displayVelocityPG = createPG(agal, vtx, displayVelocityShader);
			advectColorPG = createPG(agal, vtx, advectColorShader);
			advectVelocityPG = createPG(agal, vtx, advectVelocityShader);
			calcDivergencePG = createPG(agal, vtx, calcDivergenceShader);
			calcPressurePG = createPG(agal, vtx, calcPressureShader);
			applyPressurePG = createPG(agal, vtx, applyPressureShader);
			//
			fragmentConsts = Vector.<Number>([
				0, 10 / texSize, 1 / texSize, 0.5, // 0, vel scale, pix unit, 0.5
				2.5, 0.5, Math.pow(5 / texSize, -2), 1, // 2.5, pressure coefficient, inv square drag radius, 1
				0, 0, 0, 0, // prev mouse x, prev mouse y, mouse x, mouse y
				0.25, 0.995, 0.92, 0, // 0.25, viscosity, warm starting coefficient, mouse dragging coefficient
				4, 0, 0, 0 // velocity display coefficient, UNUSED, UNUSED, UNUSED
			]);
			//
			pmouseX = mouseX;
			pmouseY = mouseY;
			//
			var menu:ContextMenuItem = new ContextMenuItem("Reset");
			menu.addEventListener(ContextMenuEvent.MENU_ITEM_SELECT, function():void {
				reset();
			});
			contextMenu = new ContextMenu();
			contextMenu.hideBuiltInItems();
			contextMenu.customItems = [menu];
			//
			addEventListener(Event.ENTER_FRAME, frame);
		}
		
		private function createPG(agal:AGALMiniAssembler, vtx:String, frg:String):Program3D {
			agal.assemble(Context3DProgramType.VERTEX, vtx, 2);
			var codeV:ByteArray = agal.agalcode;
			agal.assemble(Context3DProgramType.FRAGMENT, frg, 2);
			var codeF:ByteArray = agal.agalcode;
			var pg:Program3D = c3d.createProgram();
			pg.upload(codeV, codeF);
			return pg;
		}
		
		private function frame(e:Event = null): void {
			fragmentConsts[8] = pmouseX / texSize;
			fragmentConsts[9] = pmouseY / texSize;
			fragmentConsts[10] = mouseX / texSize;
			fragmentConsts[11] = mouseY / texSize;
			fragmentConsts[15] = pressed ? 6 : 0;
			c3d.setProgramConstantsFromVector(Context3DProgramType.FRAGMENT, 0, fragmentConsts);
			calcDivergence();
			calcPressure();
			applyPressure();
			advect();
			display();
			//
			pmouseX = mouseX;
			pmouseY = mouseY;
			//
		}
		
		private function reset():void {
			var bmd:BitmapData = new BitmapData(texSize, texSize, false, 0xff101010);
			for (var i:int = 0; i < texSize >> 5; i++) {
				for (var j:int = 0; j < texSize >> 5; j++) {
					var bright:int = i + j & 1; // 市松模様の色テクスチャで初期化する
					bmd.fillRect(new Rectangle(i * 32, j * 32, 32, 32), 0xff000000 | (
						int((0.05 + Math.random() * 0.4 + bright * 0.5) * 0xff + 0.5) << 16 |
						int((0.05 + Math.random() * 0.4 + bright * 0.5) * 0xff + 0.5) << 8 |
						int((0.05 + Math.random() * 0.4 + bright * 0.5) * 0xff + 0.5)
					)
					);
				}
			}
			//
			color1.uploadFromBitmapData(bmd);
			color2.uploadFromBitmapData(bmd);
			bmd = new BitmapData(texSize, texSize, true, 0);
			//
			velocity1 = c3d.createTexture(texSize, texSize, Context3DTextureFormat.RGBA_HALF_FLOAT, true);
			velocity1.uploadFromBitmapData(bmd);
			velocity2 = c3d.createTexture(texSize, texSize, Context3DTextureFormat.RGBA_HALF_FLOAT, true);
			velocity2.uploadFromBitmapData(bmd);
		}
		
		private function flipColorTexture():void {
			var tmp:Texture = color2;
			color2 = color1;
			color1 = tmp;
		}
		
		private function flipVelocityTexture():void {
			var tmp:Texture = velocity2;
			velocity2 = velocity1;
			velocity1 = tmp;
		}
		
		private function calcDivergence():void {
			// 発散を計算する
			c3d.setProgram(calcDivergencePG);
			
			c3d.setRenderToTexture(velocity1); // dst
			c3d.clear();
			c3d.setTextureAt(0, velocity2); // src
			c3d.drawTriangles(idxB);
			flipVelocityTexture();
		}
		
		private function calcPressure():void {
			// 圧力を計算する
			c3d.setProgram(calcPressurePG);
			
			// 反復法でポアソン方程式を解く
			for (var i:int = 0; i < 4; i++) {
				c3d.setRenderToTexture(velocity1); // dst
				c3d.clear();
				c3d.setTextureAt(0, velocity2); // src
				c3d.drawTriangles(idxB);
				flipVelocityTexture();
			}
		}
		
		private function applyPressure():void {
			// 圧力を適用する
			c3d.setProgram(applyPressurePG);
			
			c3d.setRenderToTexture(velocity1); // dst
			c3d.clear();
			c3d.setTextureAt(0, velocity2); // src
			c3d.drawTriangles(idxB);
			flipVelocityTexture();
		}
		
		private function advect():void {
			// 色と速度を移流させる
			
			// advect color
			c3d.setProgram(advectColorPG);
			
			c3d.setRenderToTexture(color1);
			c3d.clear();
			c3d.setTextureAt(0, color2);
			c3d.setTextureAt(1, velocity2);
			c3d.drawTriangles(idxB);
			c3d.setTextureAt(1, null);
			flipColorTexture();
			
			// advect velocity
			c3d.setProgram(advectVelocityPG);
			
			c3d.setRenderToTexture(velocity1);
			c3d.clear();
			c3d.setTextureAt(0, velocity2);
			c3d.drawTriangles(idxB);
			c3d.setTextureAt(1, null);
			flipVelocityTexture();
		}
		
		private function display():void {
			c3d.setRenderToBackBuffer();
			c3d.clear();
			switch (displayType) {
			case 0:
				c3d.setProgram(displayColorPG);
				c3d.setTextureAt(0, color2);
				break;
			case 1:
				c3d.setProgram(displayVelocityPG);
				c3d.setTextureAt(0, velocity2);
				break;
			}
			c3d.drawTriangles(idxB);
			c3d.present();
		}
	}
}

var displayColorShader:String = <![CDATA[
tex ft0, v0, fs0, <2d, nearest, repeat>; // color
mov oc, ft0;
]]>;

var displayVelocityShader:String = <![CDATA[
tex ft0, v0, fs0, <2d, nearest, repeat>; // velocity
mul ft0.xyz, ft0.xyw, fc4.xxx; // vel *= coefficient
add ft0, ft0, fc0.www; // 0-centered -> 0.5-centered
mov oc, ft0;
]]>;

var advectColorShader:String = <![CDATA[
tex ft1, v0, fs1, <2d, linear, repeat>; // velocity
mul ft1.xy, ft1.xy, fc0.yy; // unit(vel) = [pix/frame]
sub ft0.xy, v0.xy, ft1.xy; // uv -= vel
tex ft3, ft0.xy, fs0, <2d, linear, repeat>; // advected color
mov oc, ft3; // output color
]]>;

var advectVelocityShader:String = <![CDATA[
tex ft1, v0, fs0, <2d, linear, repeat>; // velocity
mul ft1.xy, ft1.xy, fc0.yy; // unit(vel) = [pix/frame]
sub ft0.xy, v0.xy, ft1.xy; // uv -= vel
tex ft4, ft0.xy, fs0, <2d, linear, repeat>; // advected velocity
mov oc, ft4; // output velocity
]]>;

var calcDivergenceShader:String = <![CDATA[
sub ft1, v0.xy, fc0.zx; // u-1 v
add ft2, v0.xy, fc0.zx; // u+1 v
sub ft3, v0.xy, fc0.xz; // u v-1
add ft4, v0.xy, fc0.xz; // u v+1
tex ft0, v0, fs0, <2d, linear, repeat>; // center velocity
tex ft1, ft1, fs0, <2d, linear, repeat>; // left velocity
tex ft2, ft2, fs0, <2d, linear, repeat>; // right velocity
tex ft3, ft3, fs0, <2d, linear, repeat>; // top velocity
tex ft4, ft4, fs0, <2d, linear, repeat>; // bottom velocity

sub ft5.x, ft1.x, ft2.x; // divergence = (l.vx - r.vx) + (t.vy - b.vy)
sub ft5.y, ft3.y, ft4.y;
add ft0.z, ft5.x, ft5.y;

// 前回の値をそのまま用いるとやや不安定になるので適当に減衰させた値を用いる
mul ft0.w, ft0.w, fc3.z; // warm starting
mov oc, ft0;
]]>;

var calcPressureShader:String = <![CDATA[
sub ft1, v0.xy, fc0.zx; // u-1 v
add ft2, v0.xy, fc0.zx; // u+1 v
sub ft3, v0.xy, fc0.xz; // u v-1
add ft4, v0.xy, fc0.xz; // u v+1
tex ft0, v0, fs0, <2d, linear, repeat>; // center velocity
tex ft1, ft1, fs0, <2d, linear, repeat>; // left velocity
tex ft2, ft2, fs0, <2d, linear, repeat>; // right velocity
tex ft3, ft3, fs0, <2d, linear, repeat>; // top velocity
tex ft4, ft4, fs0, <2d, linear, repeat>; // bottom velocity
add ft5.x, ft1.w, ft2.w; // pressure = (l.pls + r.pls + t.pls + b.pls + divergence) * 0.25
add ft5.y, ft3.w, ft4.w;
add ft5.z, ft5.x, ft5.y;
add ft0.w, ft5.z, ft0.z;
mul ft0.w, ft0.w, fc3.x;
mov oc, ft0;
]]>;

var applyPressureShader:String = <![CDATA[
sub ft1, v0.xy, fc0.zx; // u-1 v
add ft2, v0.xy, fc0.zx; // u+1 v
sub ft3, v0.xy, fc0.xz; // u v-1
add ft4, v0.xy, fc0.xz; // u v+1
tex ft0, v0, fs0, <2d, linear, repeat>; // center velocity
tex ft1, ft1, fs0, <2d, linear, repeat>; // left velocity
tex ft2, ft2, fs0, <2d, linear, repeat>; // right velocity
tex ft3, ft3, fs0, <2d, linear, repeat>; // top velocity
tex ft4, ft4, fs0, <2d, linear, repeat>; // bottom velocity
sub ft5.x, ft1.w, ft2.w; // force.x = (l.pls - r.pls) * 0.5
sub ft5.y, ft3.w, ft4.w; // force.y = (t.pls - b.pls) * 0.5
mul ft5.xy, ft5.xy, fc1.yy;

// apply dragging force
mov ft1.xy, fc2.zw;
sub ft1.xy, ft1.xy, fc2.xy; // ft1.xy = vector(pmouse -> mouse)
dp3 ft1.z, ft1.xy, ft1.xy;  // ft1.z = |ft1.xy|^2 
ine ft1.z, fc0.x;           // if (ft1.z != 0) {
sub ft2.xy, v0.xy, fc2.xy;  //   ft2.xy = vector(pmouse -> pixel)
dp3 ft2.z, ft1.xy, ft2.xy;  //   ft2.z = ft1.xy dot ft2.xy ("dp3" opcode can be used for 2D vector too)
div ft2.z, ft2.z, ft1.z;    //   ft2.z /= ft1.z
sat ft2.z, ft2.z;           //   ft2.z = clamp(ft2.z, 0, 1)
mul ft3.xy, ft1.xy, ft2.zz; //   ft3.xy = ft1.xy * ft2.z
add ft2.xy, fc2.xy, ft3.xy; //   ft2.xy = fc2.xy + ft3.xy (ft2.xy is the closest point)
mul ft1.xy, ft1.xy, fc3.ww; //   ft1.xy *= dragging coefficient (ft1.xy is the dragging force)
sub ft6.xy, v0.xy, ft2.xy;  //   calculate distance from the closest point...
dp3 ft6.z, ft6.xy, ft6.xy;  //   
mul ft6.z, ft6.z, fc1.z;    //   
sub ft6.z, fc1.w, ft6.z;    //   
sat ft6.z, ft6.z;           //   
mul ft1.xy, ft1.xy, ft6.zz; //   adjust dragging force...
add ft5.xy, ft5.xy, ft1.xy; //   then add to total force
eif;                        // }

add ft0.xy, ft0.xy, ft5.xy; // apply total force
mul ft0.xy, ft0.xy, fc3.yy; // apply viscosity

mov oc, ft0;
]]>;