/* 
 *  Copyright (c) 2008  Noah Snavely (snavely (at) cs.washington.edu)
 *    and the University of Washington
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 */

/* keys2.cpp */
/* Class for SIFT keypoints */

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

#include <fcntl.h>
#include <unistd.h>
#include <sys/mman.h>
#include <sys/stat.h>

#include <zlib.h>

#include "keys2a.h"

int GetNumberOfKeysNormal(FILE *fp)
{
    int num, len;

    if (fscanf(fp, "%d %d", &num, &len) != 2) {
	printf("Invalid keypoint file.\n");
	return 0;
    }

    return num;
}

int GetNumberOfKeysGzip(gzFile fp)
{
    int num, len;

    char header[256];
    gzgets(fp, header, 256);

    if (sscanf(header, "%d %d", &num, &len) != 2) {
	printf("Invalid keypoint file.\n");
	return 0;
    }

    return num;
}

/* Returns the number of keys in a file */
int GetNumberOfKeys(const char *filename)
{
    FILE *file;

    file = fopen (filename, "r");
    if (! file) {
        /* Try to file a gzipped keyfile */
        char buf[1024];
        sprintf(buf, "%s.gz", filename);
        gzFile gzf = gzopen(buf, "rb");

        if (gzf == NULL) {
            printf("Could not open file: %s\n", filename);
            return 0;
        } else {
            int n = GetNumberOfKeysGzip(gzf);
            gzclose(gzf);
            return n;
        }
    }
    
    int n = GetNumberOfKeysNormal(file);
    fclose(file);
    return n;
}

/* This reads a keypoint file from a given filename and returns the list
 * of keypoints. */
int ReadKeyFile(const char *filename, unsigned char **keys, keypt_t **info)
{
    FILE *file;

    file = fopen (filename, "r");
    if (! file) {
        /* Try to file a gzipped keyfile */
        char buf[1024];
        sprintf(buf, "%s.gz", filename);
        gzFile gzf = gzopen(buf, "rb");

        if (gzf == NULL) {
            printf("Could not open file: %s\n", filename);
            return 0;
        } else {
            int n = ReadKeysGzip(gzf, keys, info);
            gzclose(gzf);
            return n;
        }
    }
    
    int n = ReadKeys(file, keys, info);
    fclose(file);
    return n;

    // return ReadKeysMMAP(file);
}

#if 0
/* Read keys using MMAP to speed things up */
std::vector<Keypoint *> ReadKeysMMAP(FILE *fp) 
{    
    int i, j, num, len, val, n;

    std::vector<Keypoint *> kps;

    struct stat sb;

    /* Stat the file */
    if (fstat(fileno(fp), &sb) < 0) {
	printf("[ReadKeysMMAP] Error: could not stat file\n");
	return kps;
    }

    char *file = (char *)mmap(NULL, sb.st_size, PROT_READ, MAP_SHARED, 
			      fileno(fp), 0);

    char *file_start = file;

    if (sscanf(file, "%d %d%n", &num, &len, &n) != 2) {
	printf("[ReadKeysMMAP] Invalid keypoint file beginning.");
	return kps;
    }

    file += n;

    if (len != 128) {
	printf("[ReadKeysMMAP] Keypoint descriptor length invalid "
	       "(should be 128).");
	return kps;
    }

    for (i = 0; i < num; i++) {
	/* Allocate memory for the keypoint. */
	unsigned char *d = new unsigned char[len];
	float x, y, scale, ori;

	if (sscanf(file, "%f %f %f %f%n", &y, &x, &scale, &ori, &n) != 4) {
	    printf("[ReadKeysMMAP] Invalid keypoint file format.");
	    return kps;
	}

	file += n;

	for (j = 0; j < len; j++) {
	    if (sscanf(file, "%d%n", &val, &n) != 1 || val < 0 || val > 255) {
		printf("[ReadKeysMMAP] Invalid keypoint file value.");
		return kps;
	    }
	    d[j] = (unsigned char) val;
	    file += n;
	}

	kps.push_back(new Keypoint(x, y, scale, ori, d));
    }

    /* Unmap */
    if (munmap(file_start, sb.st_size) < 0) {
	printf("[ReadKeysMMAP] Error: could not unmap memory\n");
	return kps;
    }

    return kps;    
}
#endif

