#include <cstdio>
#include <cstdlib>
#include "image.h"
#include "misc.h"
#include "pnmfile.h"
#include <math.h>

#include <time.h>

#include <vector>
#include "segment-image.h"


using std::vector;

struct vec{
	float x, y;
	vec():x(0), y(0){}
	vec(const vec& a){ *this = a;}
	vec(float x, float y):x(x), y(y){}
	vec operator +=(vec a){ x += a.x; y += a.y; }
	vec operator -=(vec a){ x -= a.x; y -= a.y; }
	vec operator *=(float a){ x *= a; y *= a; }	
	vec operator /=(float a){ x /= a; y /= a; }
	bool isZero(){ return x == 0 && y == 0;}
};

vec operator +(vec a, vec b){ vec n = a; n+=b; return n; }
vec operator -(vec a, vec b){ vec n = a; n-=b; return n; }
vec operator *(vec a, float b){ vec n = a; n*=b; return n; }

float dot(vec a, vec b)
{
	return a.x * b.x + a.y * b.y;
}

float length(vec a)
{
	return sqrt(dot(a, a));
}

vec perp(vec a)
{
	return vec(a.y, -a.x);
}

rgb toRGB(float r, float g, float b)
{
	if (r < 0) r = 0;
	if (g < 0) g = 0;
	if (b < 0) b = 0;
	
	if (r > 255) r = 255;
	if (g > 255) g = 255;
	if (b > 255) b = 255;
	rgb o;
	
	o.r = (uchar)r;
	o.g = (uchar)g;
	o.b = (uchar)b;
	return o;
}

struct rect{
	int minX, minY;
	int maxX, maxY;
	rect()
	{
		minX = 99999999;
		minY = minX;
		maxX = -minX;
		maxY = -minY;
	}
	
	void addpoint(int x,int y)
	{
		if (x < minX)
		{minX = x;}
		if (x > maxX)
		{maxX = x;}
		if (y < minY)
		{minY = y;}
		if (y> maxY)
		{maxY = y;}
	}
	
	bool isLong()
	{
		int w = maxX - minX;
		int h = maxY - minY;

		return w > h+70;
	}
	void print()
	{
		printf("(%i, %i), (%i, %i)\n", minX, minY, maxX, maxY);
	}
};

int wrap(int i, int max)
{
	i = i % max;
	if (i < 0) i += max;
	return i;
}

float lerp(float a, float b, float t)
{
	return a * (1 - t) + b * t;
}

float fetch(image<rgb> *im, vec at)
{
	int w = im->width();
	int h = im->height();

	int x = (int)at.x;
	int y = (int)at.y;
	
	float u = at.x - x;
	float v = at.y - y;
	
	
	
	float a = imRef(im, wrap(x, w), wrap(y, h)).r;
	float b = imRef(im, wrap(x + 1, w), wrap(y, h)).r;
	//
	float c = imRef(im, wrap(x, w), wrap(y+1, h)).r;
	float d = imRef(im, wrap(x + 1, w), wrap(y+1, h)).r;
	
	return lerp(lerp(a,b, u), lerp(c, d, u), v);
}

void castRays(image<rgb> *input, vector<vec> points[256], int direction = 1)
{
	int width = input->width();
	int height = input->height();
	
	int start = 1;
	int end = width - 1;
	if(direction == -1)
	{
		start = width - 2;
		end = 0;
	}
	
	for (int y = 0; y < height; y+=10)
	{
		rgb last = imRef(input, 0, y);
		
		bool visited[256];
		for(int i = 0; i < 256; i++)
			visited[i] = false;
		
		for (int x = start; x != end; x+=direction)
		{
			rgb c = imRef(input, x, y);

			if (!(c == last) && !visited[c.r]){
				points[c.r].push_back( vec(x, y) );
				visited[c.r] = true;
			}
			last = c;
		}
  }
}

