"""
Tests of the linear solver
"""

from . import multigrid

import numpy as np

###### Laplacian in 2D

def test_direct_laplacian():
    nx = 65
    ny = 65

    Lx = 2.
    Ly = 3.
    
    dx = Lx/(nx - 1)
    dy = Ly/(ny - 1)

    xx, yy = np.meshgrid(np.linspace(0,Lx,nx), np.linspace(0,Ly,ny))

    solution = np.exp( - ( (xx - 0.5)**2 + (yy - 0.5)**2 ) / 0.4**2 )
    
    A = multigrid.LaplacianOp()
    rhs = A(solution, dx, dy)

    # Copy boundaries
    rhs[:,0] = solution[:,0]
    rhs[0,:] = solution[0,:]
    rhs[:,-1] = solution[:,-1]
    rhs[-1,:] = solution[-1,:]
    
    solve = multigrid.MGDirect(multigrid.LaplaceSparse(Lx, Ly)(nx, ny))
    x = solve(None, rhs)
    
    assert np.allclose(x, solution)


def test_multigrid_laplacian():
    nx = 65
    ny = 65

    Lx = 2.
    Ly = 3.
    
    dx = Lx/(nx - 1)
    dy = Ly/(ny - 1)

    xx, yy = np.meshgrid(np.linspace(0,Lx,nx), np.linspace(0,Ly,ny))

    solution = np.exp( - ( (xx - 0.5)**2 + (yy - 0.5)**2 ) / 0.4**2 )
    
    A = multigrid.LaplacianOp()
    rhs = A(solution, dx, dy)

    # Copy boundaries
    rhs[:,0] = solution[:,0]
    rhs[0,:] = solution[0,:]
    rhs[:,-1] = solution[:,-1]
    rhs[-1,:] = solution[-1,:]
    
    solve = multigrid.createVcycle(nx,ny, multigrid.LaplaceSparse(Lx, Ly),
                                   ncycle=50, niter=50, nlevels=4, direct=True)

    xinput = np.zeros(rhs.shape)
    xinput[:,0] = solution[:,0]
    xinput[0,:] = solution[0,:]
    xinput[:,-1] = solution[:,-1]
    xinput[-1,:] = solution[-1,:]
    x = solve(xinput, rhs)
    
    assert np.allclose(x, solution, atol=1e-6)
