#include "cutil_math.h"

#define tex3Dv(t,p) tex3D(t, p.x, p.y, p.z)

// a bit of matrix algebra
typedef struct {
    float4 m[3];
} float3x4;
__device__ float3 operator*(const float3x4 &M, const float3 &v) { return make_float3(dot(v, make_float3(M.m[0])), dot(v, make_float3(M.m[1])), dot(v, make_float3(M.m[2]))); }
__device__ float4 operator*(const float3x4 &M, const float4 &v) { return make_float4(dot(v, M.m[0]), dot(v, M.m[1]), dot(v, M.m[2]), 1.0f); }

struct Ray {
	float3 o;	// origin
	float3 d;	// direction
};

// intersect ray with a box
// http://www.siggraph.org/education/materials/HyperGraph/raytrace/rtinter3.htm
__device__ int intersectBox(Ray r, float3 boxmin, float3 boxmax, float *tnear, float *tfar)
{
    // compute intersection of ray with all six bbox planes
    float3 invR = make_float3(1.0f) / r.d;
    float3 tbot = invR * (boxmin - r.o);
    float3 ttop = invR * (boxmax - r.o);

    // re-order intersections to find smallest and largest on each axis
    float3 tmin = fminf(ttop, tbot);
    float3 tmax = fmaxf(ttop, tbot);

    // find the largest tmin and the smallest tmax
    float largest_tmin = fmaxf(fmaxf(tmin.x, tmin.y), fmaxf(tmin.x, tmin.z));
    float smallest_tmax = fminf(fminf(tmax.x, tmax.y), fminf(tmax.x, tmax.z));

    *tnear = largest_tmin;
    *tfar = smallest_tmax;

    return smallest_tmax > largest_tmin;
}

// color conversion functions
__device__ uint rgbaFloatToInt(float4 rgba) { return (uint(__saturatef(rgba.w)*255)<<24) | (uint(__saturatef(rgba.z)*255)<<16) | (uint(__saturatef(rgba.y)*255)<<8) | uint(__saturatef(rgba.x)*255); }
__device__ uint rgbFloatToInt(float3 rgba) { return (uint(__saturatef(rgba.z)*255)<<16) | (uint(__saturatef(rgba.y)*255)<<8) | uint(__saturatef(rgba.x)*255); }
__device__ uint rFloatToInt(float r)
{
    r = __saturatef(r);
    return (uint(r*255)<<24) | (uint(r*255)<<16) | (uint(r*255)<<8) | uint(r*255);
}

////////////////////////////////////////////////////////////////////////////////
// constants

#define RENDER_COLOR_MASK        0x00000018
#define RENDER_COLOR_INCREMENT   0x00000008
#define RENDER_COLOR_VELOCITY    0x00000000
#define RENDER_COLOR_VEL_SPECT   0x00000008
#define RENDER_COLOR_POS_SPECT   0x00000010
#define RENDER_COLOR_GREEN       0x00000018
#define RENDER_FOG               0x00000020
#define RENDER_ISOMETRIC         0x00000040
#define RENDER_RAND_LOC          0x00000080
#define DEFAULT_RENDER_FLAGS     0x00000060

// zoom factor for isometric rendering
#define ISOMETRIC_ZOOM      1.5f

// boundary condition velocities
#define BOUNDARY_VELOCITY   5.0f
#define TOP_VELOCITY        15.0f

#define PRESSURE_FOG_CONSTANT   0.025f
#define VELOCITY_FOG_CONSTANT   0.0125f

// the inverse view matrix (needed for raytracing)
__constant__ float3x4 dInvViewMatrix;

// the forward view matrix (needed for rendering particles)
__constant__ float3x4 dViewMatrix;

// these are set by host code before the kernels are run
__constant__ int dFieldDimension; // dimension of vector / scalar fields
__constant__ int dRenderFlags; // parameters
__constant__ float dTransparencyFactor; // factor by which particles should be transparent (larger = dimmer particles)

// the velocity and pressure textures (from the previous frame)
//// apparently float3 textures can't be done, so I'm making this a float4
//// texture. the extra coordinate can be used for....?
texture<float4, 3, cudaReadModeElementType> velocityTex;
texture<float, 3, cudaReadModeElementType> pressureTex;