/* Read keypoints from the given file pointer and return the list of
 * keypoints.  The file format starts with 2 integers giving the total
 * number of keypoints and the size of descriptor vector for each
 * keypoint (currently assumed to be 128). Then each keypoint is
 * specified by 4 floating point numbers giving subpixel row and
 * column location, scale, and orientation (in radians from -PI to
 * PI).  Then the descriptor vector for each keypoint is given as a
 * list of integers in range [0,255]. */
int ReadKeys(FILE *fp, unsigned char **keys, keypt_t **info)
{
    int i, num, len;

    std::vector<Keypoint *> kps;

    if (fscanf(fp, "%d %d", &num, &len) != 2) {
	printf("Invalid keypoint file\n");
	return 0;
    }

    if (len != 128) {
	printf("Keypoint descriptor length invalid (should be 128).");
	return 0;
    }

    *keys = new unsigned char[128 * num];

    if (info != NULL) 
        *info = new keypt_t[num];

    unsigned char *p = *keys;
    for (i = 0; i < num; i++) {
	/* Allocate memory for the keypoint. */
	// short int *d = new short int[128];
	float x, y, scale, ori;

	if (fscanf(fp, "%f %f %f %f\n", &y, &x, &scale, &ori) != 4) {
	    printf("Invalid keypoint file format.");
	    return 0;
	}

        if (info != NULL) {
            (*info)[i].x = x;
            (*info)[i].y = y;
            (*info)[i].scale = scale;
            (*info)[i].orient = ori;
        }
        
	char buf[1024];
	for (int line = 0; line < 7; line++) {
	    fgets(buf, 1024, fp);

	    if (line < 6) {
		sscanf(buf, 
		       "%hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu "
		       "%hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu", 
		       p+0, p+1, p+2, p+3, p+4, p+5, p+6, p+7, p+8, p+9, 
		       p+10, p+11, p+12, p+13, p+14, 
		       p+15, p+16, p+17, p+18, p+19);

		p += 20;
	    } else {
		sscanf(buf, 
		       "%hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu",
		       p+0, p+1, p+2, p+3, p+4, p+5, p+6, p+7);
		p += 8;
	    }
	}
    }

    return num; // kps;
}

int ReadKeysGzip(gzFile fp, unsigned char **keys, keypt_t **info)
{
    int i, num, len;

    std::vector<Keypoint *> kps;
    char header[256];
    gzgets(fp, header, 256);

    if (sscanf(header, "%d %d", &num, &len) != 2) {
	printf("Invalid keypoint file.\n");
	return 0;
    }

    if (len != 128) {
	printf("Keypoint descriptor length invalid (should be 128).");
	return 0;
    }

    *keys = new unsigned char[128 * num];

    if (info != NULL) 
        *info = new keypt_t[num];

    unsigned char *p = *keys;
    for (i = 0; i < num; i++) {
	/* Allocate memory for the keypoint. */
	// short int *d = new short int[128];
	float x, y, scale, ori;
        char buf[1024];
        gzgets(fp, buf, 1024);

	if (sscanf(buf, "%f %f %f %f\n", &y, &x, &scale, &ori) != 4) {
	    printf("Invalid keypoint file format.");
	    return 0;
	}

        if (info != NULL) {
            (*info)[i].x = x;
            (*info)[i].y = y;
            (*info)[i].scale = scale;
            (*info)[i].orient = ori;
        }
        
	for (int line = 0; line < 7; line++) {
	    gzgets(fp, buf, 1024);

	    if (line < 6) {
		sscanf(buf, 
		       "%hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu "
		       "%hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu", 
		       p+0, p+1, p+2, p+3, p+4, p+5, p+6, p+7, p+8, p+9, 
		       p+10, p+11, p+12, p+13, p+14, 
		       p+15, p+16, p+17, p+18, p+19);

		p += 20;
	    } else {
		sscanf(buf, 
		       "%hhu %hhu %hhu %hhu %hhu %hhu %hhu %hhu",
		       p+0, p+1, p+2, p+3, p+4, p+5, p+6, p+7);
		p += 8;
	    }
	}
    }

    return num; // kps;
}

