//
// Simple ambient occlusion renderer with ActionScript3
//
// Compile: $ mxmlc AmbientOcclusion.as
//
package {
    import flash.display.*;
    import flash.text.*;
    import flash.utils.getTimer;

    public class AmbientOcclusion extends Sprite
    {
        private var bitmap:Bitmap;
        private var bitmapData:BitmapData;

        public static const NAO_SAMPLES:int = 8;

        private var spheres:Array;
        private var plane:Plane;

        public function initScene():void
        {
            spheres    = new Array(3);
            spheres[0] = new Sphere(new Vec3(-2.0, 0.0, -3.5), 0.5);
            spheres[1] = new Sphere(new Vec3(-0.5, 0.0, -3.0), 0.5);
            spheres[2] = new Sphere(new Vec3(1.0, 0.0, -2.2), 0.5);
            plane  = new Plane(new Vec3(0.0, -0.5, 0.0), new Vec3(0.0, 1.0, 0.0));
        }


        public function mkColor(r:Number, g:Number, b:Number):uint
        {
            var ri:uint = uint(r * 255.5);
            if (ri < 0) ri = 0;
            if (ri > 255) ri = 255;

            var gi:uint = uint(g * 255.5);
            if (gi < 0)   gi = 0;
            if (gi > 255) gi = 255;

            var bi:uint = uint(b * 255.5);
            if (bi < 0)   bi = 0;
            if (bi > 255) bi = 255;

            return (ri << 16) | (gi << 8) | bi;

        }

        public function orthoBasis(basis:Array, n:Vec3):void
        {
            basis[2] = new Vec3(n.x, n.y, n.z);
            basis[1] = new Vec3(0.0, 0.0, 0.0);

            if ((n.x < 0.6) && (n.x > -0.6)) {
              basis[1].x = 1.0;
            } else if ((n.y < 0.6) && (n.y > -0.6)) {
              basis[1].y = 1.0;
            } else if ((n.z < 0.6) && (n.z > -0.6)) {
              basis[1].z = 1.0;
            } else {
              basis[1].x = 1.0;
            }

            basis[0] = basis[1].cross(basis[2]);
            basis[0].normalize();

            basis[1] = basis[2].cross(basis[0]);
            basis[1].normalize();

        }

        public function computeAO(isect:Intersection):Vec3
        {
            var i:int, j:int;

            var ntheta:int = NAO_SAMPLES;
            var nphi:int   = NAO_SAMPLES;
            var eps:Number = 0.0001;

            // Slightly move ray org towards ray dir to avoid numerical probrem.
            var p:Vec3 = new Vec3(isect.p.x + eps * isect.n.x,
                             isect.p.y + eps * isect.n.y,
                             isect.p.z + eps * isect.n.z);

            // Calculate orthogonal basis.
            var basis:Array; basis = new Array(3);
            orthoBasis(basis, isect.n);

            var occlusion:Number = 0.0;


            for (j = 0; j < ntheta; j++) {
                for (i = 0; i < nphi; i++) {

                  // Pick a random ray direction with importance sampling.
                  // p = cos(theta) / PI
                  var r:Number     = Math.random();
                  var phi:Number   = 2.0 * Math.PI * Math.random();

                  var x:Number = Math.cos(phi) * Math.sqrt(1.0 - r);
                  var y:Number = Math.sin(phi) * Math.sqrt(1.0 - r);
                  var z:Number = Math.sqrt(r);

                  // local -> global
                  var rx:Number = x * basis[0].x + y * basis[1].x + z * basis[2].x;
                  var ry:Number = x * basis[0].y + y * basis[1].y + z * basis[2].y;
                  var rz:Number = x * basis[0].z + y * basis[1].z + z * basis[2].z;

                  var raydir:Vec3 = new Vec3(rx, ry, rz);

                  var ray:Ray = new Ray(p, raydir);

                  var occIsect:Intersection = new Intersection();
                  spheres[0].intersect(occIsect, ray);
                  spheres[1].intersect(occIsect, ray);
                  spheres[2].intersect(occIsect, ray);
                  plane.intersect(occIsect, ray);

                  if (occIsect.hit) occlusion += 1.0;

                }
            }

            // [0.0, 1.0]
            occlusion = (ntheta * nphi - occlusion) / (ntheta * nphi);

            return new Vec3(occlusion, occlusion, occlusion);

        }



        public function AmbientOcclusion()
        {

            initScene();

            var w:int = 256;
            var h:int = 256;
            var nsubsamples:int = 2;

            bitmapData = new BitmapData(w, h);
            bitmap = new Bitmap(bitmapData);
            addChild(bitmap);

            var startTime:uint, endTime:uint;
            startTime = getTimer();

            for (var i:int = 0; i < w; i++) {
                for (var j:int = 0; j < h; j++) {

                    var fcol:Vec3 = new Vec3(0.0, 0.0, 0.0);

                    // subsampling
                    for (var v:int = 0; v < nsubsamples; v++) {
                        for (var u:int = 0; u < nsubsamples; u++) {
                            var px:Number = (i + (u / Number(nsubsamples)) - (w / 2.0))/(w / 2.0);
                            var py:Number = (j + (v / Number(nsubsamples)) - (h / 2.0))/(h / 2.0);
                            py = -py;   // flip Y

                            trace(py);

                            var t:Number = 10000.0;
                            var eye:Vec3 = new Vec3(px, py, -1.0);
                            eye.normalize();

                            var ray:Ray = new Ray(new Vec3(0.0, 0.0, 0.0), new Vec3(eye.x, eye.y, eye.z));

                            var isect:Intersection = new Intersection();
                            spheres[0].intersect(isect, ray);
                            spheres[1].intersect(isect, ray);
                            spheres[2].intersect(isect, ray);
                            plane.intersect(isect, ray);

                            if (isect.hit) {

                              var col:Vec3 = computeAO(isect);

                              fcol.x += col.x;
                              fcol.y += col.y;
                              fcol.z += col.z;
                            }
                        }
                    }

                    fcol.x /= Number(nsubsamples * nsubsamples);
                    fcol.y /= Number(nsubsamples * nsubsamples);
                    fcol.z /= Number(nsubsamples * nsubsamples);

                    bitmapData.setPixel(i, j, mkColor(fcol.x, fcol.y, fcol.z));
                    //bitmapData.setPixel(i, j, mkColor(eye.x, eye.y, eye.z));
                }
            }

            endTime = getTimer();

            var elapsed:uint = (endTime - startTime) / 1000; // msec -> sec

            // display render time
            var text:TextField = new TextField();
            text.text= elapsed + " sec";
            text.x = 10;
            text.y = 10;
            text.textColor = 0xFFFFFF;
            addChild(text);
        }
    }
}