////////////////////////////////////////////////////////////////////////////////
// display kernel

// calculates a normal vector from the pressure texture (this was used for
// isosurface rendering, which I later removed)
/*__device__ float3 PressureTexNormal(float3 pos)
{
    float3 normal;
    float d = 0.5f / dFieldDimension;
    normal.x = (tex3D(pressureTex, pos.x+d, pos.y, pos.z) - tex3D(pressureTex, pos.x-d, pos.y, pos.z));
    normal.y = (tex3D(pressureTex, pos.x, pos.y+d, pos.z) - tex3D(pressureTex, pos.x, pos.y-d, pos.z));
    normal.z = (tex3D(pressureTex, pos.x, pos.y, pos.z+d) - tex3D(pressureTex, pos.x, pos.y, pos.z-d));
    return normalize(normal);
} */

// renders the pressure and (magnitude of) velocity fields as fog. pressure is
// red, velocity is blue.
__global__ void RenderParameterFogK(unsigned long* output, unsigned long imageW, unsigned long imageH)
{
    // amount to step by
    float tstep = 0.003f;
    int maxSteps = 2000;

    // pixel coordinates
    unsigned long x = blockIdx.x * blockDim.x + threadIdx.x;
    unsigned long y = blockIdx.y * blockDim.y + threadIdx.y;
    float u = (x / (float) imageW) * 2.0f - 1.0f;
    float v = (y / (float) imageH) * 2.0f - 1.0f;

    // the index of the pixel to be written by this call
    long outputColorIndex = y * imageW + x;

    // calculate eye ray in world space
    Ray eyeRay;
    if (dRenderFlags & RENDER_ISOMETRIC)
    {
        float3 offset = dInvViewMatrix * make_float3(u, v, 0.0f);
        eyeRay.o = 4.0f / ISOMETRIC_ZOOM * (make_float3(dInvViewMatrix * make_float4(0.0f, 0.0f, 0.0f, 1.0f)) + offset);
        eyeRay.d = dInvViewMatrix * make_float3(0.0f, 0.0f, -1.0f);
    } else {
        eyeRay.o = make_float3(dInvViewMatrix * make_float4(0.0f, 0.0f, 0.0f, 1.0f));
        eyeRay.d = dInvViewMatrix * normalize(make_float3(u, v, -2.0f));
    }

    // find intersection with box
    float tnear, tfar;

    // return if not intersecting
    if (!intersectBox(eyeRay, make_float3(-2.0f,-2.0f,-2.0f), make_float3(2.0f,2.0f,2.0f), &tnear, &tfar))
    {
        output[outputColorIndex] = 0;
        return;
    }

    // clamp to near plane
    if (tnear < 0.0f) tnear = 0.0f;

    // start stepping through space
    float t = tnear;
    float3 texPos;
    float counted = 0.0f;
    float velCounted = 0.0f;
    for (int i = 0; (i < maxSteps) && (t <= tfar); i++)
    {
        texPos = (eyeRay.o + eyeRay.d * t) * 0.25f + 0.5f;
        counted += tstep * tex3Dv(pressureTex, texPos) * PRESSURE_FOG_CONSTANT;
        velCounted += tstep * length(tex3Dv(velocityTex, texPos)) * VELOCITY_FOG_CONSTANT;
        t += tstep;
    }

    // clamp velCounted to 0..1
    if (velCounted > 1.0f) velCounted = 1.0f;

    // finally, red = pressure, blue = velocity
    float4 color = make_float4(counted, 0.0f, velCounted, 0.0f);
    output[outputColorIndex] = rgbaFloatToInt(color);

    return;
}