/* Create a search tree for the given set of keypoints */
ANNkd_tree *CreateSearchTree(int num_keys, unsigned char *keys)
{
    // clock_t start = clock();

    /* Create a new array of points */
    ANNpointArray pts = annAllocPts(num_keys, 128);

    for (int i = 0; i < num_keys; i++) {
	memcpy(pts[i], keys + 128 * i, sizeof(unsigned char) * 128);
    }

    /* Create a search tree for k2 */
    ANNkd_tree *tree = new ANNkd_tree(pts, num_keys, 128, 16);
    // clock_t end = clock();

    // printf("Building tree took %0.3fs\n", 
    //        (end - start) / ((double) CLOCKS_PER_SEC));    

    return tree;
}

std::vector<KeypointMatch> MatchKeys(int num_keys1, unsigned char *k1, 
                                     ANNkd_tree *tree2,
				     double ratio, int max_pts_visit)
{
    annMaxPtsVisit(max_pts_visit);
    std::vector<KeypointMatch> matches;

    /* Now do the search */
    // clock_t start = clock();
    for (int i = 0; i < num_keys1; i++) {
	ANNidx nn_idx[2];
	ANNdist dist[2];

	tree2->annkPriSearch(k1 + 128 * i, 2, nn_idx, dist, 0.0);

	if (((double) dist[0]) < ratio * ratio * ((double) dist[1])) {
	    matches.push_back(KeypointMatch(i, nn_idx[0]));
	}
    }
    // clock_t end = clock();

    // printf("Searching tree took %0.3fs\n", 
    //        (end - start) / ((double) CLOCKS_PER_SEC));

    return matches;    
}

/* Compute likely matches between two sets of keypoints */
std::vector<KeypointMatch> MatchKeys(int num_keys1, unsigned char *k1, 
				     int num_keys2, unsigned char *k2, 
				     double ratio, int max_pts_visit) 
{
    annMaxPtsVisit(max_pts_visit);

    int num_pts = 0;
    std::vector<KeypointMatch> matches;

    num_pts = num_keys2;
    clock_t start = clock();

    /* Create a new array of points */
    ANNpointArray pts = annAllocPts(num_pts, 128);

    for (int i = 0; i < num_pts; i++) {
	memcpy(pts[i], k2 + 128 * i, sizeof(unsigned char) * 128);
    }

    /* Create a search tree for k2 */
    ANNkd_tree *tree = new ANNkd_tree(pts, num_pts, 128, 16);
    clock_t end = clock();
    
    // printf("Building tree took %0.3fs\n", 
    //	      (end - start) / ((double) CLOCKS_PER_SEC));

    /* Now do the search */
    start = clock();
    for (int i = 0; i < num_keys1; i++) {
	ANNidx nn_idx[2];
	ANNdist dist[2];

	tree->annkPriSearch(k1 + 128 * i, 2, nn_idx, dist, 0.0);

	if (((double) dist[0]) < ratio * ratio * ((double) dist[1])) {
	    matches.push_back(KeypointMatch(i, nn_idx[0]));
	}
    }
    end = clock();
    // printf("Searching tree took %0.3fs\n", 
    //        (end - start) / ((double) CLOCKS_PER_SEC));

    /* Cleanup */
    annDeallocPts(pts);
    // annDeallocPt(axis_weights);

    delete tree;

    return matches;
}
