"""Initializer invariants on a host field; no CUDA context is required."""
from types import SimpleNamespace

import numpy as np
import pytest
from scipy.special import expi

from CelerisLab.lbm.descriptors import D2Q9_EX, D2Q9_EY
from CelerisLab.lbm.equilibrium import compute_feq_d2q9
from CelerisLab.lbm.initializers import _oseen_vortex, add_vortex


class HostField:
    def __init__(self):
        self.nx = self.ny = 7
        self.nq = 9
        self.cfg = SimpleNamespace(is_d2q9=True, viscosity=.02,
                                   streaming="double_buffer")
        self.ddf = np.repeat(compute_feq_d2q9(1.2, .04, .01)[:, None],
                             49, axis=1).ravel()

    def download_ddf(self):
        pass

    def upload_ddf(self):
        pass

    def _read_lattice_vectors(self):
        return np.asarray(D2Q9_EX), np.asarray(D2Q9_EY)


@pytest.mark.parametrize("kind", ["lamb", "oseen", "taylor"])
def test_zero_strength_preserves_moving_nonunit_density_equilibrium(kind):
    field = HostField()
    before = field.ddf.copy()
    add_vortex(field, (3, 3), 2., 0., kind)
    np.testing.assert_allclose(field.ddf, before, rtol=2e-6, atol=2e-8)


def test_oseen_center_limit_and_original_model_away_from_center():
    radius, strength = 2., .02
    r = np.array([0., 1e-10, 1e-5, .2, 1., 10.])
    with np.errstate(all="raise"):
        u, v, p = _oseen_vortex(r, np.zeros_like(r), radius, strength)
    assert np.all(np.isfinite(p))
    assert u[0] == v[0] == 0
    np.testing.assert_allclose(p[:2], -4*radius**2*strength**2*np.log(2),
                               rtol=1e-12)
    z = 4*r[3:]**2/radius**2
    kappa = 2*np.pi*radius**2*strength
    original = -kappa**2/(8*np.pi**2*r[3:]**2) * (
        -2*z*(expi(-z)-expi(-2*z)) + (1-np.exp(-z))**2)
    np.testing.assert_allclose(p[3:], original, rtol=2e-13)


def test_oseen_grid_center_stays_finite():
    field = HostField()
    add_vortex(field, (3, 3), 2., .02, "oseen")
    assert np.isfinite(field.ddf).all()


def test_esopull_is_rejected_before_editing_raw_storage():
    field = HostField()
    field.cfg.streaming = "esopull"
    before = field.ddf.copy()
    with pytest.raises(NotImplementedError, match="time/parity"):
        add_vortex(field, (3, 3), 2., .02)
    np.testing.assert_array_equal(field.ddf, before)