// draws particles (with transparency!) into the pixel buffer
__global__ void RenderParticlesK(unsigned long* output, float4* particles, unsigned long imageW, unsigned long imageH)
{
    // indices and locations
    unsigned int index = blockIdx.x * blockDim.x + threadIdx.x;
    float4 worldLoc = particles[index];
    float4 loc;

    // transform to camera space
    if (dRenderFlags & RENDER_ISOMETRIC)
    {
        loc = ISOMETRIC_ZOOM * (dViewMatrix * (worldLoc - make_float4(0.5f, 0.5f, 0.5f, 0.0f)));
    } else {
        float3 camera = make_float3(dInvViewMatrix * make_float4(0.0f, 0.0f, 0.0f, 1.0f));
        float4 direction = make_float4(make_float3(worldLoc - make_float4(0.5f, 0.5f, 0.5f, 0.0f)) - camera, 1.0f);
        loc = dViewMatrix * direction;
        loc /= (-2.0f / loc.z) * 4.0f;
    }
    int xScreen = ((loc.x + 1.0f) / 2.0f) * imageW;
    int yScreen = ((loc.y + 1.0f) / 2.0f) * imageH;

    // offscreen? then forget about it
    if ((xScreen < 0) || (xScreen >= imageW) || (yScreen < 0) || (yScreen >= imageH)) return;
    long outputColorIndex = yScreen * imageW + xScreen;

    // variables for calculating color
    float3 lowColor = make_float3(0.0f, 0.0f, 1.0f);
    float3 highColor = make_float3(1.0f, 0.0f, 0.0f);
    float3 vel, rgba;
    float velm;

    // figure out particle color and store it in rgba
    switch (dRenderFlags & RENDER_COLOR_MASK)
    {
      case RENDER_COLOR_VELOCITY:
        // velocity "spectrum" from 0 - 10
        velm = length(make_float3(tex3Dv(velocityTex, worldLoc))) / 10.0f;
        if (velm > 1.0f) velm = 1.0f;
        rgba = lowColor * velm + highColor * (1 - velm);
        break;

      case RENDER_COLOR_VEL_SPECT:
        // velocity by coordinates (x is red, etc)
        vel = make_float3(tex3Dv(velocityTex, worldLoc));
        rgba = normalize(vel / 2.0f + 0.5f);
        break;

      case RENDER_COLOR_POS_SPECT:
        // position by coordinates
        rgba = make_float3(worldLoc);
        break;

      case RENDER_COLOR_GREEN:
        // just green
        rgba = make_float3(0.0f, 1.0f, 0.0f);
    }

    // transparency!
    rgba /= dTransparencyFactor;

    // extract color components, and add this particle's color
    // YES, if two particles map to the same pixel, then it's possible that only
    // one of them will be drawn (since two threads might try to write to the
    // same location at the same time, and one will "get there" first). In a
    // system of hundreds of thousands of particles, though, it's highly
    // unlikely that two threads will be running at exactly the same time and
    // that this will happen... I'm sure it *does* happen, but the effect is
    // unnoticeable.
    long icolor = output[outputColorIndex];
    unsigned long b = (unsigned long)(__saturatef(rgba.z) * 255) + ((icolor >> 16) & 0x000000FF);
    unsigned long g = (unsigned long)(__saturatef(rgba.y) * 255) + ((icolor >> 8) & 0x000000FF);
    unsigned long r = (unsigned long)(__saturatef(rgba.x) * 255) + (icolor & 0x000000FF);
    if (b > 0xFF) b = 0xFF;
    if (g > 0xFF) g = 0xFF;
    if (r > 0xFF) r = 0xFF;
    output[outputColorIndex] = ((b & 0xFF) << 16) | ((g & 0xFF) << 8) | (r & 0xFF);
}

////////////////////////////////////////////////////////////////////////////////
// computation kernels

#define idx(x,y,z)          (z * dFieldDimension * dFieldDimension + y * dFieldDimension + x)
#define _fetch_xyz          int x = threadIdx.x; \
                            int y = blockIdx.x; \
                            int z = blockIdx.y
#define _fetch_xyz_float    float xf = (float)(x + 0.5f) / dFieldDimension; \
                            float yf = (float)(y + 0.5f) / dFieldDimension; \
                            float zf = (float)(z + 0.5f) / dFieldDimension
#define tex3Di(t,x,y,z) tex3D(t, x, y, z)

__device__ float3 _transport_integrator(float3 start, float dt)
{
    // TODO: make this not suck (use Runge-Kutta or something)
    return start + make_float3(tex3Dv(velocityTex, start)) * dt;
}

