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

#include "util.h"

/* This reads Msieve's matrix file format, and writes Msieve's
dependency file format. Some of the I/O code is adapted from Msieve,
due of course to Jason Papadopoulos.

The method of solving the system is a recent innovation discovered by
Prasad Raghavendra. You can read about his method and see his paper here:
http://rjlipton.wordpress.com/2012/08/09/a-new-way-to-solve-linear-equations/
*/

uint32 seed1, seed2; // used for random number generation

typedef struct {
	uint32* cols;	/* The array of occupied columns in this row */
	uint32 weight;	/* Length of data */
	uint32 size; 	/* Mem alloc'd (should always be >= weight) */
} sparse_row_t;

typedef struct {
	uint32 nrows, ncols, n_dense_rows, n_sparse_rows;
	sparse_row_t* row_list; // A list of the sparse rows
	uint32** dense_row_list, dense_row_length;
	/* Store the dense rows as bit arrays each (ncols+31)/32 words long */
} matrix_t;

void sparse_row_add_entry(sparse_row_t* row, uint32 col) {
// We have to be careful since sparse_row_t.cols is built on the fly
	if(row->weight > row->size) {
		fprintf(stderr, "...Uh, Houston, we have a problem. "
		"There's more data than array holding the data.\n");
		exit(-1);
	} else if(row->weight == row->size)
		row->cols = xrealloc(row->cols, (size_t)((row->size += 5) * sizeof(uint32)));
	row->cols[row->weight++] = col;
}

void matrix_init(uint32 n_sparse_rows, uint32 n_dense_rows, uint32 dense_row_length,
		sparse_row_t** sparse_out, uint32*** dense_out)
{
	uint32 i;
	sparse_row_t* row_list;
	uint32** dense_row_list;

	printf("calloc sparse row list\n");
	row_list = (sparse_row_t*)xcalloc(n_sparse_rows, sizeof(sparse_row_t));
	// A list of the sparse rows
	for(i = 0; i < n_sparse_rows; i++) {
		//printf("calloc one sparse row\n");
		row_list[i].cols = (uint32*)xcalloc(20, sizeof(uint32));
		row_list[i].size = 20;
		row_list[i].weight = 0;
	}

	printf("calloc dense list\n");
	dense_row_list = (uint32**)xcalloc(n_dense_rows, sizeof(uint32*));
	/* Store the dense rows as bit arrays each (ncols+31)/32 words long */
	for(i = 0; i < n_dense_rows; i++) {
		//printf("calloc one dense row\n");
		dense_row_list[i] = (uint32*)xcalloc(dense_row_length, sizeof(uint32));
	}

	*sparse_out = row_list;
	*dense_out = dense_row_list;
}

matrix_t* read_matrix(char* file) {
	uint32 i, j, k;
	uint32 nrows, ncols, n_dense_rows, n_sparse_rows, col_dense_words,
		dense_row_length;
	FILE* f;
	sparse_row_t* row_list;
	uint32** dense_row_list;
	matrix_t* matrix = (matrix_t*)xcalloc(1, sizeof(matrix_t));

	printf("Attempting to read matrix from %s\n", file);

	if((f = fopen(file, "rb")) == NULL) {
		fprintf(stderr, "Couldn't open matrix file\n");
		exit(-1);
	}

	fread(&nrows, sizeof(uint32), (size_t)1, f);
	fread(&n_dense_rows, sizeof(uint32), (size_t)1, f);
	fread(&ncols, sizeof(uint32), (size_t)1, f);
	col_dense_words = (n_dense_rows + 31) / 32;
	n_sparse_rows = nrows - n_dense_rows;
	dense_row_length = (ncols + 31) / 32;
	printf("Matrix is %u x %u with %u dense rows\n", nrows, ncols, n_dense_rows);

	matrix_init(n_sparse_rows, n_dense_rows, dense_row_length,
		&row_list, &dense_row_list);

	for(i = 0; i < ncols; i++) {
		uint32 sparse_count, word, d_row; // word/d_row used in second loop

		fread(&sparse_count, sizeof(uint32), (size_t)1, f);

		for(j = 0; j < sparse_count; j++) {
		// read in the sparse row data
			uint32 row_entry;
			sparse_row_t* row;

			fread(&row_entry, sizeof(uint32), (size_t)1, f);
			// Read in which row has an entry in this col

			row = row_list + (row_entry - n_dense_rows);
			// Find the data for that row (as a pointer)

			sparse_row_add_entry(row, i);
			// Add this col (i) to that row's data
		}

		for(j = 0; j < col_dense_words - 1; j++) {
		// now read in the dense row data
		// do the last word separately
			fread(&word, sizeof(uint32), (size_t)1, f);

			for(k = 0; k < 32; k++) {
				//uint32* row = dense_row_list[32*j + k];
				//uint32 row_word = i / 32; // i is the column...
				//uint32 word_bit = i % 32;
				//row[row_word] |= ((word >> k) & 1) << word_bit;
				dense_row_list[(j << 5) + k][i >> 5] |= (word & 1) << (i & 0x1f);
				word >>= 1;
				/* Take the kth bit of the read-in word (which
				is the 32*j+k'th row's entry in this col) and set
				it as the i'th bit of the row's bitfield...
				I would really appreciate if someone double
				checked this... */
			}
		}

		fread(&word, sizeof(uint32), (size_t)1, f);
		d_row = n_dense_rows & 0x1f; if(d_row == 0) d_row = 32;
		for(k = 0; k < d_row; k++) {
			//uint32* row = dense_row_list[32*j + k];
			//uint32 row_word = i / 32; // i is the column...
			//uint32 word_bit = i % 32;
			//row[row_word] |= (word & 1) << word_bit;
			dense_row_list[(j << 5) + k][i >> 5] |= (word & 1) << (i & 0x1f);
			word >>= 1;
		}
	}
	fclose(f);
	matrix->nrows = nrows; // too lazy to redo the above code
	matrix->ncols = ncols;
	matrix->n_dense_rows = n_dense_rows;
	matrix->n_sparse_rows = n_sparse_rows;
	matrix->dense_row_length = dense_row_length;
	matrix->row_list = row_list;
	matrix->dense_row_list = dense_row_list;
	return matrix;
}