void findDirections(vector<vec> points[256], vec dir[256])
{
	for(int i = 0; i < 255; i++) {
  		if (points[i].size() < 2) continue;
  		vec d;
  		
  		int num = points[i].size();
  		
  		float bestAmount = 0;
  		for(int ransacIt = 0; ransacIt < 10; ransacIt++){
  		
	  		int r = rand() % (num - 1);
	  		vec rDir = points[i][r + 1] - points[i][r];
	  		rDir /= length(rDir);
	  		
	  		float amount = 0;
	  		for(int j = 0; j < num - 1; j++){
	  			vec vDir = points[i][j + 1] - points[i][j];
	  			vDir/=length(vDir);
	  			 amount += dot(vDir, rDir); 
	  		}
	  		if (amount > bestAmount)
	  		{
	  			bestAmount = amount;
	  			d = rDir;
	  		}
	  		
	  		/*
	  		for(int j = 1; j < num; j++){
	  			vec d2 = points[i][j] - points[i][0];
	  			d += d2;
	  		}*/
	  		
	  		
	  		if (points[i].size() != 0)
	  		{
	  			d /= sqrt(d.x*d.x + d.y*d.y);
	  			
	  			dir[i] = d;
	  		//printf("(%f, %f)\n", d.x, d.y);
	  		}
	  		else dir[i] = vec(0, 0);
  		}
  	}
}
void getBounds(image<rgb> *input, vector<vec> points[256], rect rc[256])
{
	for (int y = 0; y < input->height(); y++) {
    	for (int x = 0; x < input->width(); x++) {
    		rgb d = imRef(input, x, y);	
    		rc[d.r].addpoint(x,y);
    	}
    }
}


image<float>* GetEdges(image<rgb> *in )
{
  int xsize = in->width();
  int ysize = in->height();
  image<float> *input = new image<float>(xsize, ysize);
  	  for(int x =0; x<xsize; x++){
		for(int y=0; y<ysize;y++){
			rgb p = imRef(in, x, y);
  			imRef(input,x,y) = (p.r+p.g+p.b)/3.0;
  		}
  		}
	
  printf("processing\n");

  int hy[3][3] = {{-1,-2,-1},{0,0,0},{1,2,1}};
  int hx[3][3] = {{-1,0,1},{-2,0,2},{-1,0,1}};

 
  image<float> *sbx = new image<float>(xsize, ysize);
  image<float> *sby = new image<float>(xsize, ysize);
  image<float> *sbout = new image<float>(xsize, ysize);

  for(int x = 1; x<xsize-1; x++){
	for(int y=1; y<ysize-1;y++){
		float sum = 0;
	for (int j = -1; j<=1; j++){
		for (int i = -1; i <= 1; i++) {
		sum += hx[j+1][i+1]*imRef(input,x+i,y+i);
		}
	}
	imRef(sbx, x, y ) = sum;
  }
  }
  
    for(int x = 1; x<xsize-1; x++){
	for(int y=1; y<ysize-1;y++){
		float sum = 0;
	for (int j = -1; j<=1; j++){
		for (int i = -1; i <= 1; i++) {
		sum += hy[j+1][i+1]*imRef(input,x+i,y+i);
		}
	}
	imRef(sby, x, y ) = sum;
  }
  }
  
 	for(int x = 1; x<xsize-1; x++){
		for(int y=1; y<ysize-1;y++){
			float gsum = sqrt(pow(imRef(sbx,x,y),2)+pow(imRef(sby,x,y),2));
			imRef(sbout,x,y) = gsum;
		}
	}
	
	delete sbx;
	delete sby;
	return sbout;
}


