#include "uintcore.h"
#include "uintarith.h"
#include "uintarithmod.h"
#include "polycore.h"
#include "polyarith.h"

using namespace std;

namespace seal
{
    namespace util
    {
        void multiply_poly_poly(const uint64_t *operand1, 
            int operand1_coeff_count, int operand1_coeff_uint64_count, 
            const uint64_t *operand2, int operand2_coeff_count, 
            int operand2_coeff_uint64_count, int result_coeff_count, 
            int result_coeff_uint64_count, uint64_t *result, MemoryPool &pool)
        {
#ifdef SEAL_DEBUG
            if (operand1 == nullptr && operand1_coeff_count > 0 && 
                operand1_coeff_uint64_count > 0)
            {
                throw invalid_argument("operand1");
            }
            if (operand1_coeff_count < 0)
            {
                throw invalid_argument("operand1_coeff_count");
            }
            if (operand1_coeff_uint64_count < 0)
            {
                throw invalid_argument("operand1_coeff_uint64_count");
            }
            if (operand2 == nullptr && operand2_coeff_count > 0 && 
                operand2_coeff_uint64_count > 0)
            {
                throw invalid_argument("operand2");
            }
            if (operand2_coeff_count < 0)
            {
                throw invalid_argument("operand2_coeff_count");
            }
            if (operand2_coeff_uint64_count < 0)
            {
                throw invalid_argument("operand2_coeff_uint64_count");
            }
            if (result_coeff_count < 0)
            {
                throw invalid_argument("result_coeff_count");
            }
            if (result_coeff_uint64_count < 0)
            {
                throw invalid_argument("result_coeff_uint64_count");
            }
            if (result == nullptr && result_coeff_count > 0 && 
                result_coeff_uint64_count > 0)
            {
                throw invalid_argument("result");
            }
            if (result != nullptr && 
                (operand1 == result || operand2 == result))
            {
                throw invalid_argument("result cannot point to the same value as operand1 or operand2");
            }
#endif
            Pointer intermediate(allocate_uint(result_coeff_uint64_count, pool));

            // Clear product.
            set_zero_poly(result_coeff_count, result_coeff_uint64_count, result);

            operand1_coeff_count = get_significant_coeff_count_poly(
                operand1, operand1_coeff_count, operand1_coeff_uint64_count);
            operand2_coeff_count = get_significant_coeff_count_poly(
                operand2, operand2_coeff_count, operand2_coeff_uint64_count);
            for (int operand1_index = 0; operand1_index < operand1_coeff_count; operand1_index++)
            {
                const uint64_t *operand1_coeff = get_poly_coeff(
                    operand1, operand1_index, operand1_coeff_uint64_count);
                for (int operand2_index = 0; operand2_index < operand2_coeff_count; operand2_index++)
                {
                    int product_coeff_index = operand1_index + operand2_index;
                    if (product_coeff_index >= result_coeff_count)
                    {
                        break;
                    }

                    const uint64_t *operand2_coeff = get_poly_coeff(
                        operand2, operand2_index, operand2_coeff_uint64_count);
                    multiply_uint_uint(operand1_coeff, operand1_coeff_uint64_count, 
                        operand2_coeff, operand2_coeff_uint64_count, result_coeff_uint64_count, intermediate.get());
                    uint64_t *result_coeff = get_poly_coeff(
                        result, product_coeff_index, result_coeff_uint64_count);
                    add_uint_uint(result_coeff, intermediate.get(), result_coeff_uint64_count, result_coeff);
                }
            }
        }
        void poly_eval_poly(const uint64_t *poly_to_eval, 
            int poly_to_eval_coeff_count, 
            int poly_to_eval_coeff_uint64_count, 
            const uint64_t *value, int value_coeff_count, 
            int value_coeff_uint64_count, int result_coeff_count, 
            int result_coeff_uint64_count, uint64_t *result, MemoryPool &pool)
        {
#ifdef SEAL_DEBUG
            if (poly_to_eval == nullptr)
            {
                throw invalid_argument("poly_to_eval");
            }
            if (value == nullptr)
            {
                throw invalid_argument("value");
            }
            if (result == nullptr)
            {
                throw invalid_argument("result");
            }
            if (poly_to_eval_coeff_count <= 0)
            {
                throw invalid_argument("poly_to_eval_coeff_count");
            }
            if (poly_to_eval_coeff_uint64_count <= 0)
            {
                throw invalid_argument("poly_to_eval_coeff_uint64_count");
            }
            if (value_coeff_count <= 0)
            {
                throw invalid_argument("value_coeff_count");
            }
            if (value_coeff_uint64_count <= 0)
            {
                throw invalid_argument("value_coeff_uint64_count");
            }
            if (result_coeff_count <= 0)
            {
                throw invalid_argument("result_coeff_count");
            }
            if (result_coeff_uint64_count <= 0)
            {
                throw invalid_argument("result_coeff_uint64_count");
            }
#endif
            // Evaluate poly at value using Horner's method
            Pointer temp1(allocate_poly(result_coeff_count, result_coeff_uint64_count, pool));
            Pointer temp2(allocate_zero_poly(result_coeff_count, result_coeff_uint64_count, pool));
            uint64_t *productptr = temp1.get();
            uint64_t *intermediateptr = temp2.get();

            for (int coeff_index = poly_to_eval_coeff_count - 1; coeff_index >= 0; coeff_index--)
            {
                multiply_poly_poly(intermediateptr, result_coeff_count, result_coeff_uint64_count, value, value_coeff_count, value_coeff_uint64_count, result_coeff_count, result_coeff_uint64_count, productptr, pool);
                const uint64_t *curr_coeff = get_poly_coeff(poly_to_eval, coeff_index, poly_to_eval_coeff_uint64_count);
                add_uint_uint(productptr, result_coeff_uint64_count, curr_coeff, poly_to_eval_coeff_uint64_count, false, result_coeff_uint64_count, productptr);
                swap(productptr, intermediateptr);
            }
            set_poly_poly(intermediateptr, result_coeff_count, result_coeff_uint64_count, result);
        }

