#include <u.h>
#include <libc.h>

typedef struct { double x, y, z; } Vec;

/* when V(n) then all n. when V(x,y) then z = 0 */
Vec V(double x, double y, double z) { return (Vec){x,y,z}; }
Vec vadd(Vec a, Vec b) { return V(a.x+b.x,a.y+b.y,a.z+b.z); }
Vec vmul(Vec v, double s) { return V(v.x*s,v.y*s,v.z*s); }
Vec vmulv(Vec a, Vec b) { return V(a.x*b.x,a.y*b.y,a.z*b.z); }
double vdot(Vec a, Vec b) { return a.x*b.x + a.y*b.y + a.z*b.z; }
Vec vnorm(Vec v) { return vmul(v, 1/sqrt(vdot(v, v))); }

double
min(double a, double b)
{
	return a < b? a: b;
}

double
randomVal(void)
{
	return frand();
}

/*
 * Rectangle CSG equation. Returns minimum signed distance from
 * space carved by
 * lowerLeft vertex and opposite rectangle vertex upperRight.
 */
double
BoxTest(Vec p, Vec lowerLeft, Vec upperRight)
{
	lowerLeft = vadd(p, vmul(lowerLeft, -1));
	upperRight = vadd(upperRight, vmul(p, -1));
	return -min(
		min(
			min(lowerLeft.x, upperRight.x),
			min(lowerLeft.y, upperRight.y)
		),
		min(lowerLeft.z, upperRight.z)
	);
}

enum {
	HIT_NONE,
	HIT_LETTER,
	HIT_WALL,
	HIT_SUN
};

/* Sample the world using Signed Distance Fields. */
double
QueryDatabase(Vec position, int *hitType)
{
	Vec f;
	double distance, roomDist, sun;
	int i;
	char letters[15*4+1] =               /* 15 two points lines */
		"5O5_" "5W9W" "5_9_"         /* P (without curve) */
		"AOEO" "COC_" "A_E_"         /* I */
		"IOQ_" "I_QO"                /* X */
		"UOY_" "Y_]O" "WW[W"         /* A */
		"aOa_" "aWeW" "a_e_" "cWiO"; /* R (without curve) */

	distance = 1e9;
	/* flattened position (z = 0) */
	f = position;
	f.z = 0;

	for(i = 0; i < sizeof letters; i += 4){
		Vec begin, e, o;

		begin = vmul(V(letters[i] - 79,letters[i + 1] - 79,0), 0.5);
		e = vadd(vmul(V(letters[i + 2] - 79,letters[i + 3] - 79,0), 0.5), vmul(begin, -1));
		o = vmul(vadd(f, vadd(begin, vmul(e, min(-min(vdot(vadd(begin, vmul(f, -1)), e) / vdot(e, e), 0), 1)))), -1);
		distance = min(distance, vdot(e, e));
	}

	/* actual distance */
	distance = sqrt(distance);

	/* Two curves (for P and R in PixaR) with hard-coded locations. */
	Vec curves[] = {
		V(-11,6,0),
		V(11,6,0)
	};
	for(i = 2; i--;){
		Vec o;

		o = vadd(f, vmul(curves[i], -1));
		distance = min(distance, o.x > 0? fabs(sqrt(vdot(o, o)) - 2): (o.y += o.y > 0? -2: 2, sqrt(vdot(o, o))));
	}
	distance = pow(pow(distance, 8) + pow(position.z, 8), 0.125) - 0.5;
	*hitType = HIT_LETTER;

	/*
	 * min(A,B) = Union with Constructive solid geometry
	 * -min carves an empty space
	 */
	roomDist = min(
		-min(	/* Lower room */
			BoxTest(position, V(-30,-0.5,-30), V(30,18,30)),
			/* Upper room */
			BoxTest(position, V(-25,17,-25), V(25,20,25))
		),
		BoxTest(/* Ceiling "planks" spaced 8 units apart. */
			V(fmod(fabs(position.x), 8),
				position.y,
				position.z),
			V(1.5,18.5,-25),
			V(6.5,20,25)
		)
	);

	if(roomDist < distance){
		distance = roomDist;
		*hitType = HIT_WALL;
	}

	sun = 19.9 - position.y ; /* Everything above 19.9 is light source. */
	if(sun < distance){
		distance = sun;
		*hitType = HIT_SUN;
	}

	return distance;
}

