/* SPDX-License-Identifier: GPL-2.0-only */
#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <stdbool.h>
#include <gmp.h>
#include <time.h>

#define square( x ) ( x * x )

// If endlessly looping, this is how many loops are made between allocing for a
//  bigger array
#define DIMENSION_PER_ALLOC 4096

// Define this for debugging.
//#define DEBUG

// Define this for showing time for each dimension.
// Not useful for small inputs.
#define TIME

// Which clock to use for measuring time used to calculate.
// CLOCK_PROCESS_CPUTIME_ID to counts CPU time of the current process.
// This is preferred as this makes optimisation easier
#define CLOCK CLOCK_MONOTONIC

// Returns  1 if the root of length2 is larger than the length of the vector.
// Returns  0 if they're equal.
// Returns -1 if the root of length2 is shorter than the length of the vector.
static int vector_length2_cmp(
	const int length2,
	const int *const vector,
	const int dimension
){
	int vector_length2 = 0;
	for( int i=0; i < dimension && vector[i] != 0; i++ ){
		vector_length2 += square( vector[i] );
		if( vector_length2 > length2 ) return -1;
	}
	if( vector_length2 == length2 )
		return 0;
	else
		return 1;
}

// Amount of unique orderings given a sorted array, including flipping of signs.
// A table of factorial values is used to not need to compute the same stuff
//  over and over again.
static void unique_array_orders(
	mpz_t rop,
	const int *const intptr,
	const int length,
	mpz_t *const factorial_table
){
	// rop = 1
	mpz_set_ui( rop, 1 );
	int n = 1;
	int zeros = 0;
	// Only one order foe all 0
	if( intptr[0] == 0 ) return;
	for( int i = 1; i < length; i++ ){
		if( intptr[i-1] == intptr[i] ){
			n++;
		}else{
			// rop = rop * n!
			mpz_mul( rop, rop, factorial_table[n] );
			n = 1;
			if( intptr[i] == 0 ){
				zeros = length - i;
				n = zeros;
				break;
			}
		}
	}
	if( n != 1 ){
		// rop = rop * n!
		mpz_mul( rop, rop, factorial_table[n] );
	}
	// rop = 2^( length - zeros ) * ( length! / rop )
	mpz_divexact( rop, factorial_table[length], rop );
	if(zeros != length)
		mpz_mul_2exp( rop, rop, length - zeros );
}

#ifdef TIME
static void print_elapsed_time(
	const struct timespec t_start,
	const struct timespec t_finish
){
	const int nsec_elapsed =
		  ( t_finish.tv_nsec - t_start.tv_nsec )
		+ ( t_finish.tv_nsec < t_start.tv_nsec ) * 1000000000;
	printf( "Time elapsed: " );
	if( nsec_elapsed < 1000000 ){
		printf( "%7.3fus\n", nsec_elapsed / 1000.0 );
		return;
	}
	int seconds_elapsed =
		  ( t_finish.tv_sec  - t_start.tv_sec  )
		- ( t_finish.tv_nsec < t_start.tv_nsec );
	if( seconds_elapsed < 1 ){
		printf( "%7.3fms\n",
			nsec_elapsed / 1000000.0 );
		return;
	}
	int msec_elapsed = nsec_elapsed / 1000000;
	if( seconds_elapsed < 60 ){
		printf( "%2d.%03ds\n", seconds_elapsed, msec_elapsed );
		return;
	}
	int minutes_elapsed = seconds_elapsed / 60;
	    seconds_elapsed = seconds_elapsed % 60;
	if( minutes_elapsed < 60 ){
		printf( "%2dm %2d.%03ds\n", minutes_elapsed, seconds_elapsed, msec_elapsed );
		return;
	}
	int hours_elapsed = minutes_elapsed / 60;
	  minutes_elapsed = minutes_elapsed % 60;
	printf( "%2dh %2dm %2ds\n", hours_elapsed, minutes_elapsed, seconds_elapsed );
	// Not bothering making more than this,
	//  though I wonder how big the output would get
}
#endif /* TIME */

