import numpy as np
import sympy as sp
from pystencils import Field, Assignment, create_kernel
from pystencils.bit_masks import flag_cond
from pystencils import TypedSymbol


def test_flag_condition():
    f_arr = np.zeros((2,2,2), dtype=np.float64)
    mask_arr = np.zeros((2,2), dtype=np.uint64)

    mask_arr[0,1] = (1<<3)
    mask_arr[1,0] = (1<<5)
    mask_arr[1,1] = (1<<3) + (1 << 5)

    f = Field.create_from_numpy_array('f', f_arr, index_dimensions=1)
    mask = Field.create_from_numpy_array('mask', mask_arr)

    v1 = 42.3
    v2 = 39.7
    v3 = 119.87

    assignments = [
        Assignment(f(0), flag_cond(3, mask(0), v1)),
        Assignment(f(1), flag_cond(5, mask(0), v2, v3))
    ]

    kernel = create_kernel(assignments).compile()
    kernel(f=f_arr, mask=mask_arr)

    reference = np.zeros((2,2,2), dtype=np.float64)
    reference[0,1,0] = v1
    reference[1,1,0] = v1

    reference[0,0,1] = v3
    reference[0,1,1] = v3

    reference[1,0,1] = v2
    reference[1,1,1] = v2

    np.testing.assert_array_equal(f_arr, reference)