/*
    Copyright (C) 2025, Vincent Neiger, Éric Schost
    Copyright (C) 2025, Mael Hostettler

    This file is part of FLINT.

    FLINT is free software: you can redistribute it and/or modify it under
    the terms of the GNU Lesser General Public License (LGPL) as published
    by the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.  See <https://www.gnu.org/licenses/>.
*/

#include "flint.h"
#include "test_helpers.h"
#include "nmod_vec.h"
#include "nmod_poly.h"

TEST_FUNCTION_START(nmod_poly_interpolate_geometric_nmod_vec_fast, state)
{
    int i, result = 1;

    for (i = 0; i < 1000 * flint_test_multiplier(); i++)
    {
        nmod_poly_t P, Q;
        nn_ptr y;
        ulong mod, r;
        slong len, npoints;

        /* npoints = 1 + (i == 0 ? 1 : n_randint(state, 200)); */
        npoints = 1 + n_randint(state, 200);
        len = 1 + n_randint(state, npoints);
        do 
        { 
            mod = n_randtest_prime(state, 1); 
        }
        while (mod <= (ulong)(2*FLINT_MAX(npoints, len) + 1)); // minimum limit for maximum order r

        r = n_primitive_root_prime(mod);
        y = _nmod_vec_init(npoints);

        /* use full `npoints` points */
        {
            nmod_poly_init(P, mod);
            nmod_poly_init(Q, mod);
            nmod_poly_randtest(P, state, len);

            nmod_poly_evaluate_geometric_nmod_vec_fast(y, P, r, npoints);
            nmod_poly_interpolate_geometric_nmod_vec_fast(Q, r, y, npoints);
            result = nmod_poly_equal(P, Q);
            if (!result)
            {
                flint_printf("FAIL (all points):\n");
                flint_printf("mod=%wu, len=%wd, npoints=%wd\n\n", mod, len, npoints);
                nmod_poly_print_pretty(P, "x"), flint_printf("\n\n");
                nmod_poly_print_pretty(Q, "x"), flint_printf("\n\n");
                fflush(stdout);
                flint_abort();
            }

            nmod_poly_clear(P);
            nmod_poly_clear(Q);
        }

        /* use only `len` points */
        {
            nmod_poly_init(P, mod);
            nmod_poly_init(Q, mod);
            nmod_poly_randtest(P, state, len);

            nmod_poly_evaluate_geometric_nmod_vec_fast(y, P, r, npoints);
            nmod_poly_interpolate_geometric_nmod_vec_fast(Q, r, y, len);
            result = nmod_poly_equal(P, Q);
            if (!result)
            {
                flint_printf("FAIL (`len` points):\n");
                flint_printf("mod=%wu, len=%wd, npoints=%wd\n\n", mod, len, npoints);
                nmod_poly_print_pretty(P, "x"), flint_printf("\n\n");
                nmod_poly_print_pretty(Q, "x"), flint_printf("\n\n");
                fflush(stdout);
                flint_abort();
            }

            nmod_poly_clear(P);
            nmod_poly_clear(Q);
        }

        /* use only `len` points, interpolate then evaluate */
        {
            nmod_poly_init(P, mod);
            nn_ptr val = FLINT_ARRAY_ALLOC(len, ulong);
            _nmod_vec_randtest(y, state, len, P->mod);

            nmod_poly_interpolate_geometric_nmod_vec_fast(P, r, y, len);
            nmod_poly_evaluate_geometric_nmod_vec_fast(val, P, r, len);
            result = _nmod_vec_equal(val, y, len);
            if (!result)
            {
                flint_printf("FAIL (`len` points):\n");
                flint_printf("mod=%wu, len=%wd, npoints=%wd\n\n", mod, len, npoints);
                flint_printf("val = %{slong*}...\n", val, len);
                flint_printf("y = %{slong*}...\n", y, len);
                fflush(stdout);
                flint_abort();
            }

            flint_free(val);
            nmod_poly_clear(P);
        }

        /* use variant with given precomputation */
        {
            nmod_t mod2;
            nmod_init(&mod2, mod);
            nmod_geometric_progression_t G;
            nmod_geometric_progression_init(G, r, npoints, mod2);
            nmod_poly_init(P, mod);
            nmod_poly_init(Q, mod);
            nmod_poly_randtest(P, state, len);

            _nmod_poly_evaluate_geometric_nmod_vec_fast_precomp(y, P->coeffs, P->length, G, npoints, mod2);
            nmod_poly_fit_length(Q, len);
            _nmod_poly_interpolate_geometric_nmod_vec_fast_precomp(Q->coeffs, y, G, len, mod2);
            _nmod_poly_set_length(Q, len);
            _nmod_poly_normalise(Q);
            result = nmod_poly_equal(P, Q);
            if (!result)
            {
                flint_printf("FAIL (precomp):\n");
                flint_printf("mod=%wu, len=%wd, npoints=%wd\n\n", mod, len, npoints);
                nmod_poly_print_pretty(P, "x"), flint_printf("\n\n");
                nmod_poly_print_pretty(Q, "x"), flint_printf("\n\n");
                fflush(stdout);
                flint_abort();
            }

            nmod_poly_clear(P);
            nmod_poly_clear(Q);
            nmod_geometric_progression_clear(G);
        }

        flint_free(y);
    }

    TEST_FUNCTION_END(state);
}
