"""Public API demo: small periodic shear, one fixed 2:1 patch, eight coarse steps.

From this checkout, after checking the selected physical GPU is idle:
    PYTHONPATH="$PWD/src" CUDA_VISIBLE_DEVICES=0 python examples/fixed_periodic_patch.py

This demonstrates initialization/stepping/leaf output, not interface accuracy.
There are no walls, objects, forces, checkpoint files or plots.
"""
from dataclasses import replace
from pathlib import Path

import numpy as np

import CelerisLab
from CelerisLab.config import load_lbm_config
from CelerisLab.cuda.context import CudaContext
from CelerisLab.lbm.equilibrium import compute_feq_d2q9, compute_macro_from_ddf
from CelerisLab.lbm.refinement import ACTIVE
from CelerisLab.lbm.refinement_stepper import PeriodicPatchStepper


def shear_equilibrium(shape, origin, dx, period):
    """Sample ux(y)=1e-3*sin(2*pi*y/period) at physical cell centers.

    Centers are origin + dx*(index+.5); this shear is independent of x.
    Fine origin includes patch offset AND the four-cell halo, with dx=.5.
    Equilibrium is already postcollision G, so no initial collision is needed.
    """
    y = origin[1] + dx * (np.arange(shape[0]) + .5)
    ux = 1e-3 * np.sin(2*np.pi*y/period)
    rows = np.stack([compute_feq_d2q9(1.0, float(u), 0.0) for u in ux], axis=1)
    return np.repeat(rows[:, :, None], shape[1], axis=2)


def leaf_summary(coarse, fine, coarse_leaf, fine_leaf):
    """Use the same leaf masks/areas for initial and final mass."""
    c, f = coarse[:, coarse_leaf], fine[:, fine_leaf]
    mass = c.sum(dtype=np.float64) + .25*f.sum(dtype=np.float64)
    macros = np.array([compute_macro_from_ddf(v) for v in np.concatenate((c, f), axis=1).T])
    return mass, macros.min(axis=0), macros.max(axis=0)


def main():
    root = Path(__file__).resolve().parents[1]
    if Path(CelerisLab.__file__).resolve() != root / "src/CelerisLab/__init__.py":
        raise RuntimeError(f"Set PYTHONPATH to {root / 'src'}")
    cfg = replace(load_lbm_config(str(root / "src/CelerisLab/configs/config_lbm.json")),
                  nx=16, ny=16, viscosity=.1, velocity=1e-3, rho=1.0,
                  lattice_model="D2Q9", collision="SRT", streaming="double_buffer",
                  store_precision="FP32", ddf_shifting=False, les_enabled=False)
    with CudaContext(0) as context, PeriodicPatchStepper(cfg, (4, 4, 12, 12), context=context) as run:
        p = run.patch
        coarse = shear_equilibrium(p.coarse_shape, (0, 0), 1.0, cfg.ny)
        fine = shear_equilibrium(p.fine_shape, p.fine_origin, .5, cfg.ny)
        coarse_leaf = np.ones(p.coarse_shape, dtype=bool)
        x0, y0, x1, y1 = p.bounds
        coarse_leaf[y0:y1, x0:x1] = False
        fine_leaf = p.fine_kind == ACTIVE
        coarse[:, ~coarse_leaf] = np.nan
        fine[:, ~fine_leaf] = np.nan
        run.initialize(coarse, fine)
        initial_mass, _, _ = leaf_summary(*run.download(), coarse_leaf, fine_leaf)
        run.step(8)
        coarse, fine = run.download()  # Both levels synchronized, stored postcollision G.
        assert np.isnan(coarse[:, ~coarse_leaf]).all() and np.isnan(fine[:, ~fine_leaf]).all()
        assert np.isfinite(coarse[:, coarse_leaf]).all() and np.isfinite(fine[:, fine_leaf]).all()
        mass, low, high = leaf_summary(coarse, fine, coarse_leaf, fine_leaf)
        print(f"periodic patch API demo: synchronized coarse_steps={run.step_count}, fine_steps={2*run.step_count}")
        print(f"leaf mass: initial={initial_mass:.12g} final={mass:.12g} drift={mass-initial_mass:+.6e}")
        for name, lo, hi in zip(("rho", "ux", "uy"), low, high):
            print(f"leaf {name}: [{lo:.9g}, {hi:.9g}]")
        print("covered coarse / fine halo: NaN; functional demonstration only")


if __name__ == "__main__":
    main()