"""Bounded pinned-body transfer and host lifecycle; no CUDA numerical claim."""
import unittest
from unittest.mock import patch
import numpy as np

from CelerisLab.body.geometry.circle import CircleGeometry
from CelerisLab.body.objects import SimObject, SurfaceActuation
from CelerisLab.body.jets import Jet, Tube, Parallel
from CelerisLab.lbm.synchronized_remap import RemapTopology, same_physical_boundary
from tests.unit.test_synchronized_remap import CFG, PARENTS, fields, remap
from tests.unit.test_refinement_three_level_multipatch import mock_backend
from CelerisLab.lbm.refinement import FixedPatchSet2D

A = ((16,16,32,40), (56,16,72,40))
B = ((16,16,32,40), (60,16,76,40))


def body(spec=None, center=(12.,14.)):
    return SimObject(71, geometry=CircleGeometry(*center,.8), center=center,
                     radius=.8, surface_actuation=spec)


def jet():
    return Jet(Tube((10.5,14.),0.,3.,.3), Parallel(0.), 'parabolic', .003)


def mesh(rects, objects, jets=()):
    return RemapTopology.from_config(CFG, FixedPatchSet2D.from_config(CFG,PARENTS),
                                    rects, child_objects=objects, child_jets=jets)


def ready(objects=None, jets=None):
    objects = (body(),) if objects is None else objects
    jets = (jet(),) if jets is None else jets
    with patch('tests.unit.test_refinement_three_level_multipatch.bodies_and_jets',
               return_value=(objects,jets)):
        run, module, allocations = 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(*fields(mesh(A,objects,jets),np.float32))
    run._step_count = 7
    return run,module,allocations,storage,objects


class PinnedBodyRemap(unittest.TestCase):
    def test_fluid_only_transfer_poison_and_round_trip(self):
        old,new = mesh(A,(body(),)),mesh(B,(body(),))
        arrays=fields(old,np.float32)
        for a,leaf in zip(arrays,old.leaf_masks): a[:,~leaf]=12345
        out=remap(old,new,arrays)
        for a,b,before,after in zip(arrays,out.g,old.leaf_masks,new.leaf_masks):
            np.testing.assert_array_equal(a[:,before & after],b[:,before & after])
            self.assertTrue(np.isnan(b[:,~after]).all())
        self.assertTrue(same_physical_boundary(old.boundary,new.boundary))
        back=remap(new,old,out.g)
        self.assertTrue(np.isnan(back.g[2][:,~old.leaf_masks[2]]).all())
        self.assertEqual(old.pinned_rectangles,(A[0],))

    def test_current_jets_and_source_independence(self):
        run,module,allocations,storage,objects=ready()
        run.set_jet_amplitude(0,.017,level=2)
        before=run.download()
        clone=run.remap_finest(B,objects=objects)
        self.assertEqual(clone._level_boundaries[2].jets[0].amplitude,.017)
        np.testing.assert_array_equal(clone._level_boundaries[2].tables['surface_velocity'],
                                      run._level_boundaries[2].tables['surface_velocity'])
        self.assertTrue(same_physical_boundary(run.boundaries[2],clone.boundaries[2]))
        clone.step()
        self.assertEqual((run.step_count,clone.step_count),(7,8))
        for got,want in zip(run.download(),before): np.testing.assert_array_equal(got,want)
        clone.close();run.close()

    def test_static_body_without_actuation(self):
        run,module,allocations,storage,objects=ready(jets=())
        clone=run.remap_finest(A,objects=objects)
        self.assertFalse(clone._level_boundaries[2].actuated)
        self.assertNotIn('surface_velocity',clone._level_boundaries[2].tables)
        for a,b in zip(run.download(),clone.download()): np.testing.assert_array_equal(a,b)
        clone.close();run.close()

    def test_reject_before_device_reads_or_allocation(self):
        run,module,allocations,storage,objects=ready()
        start=len(allocations)
        cases=[(B,()),(B,run.boundaries[2].objects),(B,(body(center=(12.1,14.)),)),
               (((24,16,36,40),B[1]),objects),(((16,28,32,40),B[1]),objects),
               ((B[1],),objects)]
        module.cuda.memcpy_dtoh.reset_mock()
        for rects,objs in cases:
            with self.subTest(rects=rects),self.assertRaises(ValueError):
                run.remap_finest(rects,objects=objs)
        self.assertEqual(len(allocations),start)
        module.cuda.memcpy_dtoh.assert_not_called()
        self.assertEqual((run._status,run.step_count),('ready',7))
        run.close()

    def test_legacy_velocity_and_zero_velocity_ownership_rejected(self):
        for speed in (.002,0.):
            run,module,allocations,storage,objects=ready(
                (body(SurfaceActuation(0.,.7,speed)),),())
            for spec in (SurfaceActuation(0.,.7,.003),SurfaceActuation(2.,.7,speed)):
                module.cuda.memcpy_dtoh.reset_mock()
                with self.assertRaises(ValueError): run.remap_finest(B,objects=(body(spec),))
                module.cuda.memcpy_dtoh.assert_not_called()
            clone=run.remap_finest(B,objects=objects)
            clone.close();run.close()

    def test_allocation_failure_keeps_old_ready_and_count(self):
        run,module,allocations,storage,objects=ready()
        before=run.download(); saved=list(allocations)
        allocate=module.cuda.mem_alloc.side_effect
        calls=0
        def fail(size):
            nonlocal calls
            calls+=1
            if calls == 3: raise RuntimeError('candidate allocation')
            return allocate(size)
        with patch.object(module.cuda,'mem_alloc',side_effect=fail):
            with self.assertRaisesRegex(RuntimeError,'allocation'):
                run.remap_finest(B,objects=objects)
        self.assertEqual((run._status,run.step_count),('ready',7))
        for got,want in zip(run.download(),before): np.testing.assert_array_equal(got,want)
        for a in saved: a.free.assert_not_called()
        for a in allocations[len(saved):]: a.free.assert_called_once()
        run.close()

    def test_direct_helper_allows_changed_rectangle_with_same_boundary(self):
        old=mesh(A,(body(),))
        new=mesh(((12,16,32,40),A[1]),(body(),))
        out = remap(old,new,fields(old))
        self.assertTrue(same_physical_boundary(old.boundary,new.boundary))
        self.assertIs(out.topology,new)