'''
.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
'''

import numpy as np

import pytest

from ... import UnitGrid, CartesianGrid, ScalarField
from .. import solve_laplace_equation, solve_poisson_equation



def test_pde_poisson_solver_1d():
    """ test the poisson solver on 1d grids """
    # solve Laplace's equation
    grid = UnitGrid([4])
    res = solve_laplace_equation(grid, bc=[{'value': -1}, {'value': 3}])
    np.testing.assert_allclose(res.data, grid.axes_coords[0] - 1)
    
    res = solve_laplace_equation(grid, bc=[{'value': -1}, {'derivative': 1}])
    np.testing.assert_allclose(res.data, grid.axes_coords[0] - 1)

    # test Poisson equation with 2nd Order BC
    res = solve_laplace_equation(grid, bc=[{'value': -1}, 'extrapolate'])
    
    # solve Poisson's equation
    grid = CartesianGrid([[0, 1]], 4)
    field = ScalarField(grid, data=1)
    
    res = solve_poisson_equation(field, bc=[{'value': 1}, {'derivative': 1}])
    xs = grid.axes_coords[0]
    np.testing.assert_allclose(res.data, 1 + 0.5 * xs**2, rtol=1e-2)
    
    # test inconsistent problem
    field.data = 1
    with pytest.raises(RuntimeError, match='Neumann'):
        solve_poisson_equation(field, {'derivative': 0})
    


def test_pde_poisson_solver_2d():
    """ test the poisson solver on 2d grids """
    grid = CartesianGrid([[0, 2 * np.pi]] * 2, 16)
    bcs = [{'value': 'sin(y)'}, {'value': 'sin(x)'}]
    
    # solve Laplace's equation
    res = solve_laplace_equation(grid, bcs)
    xs = grid.cell_coords[..., 0]
    ys = grid.cell_coords[..., 1]
    
    # analytical solution was obtained with Mathematica
    expect = (np.cosh(np.pi - ys) * np.sin(xs) +
              np.cosh(np.pi - xs) * np.sin(ys)) / np.cosh(np.pi)
    np.testing.assert_allclose(res.data, expect, atol=1e-2, rtol=1e-2)
    
    # test more complex case for exceptions
    bcs = [{'value': 'sin(y)'}, {'curvature': 'sin(x)'}]
    res = solve_laplace_equation(grid, bc=bcs)
    
    
    