/* Perform signed sphere marching */
int
RayMarching(Vec origin, Vec direction, Vec *hitPos, Vec *hitNorm)
{
	int hitType, noHitCount;
	double d; /* distance from closest object in world. */
	double total_d;

	hitType = HIT_NONE;
	noHitCount = 0;

	/* Signed distance marching */
	for (total_d = 0; total_d < 100; total_d += d){
		*hitPos = vadd(origin, vmul(direction, total_d));
		d = QueryDatabase(*hitPos, &hitType);
		if(d < 0.01 || ++noHitCount > 99){
			*hitNorm = vnorm(V(QueryDatabase(vadd(*hitPos, V(0.01,0,0)), &noHitCount) - d,
				QueryDatabase(vadd(*hitPos, V(0,0.01,0)), &noHitCount) - d,
				QueryDatabase(vadd(*hitPos, V(0,0,0.01)), &noHitCount) - d));
			return hitType;
		}
	}
	return 0;
}

Vec
Trace(Vec origin, Vec direction)
{
	Vec sampledPosition, normal, color, attenuation, lightDirection;
	int bounceCount;
	double incidence, p, c, s, g, u, v;

	attenuation = V(1,1,1);
	lightDirection = vnorm(V(0.6,0.6,1)); /* Directional light */

	for(bounceCount = 3; bounceCount--;){
		int hitType;

		hitType = RayMarching(origin, direction, &sampledPosition, &normal);
		switch(hitType){
		case HIT_LETTER: /* Specular bounce on a letter. No color acc. */
			direction = vadd(direction, vmul(normal, vdot(normal, direction) * -2));
			origin = vadd(sampledPosition, vmul(direction, 0.1));
			attenuation = vmul(attenuation, 0.2); /* Attenuation via distance traveled. */
			break;
		case HIT_WALL:
			incidence = vdot(normal, lightDirection);
			p = 6.283185 * randomVal();
			c = randomVal();
			s = sqrt(1 - c);
			g = normal.z < 0? -1: 1;
			u = -1/(g + normal.z);
			v = normal.x * normal.y * u;
			direction = vadd(vadd(vmul(V(v,g + normal.y * normal.y * u,-normal.y), cos(p)*s),vmul(V(1 + g * normal.x * normal.x * u,g * v,-g * normal.x), sin(p)*s)), vmul(normal, sqrt(c)));
			origin = vadd(sampledPosition, vmul(direction, 0.1));
			attenuation = vmul(attenuation, 0.2);
			if(incidence > 0 &&
				RayMarching(vadd(sampledPosition, vmul(normal, 0.1)),
					lightDirection,
					&sampledPosition,
					&normal) == HIT_SUN)
				color = vadd(color, vmul(vmulv(attenuation, V(500,400,100)), incidence));
			break;
		case HIT_SUN:
			color = vadd(color, vmulv(attenuation, V(50,80,100)));
			break;
		default:
			goto Out;
		}
	}
Out:
	return color;
}

void
usage(void)
{
	fprint(2, "usage: %s [-s nsamples]\n", argv0);
	exits("usage");
}

void
main(int argc, char *argv[])
{
	Vec position, goal, left, up;
	int w, h, samplesCount, x, y, p;

	w = 960;
	h = 540;
	samplesCount = 1;
	ARGBEGIN{
	case 's':
		samplesCount = strtol(EARGF(usage()), nil, 10);
		break;
	}ARGEND;

	position = V(-22,5,25);
	goal = vnorm(vadd(V(-3,4,0), vmul(position, -1)));
	left = vmul(vnorm(V(goal.z,0,-goal.x)), (1.0/w));

	/* Cross-product to get the up vector */
	up = V(goal.y * left.z - goal.z * left.y,
		goal.z * left.x - goal.x * left.z,
		goal.x * left.y - goal.y * left.x);

	print("P3 %d %d 255\n", w, h);
	for(y = h; y--;)
		for(x = w; x--;){
			Vec color, o;

			for (p = samplesCount; p--;)
				color = vadd(color, Trace(position, vnorm(vadd(vadd(goal, vmul(left, x - w/2 + randomVal())), vmul(up, y - h/2 + randomVal())))));

			/* Reinhard tone mapping */
			color = vadd(vmul(color, 1.0/samplesCount), V(14.0/241,14.0/241,14.0/241));
			o = vadd(color, V(1,1,1));
			color = vmul(V(color.x/o.x,color.y/o.y,color.z/o.z), 255);
			print("%d %d %d\n", (int)color.x, (int)color.y, (int)color.z);
		}
}