        void exponentiate_poly(const std::uint64_t *poly, 
            int poly_coeff_count, int poly_coeff_uint64_count, 
            const uint64_t *exponent, int exponent_uint64_count, 
            int result_coeff_count, int result_coeff_uint64_count, 
            std::uint64_t *result, MemoryPool &pool)
        {
#ifdef SEAL_DEBUG
            if (poly == nullptr)
            {
                throw invalid_argument("poly");
            }
            if (poly_coeff_count <= 0)
            {
                throw invalid_argument("poly_coeff_count");
            }
            if (poly_coeff_count <= 0)
            {
                throw invalid_argument("poly_coeff_uint64_count");
            }
            if (exponent == nullptr)
            {
                throw invalid_argument("exponent");
            }
            if (exponent_uint64_count <= 0)
            {
                throw invalid_argument("exponent_uint64_count");
            }
            if (result == nullptr)
            {
                throw invalid_argument("result");
            }
            if (result_coeff_count <= 0)
            {
                throw invalid_argument("result_coeff_count");
            }
            if (result_coeff_uint64_count <= 0)
            {
                throw invalid_argument("result_coeff_uint64_count");
            }
#endif
            // Fast cases
            if (is_zero_uint(exponent, exponent_uint64_count))
            {
                set_zero_poly(result_coeff_count, result_coeff_uint64_count, result);
                *result = 1;
                return;
            }
            if (is_equal_uint(exponent, exponent_uint64_count, 1))
            {
                set_poly_poly(poly, poly_coeff_count, poly_coeff_uint64_count, 
                    result_coeff_count, result_coeff_uint64_count, result);
                return;
            }

            // Need to make a copy of exponent
            Pointer exponent_copy(allocate_uint(exponent_uint64_count, pool));
            set_uint_uint(exponent, exponent_uint64_count, exponent_copy.get());

            // Perform binary exponentiation.
            Pointer big_alloc(allocate_uint((static_cast<int64_t>(result_coeff_count) + 
                result_coeff_count + result_coeff_count) * result_coeff_uint64_count, pool));

            uint64_t *powerptr = big_alloc.get();
            uint64_t *productptr = get_poly_coeff(powerptr, result_coeff_count, result_coeff_uint64_count);
            uint64_t *intermediateptr = get_poly_coeff(productptr, result_coeff_count, result_coeff_uint64_count);

            set_poly_poly(poly, poly_coeff_count, poly_coeff_uint64_count, result_coeff_count, result_coeff_uint64_count, powerptr);
            set_zero_poly(result_coeff_count, result_coeff_uint64_count, intermediateptr);
            *intermediateptr = 1;

            // Initially: power = operand and intermediate = 1, product is not initialized.
            while (true)
            {
                if ((*exponent_copy.get() % 2) == 1)
                {
                    multiply_poly_poly(powerptr, result_coeff_count, result_coeff_uint64_count, intermediateptr, result_coeff_count, result_coeff_uint64_count,
                        result_coeff_count, result_coeff_uint64_count, productptr, pool);
                    swap(productptr, intermediateptr);
                }
                right_shift_uint(exponent_copy.get(), 1, exponent_uint64_count, exponent_copy.get());
                if (is_zero_uint(exponent_copy.get(), exponent_uint64_count))
                {
                    break;
                }
                multiply_poly_poly(powerptr, result_coeff_count, result_coeff_uint64_count, powerptr, result_coeff_count, result_coeff_uint64_count,
                    result_coeff_count, result_coeff_uint64_count, productptr, pool);
                swap(productptr, powerptr);
            }
            set_poly_poly(intermediateptr, result_coeff_count, result_coeff_uint64_count, result);
        }
    }
}