/* add force kernels (which I didn't end up using)

__global__ void VelocityAddCircularForceK(float4* newVelocity, float power, float dt)
{
    _fetch_xyz;
    _fetch_xyz_float;
    float2 m;
    if ((0.5f - zf == 0) && (xf - 0.5f == 0.0f)) m = make_float2(0.0f, 0.0f);
    else m = normalize(make_float2(0.5f - zf, xf - 0.5f));
    newVelocity[idx(x, y, z)] = tex3Di(velocityTex, xf, yf, zf) + make_float4(m.x, 0.0f, m.y, 0.0f);
}

__global__ void VelocityAddConstForceK(float4* newVelocity, float fx, float fy, float fz, float dt)
{
    _fetch_xyz;
    _fetch_xyz_float;
    newVelocity[idx(x, y, z)] = tex3Di(velocityTex, xf, yf, zf) + make_float4(fx, fy, fz, 0.0f);
} */

__global__ void VelocityTransportK(float4* newVelocity, float dt)
{
    _fetch_xyz;
    _fetch_xyz_float;
    float3 start = make_float3(xf, yf, zf);
    float3 end = _transport_integrator(start, -dt);
    float3 texSampleCoord = end;// + make_float3(0.5f / dFieldDimension);
    newVelocity[idx(x, y, z)] = tex3Dv(velocityTex, texSampleCoord);
}

__global__ void ParticlesTransportK(float4* pos, float dt, float rand)
{
    unsigned int index = blockIdx.x * blockDim.x + threadIdx.x;

    // first, we get the current position, velocity, and particle life
    float4 oldPosition = pos[index];
    float3 newPosition;

    // has the particle died? (left the cube?)
    if ((oldPosition.x >= 0.999f) || (oldPosition.y >= 0.999f) || (oldPosition.z >= 0.999f) ||
        (oldPosition.z < 0.0f) || (oldPosition.y < 0.0f) || (oldPosition.z < 0.0f))
    {
        // if so, it needs to be reborn somewhere. we calculate a few parameters
        // before recreating the particle:
        float w8 = index * rand;
        float theta = (w8 - floor(w8));
        float phi = abs(sin((index * 489.0f - oldPosition.x * 237.0f) * rand));
        float dv = abs(sin((index + oldPosition.x + oldPosition.z) * rand));

        if (dRenderFlags & RENDER_RAND_LOC) newPosition = make_float3(phi, theta, dv);
        else {
            // the value of phi specifies the particle's start point
            if (phi < 0.25f)
            {
                // the particle will come from the x = 0 plane
                newPosition = make_float3(1.0f / dFieldDimension, theta, dv);
            } else if (phi < 0.5f)
            {
                // the particle will come from the x = 1.0f plane
                newPosition = make_float3(1.0f - 1.0f / dFieldDimension, theta, dv);
            } else if (phi < 0.75f)
            {
                // the particle will come from the z = 0 plane
                newPosition = make_float3(dv, theta, 1.0f / dFieldDimension);
            } else {
                // the particle will come from the z = 1.0f plane
                newPosition = make_float3(dv, theta, 1.0f - 1.0f / dFieldDimension);
            }
        }
    } else {
        // otherwise, we just continue its trajectory. a simple Euler
        // approximation will suffice here.
        newPosition = _transport_integrator(make_float3(oldPosition), dt);
    }

    // finally, assign the calculated position, velocity, and life to this particle
    pos[index] = make_float4(newPosition, 1.0f);
}

////////////////////////////////////////////////////////////////////////////////

__device__ int IsSolidCell(int x, int y, int z)
{
    if (x < 0) return 1;
    if (y < 0) return 1;
    if (z < 0) return 1;
    if (x >= dFieldDimension) return 1;
    if (y >= dFieldDimension) return 1;
    if (z >= dFieldDimension) return 1;
    return 0;
}

__device__ int IsBoundaryCell(int x, int y, int z)
{
    if (x < 0) return 1;
    if (y < 0) return 1;
    if (z < 0) return 1;
    if (x >= dFieldDimension) return 1;
    if (y >= dFieldDimension) return 1;
    if (z >= dFieldDimension) return 1;
    return 0;
}

