"""Direct Simulation Jet control upload; boundary-law accuracy is tested elsewhere."""
import numpy as np
import pycuda.driver as cuda

from CelerisLab.body.jets import Jet, Radial, Sector
from CelerisLab.lbm.descriptors import FLUID
from ._cases import make_sim


def test_simulation_jet_amplitude_upload_and_continued_steps():
    # Same grid/circle/sector as the existing uniform surface-actuation fixture.
    center = (16., 12.)
    with make_sim(nx=32, ny=24, velocity=0) as sim:
        sim.add_body("circle", center=center, radius=4.3)
        index = sim.add_jet(Jet(Sector(center, .7, .85, 4.4),
                                Radial(center, inward=False), "quartic", .002))
        sim.initialize()
        curved = sim.field.curved
        shape, jet_id = curved.jet_shape.copy(), curved.jet_id.copy()
        flags = sim.get_flags()
        driven = jet_id == index
        assert driven.any() and (~driven).any()
        velocity_pointer = int(curved.surface_velocity_gpu)
        previous = None
        for interval, amplitude in enumerate((.002, -.001)):
            if interval:
                count = sim.stepper.step_count
                sim.set_jet_amplitude(index, amplitude)
                assert sim.stepper.step_count == count  # update does not step
            expected = np.zeros_like(curved.surface_velocity)
            expected[driven] = (amplitude * shape[driven]).astype(np.float32)
            actual = np.empty_like(expected)
            with sim.ctx.activate():
                cuda.memcpy_dtoh(actual, curved.surface_velocity_gpu)
            np.testing.assert_array_equal(actual, expected)  # actual device column
            np.testing.assert_array_equal(curved.surface_velocity, expected)
            assert np.any(actual[driven] != 0)
            if previous is not None:
                assert not np.array_equal(actual[driven], previous[driven])
            previous = actual.copy()
            assert sim.bodies.jets[index].amplitude == amplitude
            assert curved.jets[index].amplitude == amplitude
            assert int(curved.surface_velocity_gpu) == velocity_pointer
            np.testing.assert_array_equal(curved.jet_shape, shape)
            np.testing.assert_array_equal(curved.jet_id, jet_id)
            np.testing.assert_array_equal(sim.get_flags(), flags)
            sim.run(2)
            assert sim.stepper.step_count == 2 * (interval + 1)
            assert sim._failed is None
            macro = sim.get_macroscopic()
            fluid = (flags.reshape(24, 32) & FLUID) != 0
            assert all(np.isfinite(value[fluid]).all() for value in macro.values())
            assert np.all(macro["rho"][fluid] > 0)