"""Manual L2 continuation gates; run explicitly on an approved idle GPU.

Frozen numerical gate: max |G - independent FP64 replay| <= 3e-6 on all
ACTIVE leaves, with exact masks and exact same-level overlap at remap. The
separate independent capture/replay acceptance harness supplies that numerical
comparison; these tests cover lifecycle, exact transfer and GPU continuation.
Remap mass/momentum must agree to FP32 roundoff, not per-q conservation during
advancement (collision changes individual populations). No GPU run is implied.
"""
from contextlib import ExitStack
import unittest

import numpy as np
import pycuda.driver as cuda

from CelerisLab.cuda.context import CudaContext
from CelerisLab.lbm.equilibrium import compute_feq_d2q9
from CelerisLab.lbm.refinement_stepper import FixedPatchStepper
from CelerisLab.lbm.synchronized_remap import RemapTopology, remap_synchronized
from test_refinement_blocked_base import exact
from test_refinement_three_level_multipatch_gpu import CFG, PARENTS, frames, physical_state

A = ((16, 16, 32, 40),)
B = ((20, 16, 36, 40),)
FP64_ATOL = 3e-6


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


def initial(run, uniform):
    equilibrium = np.asarray(compute_feq_d2q9(1., .013, -.007), np.float32)[:, None, None]
    return tuple(np.broadcast_to(equilibrium, level.shape).copy() if uniform else
                 physical_state(level.shape[1:], *frame)
                 for level, frame in zip(run._levels, frames(run)))


class RuntimeChildRemapGPU(unittest.TestCase):
    def check_public(self, run):
        for array, leaf in zip(run.download(), (~run.hierarchy.pairs[0].covered, run.leaf[1], run.leaf[2])):
            self.assertTrue(np.isfinite(array[:, leaf]).all())
            self.assertTrue(np.isnan(array[:, ~leaf]).all())

    def check_reset(self, run, expected):
        for level, want in zip(run._levels, expected):
            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)
                exact(self, actual, blocked, 'fresh current and spare, including halos')

    def test_identity_next_step_bitwise_and_old_still_usable(self):
        with CudaContext(0) as ctx, ExitStack() as owned:
            old = owned.enter_context(make(ctx))
            old.initialize(*initial(old, False))
            old.step(2)
            before = old.download()
            new = owned.enter_context(old.remap_finest(A))
            self.assertEqual((old.step_count, new.step_count), (2, 2))
            self.check_reset(new, before)
            for got, want in zip(old.download(), before): exact(self, got, want, 'old unchanged')
            old.step(); new.step()
            self.assertEqual((old.step_count, new.step_count), (3, 3))
            for got, want in zip(new.download(), old.download()):
                exact(self, got, want, 'identity next step')
            self.check_public(new)

    def test_shift_then_return_uniform_and_nonequilibrium(self):
        for uniform in (True, False):
            with self.subTest(uniform=uniform), CudaContext(0) as ctx, ExitStack() as owned:
                current = owned.enter_context(make(ctx))
                current.initialize(*initial(current, uniform))
                current.step(2)
                for rectangles in (B, A):
                    before, count = current.download(), current.step_count
                    old_mesh = RemapTopology.from_config(CFG, current.hierarchy.pairs[0], current.hierarchy.pairs[1].bounds)
                    new_mesh = RemapTopology.from_config(CFG, current.hierarchy.pairs[0], rectangles)
                    expected = remap_synchronized(old_mesh, new_mesh, before, base_count=count,
                        level_steps=(count, 2*count, 4*count), phase='postcollision_g').g
                    previous = current
                    current = owned.enter_context(previous.remap_finest(rectangles))
                    self.assertEqual(current.step_count, count)
                    self.check_reset(current, expected)
                    for got, want, old_g, old_leaf, new_leaf in zip(
                            current.download(), expected, before, old_mesh.leaf_masks, new_mesh.leaf_masks):
                        exact(self, got, want, 'CPU remap exact')
                        overlap = old_leaf & new_leaf
                        exact(self, got[:, overlap], old_g[:, overlap], 'same-level overlap')
                    for got, want in zip(previous.download(), before):
                        exact(self, got, want, 'source unchanged')
                    current.step(2)
                    self.assertEqual(current.step_count, count+2)
                    self.assertEqual(previous.step_count, count)
                    self.check_public(current)
                    if uniform:
                        equilibrium = np.asarray(compute_feq_d2q9(1., .013, -.007), np.float64)
                        for array in current.download():
                            leaf = np.isfinite(array[0])
                            self.assertLessEqual(np.max(np.abs(array[:, leaf]-equilibrium[:, None])), FP64_ATOL)


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