__device__ float4 GetObstacleVelocity(int x, int y, int z)
{
    // counterclockwise flows on boundaries
    float4 vel = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
    //if (x < 0) vel.z = 1.0f;
    //if (x >= dFieldDimension) vel.z = -1.0f;
    //if (z < 0) vel.x = -1.0f;
    //if (z >= dFieldDimension) vel.x = 1.0f;
    return vel * 10.0f;
}

__device__ float BoundaryCellPressure(int x, int y, int z)
{
    if (x < 0) return 5.0f;
    if (y < 0) return 0.0f;
    if (z < 0) return 5.0f;
    if (x >= dFieldDimension) return 5.0f;
    if (y >= dFieldDimension) return 0.0f;
    if (z >= dFieldDimension) return 5.0f;
    return 0.0f;
}

__device__ float VelocityTexDivergence(float3 pos, int x, int y, int z)
{
    float vr, vl, vt, vb, vu, vd;
    float d = 1.0f / dFieldDimension;

    if (IsBoundaryCell(x+1, y, z)) vr = GetObstacleVelocity(x+1, y, z).x;
    else vr = tex3Di(velocityTex, pos.x+d, pos.y, pos.z).x;
    if (IsBoundaryCell(x-1, y, z)) vl = GetObstacleVelocity(x-1, y, z).x;
    else vl = tex3Di(velocityTex, pos.x-d, pos.y, pos.z).x;

    if (IsBoundaryCell(x, y+1, z)) vt = GetObstacleVelocity(x, y+1, z).y;
    else vt = tex3Di(velocityTex, pos.x, pos.y+d, pos.z).y;
    if (IsBoundaryCell(x, y-1, z)) vb = GetObstacleVelocity(x, y-1, z).y;
    else vb = tex3Di(velocityTex, pos.x, pos.y-d, pos.z).y;

    if (IsBoundaryCell(x, y, z+1)) vu = GetObstacleVelocity(x, y, z+1).z;
    else vu = tex3Di(velocityTex, pos.x, pos.y, pos.z+d).z;
    if (IsBoundaryCell(x, y, z-1)) vd = GetObstacleVelocity(x, y, z-1).z;
    else vd = tex3Di(velocityTex, pos.x, pos.y, pos.z-d).z;

    return 0.5f * (vr-vl+vt-vb+vu-vd);
}

__global__ void PressureRecalculateK(float* newPressure)
{
    _fetch_xyz;
    _fetch_xyz_float;
    if ((x == 0) || (z == 0) || (x == dFieldDimension - 1) || (z == dFieldDimension - 1))
    {
        newPressure[idx(x, y, z)] = 5.0f;
        return;
    }
    if ((y == dFieldDimension - 1) && (sqrtf((x - dFieldDimension / 2) * (x - dFieldDimension / 2) + (z - dFieldDimension / 2) * (z - dFieldDimension / 2)) <= dFieldDimension / 8))
    {
        newPressure[idx(x, y, z)] = 0.0f;
        return;
    }
    float d = 1.0f / dFieldDimension;
    float divergence = VelocityTexDivergence(make_float3(xf, yf, zf), x, y, z);
    float pC = tex3Di(pressureTex, xf, yf, zf);
    float pR = tex3Di(pressureTex, xf+d, yf, zf);
    float pL = tex3Di(pressureTex, xf-d, yf, zf);
    float pT = tex3Di(pressureTex, xf, yf+d, zf);
    float pB = tex3Di(pressureTex, xf, yf-d, zf);
    float pU = tex3Di(pressureTex, xf, yf, zf+d);
    float pD = tex3Di(pressureTex, xf, yf, zf-d);
    if (IsSolidCell(x+1, y, z)) pR = pC;
    if (IsSolidCell(x-1, y, z)) pL = pC;
    if (IsSolidCell(x, y+1, z)) pT = pC;
    if (IsSolidCell(x, y-1, z)) pB = pC;
    if (IsSolidCell(x, y, z+1)) pU = pC;
    if (IsSolidCell(x, y, z-1)) pD = pC;
    newPressure[idx(x, y, z)] = (pL + pR + pB + pT + pU + pD - divergence) / 6.0f;
}