typedef struct {
	uint32* coeffs; /* bit field */
	uint32 len;
} vector_t;

typedef struct {
	vector_t* set;
	uint32 len, len_a;
} set_t;

vector_t* vector_init(uint32 n, uint32 rand_fill) {
	vector_t* vec = (vector_t*)xcalloc(1, sizeof(vector_t));
	uint32 i;
	vec->len = (n + 31) / 32;
	//printf("calloc one vector (rand: %u)\n", rand_fill);
	vec->coeffs = (uint32*)xcalloc(vec->len, sizeof(uint32));
	if(rand_fill)
		for(i = 0; i < vec->len; i++)
			vec->coeffs[i] = get_rand(&seed1, &seed2);
	return vec;
}

void vector_free(vector_t* v) {
	free(v->coeffs);
}

set_t* set_init(uint32 N) {
	set_t* set = (set_t*)xcalloc(1, sizeof(set_t));
	set->len = 0;
	set->len_a = 1000;
	printf("calloc a set\n");
	set->set = (vector_t*)xcalloc(set->len_a, sizeof(vector_t));
	return set;
}

void set_free(set_t* set) {
	vector_t* v;
	for(v = set->set; v < set->set+set->len; v++) vector_free(v);
	free(set->set);
}

void add_vec_to_set(vector_t* v, set_t* s) {
/* We're supposed to check for duplicates, however that would be an incredibly
monstrous task, and given the size of these vectors, I think the chance of there
being a duplicate is *incredibly* low, like N*2^(-n) with n > 1e5. */
	if(s->len > s->len_a) {
		fprintf(stderr, "Mem alloc error (adding vector to set)\n");
		exit(-1);
	} else if(s->len == s->len_a) {
		uint32 old = s->len_a;
		s->len_a += 8000;
		printf("Reallocating a set from %lu KB to %lu KB\n", (old*sizeof(vector_t))/1000, (s->len_a*sizeof(vector_t))/1000);
		s->set = (vector_t*)xrealloc(s->set, (size_t)((s->len_a) * sizeof(vector_t)));
	}
	s->set[s->len++] = *v;
}

uint32 dense_satisfies(uint32* dense_row, uint32 row_length, vector_t* vec) {
/* The bit parity method comes from (and is explained in) this awesome page:
http://graphics.stanford.edu/~seander/bithacks.html#ParityParallel */
	uint32 i; uint8 out = 0;
	//printf("Testing a dense row... row length: %u vec length: %u\n",
	//	row_length, vec->len);
	for(i = 0; i < row_length; i++) { // dot product the two bit fields
		uint32 word = dense_row[i] & vec->coeffs[i];
		// count the 1s in the word, see comment above
		word ^= word >> 16;
		word ^= word >> 8;
		word ^= word >> 4;
		word &= 0xf;
		out ^= (0x6996 >> word) & 1;
	}
	return (uint32)!out;
}

uint32 sparse_satisfies(sparse_row_t* row, vector_t* vec) {
	uint32 i; uint8 out = 0;
	for(i = 0; i < row->weight; i++) {
		uint32 bit = row->cols[i];
		out ^= vec->coeffs[bit >> 5] >> (bit & 0xff);
	}
	return (uint32)!out;
}

