# CelerisLab/lbm/stepper.py
"""
LBMStepper — time-advance driver for LBM kernels.

Owns kernel function handles and manages the double-buffer swap.
Does NOT return observations by default — callers use field.get_macroscopic()
or body manager methods to pull data when they need it.
"""

import numpy as np
import pycuda.driver as cuda


class LBMStepper:
    """Drive the LBM kernel forward in time."""

    def __init__(self, field, module: cuda.Module, cfg):
        self.field = field
        self.module = module
        self.cfg = cfg

        self._esopull = (cfg.streaming == "esopull")

        if self._esopull:
            self.step_fn = module.get_function("EsoPullStep")
            self.init_fn = module.get_function("InitEsoPull")
        else:
            self.step_fn = module.get_function("OneStep")
            self.init_fn = module.get_function("InitTubeFlow_v2")

        self.curved_fn = module.get_function("CurvedBoundaryKernel")
        self.prepare_esopull_curved_fn = module.get_function("PrepareEsoPullCurvedKernel")
        self.apply_esopull_curved_fn = module.get_function("ApplyEsoPullCurvedKernel")
        self.sensor_fn = module.get_function("SensorKernel")
        self.force_region_fn = module.get_function("ForceRegionKernel")

        tpb = cfg.threads_per_block
        self.block = (tpb, 1, 1)
        if cfg.is_d2q9:
            self.grid = ((cfg.nx + tpb - 1) // tpb, cfg.ny, 1)
        else:
            self.grid = ((cfg.nx + tpb - 1) // tpb, cfg.ny, cfg.nz)

        self._step_count = 0

    def initialize(self):
        """Run the init kernel to set up channel flow + flags."""
        f = self.field
        self.init_fn(
            f.flag_gpu, f.ddf_gpu,
            block=self.block, grid=self.grid,
        )
        if not self._esopull:
            _gpu_bytes = f.n * f.nq * f.store_bytes
            cuda.memcpy_dtod(f.temp_gpu, f.ddf_gpu, _gpu_bytes)
        cuda.memcpy_dtoh(f.flag, f.flag_gpu)
        f.download_ddf(step_id=self._step_count, force=True)

    def step(
        self,
        n: int = 1,
        *,
        action_gpu: cuda.DeviceAllocation,
        obs_gpu: cuda.DeviceAllocation,
        stream: cuda.Stream | None = None,
    ):
        """Advance *n* time steps."""
        f = self.field
        launch_kw = {}
        if stream is not None:
            launch_kw["stream"] = stream

        for _ in range(n):
            self._launch_curved(action_gpu, obs_gpu, **launch_kw)

            if self._esopull:
                self.step_fn(
                    f.ddf_gpu, f.flag_gpu,
                    action_gpu, obs_gpu,
                    np.uint64(self._step_count),
                    block=self.block, grid=self.grid,
                    **launch_kw,
                )
            else:
                self.step_fn(
                    f.flag_gpu, f.ddf_gpu, f.temp_gpu,
                    block=self.block, grid=self.grid,
                    **launch_kw,
                )
                f.ddf_gpu, f.temp_gpu = f.temp_gpu, f.ddf_gpu

            self._launch_force_region(action_gpu, **launch_kw)
            self._launch_sensor(obs_gpu, **launch_kw)
            self._step_count += 1
        if n > 0:
            # Any completed GPU step invalidates host DDF mirror.
            f.invalidate_host_ddf_cache()

    def _launch_curved(self, action_gpu, obs_gpu, **launch_kw):
        f = self.field
        if f.n_curved == 0:
            return
        tpb = self.cfg.threads_per_block
        grid_c = ((f.n_curved + tpb - 1) // tpb, 1, 1)
        cl = f.curved
        assert cl.fluid_idx_gpu is not None
        assert cl.dir_gpu is not None
        assert cl.q_gpu is not None
        assert cl.rx_gpu is not None
        assert cl.ry_gpu is not None
        assert cl.rz_gpu is not None
        assert cl.fallback_class_gpu is not None
        assert cl.body_id_gpu is not None

        if self._esopull:
            eb = f.esopull_curved
            assert eb.corr_value_gpu is not None
            assert eb.target_cell_gpu is not None
            assert eb.target_slot_gpu is not None
            self.prepare_esopull_curved_fn(
                f.ddf_gpu,
                cl.fluid_idx_gpu, cl.dir_gpu, cl.q_gpu,
                cl.rx_gpu, cl.ry_gpu, cl.fallback_class_gpu, cl.body_id_gpu,
                action_gpu, obs_gpu,
                eb.corr_value_gpu, eb.target_cell_gpu, eb.target_slot_gpu,
                np.uint32(f.n_curved),
                np.uint64(self._step_count),
                block=(tpb, 1, 1), grid=grid_c,
                **launch_kw,
            )
            self.apply_esopull_curved_fn(
                f.ddf_gpu,
                eb.corr_value_gpu, eb.target_cell_gpu, eb.target_slot_gpu,
                np.uint32(f.n_curved),
                block=(tpb, 1, 1), grid=grid_c,
                **launch_kw,
            )
        else:
            assert cl.fluid_idx_gpu is not None
            self.curved_fn(
                f.ddf_gpu,
                cl.fluid_idx_gpu, cl.dir_gpu, cl.q_gpu,
                cl.rx_gpu, cl.ry_gpu, cl.rz_gpu, cl.fallback_class_gpu, cl.body_id_gpu,
                action_gpu, obs_gpu,
                np.uint32(f.n_curved),
                block=(tpb, 1, 1), grid=grid_c,
                **launch_kw,
            )

    def _launch_sensor(self, obs_gpu, **launch_kw):
        f = self.field
        if f.n_sensor == 0:
            return
        tpb = self.cfg.threads_per_block
        grid_s = ((f.n_sensor + tpb - 1) // tpb, 1, 1)
        fi_in = f.ddf_gpu
        se = f.sensors
        assert se.cells_gpu is not None and se.obj_id_gpu is not None
        t = np.uint64(self._step_count + 1 if self._esopull else 0)
        self.sensor_fn(
            fi_in, f.flag_gpu,
            se.cells_gpu, se.obj_id_gpu,
            obs_gpu,
            np.uint32(f.n_sensor),
            t,
            block=(tpb, 1, 1), grid=grid_s,
            **launch_kw,
        )

    def _launch_force_region(self, action_gpu, **launch_kw):
        f = self.field
        if f.n_force_region == 0:
            return
        tpb = self.cfg.threads_per_block
        grid_f = ((f.n_force_region + tpb - 1) // tpb, 1, 1)
        fr = f.force_regions
        assert fr.cells_gpu is not None and fr.obj_id_gpu is not None
        t = np.uint64(self._step_count + 1 if self._esopull else 0)
        self.force_region_fn(
            f.ddf_gpu, f.flag_gpu,
            fr.cells_gpu, fr.obj_id_gpu,
            action_gpu,
            np.uint32(f.n_force_region),
            t,
            block=(tpb, 1, 1), grid=grid_f,
            **launch_kw,
        )

    @property
    def step_count(self) -> int:
        return self._step_count
