from typing import List, Tuple
import dataclasses
import pytest
import numpy
from numpy.testing import assert_allclose, assert_array_equal

from .. import fdfd, vec, unvec
from .utils import assert_close, assert_fields_close, PRNG


def test_residual(sim):
    A = fdfd.operators.e_full(sim.omega, sim.dxes, vec(sim.epsilon)).tocsr()
    b = -1j * sim.omega * vec(sim.j)
    residual = A @ vec(sim.e) - b
    assert numpy.linalg.norm(residual) < 1e-10


def test_poynting_planes(sim):
    mask = (sim.j != 0).any(axis=0)
    if mask.sum() != 2:
        pytest.skip(f'test_poynting_planes will only test 2-point sources, got {mask.sum()}')
    points = numpy.where(mask)
    mask[points[0][0], points[1][0], points[2][0]] = 0

    mx = numpy.roll(mask, -1, axis=0)
    my = numpy.roll(mask, -1, axis=1)
    mz = numpy.roll(mask, -1, axis=2)

    e2h = fdfd.operators.e2h(omega=sim.omega, dxes=sim.dxes, pmc=sim.pmc)
    ev = vec(sim.e)
    hv = e2h @ ev

    exh = fdfd.operators.poynting_e_cross(e=ev, dxes=sim.dxes) @ hv.conj()
    s = unvec(exh.real / 2, sim.shape[1:])
    planes = [s[0, mask].sum(), -s[0, mx].sum(),
              s[1, mask].sum(), -s[1, my].sum(),
              s[2, mask].sum(), -s[2, mz].sum()]

    e_dot_j = sim.e * sim.j * sim.dxes[0][0][:, None, None] * sim.dxes[0][1][None, :, None] * sim.dxes[0][2][None, None, :]
    src_energy = -e_dot_j[:, mask].real / 2

    assert_close(sum(planes), src_energy.sum(), rtol=2e-5)  # TODO why is the tolerance so bad?


#####################################
#      Test fixtures
#####################################
# Also see conftest.py

@pytest.fixture(params=[1/1500])
def omega(request):
    yield request.param


@pytest.fixture(params=[None])
def pec(request):
    yield request.param


@pytest.fixture(params=[None])
def pmc(request):
    yield request.param


@pytest.fixture(params=['center', 'diag'])
def j_distribution(request, shape, j_mag):
    j = numpy.zeros(shape, dtype=complex)
    center_mask = numpy.zeros(shape, dtype=bool)
    center_mask[:, shape[1]//2, shape[2]//2, shape[3]//2] = True

    if request.param == 'center':
        j[center_mask] = j_mag
    elif request.param == 'diag':
        j[numpy.roll(center_mask, [1, 1, 1], axis=(1, 2, 3))] = j_mag
        j[numpy.roll(center_mask, [-1, -1, -1], axis=(1, 2, 3))] = -1j * j_mag
    yield j


@dataclasses.dataclass()
class SimResult:
    shape: Tuple[int]
    dxes: List[List[numpy.ndarray]]
    epsilon: numpy.ndarray
    omega: complex
    j: numpy.ndarray
    e: numpy.ndarray
    pmc: numpy.ndarray
    pec: numpy.ndarray


@pytest.fixture()
def sim(request, shape, epsilon, dxes, j_distribution, omega, pec, pmc):
#    is3d = (numpy.array(shape) == 1).sum() == 0
#    if is3d:
#        pytest.skip('Skipping dt != 0.3 because test is 3D (for speed)')

    j_vec = vec(j_distribution)
    eps_vec = vec(epsilon)
    e_vec = fdfd.solvers.generic(J=j_vec, omega=omega, dxes=dxes, epsilon=eps_vec, matrix_solver_opts={'atol': 1e-12})
    e = unvec(e_vec, shape[1:])

    sim = SimResult(
        shape=shape,
        dxes=dxes,
        epsilon=epsilon,
        j=j_distribution,
        e=e,
        pec=pec,
        pmc=pmc,
        omega=omega,
        )

    return sim