int main(int argc, char **argv) {
	if (argc != 3) {
		fprintf(stderr, "usage: %s input(ppm) output(ppm)\n", argv[0]);
		return 1;
	}
  
  	srand ( time(NULL) );

	float sigma = .5;
	int k = 500;
	int min_size = 2000;
	int num_ccs;
	

	printf("loading input image.\n");
	image<rgb> *color = loadPPM(argv[1]);
	
	printf("Segmenting...\n");
	image<rgb> *input = segment_image(color, sigma, k, min_size, &num_ccs);
	
	printf("got %i components\n", num_ccs);
	
	image<rgb> *h1 = loadPPM("h1.ppm");
	image<rgb> *h2 = loadPPM("h2.ppm");
	image<rgb> *h3 = loadPPM("h3.ppm");
	image<float> *edge = GetEdges(color);// loadPPM("edge.ppm");
	

	printf("Computing directions\n");

	int width = input->width();
	int height = input->height();
	
	
	vector<vec> points[256];
	vec dir[256];
	rect rc[256];
	
	castRays(input, points);
	findDirections(points, dir);
	getBounds(input, points, rc);
	
	vector<vec> points2[256];
	vec dir2[256];
	
	castRays(input, points2, -1);
	findDirections(points2, dir2);
	
	for(int i = 0; i < 256; i++)
	{
		if (dir[i].isZero())
			dir[i] = dir2[i];
		else
			if(!dir2[i].isZero())
				dir[i] = dir[i] * .5 + dir2[i] * .5;
		//dir[i] /= length(dir[i]);
	}
 
    
    int hw = h1->width();
	int hh = h1->height();

	printf("Rendering strokes.\n");
	for (int y = 0; y < height; y++) {
    	for (int x = 0; x < width; x++) {
    	
    		rgb i = imRef(input, x, y);
    		vec d = dir[i.r];
    		
    		d.x *= -1;
    		
    		rgb pix = imRef(input, x, y);
    		
    		//if (rc[pix.r].isLong())	d = perp(d);
    		
    		//d = vec(.7, .4);
    		
    		vec p(x*d.x + y*d.y, y * d.x + x * d.y);
    		
    		
    		d = perp(d);
    		vec p2(x*d.x + y*d.y, y * d.x + x * d.y);
    		
    	
    		rgb b = imRef(color,x , y);
    		



			float intense = 1-(b.g+b.b+b.r)/3.0 /255.0;
			float t;

			float hatchingAmount = 0;
			
			float shades = 6;
			
			
			if(intense <= 1.0/shades)
			{
				t = intense*shades;
				hatchingAmount = lerp(255, fetch(h1, p), t);
			}
			else if (intense<= 2.0/shades)
			{
				t = (intense-1.0/shades)*shades;
				hatchingAmount = lerp(fetch(h1, p), fetch(h2, p), t);
			}
			else if (intense<= 3.0 / shades)
			{
				t = (intense-2.0/shades)*shades;
				hatchingAmount = lerp(fetch(h2, p), fetch(h3, p), t);
			}
			else if(intense <= 4.0/shades)
			{
				t = (intense-3.0/shades)*shades;
				hatchingAmount = fetch(h2, p) * lerp(255, fetch(h1, p2), t)/255.0;
			}
			else if (intense<= 5.0/shades)
			{
				t = (intense-4.0/shades)*shades;
				hatchingAmount = fetch(h2, p) * lerp(fetch(h1, p2), fetch(h2, p2), t)/255.0;
			}
			else if (intense<= 1)
			{
				t = (intense-5.0/shades)*shades;
				hatchingAmount = fetch(h2, p) * lerp(fetch(h2, p2), fetch(h3, p2), t)/255.0;
			}
    		
    		t = 1 - hatchingAmount/ 255.0;
			t = t * .6;
    			
    		
    		float e = imRef(edge, x, y);
    		
    		float edgeAmount = e;
    		if(edgeAmount > 200) edgeAmount = 200;
    		
    		float R = lerp(255, b.r, t) - edgeAmount;
    		float G = lerp(255, b.g, t) - edgeAmount;
    		float B = lerp(255, b.b, t) - edgeAmount;
    		
    		
    		
    		imRef(input, x, y) = toRGB(R, G, B);
    	}
	}
	
	
  
	savePPM(input, argv[2]);
	printf("Output saved to %s\n", argv[2]);
	return 0;
}
