#include <assert.h>
#include <math.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "ts.h"


/**
 * Fill an array with linearly spaced values from x0 to x1.
 * Uses a running accumulator to avoid floating-point drift,
 * and forces the last element to exactly x1.
 */
fp_t* ts_fill(fp_t *x, uint_t n, fp_t x0, fp_t x1 )
{
	if( !x || n == 0 ) return NULL;
	if( n == 1 )
	{
		x[0] = x0;
		return x;
	}

	/* TS_FILL_ZERO_THRESHOLD avoids division by zero when both endpoints
	 * are effectively zero, and skips unnecessary arithmetic. */
	if( fabs( x0 ) < TS_FILL_ZERO_THRESHOLD && fabs( x1 ) < TS_FILL_ZERO_THRESHOLD )
	{
		for( uint_t i = 0; i < n; i++ )
			x[i] = 0.0;
	}
	else
	{
		fp_t step = ( x1 - x0 ) / (fp_t)( n - 1 );
		fp_t acc = x0;
		for( uint_t i = 0; i < n - 1; i++ )
		{
			x[i] = acc;
			acc += step;
		}
		x[n - 1] = x1;
	}

	return x;
}


/**
 * Print an array of values to a stream in bracket notation.
 * If stream is NULL, output goes to stdout.
 */
void ts_print( fp_t const *x, uint_t n, FILE *stream )
{
	assert( x != NULL );
	if( !stream ) stream = stdout;
	fputc( '[', stream );
	if( !x || n == 0 )
	{
		fputc( ']', stream );
		return;
	}

	uint_t i;
	for( i = 0; i < n - 1; i++ )
	{
		fprintf( stream, "%.17g,", x[i] );
	}

	fprintf( stream, "%.17g]", x[i] );
}


/**
 * Pretty-print a 2D array of values to a stream.
 * If stream is NULL, output goes to stdout.
 */
void ts_pretty_print( fp_t const * const *x, uint_t m, uint_t n, FILE *stream )
{
	if( !stream ) stream = stdout;
	if( !x || m == 0 || n == 0 )
	{
		return;
	}

	if( m > 1 )
	{
		for( uint_t i = 0; i < n; i++ )
		{
			for( uint_t j = 0; j < m; j++ )
			{
				if( x[j] )
					fprintf( stream, "%.17g  ", x[j][i] );
				else
					fprintf( stream, "NULL   " );
			}

			fputs( "\n", stream );
		}
	}
	else
	{
		if( !x[0] )
		{
			fputs( "[]\n", stream );
			return;
		}
		for( uint_t i = 0; i < n; i++ )
		{
			fprintf( stream, "%.17g\n", x[0][i] );
		}
	}
}


fp_t _ts_square( fp_t x )
{
	return x * x;
}


fp_t _ts_reciprocal( fp_t x )
{
	return x != 0 ? 1.0 / x : (fp_t)INFINITY;
}


fp_t _ts_triangle_number( fp_t x )
{
	return x * ( x + 1 ) * 0.5;
}


/**
 * Compute covariance from accumulated sums.
 * NOTE: This formula is susceptible to catastrophic cancellation for
 * large or closely clustered inputs. For production use with high
 * numerical sensitivity, prefer Welford's online algorithm.
 */
fp_t _ts_sum_to_covariance( fp_t sum_xy, fp_t sum_x, fp_t sum_y, uint_t k )
{
	if( k == 0 )
		return 0.0;
	fp_t result = ( sum_xy - ( sum_x * sum_y / k ) ) / k;
	if( result < 0.0 )
		result = 0.0;
	return result;
}


fp_t _ts_sum_to_variance( fp_t sum_x2, fp_t sum_x, uint_t k )
{
	return _ts_sum_to_covariance( sum_x2, sum_x, sum_x, k );
}


uint_t _ts_circle_increment( uint_t i, uint_t k )
{
	if( k == 0 )
		return i;
	return ( i + 1 ) % k;
}


/**
 * Allocate and zero-initialize a ts_constant_t structure.
 */
ts_constant_t* ts_constant( void )
{
	return calloc( 1, sizeof( ts_constant_t ) );
}


/**
 * Free a ts_constant_t structure.
 */
void ts_constant_free( ts_constant_t *c )
{
	free( c );
}


/**
 * Run a test harness: call ts_function, then compare outputs against
 * expected values with tolerance TS_TEST_EPSILON.
 */
bool ts_test(
	int ( *ts_function )(
		fp_t * const *y,
		uint_t y_n,
		fp_t const * const *x,
		uint_t x_n,
		uint_t n,
		ts_constant_t const *c
	),
	uint_t ( *ts_offset )(
		ts_constant_t const *c
	),
	fp_t const * const *x,
	fp_t const * const *e,
	uint_t n,
	ts_constant_t const *c,
	ts_parameter_t const *m )
{
	if ( !ts_function || !ts_offset )
		return false;

	if( !m ) return false;
	if( !e ) return false;
	if( !x ) return false;
	if( m->n_constants > 0 && !c ) return false;
	if( m->n_outputs > 1024 ) return false;

	uint_t offset = ts_offset( c );
	if ( offset >= n )
	{
		return false;
	}
	uint_t y_n = n - offset;
	bool result = true;

	fp_t **y = calloc( m->n_outputs, sizeof( fp_t * ) );
	if ( !y )
	{
		return false;
	}

	if ( n > UINT64_MAX / sizeof( fp_t ) )
	{
		return false;
	}

	for( uint_t i = 0; i < m->n_outputs; i++ )
	{
		fp_t *tmp = calloc( n, sizeof( fp_t ) );
		if ( !tmp )
		{
			for ( uint_t j = 0; j < i; j++ )
				free( y[j] );
			free( y );
			return false;
		}
		y[i] = tmp;
	}

	int ret = ts_function( y, m->n_outputs, x, m->n_inputs, n, c );
	if ( ret != 0 )
	{
		for( uint_t i = 0; i < m->n_outputs; i++ )
			free( y[i] );
		free( y );
		return false;
	}

	for( uint_t i = 0; i < m->n_outputs; i++ )
	{
		for( uint_t j = 0; j < y_n; j++ )
		{
			fp_t yv = y[i][j + offset];
			fp_t ev = e[i][j];
			if( isnan( yv ) || isnan( ev ) )
			{
				fprintf( stderr, "FAIL NaN y[%d][%d]=%f exp=%f\n", (int)i, (int)j, yv, ev );
				result = false;
				continue;
			}
			/* Matching infinities are equal */
			if( isinf( yv ) && isinf( ev ) && ( yv > 0 ) == ( ev > 0 ) )
				continue;
			if( isinf( yv ) || isinf( ev ) )
			{
				fprintf( stderr, "FAIL inf y[%d][%d]=%f exp=%f\n", (int)i, (int)j, yv, ev );
				result = false;
				continue;
			}
			fp_t diff = fabs( yv - ev );
			fp_t abs_e = fabs( ev );
			fp_t tol = TS_TEST_EPSILON + abs_e * TS_TEST_EPSILON;
			if( diff > tol )
				fprintf( stderr, "FAIL y[%d][%d]=%.6f exp=%.6f diff=%.8f tol=%.8f\n", (int)i, (int)j, yv, ev, diff, tol );
			if( diff > tol )
				result = false;
		}

		free( y[i] );
	}
	free( y );

	return result;
}