//
// -- Common classes for a global illumination renderer.
//
class Vec3
{
    public var x:Number;
    public var y:Number;
    public var z:Number;

    public function Vec3(x:Number = 0.0, y:Number = 0.0, z:Number = 0.0)
    {
        this.x = x;
        this.y = y;
        this.z = z;
    }

    public function length():Number
    {
        return Math.sqrt(this.x * this.x + this.y * this.y + this.z * this.z);
    }

    public function normalize():void
    {
        var dist:Number = this.length();

        var invdist:Number = 1.0
        if (Math.abs(dist) > 1.0e-6) {
            invdist = 1.0 / dist;
        }

        this.x *= invdist;
        this.y *= invdist;
        this.z *= invdist;
    }

    public function add(b:Vec3):Vec3
    {
        return new Vec3(this.x + b.x, this.y + b.y, this.z + b.z);
    }

    public function sub(b:Vec3):Vec3
    {
        return new Vec3(this.x - b.x, this.y - b.y, this.z - b.z);
    }

    public function cross(b:Vec3):Vec3
    {
        var s:Number = this.y * b.z - this.z * b.y;
        var t:Number = this.z * b.x - this.x * b.z;
        var q:Number = this.x * b.y - this.y * b.x;

        return new Vec3(s, t, q);
    }

    public function dot(b:Vec3):Number
    {
        return this.x * b.x + this.y * b.y + this.z * b.z;
    }

}

class Ray
{
    public var org:Vec3;
    public var dir:Vec3;

    public function Ray(o:Vec3, d:Vec3)
    {
        this.org = o;
        this.dir = d;
    }
}

class Intersection
{
    public var t:Number;
    public var p:Vec3;     // hit point
    public var n:Vec3;     // normal
    public var hit:Boolean;

    public function Intersection() {
      hit = false;
      t   = 1.0e+30;
      n   = new Vec3(0.0, 0.0, 0.0);
    }
}

class Sphere
{
    public var center:Vec3;
    public var radius:Number;

    public function Sphere(center:Vec3, radius:Number) {
        this.center = center;
        this.radius = radius;
    }

    public function intersect(isect:Intersection, ray:Ray):void
    {
        // rs = ray.org - sphere.center
        var rs:Vec3  = ray.org.sub(this.center);
        var B:Number = rs.dot(ray.dir);
        var C:Number = rs.dot(rs) - (this.radius * this.radius);
        var D:Number = B * B - C;

        if (D > 0.0) {
            var t:Number = -B - Math.sqrt(D);
            if ( (t > 0.0) && (t < isect.t) ) {
                isect.t = t;
                isect.hit = true;

                // calculate normal.
                var p:Vec3 = new Vec3(ray.org.x + ray.dir.x * t,
                                      ray.org.y + ray.dir.y * t,
                                      ray.org.z + ray.dir.z * t);
                var n:Vec3 = p.sub(center);
                n.normalize();
                isect.n = n;
                isect.p = p;
            }
        }
    }
}

class Plane
{
    public var p:Vec3;  // point on the plane
    public var n:Vec3;  // normal to the plane

    public function Plane(p:Vec3, n:Vec3) {
        this.p = p;
        this.n = n;
    }

    public function intersect(isect:Intersection, ray:Ray):void
    {
        // d = -(p . n)
        // t = -(ray.org . n + d) / (ray.dir . n)
        var d:Number = -p.dot(n);
        var v:Number = ray.dir.dot(n);

        if (Math.abs(v) < 1.0e-6) return; // the plane is parallel to the ray.

        var t:Number = -(ray.org.dot(n) + d) / v;

        if ( (t > 0) && (t < isect.t) ) {
          isect.hit = true;
          isect.t   = t;
          isect.n   = n;

          var p:Vec3 = new Vec3(ray.org.x + t * ray.dir.x,
                                ray.org.y + t * ray.dir.y,
                                ray.org.z + t * ray.dir.z);
          isect.p   = p;
        }
    }
}