/* This software contains source code provided by NVIDIA Corporation.
 * The program is based on the example code ray_const.cu of Chapter 6
 * of the "CUDA by Example" book.
 * 
 * A program az NVIDIA Corporation által készített programkódot tartalmaz.
 * A program a "Cuda by Example" tankönyv 6. fejezetének ray_const.cu 
 * példaprogramjára épül.
 */

#include "cuda_runtime.h"
#include "device_launch_parameters.h"

#include <stdio.h>
#include <stdlib.h>

#include "ImageIO.h"

#define N           1024
#define BLOCKDIM    16

#define rnd( x ) (x * rand() / RAND_MAX)
#define INF     2e10f

struct Sphere {
    float   r,b,g;
    float   radius;
    float   x,y,z;
    __device__ float hit( float ox, float oy, float *n ) {
        float dx = ox - x;
        float dy = oy - y;
        if (dx*dx + dy*dy < radius*radius) {
            float dz = sqrtf( radius*radius - dx*dx - dy*dy );
            *n = dz / sqrtf( radius * radius );
            return dz + z;
        }
        return -INF;
    }
};
#define SPHERES 20

__constant__ Sphere s[SPHERES];

__global__ void kernel( unsigned char *ptr ) {
    // map from threadIdx/BlockIdx to pixel position
    int x = threadIdx.x + blockIdx.x * blockDim.x;
    int y = threadIdx.y + blockIdx.y * blockDim.y;

    int offset = x + y * blockDim.x * gridDim.x;
    
    float   ox = (x - N/2);
    float   oy = (y - N/2);

    float   r=0, g=0, b=0;
    float   maxz = -INF;
    for(int i=0; i<SPHERES; i++) {
        float   n;
        float   t = s[i].hit( ox, oy, &n );
        if (t > maxz) {
            float fscale = n;
            r = s[i].r * fscale;
            g = s[i].g * fscale;
            b = s[i].b * fscale;
            maxz = t;
        }
    } 

    ptr[offset*4 + 0] = (int)(r * 255);
    ptr[offset*4 + 1] = (int)(g * 255);
    ptr[offset*4 + 2] = (int)(b * 255);
	ptr[offset*4 + 3] = 255;
}

// globals needed by the update routine
struct DataBlock {
    unsigned char   *dev_bitmap;
};

int main( void ) {
    unsigned char* bitmap = new unsigned char[N * N * 4];
    unsigned char* dev_bitmap;

    // allocate memory on the GPU for the output bitmap
    cudaMalloc( (void**)&dev_bitmap, N * N * 4 * sizeof(unsigned char));

    // allocate temp memory, initialize it, copy to constant
    // memory on the GPU, then free our temp memory
    Sphere *temp_s = (Sphere*)malloc( sizeof(Sphere) * SPHERES );
    for (int i=0; i<SPHERES; i++) {
        temp_s[i].r = rnd( 1.0f );
        temp_s[i].g = rnd( 1.0f );
        temp_s[i].b = rnd( 1.0f );
        temp_s[i].x = rnd( 1000.0f ) - 500;
        temp_s[i].y = rnd( 1000.0f ) - 500;
        temp_s[i].z = rnd( 1000.0f ) - 500;
        temp_s[i].radius = rnd( 100.0f ) + 20;
    }
    cudaMemcpyToSymbol( s, temp_s, sizeof(Sphere) * SPHERES);
    free( temp_s );

    // generate a bitmap from our sphere data
    dim3 blockDim = dim3(BLOCKDIM, BLOCKDIM, 1);
    dim3 gridDim = dim3((N + BLOCKDIM - 1) / BLOCKDIM, (N + BLOCKDIM - 1) / BLOCKDIM, 1);

    kernel<<<gridDim, blockDim >>>( dev_bitmap );

    // copy our bitmap back from the GPU for display
    cudaMemcpy( bitmap, dev_bitmap, N * N * 4 * sizeof(unsigned char), cudaMemcpyDeviceToHost );

    writeRGBImageToFile("image.png", bitmap, N, N);

    cudaFree( dev_bitmap );

    return 0;
}