"""Opt-in GPU raw shifted remap lifecycle gates; no precision/performance claim.

Run only on an authorized GPU:
  PYTHONPATH=src python tests/validation/multigrid/test_shifted_remap_gpu.py -v
Frozen gates: bitwise raw overlap/copy and identity continuation; moved-remap
matches the existing host fill loop exactly; finite owned leaves after stepping.
"""
from contextlib import ExitStack
import unittest
import numpy as np
import pycuda.driver as cuda
from CelerisLab.config import LBMConfig
from CelerisLab.cuda.context import CudaContext
from CelerisLab.lbm.population_storage import _StoredPopulations
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
from CelerisLab.lbm.synchronized_remap import RemapTopology, _remap_stored

CFG=LBMConfig(nx=48,ny=32,viscosity=.1,ddf_shifting=True)
PARENTS=((4,4,20,28),(24,4,44,28))
A=((16,16,32,40),)
B=((20,16,36,40),)


def make(ctx):
    return FixedPatchStepper(CFG,context=ctx,refinements=[PARENTS,A],
                             level_blocks=[(4,4)]*3)


def initial(run):
    # Non-unit density, nonzero velocity and nonequilibrium perturbations.
    w=np.array([4/9,*([1/9]*4),*([1/36]*4)],np.float64)
    cx=np.array([0,1,-1,0,0,1,-1,1,-1]);cy=np.array([0,0,0,1,-1,1,-1,-1,1])
    rho,u,v=.99,.013,-.007
    cu=3*(cx*u+cy*v)
    eq=w*((rho-1)+rho*(cu+.5*cu*cu-1.5*(u*u+v*v)))
    rng=np.random.default_rng(8473)
    arrays=tuple((np.broadcast_to(eq[:,None,None],level.shape)+
                  rng.normal(0,1e-7,level.shape)).astype(np.float32) for level in run._levels)
    return _StoredPopulations(arrays,True)


class ShiftedRemapGPU(unittest.TestCase):
    def exact(self,a,b,label):
        self.assertEqual(a.shape,b.shape,label)
        self.assertEqual(a.dtype,b.dtype,label)
        self.assertEqual(a.tobytes(),b.tobytes(),label)

    def snapshot_equal(self,a,b,label):
        self.assertTrue(a.shifted and b.shifted)
        self.assertEqual(len(a.values),len(b.values))
        for x,y in zip(a.values,b.values):self.exact(x,y,label)

    def check_reset(self,run,wanted):
        for level,want in zip(run._levels,wanted):
            blocked=level.layout.to_blocked(want)
            for pointer in (level.current,level.spare):
                actual=np.empty(level.layout.blocked_shape,np.float32)
                cuda.memcpy_dtoh(actual,pointer)
                self.exact(actual,blocked,'fresh current and spare')

    def check_finite(self,run):
        raw=run._download_raw()
        self.assertTrue(raw.shifted)
        for a,leaf in zip(raw.values,run._leaf):
            self.assertTrue(np.isfinite(a[:,leaf]).all())
            self.assertTrue(np.isnan(a[:,~leaf]).all())
            self.assertTrue((1+a[:,leaf].sum(axis=0,dtype=np.float64)>0).all())

    def test_identity_next_step_bitwise(self):
        with CudaContext(0) as ctx,ExitStack() as owned:
            old=owned.enter_context(make(ctx))
            old._initialize_raw(initial(old));old.step(2)
            before=old._download_raw()
            new=owned.enter_context(old.remap_child(A))
            self.check_reset(new,before.values)
            self.snapshot_equal(old._download_raw(),before,'old unchanged')
            self.snapshot_equal(new._download_raw(),before,'identity raw copy')
            old.step();new.step()
            self.assertEqual((old.step_count,new.step_count),(3,3))
            self.snapshot_equal(old._download_raw(),new._download_raw(),'identity next step')
            self.check_finite(old);self.check_finite(new)

    def test_moved_remap_matches_raw_host_and_old_remains_usable(self):
        with CudaContext(0) as ctx,ExitStack() as owned:
            old=owned.enter_context(make(ctx))
            old._initialize_raw(initial(old));old.step(2)
            before=old._download_raw();count=old.step_count
            old_mesh=RemapTopology.from_config(CFG,old.patch,A)
            new_mesh=RemapTopology.from_config(CFG,old.patch,B)
            expected=_remap_stored(old_mesh,new_mesh,before,shifted=True,
                base_count=count,level_steps=(count,2*count,4*count),phase='postcollision_g')
            new=owned.enter_context(old.remap_child(B))
            self.check_reset(new,expected.g)
            self.snapshot_equal(new._download_raw(),_StoredPopulations(expected.g,True),'host raw fill')
            for a,b,om,nm in zip(before.values,new._download_raw().values,old_mesh.leaf_masks,new_mesh.leaf_masks):
                self.exact(a[:,om&nm],b[:,om&nm],'physical same-level overlap')
            self.snapshot_equal(old._download_raw(),before,'source unchanged')
            old.step();new.step(2)
            self.assertEqual((old.step_count,new.step_count),(3,4))
            self.check_finite(old);self.check_finite(new)

    def test_rejected_remap_does_not_mutate_source(self):
        with CudaContext(0) as ctx,make(ctx) as run:
            run._initialize_raw(initial(run));run.step(2)
            before=run._download_raw()
            with self.assertRaises(ValueError):run.remap_child([])
            self.snapshot_equal(run._download_raw(),before,'rejection preserves source')
            self.assertEqual(run.step_count,2)
            run.step();self.check_finite(run)


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