/*
 * SPDX-License-Identifier: CC0-1.0
 *
 * Copyright (C) 2024 W. Kosior <koszko@koszko.org>
 */

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

#include "pqcrypto_poly.h"

/* Exponent for Mersenne prime, for testing. */
#define TEST_MERSENNE_EXPONENT 7 /* 89 */

static void marsenne_prime_init(fmpz_t prime, const ulong exponent) {
	fmpz_init(prime);

	fmpz_ui_pow_ui(prime, 2, exponent);
	fmpz_sub_ui(prime, prime, 1);
}

static void init_read_poly(fmpz_poly_t poly, FILE * file) {
	bool first = true;
	fmpz_t coef;

	fmpz_poly_init(poly);
	fmpz_init(coef);

	for (ulong exponent = 0;; exponent++) {
		int separator_char;

		if (first) {
			first = false;
		} else {
			separator_char = getc(file);

			if (separator_char == '\n')
				break;

			if (separator_char != ' ')
				goto error;
		}

		if (fmpz_fread(file, coef) < 0)
			goto error;

		fmpz_poly_set_coeff_fmpz(poly, exponent, coef);
	}

	fmpz_clear(coef);
	return;

error:
	fprintf(stderr, "Error reading polynomial.\n");
	abort();
}

int main(const int argc, const char* const* const argv) {
	fmpz_t prime; /* integer for modulo operations */
	mod_centered_0_ctx_t mod_ctx;

	(void) argc;
	(void) argv;

	/*
	 * Marsenne primes are used just for testing.  Cryptographic algorithm
	 * will use different ones.
	 */
	marsenne_prime_init(prime, TEST_MERSENNE_EXPONENT);

	printf("Prime used for modulo operations: ");
	fmpz_fprint(stdout, prime);
	putchar('\n');

	mod_c0_ctx_init(mod_ctx, prime);

	{ /* Experiment 1 — modulo addition */
		fmpz_t num1, num2, num_sum;

		fmpz_init_set_ui(num1, 55);
		fmpz_init_set_ui(num2, 31);
		fmpz_init(num_sum);

		fmpz_fprint(stdout, num1);
		printf(" + ");
		fmpz_fprint(stdout, num2);
		printf(" mod [-");
		fmpz_fprint(stdout, mod_ctx[0].range_max);
		putchar(',');
		fmpz_fprint(stdout, mod_ctx[0].range_max);
		printf("] = ");

		fmpz_add(num_sum, num1, num2);
		mod_c0(num_sum, mod_ctx);
		fmpz_fprint(stdout, num_sum);
		putchar('\n');

		fmpz_clear(num1);
		fmpz_clear(num2);
		fmpz_clear(num_sum);
	} /* End of experiment 1 */

	{ /* Experiment 2 */
		fmpz_poly_t poly1, poly2, poly_computed;
		slong divisor_degree;
		poly_ring_ctx_t poly_ring_ctx;

		printf("Give first polynomial for the experiment:\n");
		init_read_poly(poly1, stdin);

		printf("Read polynomial: ");
		fmpz_poly_print_pretty(poly1, "x");
		putchar('\n');

		printf("Give second polynomial for the experiment:\n");
		init_read_poly(poly2, stdin);

		printf("Read polynomial: ");
		fmpz_poly_print_pretty(poly2, "x");
		putchar('\n');

		fmpz_poly_init(poly_computed);

		printf("Normal product of polynomials:\n");
		fmpz_poly_mul(poly_computed, poly1, poly2);
		fmpz_poly_print_pretty(poly_computed, "x");
		putchar('\n');

		printf("Normal sum of polynomials:\n");
		fmpz_poly_add(poly_computed, poly1, poly2);
		fmpz_poly_print_pretty(poly_computed, "x");
		putchar('\n');

		printf("Give the degree m of X^m+1 polynomial to be used as ");
		printf("divisor in the ring:\n");
		if (flint_scanf("%wd", &divisor_degree) < 1 ||
		    divisor_degree < 1) {
			fprintf(stderr, "Bad divisor.\n");
			abort();
		}
		poly_ring_ctx_init(poly_ring_ctx, mod_ctx, divisor_degree);

		printf("Product of polynomials in the ring:\n");
		poly_mul_in_ring(poly_computed, poly1, poly2, poly_ring_ctx);
		fmpz_poly_print_pretty(poly_computed, "x");
		putchar('\n');

		printf("Sum of polynomials in the ring:\n");
		poly_add_in_ring(poly_computed, poly1, poly2, poly_ring_ctx);
		fmpz_poly_print_pretty(poly_computed, "x");
		putchar('\n');

		fmpz_poly_clear(poly1);
		fmpz_poly_clear(poly2);
		fmpz_poly_clear(poly_computed);

		poly_ring_ctx_clear(poly_ring_ctx);
	} /* End of experiment 2 */

	mod_c0_ctx_clear(mod_ctx);
	fmpz_clear(prime);

	flint_cleanup();

	return EXIT_SUCCESS;
}