"""Shifted velocity Jets: bounded opt-in and host lifecycle, no CUDA numerics."""
from dataclasses import replace
import unittest
from unittest.mock import patch
import numpy as np
from CelerisLab.body.jets import Jet, Tube, Parallel
from CelerisLab.body.objects import SurfaceActuation
from CelerisLab.lbm.synchronized_remap import RemapTopology
from CelerisLab.lbm.refinement import FixedPatchSet2D
from CelerisLab.lbm.population_storage import _StoredPopulations
from tests.unit import test_refinement_three_level_multipatch as backend
from tests.unit.test_pinned_body_remap import A, B, body, jet, mesh
from tests.unit.test_shifted_storage import shifted_fields

CFG=replace(backend.CFG,ddf_shifting=True)


def topology(objects, jets, *, enabled=False, cfg=CFG):
    return RemapTopology.from_config(cfg,FixedPatchSet2D.from_config(cfg,backend.PARENTS),A,
        child_objects=objects,child_jets=jets,_shifted_blocked=enabled)


def ready():
    objects,jets=(body(),),(jet(),)
    with patch.object(backend,'CFG',CFG),patch.object(backend,'bodies_and_jets',return_value=(objects,jets)):
        run,module,allocations=backend.mock_backend(children=A,bodies=True)
    storage={}
    module.cuda.memcpy_htod.side_effect=lambda p,a:storage.__setitem__(id(p),np.array(a,copy=True))
    module.cuda.memcpy_dtoh.side_effect=lambda a,p:np.copyto(a,storage[id(p)])
    run._initialize_raw(_StoredPopulations(shifted_fields(mesh(A,objects,jets)),True))
    run._step_count=7
    return run,module,allocations,storage,objects


class ShiftedJets(unittest.TestCase):
    def test_default_guard_and_storage_neutral_host_rows(self):
        with self.assertRaisesRegex(ValueError,'unshifted'):
            topology((body(),),(jet(),))
        shifted=topology((body(),),(jet(),),enabled=True)
        unshifted=topology((body(),),(jet(),),cfg=backend.CFG)
        np.testing.assert_array_equal(shifted.surface_velocity,unshifted.surface_velocity)
        self.assertEqual(shifted.surface_velocity.dtype,np.float32)
        legacy=body(SurfaceActuation(.7,.85,.002))
        with self.assertRaisesRegex(ValueError,'unshifted'):
            topology((legacy,),(),enabled=True)

    def test_rejections_before_device_allocation(self):
        module=backend.host_module()
        moving=body();moving.state.omega=.01
        cases=[((moving,),(jet(),)),((body(),),(jet(),jet())),
               ((body(),),(Jet(Tube((100.,100.),0.,1.,.1),Parallel(0.)),)),
               ((body(center=(8.,14.)),),(jet(),)),
               ((body(),body(center=(34.,14.))),
                (Jet(Tube((10.,14.),0.,26.,.3),Parallel(0.)),))]
        for objects,jets in cases:
            with self.subTest(objects=len(objects),jets=len(jets)),self.assertRaises(ValueError):
                module.FixedPatchStepper(CFG,context=None,refinements=[backend.PARENTS,A],
                    level_blocks=[(4,4)]*3,level_objects=[(),(),objects],level_jets=[(),(),jets])
        for blocks in ((None,None,None),((4,4),None,(4,4))):
            with self.assertRaises(ValueError):
                module.FixedPatchStepper(CFG,context=None,refinements=[backend.PARENTS,A],
                    level_blocks=list(blocks),level_objects=[(),(),(body(),)],
                    level_jets=[(),(),(jet(),)])
        module.cuda.mem_alloc.assert_not_called()
        module.compiler.build_module.assert_not_called()

    def test_switch_upload_only_then_raw_remap_keeps_live_velocity(self):
        run,module,allocations,storage,objects=ready();clone=None
        try:
            boundary=run._level_boundaries[2]
            before=run._download_raw()
            original=boundary.tables['surface_velocity'].copy()
            module.cuda.memcpy_htod.reset_mock()
            run.set_jet_amplitude(0,-jet().amplitude,level=2)
            module.cuda.memcpy_htod.assert_called_once()
            self.assertIs(module.cuda.memcpy_htod.call_args.args[0],boundary.device['surface_velocity'])
            np.testing.assert_array_equal(boundary.tables['surface_velocity'],-original)
            self.assertEqual(run.step_count,7)
            for a,b in zip(before.values,run._download_raw().values):self.assertEqual(a.tobytes(),b.tobytes())
            clone=run.remap_finest(B,objects=objects)
            self.assertEqual(clone.jets[2][0].amplitude,-jet().amplitude)
            np.testing.assert_array_equal(clone._level_boundaries[2].tables['surface_velocity'],
                                          boundary.tables['surface_velocity'])
            self.assertEqual(clone.step_count,7)
            saved=run._download_raw()
            with self.assertRaises(ValueError):run.remap_finest(B,objects=(body(center=(12.1,14.)),))
            for a,b in zip(saved.values,run._download_raw().values):self.assertEqual(a.tobytes(),b.tobytes())
        finally:
            if clone is not None:clone.close()
            run.close()

    def test_invalid_amplitude_does_not_upload_or_commit(self):
        run,module,_,_,_=ready()
        try:
            before=run._level_boundaries[2].tables['surface_velocity'].copy()
            module.cuda.memcpy_htod.reset_mock()
            for amplitude in (float('nan'),float('inf'),1.):
                with self.assertRaises(ValueError):run.set_jet_amplitude(0,amplitude,level=2)
            module.cuda.memcpy_htod.assert_not_called()
            np.testing.assert_array_equal(before,run._level_boundaries[2].tables['surface_velocity'])
        finally:run.close()


if __name__=='__main__':unittest.main()