import itertools
import random
import sys

import numpy as np

from multivar_horner import HornerMultivarPolynomial, MultivarPolynomial
from multivar_horner.global_settings import FLOAT_DTYPE, UINT_DTYPE
from tests.test_settings import MAX_COEFF_MAGNITUDE, MAX_NUMERICAL_ERROR, NR_COEFF_CHANGES, NR_TEST_POLYNOMIALS


def proto_test_case(data, fct):
    all_good = True
    for input, expected_output in data:
        print('\n')
        actual_output = fct(input)
        print(f'p({input[2]}) == {expected_output}')
        if actual_output != expected_output:
            print(f'ERROR: p(x) == {actual_output}')
            all_good = False
        else:
            print('OK.')

    assert all_good


def get_rnd_poly_properties(all_exponents, degree, max_abs_coeff=1.0, integer_coeffs=False, enforce_degree=False):
    """
    for meaningful test results uniformly distributed random test polynomials are required
    -> every possible polynomial should appear with equal probability
    <=> every 'sparsity' AND every possible monomial have equal probability

    sparsity: the ratio of 'active' monomials out of all possible monomials

    there must be at least 1 entry ('active' monomial)

    :param all_exponents:
    :param degree: the total degree the monomial must (at least) have
    :param max_abs_coeff:
    :param integer_coeffs:
    :return: evenly distributed settings of a polynomial with the given possible exponents
    """

    max_degree_monomial_idxs = all_exponents.sum(axis=1) >= degree
    assert np.any(max_degree_monomial_idxs)

    # every exponent can take the values in the range [0; maximal_degree]
    max_nr_exponent_vects = all_exponents.shape[0]

    # pick a random sparsity: controlled by the uniformly distributed probability a monomial is present
    activation_thres = random.random()
    row_activation = np.random.rand(max_nr_exponent_vects)
    row_idxs = row_activation > activation_thres

    if enforce_degree:
        # the settings must correspond to a polynomial of at least the requested total degree
        # -> at least one monomial of this total degree (or higher) must be active
        while not np.any(row_idxs & max_degree_monomial_idxs):
            row_activation = np.random.rand(max_nr_exponent_vects)
            row_idxs = row_activation > activation_thres
    else:
        if np.sum(row_idxs) == 0:  # there must be at least 1 entry ('active' monomial)
            row_idxs[random.randint(0, max_nr_exponent_vects - 1)] = 1

    exponents = all_exponents[row_idxs, :]
    nr_monomials = exponents.shape[0]
    coefficients = (np.random.rand(nr_monomials, 1) - 0.5) * (2 * max_abs_coeff)  # [ -max_abs_coeff; max_abs_coeff]
    if integer_coeffs:
        coefficients = np.round(coefficients)

    return coefficients, exponents


def all_possible_exponents(dim, deg):
    """
    generate a fully occupied exponent matrix for a polynomial of lp_degree = infinity
    -> the given degree is the total degree of the polynomials
    :param dim:
    :param deg:
    :return:
    """

    # def cntr2exp_vect(cntr):
    #     exp_vect = np.empty(dim, dtype=UINT_DTYPE)
    #     for d in range(dim - 1, -1, -1):
    #         divisor = (deg + 1) ** d
    #         # cntr = quotient*divisor + remainder
    #         quotient, remainder = divmod(cntr, divisor)
    #         exp_vect[d] = quotient
    #         cntr = remainder
    #     return exp_vect

    max_nr_exponent_vects = (deg + 1) ** dim
    all_possible = np.empty((max_nr_exponent_vects, dim), dtype=UINT_DTYPE)

    # for i in range(max_nr_exponent_vects):
    #     all_possible[i] = cntr2exp_vect(i)

    for i, exponents in enumerate(itertools.product(range(deg + 1), repeat=dim)):
        all_possible[i] = np.asarray(exponents)

    # there must not be duplicate exponent vectors
    assert all_possible.shape == np.unique(all_possible, axis=0).shape

    return all_possible


def rnd_settings_list(length, dim, degree, *args, **kwargs):
    all_exponent_vect = all_possible_exponents(dim, degree)
    settings_list = [get_rnd_poly_properties(all_exponent_vect, degree, *args, **kwargs) for i in range(length)]

    # # random settings should have approx. half the amount of maximal entries on average
    # num_monomial_entries = 0
    # for settings in settings_list:
    #     num_monomial_entries += settings[0].shape[0]
    #
    # avg_monomial_entries = num_monomial_entries / length
    # max_monomial_entries = int((maximal_degree + 1) ** dim)
    # print(avg_monomial_entries, max_monomial_entries)
    return settings_list


def rnd_input_list(length, dim, max_abs_val):
    return [(np.random.rand(dim) - 0.5) * (2 * max_abs_val) for i in range(length)]


def vectorize(obj):
    return lambda x: np.array([obj(np.atleast_1d(x_i)) for x_i in x])


#
# vectorised version of naive_eval() in multivar_horner.helpers_fcts_numba
def naive_eval_reference(X, exponents, coefficients):
    return np.dot(np.array([[(x ** ex).prod() for ex in exponents] for x in X]), coefficients)


def evaluate_numerical_error(dim, max_degree):
    # basic idea: evaluating a polynomial at x = all 1 should give the sum of coefficients
    # -> any deviation is the numerical error
    results = []
    x = np.ones(dim, dtype=FLOAT_DTYPE)
    max_error = 0.0
    ctr_total = 0
    ctr_total_max = NR_TEST_POLYNOMIALS * NR_COEFF_CHANGES

    print(f'evaluating numerical error: dim: {dim}, max. degree: {max_degree} ...')
    for poly_ctr, (coefficients, exponents) in enumerate(rnd_settings_list(NR_TEST_POLYNOMIALS, dim, max_degree,
                                                                           max_abs_coeff=MAX_COEFF_MAGNITUDE,
                                                                           integer_coeffs=False)):
        # debug: validate_input=True
        nr_monomials = exponents.shape[0]
        # find factorisation (expensive)
        poly_horner = HornerMultivarPolynomial(coefficients, exponents, validate_input=True)

        for coeff_ctr in range(NR_COEFF_CHANGES):
            # simply change coefficients of the found factorisation (cheap)
            coefficients = (np.random.rand(nr_monomials, 1) - 0.5) * (2 * MAX_COEFF_MAGNITUDE)
            # is testing for in_place=True at the same time
            poly_horner.change_coefficients(coefficients, validate_input=True, in_place=True)
            p_x_horner = poly_horner.eval(x)

            poly = MultivarPolynomial(coefficients, exponents)
            p_x_expected = np.sum(coefficients)
            p_x = poly.eval(x)

            result = (poly, poly_horner, p_x_expected, p_x, p_x_horner)
            results.append(result)
            abs_numerical_error = abs(p_x_horner - p_x_expected)
            max_error = max(max_error, abs_numerical_error)
            sys.stdout.write(f'\r(poly #{poly_ctr + 1} coeff #{coeff_ctr + 1}, {(ctr_total + 1) / ctr_total_max:.1%})'
                             f' max numerical error: {max_error:.2e}')
            sys.stdout.flush()
            if max_error > MAX_NUMERICAL_ERROR:
                # # DEBUG:
                # with open('coefficients.pickle', 'wb') as f:
                #     pickle.dump(coefficients, f)
                # with open('exponents.pickle', 'wb') as f:
                #     pickle.dump(exponents, f)
                raise AssertionError(f'numerical error {max_error:.2e} exceeded limit of {MAX_NUMERICAL_ERROR :.2e} ')
            ctr_total += 1

    print('\n... done.\n')
    return results
