import numpy as np
import numpy.linalg as la  # noqa: F401
import pyopencl as cl  # noqa: F401
import pyopencl.array  # noqa
import pyopencl.tools  # noqa
import pyopencl.clrandom  # noqa
import loopy as lp  # noqa


def test_weno_weight_computation(ctx_factory, flux_test_data_fixture):
    data = flux_test_data_fixture

    def weno_weights(oscillation):
        linear = np.array([0.1, 0.6, 0.3])
        eps = 1e-6

        raw_weights = np.empty((5,3))
        for i in range(5):
            for j in range(3):
                raw_weights[i,j] = linear[j]/(oscillation[i,j] + eps)**2

        weight_sum = raw_weights.sum(axis=1)
        weights = np.empty((5,3))
        for i in range(5):
            for j in range(3):
                weights[i,j] = raw_weights[i,j]/weight_sum[i]

        return weights

    prg = u.get_weno_program_with_root_kernel("weno_weights_pos")
    queue = u.get_queue(ctx_factory)

    weights_dev = u.empty_array_on_device(queue, data.nvars, 3)

    prg(queue, nvars=data.nvars,
            characteristic_fluxes=data.char_fluxes_pos,
            combined_frozen_metrics=1.0,
            w=weights_dev)

    w = weno_weights(data.oscillation_pos)
    u.compare_arrays(weights_dev.get(), w)

    prg = u.get_weno_program_with_root_kernel("weno_weights_neg")
    queue = u.get_queue(ctx_factory)

    weights_dev = u.empty_array_on_device(queue, data.nvars, 3)

    prg(queue, nvars=data.nvars,
            characteristic_fluxes=data.char_fluxes_neg,
            combined_frozen_metrics=1.0,
            w=weights_dev)

    w = weno_weights(data.oscillation_neg)
    u.compare_arrays(weights_dev.get(), w)