/* 
 *  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.
 *
 */

/* keys.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>

#ifndef WIN32_MVC
#include <ext/hash_map>
#else
#include <hash_map>
#endif

#include "keys.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. */
std::vector<Keypoint> ReadKeyFile(const char *filename, bool descriptor)
{
    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) {
            std::vector<Keypoint> empty;
            printf("Could not open file: %s\n", filename);
            return empty;
        } else {
            std::vector<Keypoint> kps = ReadKeysFastGzip(gzf, descriptor);
            gzclose(gzf);
            return kps;
        }
    }
    
    std::vector<Keypoint> kps = ReadKeysFast(file, descriptor);

    fclose(file);
    
    return kps;
}

/* This reads a keypoint file from a given filename and returns the list
 * of keypoints. */
std::vector<KeypointWithScaleRot> 
    ReadKeyFileWithScaleRot(const char *filename, bool descriptor)
{
    FILE *file;
    std::vector<Keypoint> kps;
    float *scale = NULL, *orient = NULL;

    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) {
            std::vector<KeypointWithScaleRot> empty;
            printf("Could not open file: %s\n", filename);
            return empty;
        } else {
            kps = ReadKeysFastGzip(gzf, descriptor, &scale, &orient);
            gzclose(gzf);
        }
    } else {
        kps = ReadKeysFast(file, descriptor, &scale, &orient);
        fclose(file);
    }

    std::vector<KeypointWithScaleRot> kps_w;
    int num_keys = (int) kps.size();
    kps_w.resize(num_keys);
    for (int i = 0; i < num_keys; i++) {
        kps_w[i].m_x = kps[i].m_x;
        kps_w[i].m_y = kps[i].m_y;
        kps_w[i].m_d = kps[i].m_d;
        kps_w[i].m_scale = scale[i];
        kps_w[i].m_orient = orient[i];
    }

    kps.clear();

    if (scale != NULL)
        delete [] scale;

    if (scale != NULL)
        delete [] orient;

    return kps_w;
}

static char *strchrn(char *str, int c, int n) {
    for (int i = 0; i < n; i++) {
	str = strchr(str, c) + 1;
	if (str == NULL) return NULL;
    }

    return str - 1;
}

#if 1
/* 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;

    char string_buf[1024];
    char *str = string_buf;

    /* Find the first '\n' */
    char *newline = strchr(file, '\n');
    int pos = (int) (newline - file);
	
    memcpy(str, file, pos);
    str[pos] = 0;

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

    file += (pos + 1);

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

    for (i = 0; i < num; i++) {
	str = string_buf;

	/* Allocate memory for the keypoint. */
	unsigned char *d = new unsigned char[len];
	float x, y, scale, ori;

	/* Find the first '\n' */
	newline = strchr(file, '\n');
	pos = (int) (newline - file);
	
	memcpy(str, file, pos);
	str[pos] = 0;

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

	file += (pos + 1);

	/* Find the next seven '\n's */
	str = string_buf;

	char *seventh_newline = strchrn(file, '\n', 7);
	pos = (int) (seventh_newline - file);

	memcpy(str, file, pos);
	str[pos] = 0;

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

	file += (pos + 1);

	kps.push_back(Keypoint(x, y, 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]. */
std::vector<Keypoint> ReadKeys(FILE *fp, bool descriptor)
{
    int i, j, num, len, val;

    std::vector<Keypoint> kps;

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

    if (len != 128) {
	printf("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 (fscanf(fp, "%f %f %f %f", &y, &x, &scale, &ori) != 4) {
	    printf("Invalid keypoint file format.");
	    return kps;
	}

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

	if (descriptor) {
	    kps.push_back(Keypoint(x, y, d));
	} else {
	    delete [] d;
	    kps.push_back(Keypoint(x, y, NULL));
	}
    }

    return kps;
}

/* Read keys more quickly */
std::vector<Keypoint> ReadKeysFast(FILE *fp, bool descriptor,
                                   float **scales, float **orients)
{
    int i, j, num, len;

    std::vector<Keypoint> kps;

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

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

    if (num > 0 && scales != NULL) {
        *scales = new float[num];
    }

    if (num > 0 && orients != NULL) {
        *orients = new float[num];
    }

    for (i = 0; i < num; i++) {
	/* Allocate memory for the keypoint. */
	float x, y, scale, ori;

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

        if (scales != NULL) {
            (*scales)[i] = scale;
        }
        
        if (orients != NULL) {
            (*orients)[i] = ori;
        }

	char buf[1024];

	/* Allocate memory for the keypoint. */
	unsigned char *d = NULL;
	
	if (descriptor)
	    d = new unsigned char[len];

	int start = 0;
	for (int line = 0; line < 7; line++) {
	    fgets(buf, 1024, fp);

	    if (!descriptor) continue;

	    short int p[20];

	    if (line < 6) {
		sscanf(buf, 
		       "%hu %hu %hu %hu %hu %hu %hu %hu %hu %hu "
		       "%hu %hu %hu %hu %hu %hu %hu %hu %hu %hu", 
		       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);

		for (j = 0; j < 20; j++)
		    d[start + j] = p[j];

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

		for (j = 0; j < 8; j++)
		    d[start + j] = p[j];
	    }
	}

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

    return kps;
}

std::vector<Keypoint> ReadKeysFastGzip(gzFile fp, bool descriptor,
                                       float **scales, float **orients)
{
    int i, j, 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 kps;
    }

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

    if (num > 0 && scales != NULL) {
        *scales = new float[num];
    }

    if (num > 0 && orients != NULL) {
        *orients = new float[num];
    }

    for (i = 0; i < num; i++) {
	/* Allocate memory for the keypoint. */
	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 kps;
	}

        if (scales != NULL) {
            (*scales)[i] = scale;
        }
        
        if (orients != NULL) {
            (*orients)[i] = ori;
        }

	/* Allocate memory for the keypoint. */
	unsigned char *d = NULL;
	
	if (descriptor)
	    d = new unsigned char[len];

	int start = 0;
	for (int line = 0; line < 7; line++) {
	    gzgets(fp, buf, 1024);

	    if (!descriptor) continue;

	    short int p[20];

	    if (line < 6) {
		sscanf(buf, 
		       "%hu %hu %hu %hu %hu %hu %hu %hu %hu %hu "
		       "%hu %hu %hu %hu %hu %hu %hu %hu %hu %hu", 
		       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);

		for (j = 0; j < 20; j++)
		    d[start + j] = p[j];

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

		for (j = 0; j < 8; j++)
		    d[start + j] = p[j];
	    }
	}

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

    return kps;
}


