"""
Test the Grid class.

For parallel tests, it is preferable to use a number of processes that is a divisor of
60 (2, 3, 6, etc.).
"""

import pytest

from llg3d import solvers
from llg3d.grid import Grid


@pytest.fixture
def grid_standard():
    """Fixture providing a standard Grid instance."""
    return Grid(Jx=301, Jy=11, Jz=21, dx=2.0)


@pytest.fixture
def grid_unit_spacing():
    """Fixture providing a Grid instance with unit spacing."""
    return Grid(Jx=3 * solvers.size, Jy=3, Jz=3, dx=1.0)


def test_grid_properties(grid_standard: Grid):
    """Test basic Grid properties and calculations."""
    g = grid_standard
    assert g.dx == 2.0
    assert g.dy == 2.0
    assert g.dz == 2.0
    assert g.dV == 8.0
    assert g.V == 480000.0
    assert g.ntot == 301 * 11 * 21
    assert g.ncell == (301 - 1) * (11 - 1) * (21 - 1)
    assert g.Lx == 600.0
    assert g.Ly == 20.0
    assert g.Lz == 40.0
    assert g.dims == (301 // solvers.size, 11, 21)
    assert g.inv_dx2 == 0.25
    assert g.inv_dy2 == 0.25
    assert g.inv_dz2 == 0.25
    assert g.center_coeff == -1.5

def test_get_x_coords_local(grid_unit_spacing: Grid):
    x_local = grid_unit_spacing.get_x_coords(local=True)
    assert x_local.shape == (3,)

def test_get_x_coords_global(grid_unit_spacing: Grid):
    x_global = grid_unit_spacing.get_x_coords(local=False)
    assert x_global.shape == (3 * solvers.size,)

def test_get_mesh_local(grid_unit_spacing: Grid):
    """Test get_mesh with local=True."""
    x, y, z = grid_unit_spacing.get_mesh(local=True)
    assert x.shape == (3, 3, 3)
    assert y.shape == (3, 3, 3)
    assert z.shape == (3, 3, 3)


def test_get_mesh_global(grid_unit_spacing: Grid):
    """Test get_mesh with local=False."""
    xg, yg, zg = grid_unit_spacing.get_mesh(local=False)
    assert xg.shape == (3 * solvers.size, 3, 3)
    assert yg.shape == (3 * solvers.size, 3, 3)
    assert zg.shape == (3 * solvers.size, 3, 3)


def test_grid_as_dict(grid_standard: Grid):
    """Test that as_dict returns correct dictionary representation."""
    d = grid_standard.as_dict()

    assert d["Jx"] == 301
    assert d["Jy"] == 11
    assert d["Jz"] == 21
    assert d["dx"] == 2.0
    assert d["dy"] == 2.0
    assert d["dz"] == 2.0
    assert d["dV"] == 8.0
