#include "../squiggle.h"
#include <stdint.h>
#include <math.h>
#include <stdlib.h>
#include <stdio.h>

#define N 1000 * 1000
#define PERCENTAGE_TOLERANCE_UNIFORM 1.0/1000.0
#define PERCENTAGE_TOLERANCE_NORMAL 5.0/1000.0
#define MAX_NAME_LENGTH 500

// Structs

struct array_expectations {
	double* array;
	int n;
	char* name;
	double expected_mean;
	double expected_std;
	double tolerance;
};

void test_array_expectations(struct array_expectations e){
	double mean = array_mean(e.array, e.n);
	double delta_mean = mean - e.expected_mean;

	double std = array_std(e.array, e.n);
	double delta_std = std - e.expected_std;
	

	if(fabs(delta_mean) > e.tolerance){
		printf("[-] Mean test for %s NOT passed.\n", e.name);
		printf("Mean of %s: %f, vs expected mean: %f, delta: %f\n", e.name, mean, e.expected_mean, delta_mean);
	}else {
		printf("[x] Mean test for %s PASSED\n", e.name);
	}

	if(fabs(delta_std) > e.tolerance){
		printf("[-] Std test for %s NOT passed.\n", e.name);
		printf("Std of %s: %f, vs expected std: %f, delta: %f\n", e.name, std, e.expected_std, delta_std);
	}else {
		printf("[x] Std test for %s PASSED\n", e.name);
	}
	
	printf("\n");

}

// Test unit uniform
void test_unit_uniform(uint64_t* seed){
	double* unit_uniform_array = malloc(sizeof(double) * N);
	
	for(int i=0; i<N; i++){
		unit_uniform_array[i] = sample_unit_uniform(seed);
	}

	struct array_expectations expectations = {
		.array = unit_uniform_array,
		.n = N,
		.name = "unit uniform",
		.expected_mean = 0.5,
		.expected_std = sqrt(1.0/12.0),
		.tolerance = 1 * PERCENTAGE_TOLERANCE_UNIFORM,
	};

	test_array_expectations(expectations);
}

// Test uniforms
void test_uniform(double start, double end, uint64_t* seed){
	double* uniform_array = malloc(sizeof(double) * N);
	
	for(int i=0; i<N; i++){
		uniform_array[i] = sample_uniform(start, end, seed);
	}
	
	char* name = malloc(MAX_NAME_LENGTH * sizeof(char));
	snprintf(name, MAX_NAME_LENGTH, "[%f, %f] uniform", start, end);
	struct array_expectations expectations = {
		.array = uniform_array,
		.n = N,
		.name = name,
		.expected_mean = (start + end)/2,
		.expected_std = sqrt(1.0/12.0) * fabs(end-start),
		.tolerance = fabs(end -start) * PERCENTAGE_TOLERANCE_UNIFORM,
	};

	test_array_expectations(expectations);
	free(name);
}

// Test unit normal
void test_unit_normal(uint64_t* seed){
	double* unit_normal_array = malloc(sizeof(double) * N);
	
	for(int i=0; i<N; i++){
		unit_normal_array[i] = sample_unit_normal(seed);
	}

	struct array_expectations expectations = {
		.array = unit_normal_array,
		.n = N,
		.name = "unit normal",
		.expected_mean = 0,
		.expected_std = 1, 
		.tolerance = 1 * PERCENTAGE_TOLERANCE_NORMAL,
	};

	test_array_expectations(expectations);
}

// Test normal
void test_normal(double mean, double std, uint64_t* seed){
	double* normal_array = malloc(sizeof(double) * N);
	
	for(int i=0; i<N; i++){
		normal_array[i] = sample_normal(mean, std, seed);
	}
	
	char* name = malloc(MAX_NAME_LENGTH * sizeof(char));
	snprintf(name, MAX_NAME_LENGTH, "normal(%f, %f)", mean, std);
	struct array_expectations expectations = {
		.array = normal_array,
		.n = N,
		.name = name,
		.expected_mean = mean,
		.expected_std = std,
		.tolerance = std * PERCENTAGE_TOLERANCE_NORMAL,
	};

	test_array_expectations(expectations);
	free(name);
}

// Test beta

void test_beta(double a, double b, uint64_t* seed){
	double* beta_array = malloc(sizeof(double) * N);
	
	for(int i=0; i<N; i++){
		beta_array[i] = sample_beta(a, b, seed);
	}
	
	char* name = malloc(MAX_NAME_LENGTH * sizeof(char));
	snprintf(name, MAX_NAME_LENGTH, "beta(%f, %f)", a, b);
	struct array_expectations expectations = {
		.array = beta_array,
		.n = N,
		.name = name,
		.expected_mean = a/(a+b),
		.expected_std = sqrt((a*b)/(pow(a+b, 2) * (a + b + 1))),
		.tolerance = PERCENTAGE_TOLERANCE_UNIFORM,
	};

	test_array_expectations(expectations);
	free(name);
}

int main(){
    // set randomness seed
    uint64_t* seed = malloc(sizeof(uint64_t));
    *seed = 1000; // xorshift can't start with a seed of 0
   
		printf("Testing unit uniform\n");
		test_unit_uniform(seed);

		printf("Testing small uniforms\n");
		for(int i=0; i<100; i++){
			double start = sample_uniform(-10, 10, seed);
			double end = sample_uniform(-10, 10, seed);
			if ( end > start){
				test_uniform(start, end, seed);
			}
		}

		printf("Testing wide uniforms\n");
		for(int i=0; i<100; i++){
			double start = sample_uniform(-1000 * 1000, 1000 * 1000, seed);
			double end = sample_uniform(-1000 * 1000, 1000 * 1000, seed);
			if ( end > start){
				test_uniform(start, end, seed);
			}
		}

		printf("Testing unit normal\n");
		test_unit_normal(seed);
		
		printf("Testing small normals\n");
		for(int i=0; i<100; i++){
			double mean = sample_normal(-10, 10, seed);
			double std = sample_normal(0, 10, seed);
			if ( std > 0){
				test_normal(mean, std, seed);
			}
		}

		printf("Testing larger normals\n");
		for(int i=0; i<100; i++){
			double mean = sample_uniform(-1000 * 1000, 1000 * 1000, seed);
			double std = sample_uniform(0, 1000 * 1000, seed);
			if ( std > 0){
				test_normal(mean, std, seed);
			}
		}

		printf("Testing beta distribution\n");
		for(int i=0; i<100; i++){
			double a = sample_uniform(0, 1000, seed);
			double b = sample_uniform(0, 1000, seed);
			if ( (a > 0) && ( b >0)){
				test_beta(a, b, seed);
			}
		}

		printf("Testing larger beta distributions\n");
		for(int i=0; i<100; i++){
			double a = sample_uniform(0, 1000 * 1000, seed);
			double b = sample_uniform(0, 1000 * 1000, seed);
			if ( (a > 0) && ( b >0)){
				test_beta(a, b, seed);
			}
		}

		free(seed);
}