#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <limits.h>

#include <opencv2/imgproc/imgproc_c.h>

#include <opencv2/highgui/highgui_c.h>

#include "gck.h"

#include "kdtree.h"

#define XY_TO_INT(x, y) (((y) << 16) | (x))

#define PACK_SCOREIDX(s, i) ((uint64_t)((uint64_t)s << 32 | (i)))

#define UNPACK_SCORE(a) ((a) >> 32)

#define UNPACK_IDX(a) ((a)&(0xFFFFFFFF))

static int64_t match_score(int *coeffs, kd_node *n, int k)

{

int i, j, *u, *v, **p = n->value, best = INT_MAX, idx = -1;

for (i = 0; i < n->nb; i++) {

int dist = 0;

u = coeffs;

v = *p++;

for (j = 0; j < k; j++) {

int a = *u++;

int b = *v++;

dist += (a - b)*(a - b);

}

if (dist < best) {

best = dist;

idx = i;

}

}

return PACK_SCOREIDX(best, idx);

}

static int patch_score(int *c1, int *c2, int k)

{

int i, dist = 0;

for (i = 0; i < k; i++) {

int a = *c1++;

int b = *c2++;

dist += (a-b)*(a-b);

}

return dist;

}

static void swap2(int *scores, int *index)

{

int t = scores[0];

scores[0] = scores[1];

scores[1] = t;

int u = index[0];

index[0] = index[1];

index[1] = u;

}

static inline void check_guide(kd_tree *t, int *coeffs, int off,

int *scores, int *pos)

{

if (t->start + off >= t->end) return;

int *points = t->start + off;

int attempt = patch_score(coeffs, points, t->k);

if (attempt < scores[0]) {

pos[0] = off;

scores[0] = attempt;

if (scores[0] < scores[1]) swap2(scores, pos);

}

}

static unsigned match_enrich(kd_tree *t, int *coeffs, int x, int y,

int *prev)

{

int k = t->k, *start = t->start;

kd_node *n = kdt_query(t, coeffs);

// set results of query

int64_t res = match_score(coeffs, n, k);

int best[] = {INT_MAX, UNPACK_SCORE(res)};

int pos[] = {INT_MAX, n->value[UNPACK_IDX(res)] - start};

pos[0] = pos[1]; // hack for (x,y) == (0,0)

// now check 2 best matches for the left

if (x) {

check_guide(t, coeffs, prev[-1], best, pos);

check_guide(t, coeffs, prev[-2], best, pos);

}

// check 2 best matches for top

if (y) {

check_guide(t, coeffs, prev[0], best, pos);

check_guide(t, coeffs, prev[1], best, pos);

}

// set prev to best matches

prev[0] = pos[0];

prev[1] = pos[1];

return pos[1];

}

static IplImage* match(kd_tree *t, int *coeffs, IplImage *src,

CvSize dst_size)

{

IplImage *xy = cvCreateImage(dst_size, IPL_DEPTH_32S, 1);

int w = dst_size.width - 8 + 1, h = dst_size.height - 8 + 1;

int k = t->k, sz = w*h, i, sw = src->width - 8 + 1;

int *prevs = malloc(w*sizeof(int)*2), *prev = prevs;

int *xydata = (int*)xy->imageData;

for (i = 0; i < sz; i++) {

int x = i % w, y = i/w, sx, sy, sxy;

if (!x) {

prev = prevs;

xydata = (int*)xy->imageData + y*(xy->widthStep/sizeof(int));

}

sxy = match_enrich(t, coeffs, x, y, prev) / k;

sx = sxy % sw; sy = sxy / sw;

*xydata++ = XY_TO_INT(sx, sy);

if (sx >= src->width || sy >= src->height) {

printf("grievous error: got %d,%d but dims %d,%d sxy %d\n", sx, sy, src->width, src->height, sxy);

}

coeffs += k;

prev += 2;

}

free(prevs);

return xy;

}

static void interleave_data(int *data, int w, int h, int kern_size,

int bases, int *a, int aw)

{

// takes XXXAAAABBBBCCCCXXX -> ABCABCABCABC

// eg, remove the results that incorporate padding and interleave

int kw = w + kern_size - 1, i, j, k, kh = h + kern_size - 1;

int rw = w - kern_size + 1, rh = h - kern_size + 1;

int *src = data;

for (k = 0; k < bases; k++) {

src = data + kw*kh*k;

src += kw * (kern_size - 1) + kern_size - 1;

for (i = 0; i < rh; i++) {

int *s = src;

for (j = 0; j < rw; j++) {

int n = i*rw+j;

a[n*aw + k] = *s++;

}

src += kw;

}

}

}

static void coeffs_i(IplImage *img, int bases, int total_b, int *data)

{

CvSize s = cvGetSize(img);

int w = s.width, h = s.height;

if (w != img->widthStep) {

fprintf(stderr, "image not aligned uh oh\n");

exit(1);

}

int *res = gck_calc_2d((uint8_t*)img->imageData, w, h, 8, bases);

interleave_data(res, w, h, 8, bases, data, total_b);

free(res);

}

static void coeffs(IplImage *img, int dim, int *pc, int **in) {

CvSize size = cvGetSize(img);

IplImage *lab = cvCreateImage(size, IPL_DEPTH_8U, 3);

IplImage *l = cvCreateImage(size, IPL_DEPTH_8U, 1);

IplImage *a = cvCreateImage(size, IPL_DEPTH_8U, 1);

IplImage *b = cvCreateImage(size, IPL_DEPTH_8U, 1);

int w = size.width, h = size.height;

int *interleaved = gck_alloc_buffer(w, h, 8, dim);

cvCvtColor(img, lab, CV_BGR2YCrCb);

cvSplit(img, l, a, b, NULL);

coeffs_i(l, pc[0], dim, interleaved);

coeffs_i(a, pc[1], dim, interleaved+pc[0]);

coeffs_i(b, pc[2], dim, interleaved+pc[0]+pc[1]);

cvReleaseImage(&lab);

cvReleaseImage(&l);

cvReleaseImage(&a);

cvReleaseImage(&b);

*in = interleaved;

}

IplImage* prop_match(IplImage *src, IplImage *dst)

{

int *srcdata, *dstdata;

CvSize src_size = cvGetSize(src), dst_size = cvGetSize(dst);

int w1 = src_size.width - 8 + 1, h1 = src_size.height - 8 + 1;

int sz = w1*h1, plane_coeffs[] = {2, 9, 5};

int dim = plane_coeffs[0] + plane_coeffs[1] + plane_coeffs[2];

kd_tree kdt;

IplImage *matched;

coeffs(src, dim, plane_coeffs, &srcdata);

coeffs(dst, dim, plane_coeffs, &dstdata);

memset(&kdt, 0, sizeof(kdt));

kdt_new(&kdt, srcdata, sz, dim);

matched = match(&kdt, dstdata, src, dst_size);

free(srcdata);

free(dstdata);

kdt_free(&kdt);

return matched;

}

void prop_coeffs(IplImage *src, int *plane_coeffs, int **data)

{

int dim = plane_coeffs[0] + plane_coeffs[1] + plane_coeffs[2];

return coeffs(src, dim, plane_coeffs, data);

}

IplImage *prop_match_complete(kd_tree *kdt, int *data, IplImage *src,

CvSize dst_size)

{

return match(kdt, data, src, dst_size);

}

unsigned prop_enrich(kd_tree *t, int *coeffs, int x, int y, int *prev)

{

return match_enrich(t, coeffs, x, y, prev);

}