package {
    import com.derschmale.fluids.AddForcesShader;
    import com.derschmale.fluids.AdvectionKernel;
    import com.derschmale.fluids.ApplyProjectionShader;
    import com.derschmale.fluids.CalculateDivergenceShader;
    import com.derschmale.fluids.DiffuseShader;
    import com.derschmale.fluids.ProjectionShader;
    import com.derschmale.fluids.RenderBitmapDataShader;
    
    import flash.display.Bitmap;
    import flash.display.BitmapData;
    import flash.display.Sprite;
    import flash.display.StageAlign;
    import flash.display.StageQuality;
    import flash.display.StageScaleMode;
    import flash.events.Event;
    import flash.events.KeyboardEvent;
    import flash.events.MouseEvent;
    import flash.geom.Vector3D;
    import flash.text.TextField;
    import flash.ui.Keyboard;
    
    import net.hires.debug.Stats;
    
    /**
     * 3D Smoke simulation
     * 
     * @author David Lenaerts
     * http://www.derschmale.com
     */
    
    [SWF(width="400", height="400", frameRate="30", backgroundColor="0x000000")]
    public class Fluids3DPB extends Sprite
    {
        /**
         * grid dimensions
         */
        private static const GRID_X : int = 18;
        private static const GRID_Y : int = 32;
        private static const GRID_Z : int = 12;
        
        private static const GRID_SPACING_X : int = 10;
        private static const GRID_SPACING_Y : int = 10;
        private static const GRID_SPACING_Z : int = 13;
        
        private const TOTAL_SIZE : int = GRID_X*GRID_Y*GRID_Z;
        private const BUFFER_SIZE_3 : int = TOTAL_SIZE * 3;
        private const BUFFER_SIZE_4 : int = TOTAL_SIZE << 2;
        
        /**
         * grid dependents
         */
        public static const GRID_X_3 : int = GRID_X*3;
        public static const GRID_X_4 : int = GRID_X << 2;
        
        public static const GRID_X_1 : int = GRID_X-1;
        public static const GRID_Y_1 : int = GRID_Y-1;
        public static const GRID_Z_1 : int = GRID_Z-1;
        
        public static const TOTAL_W : int = GRID_X*GRID_Z;
        public static const TOTAL_W_3 : int = TOTAL_W*3;
        public static const TOTAL_W_4 : int = TOTAL_W << 2;
        
        /**
         * Properties of the fluid solver
         */
        public static const DT : Number = 2;
        public static const BUOYANCY : Number = 0.2;
        public static const GRAVITY : Number = 0.005;
        public static const KINEMATIC_VISCOSITY : Number = .001;
        public static const DENSITY_VISCOSITY : Number = .001;
        public static const DIFFUSION_ITERATIONS : int = 3;
        public static const PROJECTION_ITERATIONS : int = 7;
        
        /**
         * Vector and scalar fields
         */
        private var _velocityDensityField : Vector.<Number>;
        private var _divergencePressureField : Vector.<Number>;
        
        /**
         * Shaders used in the fluid solver
         */
        private var _addForces : AddForcesShader;
        private var _diffusion : DiffuseShader;
        private var _calculateDivergence : CalculateDivergenceShader;
        private var _projection : ProjectionShader;
        private var _applyProjection : ApplyProjectionShader;
        private var _advection : AdvectionKernel;
        private var _renderBitmapData : RenderBitmapDataShader;
        
        private var _startX : Number;
        private var _startY : Number;
        
        /**
         * Display
         */
        private var _container : Sprite;
        
        // axis-aligned to Y/Z arranged into X
        private var _bitmapDatasX : Vector.<BitmapData>;
        private var _bitmapsX : Vector.<Bitmap>;
        
        // axis-aligned to X/Z arranged into Y
        private var _bitmapDatasY : Vector.<BitmapData>;
        private var _bitmapsY : Vector.<Bitmap>;
        
        // axis-aligned to X/Y arranged into Z
        private var _bitmapDatasZ : Vector.<BitmapData>;
        private var _bitmapsZ : Vector.<Bitmap>;
        
        private var _activeBitmapDatas : Vector.<BitmapData>;
        private var _activeBitmaps : Vector.<Bitmap>;
        
        [Embed(source="bg.jpg")]
        private var BackgroundAsset : Class;
        
        public function Fluids3DPB()
        {
            stage.quality = StageQuality.MEDIUM;
            stage.scaleMode = StageScaleMode.NO_SCALE;
            stage.align = StageAlign.TOP_LEFT;
            
            init();
            initViewSlices();
            initShaders();
            
            // init listeners
            addEventListener(Event.ENTER_FRAME, onEnterFrame);
            stage.addEventListener(MouseEvent.MOUSE_MOVE, onMouseMove);
            stage.addEventListener(MouseEvent.MOUSE_DOWN, onMouseDown);
            stage.addEventListener(KeyboardEvent.KEY_DOWN, onKeyDown);
        }
        
        /**
         * Reset the velocity field
         */
        private function onKeyDown(event : KeyboardEvent) : void
        {
            switch (event.keyCode) {
                case Keyboard.SPACE:
                    for (var i : int = 0; i < BUFFER_SIZE_4; ++i)
                        _velocityDensityField[i] = 0;
                    break;
                case Keyboard.ENTER:
                    _container.rotationX = _container.rotationY = _container.rotationZ;
                    break;
            }
        }
        
        /**
         * Set rotation reference coordinates
         */
        private function onMouseDown(event : MouseEvent) : void
        {
            _startX = mouseX;
            _startY = mouseY;
        }
        
        /**
         * Rotate relative to reference coords and update ref coords
         */
        private function onMouseMove(event : MouseEvent) : void
        {
            if (event.buttonDown) {
                _container.rotationY += (_startX-mouseX);
                _container.rotationX += (_startY-mouseY);
                _startX = mouseX;
                _startY = mouseY;
            }
        }
        
        
        private const X_BMP_POS : Number = -GRID_X*GRID_SPACING_X*.5;
        private const Y_BMP_POS : Number = -GRID_Y*GRID_SPACING_Y*.5;
        private const Z_BMP_POS : Number = -GRID_Z*GRID_SPACING_Z*.5;
        private const GRID_X_HALF : Number = (GRID_X-1)*.5;
        private const GRID_Y_HALF : Number = (GRID_Y-1)*.5;
        private const GRID_Z_HALF : Number = (GRID_Z-1)*.5;
        
        /**
         * Create axis aligned slices.
         */
        private function initViewSlices() : void
        {
            var i : uint;
            var bmp : Bitmap;
            _bitmapDatasX = new Vector.<BitmapData>(GRID_X);
            _bitmapDatasY = new Vector.<BitmapData>(GRID_Y);
            _bitmapDatasZ = new Vector.<BitmapData>(GRID_Z);
            _bitmapsX = new Vector.<Bitmap>(GRID_X);
            _bitmapsY = new Vector.<Bitmap>(GRID_Y);
            _bitmapsZ = new Vector.<Bitmap>(GRID_Z);
            _container = new Sprite();
            
            do {
                _bitmapDatasX[i] = new BitmapData(GRID_Z, GRID_Y, true, 0xffffffff);
                bmp = new Bitmap(_bitmapDatasX[i]);
                bmp.smoothing = true;
                bmp.scaleX = GRID_SPACING_Z;
                bmp.scaleY = GRID_SPACING_Y;
                bmp.x = GRID_SPACING_X*(i-GRID_X_HALF);
                bmp.y = Y_BMP_POS;
                bmp.z = Z_BMP_POS;
                bmp.rotationY = -90;
                _bitmapsX[i] = bmp;
            } while(++i < GRID_X);
            
            i = 0;
            do {
                _bitmapDatasY[i] = new BitmapData(GRID_X, GRID_Z, true, 0xffffffff);
                bmp = new Bitmap(_bitmapDatasY[i]);
                bmp.smoothing = true;
                bmp.scaleX = GRID_SPACING_X;
                bmp.scaleY = GRID_SPACING_Z;
                bmp.x = X_BMP_POS;
                bmp.y = GRID_SPACING_Y*(i-GRID_Y_HALF);
                bmp.z = Z_BMP_POS;
                _bitmapsY[i] = bmp;
                bmp.rotationX = 90;
            } while(++i < GRID_Y);
            
            i = 0;
            do {
                _bitmapDatasZ[i] = new BitmapData(GRID_X, GRID_Y, true, 0xffffffff);
                bmp = new Bitmap(_bitmapDatasZ[i]);
                bmp.smoothing = true;
                bmp.scaleX = GRID_SPACING_X;
                bmp.scaleY = GRID_SPACING_Y;
                bmp.x = X_BMP_POS;
                bmp.y = Y_BMP_POS;
                bmp.z = GRID_SPACING_Z*(i-GRID_Z_HALF);
                _bitmapsZ[i] = bmp;
            } while(++i < GRID_Z);
            
            _container.x = stage.stageWidth*.5;
            _container.y = stage.stageHeight*.5;
            _container.z = 100;
            addChild(_container);
        }
        
        /**
         * Create fields and display
         */
        private function init() : void
        {
            var i : int;
            var text : TextField = new TextField();
            addChild(new BackgroundAsset());
            text.textColor = 0xffffff;
            text.multiline = true;
            text.text = "Move mouse: add smoke\n" + 
                        "Click & drag: rotate camera\n" + 
                        "Space: clear smoke\n" + 
                        "Enter: reset rotation";
            text.width = text.textWidth+10;
            text.height = text.textHeight+10;
            text.x = stage.stageWidth-text.width;
            addChild(text);
                        
            // contains 4 values: velocity(x,y,z) and density
            _velocityDensityField = new Vector.<Number>(BUFFER_SIZE_4);
            _divergencePressureField = new Vector.<Number>(BUFFER_SIZE_3);
            addChild(new Stats());
            
            for (i = 0; i < BUFFER_SIZE_4; ++i)
                _velocityDensityField[i] = 0;
            
            for (i = 0; i < BUFFER_SIZE_3; ++i)
                _divergencePressureField[i] = 0;
        }
        
        /**
         * initialize shaders and assign properties
         */
        private function initShaders() : void
        {
            _addForces = new AddForcesShader(_velocityDensityField, GRID_X, GRID_Y, GRID_Z);
            _addForces.dt = DT;
            _addForces.buoyancy = BUOYANCY;
            _addForces.setGlobalForce(0, GRAVITY, 0);
            _diffusion = new DiffuseShader(_velocityDensityField, GRID_X, GRID_Y, GRID_Z);
            _diffusion.dt = DT;
            _diffusion.kinematicViscosity = KINEMATIC_VISCOSITY;
            _diffusion.densityViscosity = DENSITY_VISCOSITY;
            _calculateDivergence = new CalculateDivergenceShader(_velocityDensityField, GRID_X, GRID_Y, GRID_Z);
            _projection = new ProjectionShader(_divergencePressureField, GRID_X, GRID_Y, GRID_Z);
            _applyProjection = new ApplyProjectionShader(_divergencePressureField, _velocityDensityField, GRID_X, GRID_Y, GRID_Z);
            _advection = new AdvectionKernel(_velocityDensityField, GRID_X, GRID_Y, GRID_Z);
            _advection.dt = DT;
            _renderBitmapData = new RenderBitmapDataShader(_velocityDensityField, GRID_X, GRID_Y, GRID_Z);
            _renderBitmapData.color = 0xddeeff;
        }
        
        /**
         * Update the simulation
         */
        private function onEnterFrame(event : Event) : void
        {
            var i : int;
            var x : Number = (0.5+_container.mouseX/_container.width)*GRID_X;
            var y : Number = (0.5+_container.mouseY/_container.height)*GRID_Y;
            var z : Number = GRID_Z*.5; //int(_bitmap.mouseX/GRID_X);
            var l : int = x-1;
            var r : int = x+1;
            var t : int = y-1;
            var b : int = y+1;
            var n : int = z-1;
            var f : int = z+1;
            
            // add some random global wind
            _addForces.setGlobalForce(Math.random()*0.02-0.01, GRAVITY, Math.random()*0.02-0.01);
            
            // add smoke:
            // going gung-ho on method calls - too lazy to fix
            // doesn't seem to impact performance tho
            addSmoke(x, y, z, 1);
            addSmoke(r, y, z, .5);
            addSmoke(x, t, z, .5);
            addSmoke(x, y, f, .5);
            addSmoke(x, y, n, .5);
            
            // do fluid solver stuff
            updateVelocityDensityBounds(_velocityDensityField);
            
            _addForces.execute();
            updateVelocityDensityBounds(_velocityDensityField);
            
            for (i = 0; i < DIFFUSION_ITERATIONS; ++i) {
                _diffusion.execute();
                updateVelocityDensityBounds(_velocityDensityField);
            }
            
            _calculateDivergence.execute(_divergencePressureField);
            updateScalarBounds3(_divergencePressureField, 0);
            for (i = 0; i < PROJECTION_ITERATIONS; ++i) {
                _projection.execute();
                updateScalarBounds3(_divergencePressureField, 1);
            }
            _applyProjection.execute();
            
            _advection.execute();
            updateVelocityDensityBounds(_velocityDensityField);
            
            _calculateDivergence.execute(_divergencePressureField);
            updateScalarBounds3(_divergencePressureField, 0);
            for (i = 0; i < PROJECTION_ITERATIONS; ++i) {
                _projection.execute();
                updateScalarBounds3(_divergencePressureField, 1);
            }
            _applyProjection.execute();
            
            // render
            updateView();
        }
        
        /**
         * precomputed constants, are constant and calculated from grid sizes
         */
        private static const TW_4 : int = (TOTAL_W-GRID_X) << 2;
        private static const FRONT_BACK_XBT_4 : int = TW_4-GRID_X_4;
        private static const FRONT_BACK_L_4 : int = ((TOTAL_W*GRID_Y) << 2) - TW_4;
        
        private static const TOP_BOTTOM_XB_4 : int = (TOTAL_W*(GRID_Y_1)) << 2;
        private static const TOP_BOTTOM_XBT_4 : int = TOP_BOTTOM_XB_4-TOTAL_W_4;
        
        private static const SIDE_XB_4 : int = (GRID_X_1) << 2;
        private static const SIDE_XBT_4 : int = SIDE_XB_4-4;
        private static const SIDE_BOUND_4 : int = TW_4 + (((GRID_Y_1)*TOTAL_W) << 2);

        /**
         * Set boundary conditions for the velocity and density fields
         * Aweful lot of looping
         */
        private function updateVelocityDensityBounds(field : Vector.<Number>) : void
        {
            var tw : int;
            var xt : int;
            var xb : int;
            var xtb : int;
            var xbt : int;
            var l : int;
            var bound : int; 
            
            xt = 0;
            tw = xb = TW_4;
            xtb = bound = GRID_X_4;
            l = FRONT_BACK_L_4;
            xbt = FRONT_BACK_XBT_4;
            
            // front and back border
            do {
                // x velocity
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
                // y velocity
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
                // z velocity
                field[xt++] = -field[xtb++];
                field[xb++] = -field[xbt++];
                
                // density
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
                
                if (xt >= bound) {
                    xt += tw;
                    xb += tw;
                    xtb += tw;
                    xbt += tw;
                    bound += TOTAL_W_4;
                }
            } while (xt < l);
            
            
            // top and bottom border over all depths
            xt = 0;
            xb = TOP_BOTTOM_XB_4;
            bound = xtb = TOTAL_W_4;
            xbt = TOP_BOTTOM_XBT_4;
            
            do {
                // x velocity
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
                // y velocity
                field[xt++] = -field[xtb++];
                field[xb++] = -field[xbt++];
                // z velocity
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
                
                // density
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
            } while (xt < bound);
            
            
            // side borders
            xt = 0;
            l = xb = SIDE_XB_4;
            xtb = 4;
            xbt = SIDE_XBT_4;
            bound = SIDE_BOUND_4;
            do {
                // x velocity
                field[xt++] = -field[xtb++];
                field[xb++] = -field[xbt++];
                // y velocity
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
                // z velocity
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
                
                // density
                field[xt++] = field[xtb++];
                field[xb++] = field[xbt++];
                
                xt += l;
                xb += l;
                xtb += l;
                xbt += l;
            } while (xt < bound);
        }
        
        /**
         * precomputed constants, are constant and calculated from grid sizes
         */
        private static const TW_3 : int = (TOTAL_W-GRID_X) * 3;
        private static const FRONT_BACK_L_3 : int = (TOTAL_W*GRID_Y)*3 - TW_3;
        
        private static const TOP_BOTTOM_XB_3 : int = GRID_X_3+TOTAL_W*(GRID_Y_1)* 3;
        
        private static const SIDE_XB_3 : int = (GRID_X_1)*3;
        private static const SIDE_XBT_3 : int = SIDE_XB_3-3;
        private static const SIDE_BOUND_3 : int = TW_3 + ((GRID_Y_1)*TOTAL_W)*3;

        /**
         * Set boundary conditions for the scalar fields (ie: divergence and pressure)
         * More looping.
         */
        private function updateScalarBounds3(field : Vector.<Number>, offset : int) : void
        {
            var tw : int = TW_3;
            var xt : int;
            var xb : int;
            var xtb : int;
            var xbt : int;
            var l : int;
            var bound : int; 
            
            // front and back border
            
            xbt = int((xb = int(TW_3 + (xt = offset))) - GRID_X_3);
            xtb = bound = int(GRID_X_3 + offset);
            l = int(FRONT_BACK_L_3 + offset);
            
            do {
                field[xt] = field[xtb];
                field[xb] = field[xbt];
                
                xtb += 3;
                xbt += 3;
                
                if ((xb += 3) >= bound) {
                    xt += tw;
                    xb += tw;
                    xtb += tw;
                    xbt += tw;
                    bound += TOTAL_W_3;
                }
            } while ((xt += 3) < l);
        
            // side borders
            xtb = int(3+ (xt = offset));
            xbt = int((xb = int(SIDE_XB_3 + offset)) - 3);
            bound = int(SIDE_BOUND_3 + offset);
            l = GRID_X_3;
            do {
                field[xt] = field[xtb];
                field[xb] = field[xbt];
                
                xb += l;
                xtb += l;
                xbt += l;
            } while ((xt += l) < bound);
            
            xtb = int(TOTAL_W_3 + (xt = int(GRID_X_3 + offset)) );
            xbt = int((xb = int(TOP_BOTTOM_XB_3 + offset)) - TOTAL_W_3);
            bound = int(tw + offset);
            
            // top and bottom border over all depths
            do {
                field[xt] = 0;
                field[xb] = 0;
                
                xb += 3;
                xtb += 3;
                xbt += 3;
            } while ((xt += 3) < bound);
        }
        
        private static const NODE_BRN : int = TOTAL_W_4+4;
        private static const NODE_TRF : int = GRID_X_4+4;
        private static const NODE_BLF : int = GRID_X_4+TOTAL_W_4;
        private static const NODE_BRF : int = GRID_X_4+TOTAL_W_4+4;
        
        public function addSmoke(x : Number, y : Number, z : Number, amount : Number) : void
        {
            var xi : int, yi : int, zi : int;
            var xr : Number, yr : Number, zr : Number;
            var index : int;
            var qtln : Number;
            var qtrn : Number;
            var qbln : Number;
            var qbrn : Number;
            var qtlf : Number;
            var qtrf : Number;
            var qblf : Number;
            var qbrf : Number;
            
            if (x >= 1 && x < GRID_X_1 && y >= 1 && y < GRID_Y_1 && z >= 1 && z < GRID_Z_1) {
                xi = int(x);
                yi = int(y);
                zi = int(z);
                xr = x-xi;
                yr = y-yi;
                zr = z-zi;
                qtlf = zr*(1.0-yr)*(1.0-xr)*amount;
                qtrf = zr*(1.0-yr)*xr*amount;
                qblf = zr*yr*(1.0-xr)*amount;
                qbrf = zr*yr*xr*amount;
                qtln = (1.0-zr)*(1.0-yr)*(1.0-xr)*amount;
                qtrn = (1.0-zr)*(1.0-yr)*xr*amount;
                qbln = (1.0-zr)*yr*(1.0-xr)*amount;
                qbrn = (1.0-zr)*yr*xr*amount;
                
                index = int(((xi+yi*TOTAL_W+zi*GRID_X)<<2)+3);
                
                _velocityDensityField[index] += qtln;
                _velocityDensityField[index+4] += qtrn;
                _velocityDensityField[index+TOTAL_W_4] += qbln;
                _velocityDensityField[index+NODE_BRN] += qbrn;
                
                _velocityDensityField[index+GRID_X_4] += qtlf;
                _velocityDensityField[index+NODE_TRF] += qtrf;
                _velocityDensityField[index+NODE_BLF] += qblf;
                _velocityDensityField[index+NODE_BRF] += qbrf;
            }
        }
        
        private var _chosenGridDim : int;
        private var _offsetStep : int;
        
        private const AXIS_X : int = 0;
        private const AXIS_Y : int = 1;
        private const AXIS_Z : int = 2;
        
        private var _axis : int = -1;
        
        private const VEC : Vector3D = new Vector3D(0, 0, 1); 
        private const VEC_X_AXIS : Vector3D = new Vector3D(1, 0, 0);
        private const VEC_Y_AXIS : Vector3D = new Vector3D(0, 1, 0);
        private const VEC_Z_AXIS : Vector3D = new Vector3D(0, 0, 1);
        
        private const HALF_PI : Number = Math.PI * .5;
        private const PI : Number = Math.PI;
        
        /**
         * Renders the grid to the axis-aligned slices
         */
        private function updateView() : void
        {
            var offset : int;
            var i : int;
            var closestAxis : int;
            var changeAxis : Boolean;
            var l : int;
            
            // determine axis closest to the view vector 
            var v : Vector3D = _container.transform.matrix3D.deltaTransformVector(VEC);
            var angleX : Number = Math.acos(v.dotProduct(VEC_X_AXIS));
            var angleY : Number = Math.acos(v.dotProduct(VEC_Y_AXIS));
            var angleZ : Number = Math.acos(v.dotProduct(VEC_Z_AXIS));
            
            // correct to ignore direction of axis vectors
            if (angleX > HALF_PI) angleX = PI-angleX;
            if (angleY > HALF_PI) angleY = PI-angleY;
            if (angleZ > HALF_PI) angleZ = PI-angleZ;
            
            if (angleX < angleY && angleX < angleZ)
                closestAxis = AXIS_X;
            else if (angleY < angleX && angleY < angleZ)
                closestAxis = AXIS_Y;
            else if (angleZ < angleY && angleZ < angleX)
                closestAxis = AXIS_Z;
            
            // if axis is different, update settings for rendering
            if (closestAxis != _axis) {
                if (_axis != -1) {
                    l = _activeBitmaps.length;
                    
                    // remove all slices from the old axis
                    for (i = 0; i < l; i++)
                        _container.removeChild(_activeBitmaps[i]);
                }
                
                changeAxis = true;
                _axis = closestAxis;
                
                if (closestAxis == AXIS_X) {
                    _chosenGridDim = GRID_X;
                    _offsetStep = 1;
                    _activeBitmaps = _bitmapsX;
                    _activeBitmapDatas = _bitmapDatasX;
                }
                else if (closestAxis == AXIS_Y) {
                    _chosenGridDim = GRID_Y;
                    _offsetStep = 1;
                    _activeBitmaps = _bitmapsY;
                    _activeBitmapDatas = _bitmapDatasY;
                }
                else if (closestAxis == AXIS_Z) {
                    _chosenGridDim = GRID_Z;
                    _offsetStep = GRID_X;
                    _activeBitmaps = _bitmapsZ;
                    _activeBitmapDatas = _bitmapDatasZ;
                }
                _renderBitmapData.axis = closestAxis;
            }
            
            // render all grid slices
            i = 0;
            do { 
                _renderBitmapData.execute(_activeBitmapDatas[i], offset);
                offset += _offsetStep;
                
                // if axis changed, add children of current axis
                if (changeAxis)
                    _container.addChild(_activeBitmaps[i]);
            } while (++i < _chosenGridDim);
        }
    }
}