""" Test the FDTD grid """


## Imports

# 3rd party
import pytest
import numpy as np

# fdtd
import fdtd
from fdtd.grid import curl_E, curl_H

## Tests


def test_grid_shape_of_ints():
    grid = fdtd.Grid(shape=(3, 3, 3))
    assert (grid.Nx, grid.Ny, grid.Nz) == (3, 3, 3)


def test_grid_shape_of_floats():
    grid = fdtd.Grid(shape=(10.0e-9, 10.0e-9, 10.0e-9), grid_spacing=5.0e-9)
    assert (grid.Nx, grid.Ny, grid.Nz) == (2, 2, 2)


def test_grid_shape_mix_of_floats_and_ints():
    grid = fdtd.Grid(shape=(10.0e-9, 10.0e-9, 3), grid_spacing=5.0e-9)
    assert (grid.Nx, grid.Ny, grid.Nz) == (2, 2, 3)


def test_default_courant_number_1d():
    grid = fdtd.Grid(shape=(3, 1, 1))
    assert grid.courant_number == pytest.approx(1.0, rel=0.02)


def test_default_courant_number_2d():
    grid = fdtd.Grid(shape=(3, 3, 1))
    assert grid.courant_number == pytest.approx((1.0 / 2.0) ** 0.5, rel=0.02)


def test_default_courant_number_3d():
    grid = fdtd.Grid(shape=(3, 3, 3))
    assert grid.courant_number == pytest.approx((1.0 / 3.0) ** 0.5, rel=0.02)


def test_curl_E():
    E = np.random.RandomState(0).randn(3, 3, 3, 3)
    curl = np.array(
        [
            [
                [
                    [-0.99186526, -0.01377993, 2.48607585],
                    [3.44005631, -1.38029691, -0.00954],
                    [-0.10193941, -0.25956782, 0.41985915],
                ],
                [
                    [-3.98489477, 2.19203955, 1.15586708],
                    [-2.55843231, 1.43651777, -1.06280885],
                    [0.25091678, 0.30449392, -1.47938397],
                ],
                [
                    [-1.71853194, -2.98992928, -0.39869962],
                    [2.31880187, 1.76943757, -1.07717648],
                    [0.0, 1.2263907, 0.94356054],
                ],
            ],
            [
                [
                    [-0.70206889, 1.74497282, -2.52192187],
                    [0.19384248, -2.66078848, -0.55523444],
                    [-0.66600115, 0.96949525, -0.28616729],
                ],
                [
                    [3.41574999, -1.51813931, 0.64341128],
                    [-2.47477698, -1.91668701, 2.82293032],
                    [-0.67098, -0.63863509, -3.31485752],
                ],
                [
                    [-1.04005508, -0.80067483, 0.01796954],
                    [0.29806486, 1.68418345, -0.65805687],
                    [0.0, -0.01548234, 0.56697048],
                ],
            ],
            [
                [
                    [0.59435581, 0.33065413, 1.69810037],
                    [-0.81932613, -0.97493235, 1.93267024],
                    [0.94212919, 0.0, -0.72440584],
                ],
                [
                    [0.51876646, 0.09608426, -2.86568329],
                    [0.06214039, 1.68214374, -0.94538826],
                    [-1.29413275, 0.0, 0.36349793],
                ],
                [
                    [-0.36402867, -1.82421078, 0.0],
                    [-0.92696249, 0.37325756, 0.0],
                    [0.0, 0.0, 0.0],
                ],
            ],
        ]
    )
    assert curl_E(E) == pytest.approx(curl)


def test_curl_H():
    H = np.random.RandomState(1).randn(3, 3, 3, 3)
    curl = np.array(
        [
            [
                [
                    [0.0, 0.0, 0.0],
                    [-1.47716404, -2.69731399, 0.0],
                    [1.62661453, 2.81778039, 0.0],
                ],
                [
                    [-1.53196896, 0.0, 1.87371574],
                    [5.28147043, -0.07304683, -0.75055142],
                    [-1.40852366, -0.77747406, 2.84470303],
                ],
                [
                    [0.95952153, 0.0, -0.29158412],
                    [-0.95005061, 1.10250996, -1.46714091],
                    [2.34028677, -0.24386776, -2.00074722],
                ],
            ],
            [
                [
                    [0.0, -1.05852722, 0.34386833],
                    [0.12886545, -1.37025731, -1.26216116],
                    [0.2744926, 0.17815881, 0.08996077],
                ],
                [
                    [1.12944671, -3.71994289, -1.04615132],
                    [0.22579499, 3.88075292, -1.24148611],
                    [-1.8208178, -2.41786863, 1.76683546],
                ],
                [
                    [0.44045296, -3.20087431, -0.87221443],
                    [0.76151166, 0.95947862, 0.3374976],
                    [1.35957083, -0.2459563, -0.85369879],
                ],
            ],
            [
                [
                    [0.0, -0.30862795, 0.85451127],
                    [0.30103587, -0.43231755, 0.68234085],
                    [-0.96728083, -2.21109444, 1.92411429],
                ],
                [
                    [-0.91455513, 1.73537389, 0.34330105],
                    [0.26913557, -1.64348215, 1.51112506],
                    [1.94653494, -1.97285211, -2.49447001],
                ],
                [
                    [0.39120666, 1.78462019, 0.22712353],
                    [-0.17522783, -2.71004261, 2.23042348],
                    [-0.78607747, 2.12528137, 0.27793831],
                ],
            ],
        ]
    )
    assert curl_H(H) == pytest.approx(curl)