__global__ void VelocityProjectK(float4* newVelocity)
{
    _fetch_xyz;
    _fetch_xyz_float;
    float d = 1.0f / dFieldDimension;

    if (IsSolidCell(x, y, z))
    {
        newVelocity[idx(x, y, z)] = GetObstacleVelocity(x, y, z);
        return;
    }

    if (x == 0)
    {
        newVelocity[idx(x, y, z)] = make_float4(0.0f, 0.0f, BOUNDARY_VELOCITY, 0.0f);
        return;
    }
    if (x == dFieldDimension - 1)
    {
        newVelocity[idx(x, y, z)] = make_float4(0.0f, 0.0f, -BOUNDARY_VELOCITY, 0.0f);
        return;
    }
    if (z == 0)
    {
        newVelocity[idx(x, y, z)] = make_float4(-BOUNDARY_VELOCITY, 0.0f, 0.0f, 0.0f);
        return;
    }
    if (z == dFieldDimension - 1)
    {
        newVelocity[idx(x, y, z)] = make_float4(BOUNDARY_VELOCITY, 0.0f, 0.0f, 0.0f);
        return;
    }
    if ((y == dFieldDimension - 1) && (sqrtf((x - dFieldDimension / 2) * (x - dFieldDimension / 2) + (z - dFieldDimension / 2) * (z - dFieldDimension / 2)) <= dFieldDimension / 4))
    {
        newVelocity[idx(x, y, z)] = make_float4(0.0f, TOP_VELOCITY, 0.0f, 0.0f);
        return;
    }

    float pC = tex3Di(pressureTex, xf, yf, zf);
    float pR = tex3Di(pressureTex, xf+d, yf, zf);
    float pL = tex3Di(pressureTex, xf-d, yf, zf);
    float pT = tex3Di(pressureTex, xf, yf+d, zf);
    float pB = tex3Di(pressureTex, xf, yf-d, zf);
    float pU = tex3Di(pressureTex, xf, yf, zf+d);
    float pD = tex3Di(pressureTex, xf, yf, zf-d);

    float3 vObst = make_float3(0.0f, 0.0f, 0.0f);
    float3 vMask = make_float3(1.0f, 1.0f, 1.0f);

    if (IsBoundaryCell(x+1, y, z))
    {
        pR = BoundaryCellPressure(x+1, y, z);
        vObst.x = GetObstacleVelocity(x+1, y, z).x;//vR.x;
        vMask.x = 0.0f;
    }
    if (IsBoundaryCell(x-1, y, z))
    {
        pL = BoundaryCellPressure(x-1, y, z);
        vObst.x = GetObstacleVelocity(x-1, y, z).x;//vL.x;
        vMask.x = 0.0f;
    }
    if (IsBoundaryCell(x, y+1, z))
    {
        pT = BoundaryCellPressure(x, y+1, z);
        vObst.y = GetObstacleVelocity(x, y+1, z).y;//vT.y;
        vMask.y = 0.0f;
    }
    if (IsBoundaryCell(x, y-1, z))
    {
        pB = BoundaryCellPressure(x, y-1, z);
        vObst.y = GetObstacleVelocity(x, y-1, z).y;//vB.y;
        vMask.y = 0.0f;
    }
    if (IsBoundaryCell(x, y, z+1))
    {
        pU = BoundaryCellPressure(x, y, z+1);
        vObst.z = GetObstacleVelocity(x, y, z+1).z;//vU.z;
        vMask.z = 0.0f;
    }
    if (IsBoundaryCell(x, y, z-1))
    {
        pD = BoundaryCellPressure(x, y, z-1);
        vObst.z = GetObstacleVelocity(x, y, z-1).z;//vD.z;
        vMask.z = 0.0f;
    }

    float3 grad = 0.5f * make_float3(pR - pL, pT - pB, pU - pD);
    float3 vNew = make_float3(tex3Di(velocityTex, xf, yf, zf)) - grad;

    newVelocity[idx(x, y, z)] = make_float4((vMask * vNew) + vObst, 0.0f);
}