// Less slow algorithm
int main( const int argc, const char *argv[] ){
	if( argc > 4 || argc < 2 ){
		fprintf( stderr,
"Usage: \n\t%s radius [dimension] [dimension_end]\n"
"Radius is needed, dimension is optional.\n"
"If no dimension given, this program will loop starting from 1 dimension and\n"
" increment the dimension each loop until the program is termnated.\n"
"If dimension_end is given, it loops from dimension to dimension_end.\n"
			,argv[0]
		);
		return -1;
	}
	const bool loop = ( argc == 2 );
	const int radius = atoi( argv[1] );
	if( radius < 1 ){
		fprintf( stderr, "Please input a positive number for radius\n" );
		return -1;
	}
	if( radius > INT16_MAX ){
		fprintf( stderr,
			"Please input a number for radius that's less than %d \n", INT16_MAX
		);
		return -1;
	}
	const int r2 = square( radius );
	int dimension = loop ? 1 : atoi( argv[2] );
	if( dimension < 0 ){
		fprintf( stderr, "Please input a positive number for dimension\n" );
		return -1;
	}
	int dimension_end = argc < 4 ? dimension : atoi( argv[3] );
	if( dimension_end < dimension ){
		fprintf( stderr,
			"Please bigger number for dimension_end than for dimension\n"
		);
		return -1;
	}
	// Cover the edge cases because the rest of the code assumes dimension > 1
	if( dimension == 0 ){
		printf( "0: 1\n" );
		if( dimension_end == 0 && loop == 0 ) return 0;
		dimension = 1;
	}
	if( dimension == 1 ){
		printf( "1: %d\n", 2 * radius + 1 );
		if( dimension_end == 1 && loop == 0 ) return 0;
		dimension = 2;
	}
#ifdef TIME
	// Get start time
	struct timespec t_start;
	clock_gettime( CLOCK, &t_start );
#endif
	// Init mpz_t variables
	mpz_t count, z;
	mpz_init( count );
	mpz_init( z );
	// Allocate arrays
	int dimension_allocated = loop ? DIMENSION_PER_ALLOC : dimension_end + 1;
	int *vector = calloc( dimension_allocated, sizeof(int) );
	mpz_t *factorial_table = calloc( dimension_allocated + 1, sizeof(mpz_t) );
	// Generate initial lookup table
	mpz_init_set_ui( factorial_table[0], 1 );
	mpz_init_set_ui( factorial_table[1], 1 );
	printf( "Generating factorial lookup table " );
	if( dimension < 2 * r2 + 1 ){ // Generate full lookup table
		printf( "until %d\n", dimension - 1 );
		for( int i = 2; i < dimension; i++ ){
			mpz_init( factorial_table[i] );
			mpz_mul_ui( factorial_table[i], factorial_table[i-1], i );
		}
	}else{ // Generate sparse lookup table
		printf( "until %d ", r2 );
		// Generate the first part of the table
		for( int i = 2; i <= r2; i++ ){
			mpz_init( factorial_table[i] );
			mpz_mul_ui( factorial_table[i], factorial_table[i-1], i );
		}
		printf( "and from %d to %d\n",
			dimension - r2, dimension - 1 );
		// Generate the second part of the table
		mpz_init( factorial_table[ dimension - r2 - 1 ] );
		mpz_fac_ui( factorial_table[ dimension - r2 - 1 ], dimension - r2 - 1 );
		for( int i = dimension - r2; i < dimension; i++ ){
			mpz_init( factorial_table[i] );
			mpz_mul_ui( factorial_table[i], factorial_table[i-1], i );
		}
	}
#ifdef TIME
	// Get start time
	struct timespec t_finish;
	clock_gettime( CLOCK, &t_finish );
	print_elapsed_time( t_start, t_finish );
#endif
	// Actually start doing stuff
	do{
#ifdef TIME
		// Get start time
		clock_gettime( CLOCK, &t_start );
#endif
		// Check if enough allocated
		if( dimension == dimension_allocated ){
			dimension_allocated += DIMENSION_PER_ALLOC;
			factorial_table = realloc( factorial_table,
				( dimension_allocated + 1 ) * sizeof(mpz_t) );
			free( vector );
			vector = calloc( dimension_allocated, sizeof(int) );
		}
		// Initialise count to 0
		mpz_set_ui( count, 0 );
		// Generate lookup table entry for this dimension
		mpz_init(   factorial_table[ dimension ] );
		mpz_mul_ui( factorial_table[ dimension ],
		            factorial_table[ dimension - 1 ],
		            dimension );
		// Clear old lookup table entry if unused
		if( dimension > r2 * 2 + 1 )
			mpz_clear( factorial_table[ dimension - r2 - 1 ] );
		// Calculate stuff
		for( int i=0; i<dimension; i++ ) vector[i] = 0;
		for( bool is_counting = true; is_counting;){
			const int cmp = vector_length2_cmp( r2, vector, dimension );
			if( cmp >=0 ){
				unique_array_orders( z, vector, dimension, factorial_table );
				mpz_add( count, count, z );
			};
#ifdef DEBUG
			for( int i=0; i<dimension; i++ ) printf( "%2d, ", vector[i] );
			gmp_printf( "cmp = %2d, count = %Zd\n", cmp, count );
#endif
			int max_in_vector = 0;
			for( int i=0;; i++ ){
				if( i+1 < dimension ){
					if( cmp > 0 ){
						if( vector[i] < radius ){
							vector[i]++;
							break;
						} else if( vector[i+1] < radius ){
							vector[i+1]++;
							for(int j=i;j>=0;j--)vector[j]=vector[i+1];
							break;
						}
					}else{
						if( vector[i] > max_in_vector ) max_in_vector = vector[i];
						if( max_in_vector > vector[i+1] + 1 ){
							vector[i+1]++;
							for(int j=i;j>=0;j--)vector[j]=vector[i+1];
							break;
						}
					}
				}else{
					is_counting = false;
					break;
				}
			}
		}
#ifdef TIME
		// Get elapsed time
		clock_gettime( CLOCK, &t_finish );
#endif
		// Print stuff
		gmp_printf( "%d: %Zd\n", dimension, count );
#ifdef TIME
		print_elapsed_time( t_start, t_finish );
#endif
		dimension++;
	}while( loop || dimension <= dimension_end );
	return 0;
}