set_t* recombine(set_t* T, uint32 N, uint32 n) {
	printf("Recombination: T.len is %u, ratio good vectors = %f\n", T->len, (double)T->len/N);
	uint32 i, j;
	vector_t* vec, *v1, *v2, *v3;
	set_t* S = set_init(N);
	for(i = 0; i < N; i++) {
		v1 = T->set + (get_rand(&seed1, &seed2) % T->len);
		v2 = T->set + (seed2 % T->len);
		v3 = T->set + (get_rand(&seed1, &seed2) % T->len);
		vec = vector_init(n, 0);
		for(j = 0; j < vec->len; j++) // sum the three vectors
			vec->coeffs[j] = v1->coeffs[j] ^ v2->coeffs[j] ^ v3->coeffs[j];
		add_vec_to_set(vec, S);
	}
	return S;
}

set_t Raghavendra(matrix_t* mat) {
	/* The mem use is rather high, but... (m rows, n columns)*/
	matrix_t m = *mat;
	uint32 N = (uint32)(0.45 * m.ncols), i, j;
	set_t* S = set_init(N);
	printf("Matrix is %u x %u with %u dense rows\n", m.nrows, m.ncols, m.n_dense_rows);
	/* The paper chooses N = 145q^2*ln(q)*n; 145*4*ln(2)~402.025 */
	printf("Estimated memory for the initial set is %u MB (N is %u, words per vector is %u)\n", (unsigned int)(N/(double)1000000 * 4 * (m.ncols+31)/(double)32), N, (m.ncols+31)/32);
	for(i = 0; i < N; i++)
		add_vec_to_set(vector_init(m.ncols, 1), S);
	printf("Initial set constructed\n");
	for(i = 0; i < m.n_sparse_rows; i++) { // for each sparse row...
		/*if((i % 100) == 0)*/ printf("Doing the %uth sparse row...\n", i);
		set_t* T = set_init(N);
		printf("New T init'd, looping on solutions\n");
		for(j = 0; j < S->len; j++) {
			if(sparse_satisfies(m.row_list+i, &(S->set[j])))
				add_vec_to_set(&(S->set[j]), T);
		}
		set_free(S);
		S = recombine(T, N, m.ncols);
		set_free(T);
		free(m.row_list[i].cols);
		printf("Recomb complete; memory should have just dropped\n");
	}

	for(i = 0; i < m.n_dense_rows; i++) { // for each dense row...
		/*if((i % 10) == 0)*/ printf("Doing the %uth dense row...\n", i);
		set_t* T = set_init(N);
		for(j = 0; j < S->len; j++)
			if(dense_satisfies(m.dense_row_list[i], m.dense_row_length,
								&(S->set[j])))
				add_vec_to_set(&(S->set[j]), T);
		set_free(S); //clear the old S, then recombine T into S
		S = recombine(T, N, m.ncols);
		set_free(T);
		free(m.dense_row_list[i]);
		printf("Recomb complete; memory should have just dropped\n");
	}
	free(m.dense_row_list);

	printf("Found %u dependencies (writing at most 64)\n", S->len);
	return *S;
}

void write_deps(char* file, set_t S, uint32 n) {
	uint64* cols = (uint64*)xcalloc(n, sizeof(uint64));
	uint32 i, j, k, x, y;
	FILE* f;
	if(S.len > 64) S.len = 64;
	for(i = 0; i < S.len; i++) { // for each dep...
		vector_t v = S.set[i];
		for(j = 0; j < v.len - 1; j++) { // for each word in the dep's bit array
			uint32 w = v.coeffs[j];  // (except the last)
			for(k = 0; k < 32; k++) { // for each bit in the word
				if(w & 1)
					cols[(j << 5) + k] |= 1 << i;
				w >>= 1;
			}
		}
		x = v.coeffs[j]; // the last word with < 32 bits
		y = n % 32; y = y==0?32:y;
		for(k = 0; k < y; k++) {
			if(x & 1)
				cols[(j << 5) + k] |= 1 << i;
			x >>= 1;
		}
	}

	if((f = fopen(file, "wb")) == NULL) {
		fprintf(stderr, "Well this is awkward... the solution's all wrapped"
			" up but we can't deliver it...\n");
		exit(-1);
	}

	fwrite(cols, sizeof(uint64), n, f);
	fclose(f);
}

int main(int argc, char** argv) {
	matrix_t* mat;
	set_t solutions;
	printf("\nWelcome to derptown!\n");

	if(argc < 3) {
		fprintf(stderr, "Usage: %s Msieve_matrix_file dep_file", argv[0]);
		exit(-1);
	}
	get_random_seeds(&seed1, &seed2);

	mat = read_matrix(argv[1]);

	solutions = Raghavendra(mat);

	write_deps(argv[2], solutions, mat->ncols);

	return 0;
}