# -*- coding: utf-8 -*-

import os

import numpy as np
import pytest

from africanus.gridding.perleypolyhedron import (kernels,
                                                 gridder,
                                                 degridder)
from africanus.dft.kernels import im_to_vis, vis_to_im
from africanus.coordinates import radec_to_lmn
from africanus.constants import c as lightspeed


def test_construct_kernels(tmp_path_factory):
    matplotlib = pytest.importorskip("matplotlib")
    matplotlib.use("agg")
    plt = pytest.importorskip("matplotlib.pyplot")

    plt.figure()
    WIDTH = 5
    OVERSAMP = 101
    ll = kernels.uspace(WIDTH, OVERSAMP)
    sel = ll <= (WIDTH + 2) // 2
    plt.axvline(0.5, 0, 1, ls="--", c="k")
    plt.axvline(-0.5, 0, 1, ls="--", c="k")
    plt.plot(
        ll[sel] * OVERSAMP / 2 / np.pi,
        10 * np.log10(
            np.abs(
                np.fft.fftshift(
                    np.fft.fft(
                        kernels.kbsinc(WIDTH, oversample=OVERSAMP,
                                       order=0)[sel])))),
        label="kbsinc order 0")
    plt.plot(
        ll[sel] * OVERSAMP / 2 / np.pi,
        10 * np.log10(
            np.abs(
                np.fft.fftshift(
                    np.fft.fft(
                        kernels.kbsinc(
                            WIDTH, oversample=OVERSAMP, order=15)[sel])))),
        label="kbsinc order 15")
    plt.plot(ll[sel] * OVERSAMP / 2 / np.pi,
             10 * np.log10(
                    np.abs(
                        np.fft.fftshift(
                            np.fft.fft(
                                kernels.hanningsinc(
                                    WIDTH, oversample=OVERSAMP)[sel])))),
             label="hanning sinc")
    plt.plot(
        ll[sel] * OVERSAMP / 2 / np.pi,
        10 * np.log10(
            np.abs(
                np.fft.fftshift(
                    np.fft.fft(
                        kernels.sinc(WIDTH, oversample=OVERSAMP)[sel])))),
        label="sinc")
    plt.xlim(-10, 10)
    plt.legend()
    plt.ylabel("Response [dB]")
    plt.xlabel("FoV")
    plt.grid(True)
    plt.savefig(tmp_path_factory.mktemp("plots") / "aakernels.png")


def test_taps():
    oversample = 14
    W = 5
    taps = kernels.uspace(W, oversample=oversample)
    assert taps[oversample * ((W + 2) // 2)] == 0
    assert taps[0] == -((W + 2) // 2)
    assert taps[-oversample] == (W + 2) // 2


def test_packunpack():
    oversample = 4
    W = 3
    K = kernels.uspace(W, oversample=oversample)
    Kp = kernels.pack_kernel(K, W, oversample=oversample)
    Kup = kernels.unpack_kernel(Kp, W, oversample=oversample)
    assert np.all(K == Kup)
    assert np.allclose(K, [
        -2.0, -1.75, -1.5, -1.25, -1.0, -0.75, -0.5, -0.25, 0, 0.25, 0.5,
        0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.25, 2.5, 2.75
    ])
    assert np.allclose(Kp, [
        -2.0, -1.0, 0, 1.0, 2.0, -1.75, -0.75, 0.25, 1.25, 2.25, -1.5,
        -0.5, 0.5, 1.5, 2.5, -1.25, -0.25, 0.75, 1.75, 2.75
    ])


def test_facetcodepath():
    # construct kernel
    W = 5
    OS = 3
    kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS),
                               W,
                               oversample=OS)

    # offset 0
    uvw = np.array([[0, 0, 0]])
    vis = np.array([[[1.0 + 0j, 1.0 + 0j]]])
    gridder.gridder(uvw, vis, np.array([1.0]), np.array([0]), 64,
                    30, (0, 0), (0, 0), kern, W, OS, "rotate",
                    "phase_rotate", "I_FROM_XXYY",
                    "conv_1d_axisymmetric_packed_scatter")


def test_degrid_dft(tmp_path_factory):
    # construct kernel
    W = 5
    OS = 3
    kern = kernels.kbsinc(W, oversample=OS)
    uvw = np.column_stack(
        (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 1000)),
            5000.0 * np.sin(np.linspace(0, 2 * np.pi, 1000)), np.zeros(1000)))

    pxacrossbeam = 10
    frequency = np.array([1.4e9])
    wavelength = lightspeed / frequency

    cell = np.rad2deg(
        wavelength[0] /
        (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) *
            pxacrossbeam))
    npix = 512
    mod = np.zeros((1, npix, npix), dtype=np.complex64)
    mod[0, npix // 2 - 5, npix // 2 - 5] = 1.0

    ftmod = np.fft.ifftshift(np.fft.fft2(np.fft.fftshift(
        mod[0, :, :]))).reshape((1, npix, npix))
    chanmap = np.array([0])
    vis_degrid = degridder.degridder(
        uvw,
        ftmod,
        wavelength,
        chanmap,
        cell * 3600.0,
        (0, np.pi / 4.0),
        (0, np.pi / 4.0),
        kern,
        W,
        OS,
        "None",  # no faceting
        "None",  # no faceting
        "XXYY_FROM_I",
        "conv_1d_axisymmetric_unpacked_gather")

    dec, ra = np.meshgrid(
        np.arange(-npix // 2, npix // 2) * np.deg2rad(cell),
        np.arange(-npix // 2, npix // 2) * np.deg2rad(cell))
    radec = np.column_stack((ra.flatten(), dec.flatten()))

    vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, npix * npix).T.copy(),
                        uvw, radec, frequency)

    try:
        import matplotlib
    except ImportError:
        pass
    else:
        matplotlib.use("agg")
        from matplotlib import pyplot as plt
        plt.figure()
        plt.plot(vis_degrid[:, 0, 0].real, label=r"$\Re(\mathtt{degrid})$")
        plt.plot(vis_dft[:, 0, 0].real, label=r"$\Re(\mathtt{dft})$")
        plt.plot(np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real),
                 label="Error")
        plt.legend()
        plt.xlabel("sample")
        plt.ylabel("Real of predicted")
        plt.savefig(
            os.path.join(os.environ.get("TMPDIR", "/tmp"),
                         "degrid_vs_dft_re.png"))
        plt.figure()
        plt.plot(vis_degrid[:, 0, 0].imag, label=r"$\Im(\mathtt{degrid})$")
        plt.plot(vis_dft[:, 0, 0].imag, label=r"$\Im(\mathtt{dft})$")
        plt.plot(np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag),
                 label="Error")
        plt.legend()
        plt.xlabel("sample")
        plt.ylabel("Imag of predicted")
        plt.savefig(tmp_path_factory.mktemp("degrid_dft") /
                    "degrid_vs_dft_im.png")

    assert np.percentile(
        np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real),
        99.0) < 0.05
    assert np.percentile(
        np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag),
        99.0) < 0.05


def test_degrid_dft_packed(tmp_path_factory):
    # construct kernel
    W = 5
    OS = 3
    kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS),
                               W,
                               oversample=OS)
    uvw = np.column_stack(
        (5000.0 * np.cos(np.linspace(0, 2 * np.pi, 1000)),
            5000.0 * np.sin(np.linspace(0, 2 * np.pi, 1000)), np.zeros(1000)))

    pxacrossbeam = 10
    frequency = np.array([1.4e9])
    wavelength = lightspeed / frequency

    cell = np.rad2deg(
        wavelength[0] /
        (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) *
            pxacrossbeam))
    npix = 512
    mod = np.zeros((1, npix, npix), dtype=np.complex64)
    mod[0, npix // 2 - 5, npix // 2 - 5] = 1.0

    ftmod = np.fft.ifftshift(np.fft.fft2(np.fft.fftshift(
        mod[0, :, :]))).reshape((1, npix, npix))
    chanmap = np.array([0])
    vis_degrid = degridder.degridder(
        uvw,
        ftmod,
        wavelength,
        chanmap,
        cell * 3600.0,
        (0, np.pi / 4.0),
        (0, np.pi / 4.0),
        kern,
        W,
        OS,
        "None",  # no faceting
        "None",  # no faceting
        "XXYY_FROM_I",
        "conv_1d_axisymmetric_packed_gather")

    dec, ra = np.meshgrid(
        np.arange(-npix // 2, npix // 2) * np.deg2rad(cell),
        np.arange(-npix // 2, npix // 2) * np.deg2rad(cell))
    radec = np.column_stack((ra.flatten(), dec.flatten()))

    vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1, npix * npix).T.copy(),
                        uvw, radec, frequency)

    try:
        import matplotlib
    except ImportError:
        pass
    else:
        matplotlib.use("agg")
        from matplotlib import pyplot as plt
        plt.figure()
        plt.plot(vis_degrid[:, 0, 0].real, label=r"$\Re(\mathtt{degrid})$")
        plt.plot(vis_dft[:, 0, 0].real, label=r"$\Re(\mathtt{dft})$")
        plt.plot(np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real),
                 label="Error")
        plt.legend()
        plt.xlabel("sample")
        plt.ylabel("Real of predicted")
        plt.savefig(
            os.path.join(os.environ.get("TMPDIR", "/tmp"),
                         "degrid_vs_dft_re_packed.png"))
        plt.figure()
        plt.plot(vis_degrid[:, 0, 0].imag, label=r"$\Im(\mathtt{degrid})$")
        plt.plot(vis_dft[:, 0, 0].imag, label=r"$\Im(\mathtt{dft})$")
        plt.plot(np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag),
                 label="Error")
        plt.legend()
        plt.xlabel("sample")
        plt.ylabel("Imag of predicted")
        plt.savefig(tmp_path_factory.mktemp("degrid_dft_packed") /
                    "degrid_vs_dft_im_packed.png")

    assert np.percentile(
        np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real),
        99.0) < 0.05
    assert np.percentile(
        np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag),
        99.0) < 0.05


def test_detaper(tmp_path_factory):
    W = 5
    OS = 3
    K1D = kernels.kbsinc(W, oversample=OS)
    K2D = np.outer(K1D, K1D)
    detaper = kernels.compute_detaper(128, K2D, W, OS)
    detaperdft = kernels.compute_detaper_dft(128, K2D, W, OS)
    detaperdftsep = kernels.compute_detaper_dft_seperable(128, K1D, W, OS)

    try:
        import matplotlib
    except ImportError:
        pass
    else:
        matplotlib.use("agg")
        from matplotlib import pyplot as plt
        plt.figure()
        plt.subplot(131)
        plt.title("FFT detaper")
        plt.imshow(detaper)
        plt.colorbar()
        plt.subplot(132)
        plt.title("DFT detaper")
        plt.imshow(detaperdft)
        plt.colorbar()
        plt.subplot(133)
        plt.title("ABS error")
        plt.imshow(np.abs(detaper - detaperdft))
        plt.colorbar()
        plt.savefig(tmp_path_factory.mktemp("detaper") / "detaper.png")

    assert (np.percentile(np.abs(detaper - detaperdft), 99.0) < 1.0e-14)
    assert (np.max(np.abs(detaperdft - detaperdftsep)) < 1.0e-14)


def test_grid_dft(tmp_path_factory):
    # construct kernel
    W = 7
    OS = 9
    kern = kernels.kbsinc(W, oversample=OS)
    nrow = 5000
    np.random.seed(0)
    uvw = np.random.normal(scale=6000, size=(nrow, 3))
    uvw[:, 2] = 0.0  # ignore widefield effects for now

    pxacrossbeam = 10
    frequency = np.array([30.0e9])
    wavelength = lightspeed / frequency

    cell = np.rad2deg(
        wavelength[0] /
        (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) *
            pxacrossbeam))
    npix = 256
    fftpad = 1.25
    mod = np.zeros((1, npix, npix), dtype=np.complex64)
    for n in [int(n) for n in np.linspace(npix // 8, 2 * npix // 5, 5)]:
        mod[0, npix // 2 + n, npix // 2 + n] = 1.0
        mod[0, npix // 2 + n, npix // 2 - n] = 1.0
        mod[0, npix // 2 - n, npix // 2 - n] = 1.0
        mod[0, npix // 2 - n, npix // 2 + n] = 1.0
        mod[0, npix // 2, npix // 2 + n] = 1.0
        mod[0, npix // 2, npix // 2 - n] = 1.0
        mod[0, npix // 2 - n, npix // 2] = 1.0
        mod[0, npix // 2 + n, npix // 2] = 1.0

    dec, ra = np.meshgrid(
        np.arange(-npix // 2, npix // 2) * np.deg2rad(cell),
        np.arange(-npix // 2, npix // 2) * np.deg2rad(cell))
    radec = np.column_stack((ra.flatten(), dec.flatten()))

    vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1,
                                             npix * npix).T.copy(), uvw,
                        radec, frequency).repeat(2).reshape(nrow, 1, 2)
    chanmap = np.array([0])

    detaper = kernels.compute_detaper(int(npix * fftpad),
                                      np.outer(kern, kern), W, OS)
    vis_grid = gridder.gridder(
        uvw,
        vis_dft,
        wavelength,
        chanmap,
        int(npix * fftpad),
        cell * 3600.0,
        (0, np.pi / 4.0),
        (0, np.pi / 4.0),
        kern,
        W,
        OS,
        "None",  # no faceting
        "None",  # no faceting
        "I_FROM_XXYY",
        "conv_1d_axisymmetric_unpacked_scatter",
        do_normalize=True)

    ftvis = (np.fft.fftshift(
        np.fft.ifft2(np.fft.ifftshift(vis_grid[0, :, :]))).reshape(
            (1, int(npix * fftpad), int(
                npix * fftpad)))).real / detaper * int(npix * fftpad)**2
    ftvis = ftvis[:,
                  int(npix * fftpad) // 2 -
                  npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix,
                  int(npix * fftpad) // 2 -
                  npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix]
    dftvis = vis_to_im(vis_dft, uvw, radec, frequency,
                       np.zeros(vis_dft.shape,
                                dtype=np.bool)).T.copy().reshape(
                                    2, 1, npix, npix) / nrow

    try:
        import matplotlib
    except ImportError:
        pass
    else:
        matplotlib.use("agg")
        from matplotlib import pyplot as plt
        plt.figure()
        plt.subplot(131)
        plt.title("FFT")
        plt.imshow(ftvis[0, :, :])
        plt.colorbar()
        plt.subplot(132)
        plt.title("DFT")
        plt.imshow(dftvis[0, 0, :, :])
        plt.colorbar()
        plt.subplot(133)
        plt.title("ABS diff")
        plt.imshow(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :]))
        plt.colorbar()
        plt.savefig(tmp_path_factory.mktemp("grid_dft") /
                    "grid_diff_dft.png")

    assert (np.percentile(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :]),
                          95.0) < 0.15)


def test_grid_dft_packed(tmp_path_factory):
    # construct kernel
    W = 7
    OS = 1009
    kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS)
    nrow = 5000
    np.random.seed(0)
    uvw = np.random.normal(scale=6000, size=(nrow, 3))
    uvw[:, 2] = 0.0  # ignore widefield effects for now

    pxacrossbeam = 10
    frequency = np.array([30.0e9])
    wavelength = lightspeed / frequency

    cell = np.rad2deg(
        wavelength[0] /
        (2 * max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) *
            pxacrossbeam))
    npix = 256
    fftpad = 1.25
    mod = np.zeros((1, npix, npix), dtype=np.complex64)
    for n in [int(n) for n in np.linspace(npix // 8, 2 * npix // 5, 5)]:
        mod[0, npix // 2 + n, npix // 2 + n] = 1.0
        mod[0, npix // 2 + n, npix // 2 - n] = 1.0
        mod[0, npix // 2 - n, npix // 2 - n] = 1.0
        mod[0, npix // 2 - n, npix // 2 + n] = 1.0
        mod[0, npix // 2, npix // 2 + n] = 1.0
        mod[0, npix // 2, npix // 2 - n] = 1.0
        mod[0, npix // 2 - n, npix // 2] = 1.0
        mod[0, npix // 2 + n, npix // 2] = 1.0

    dec, ra = np.meshgrid(
        np.arange(-npix // 2, npix // 2) * np.deg2rad(cell),
        np.arange(-npix // 2, npix // 2) * np.deg2rad(cell))
    radec = np.column_stack((ra.flatten(), dec.flatten()))

    vis_dft = im_to_vis(mod[0, :, :].reshape(1, 1,
                                             npix * npix).T.copy(), uvw,
                        radec, frequency).repeat(2).reshape(nrow, 1, 2)
    chanmap = np.array([0])
    detaper = kernels.compute_detaper_dft_seperable(
        int(npix * fftpad), kernels.unpack_kernel(kern, W, OS), W, OS)
    vis_grid = gridder.gridder(
        uvw,
        vis_dft,
        wavelength,
        chanmap,
        int(npix * fftpad),
        cell * 3600.0,
        (0, np.pi / 4.0),
        (0, np.pi / 4.0),
        kern,
        W,
        OS,
        "None",  # no faceting
        "None",  # no faceting
        "I_FROM_XXYY",
        "conv_1d_axisymmetric_packed_scatter",
        do_normalize=True)

    ftvis = (np.fft.fftshift(
        np.fft.ifft2(np.fft.ifftshift(vis_grid[0, :, :]))).reshape(
            (1, int(npix * fftpad), int(
                npix * fftpad)))).real / detaper * int(npix * fftpad)**2
    ftvis = ftvis[:,
                  int(npix * fftpad) // 2 -
                  npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix,
                  int(npix * fftpad) // 2 -
                  npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix]
    dftvis = vis_to_im(vis_dft, uvw, radec, frequency,
                       np.zeros(vis_dft.shape,
                                dtype=np.bool)).T.copy().reshape(
                                    2, 1, npix, npix) / nrow

    try:
        import matplotlib
    except ImportError:
        pass
    else:
        matplotlib.use("agg")
        from matplotlib import pyplot as plt
        plt.figure()
        plt.subplot(131)
        plt.title("FFT")
        plt.imshow(ftvis[0, :, :])
        plt.colorbar()
        plt.subplot(132)
        plt.title("DFT")
        plt.imshow(dftvis[0, 0, :, :])
        plt.colorbar()
        plt.subplot(133)
        plt.title("ABS diff")
        plt.imshow(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :]))
        plt.colorbar()
        plt.savefig(tmp_path_factory.mktemp("grid_dft_packed") /
                    "grid_diff_dft_packed.png")

    assert (np.percentile(np.abs(ftvis[0, :, :] - dftvis[0, 0, :, :]),
                          95.0) < 0.15)


def test_wcorrection_faceting_backward(tmp_path_factory):
    # construct kernel
    W = 5
    OS = 9
    kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS)
    nrow = 5000
    np.random.seed(0)
    # simulate some ficticious baselines rotated by an hour angle
    uvw = np.zeros((nrow, 3), dtype=np.float64)
    blpos = np.random.uniform(26, 10000, size=(25, 3))
    ntime = int(nrow / 25.0)
    d0 = np.pi / 4.0
    for n in range(25):
        for ih0, h0 in enumerate(
                np.linspace(np.deg2rad(-20), np.deg2rad(20), ntime)):
            s = np.sin
            c = np.cos
            R = np.array([[s(h0), c(h0), 0],
                          [-s(d0) * c(h0),
                           s(d0) * s(h0),
                           c(d0)], [c(d0) * c(h0), -c(d0) * s(h0),
                                    s(d0)]])
            uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T)

    pxacrossbeam = 5
    frequency = np.array([1.4e9])
    wavelength = lightspeed / frequency

    cell = np.rad2deg(
        wavelength[0] /
        (max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) *
            pxacrossbeam))
    npix = 2048
    npixfacet = 100
    fftpad = 1.1
    mod = np.ones((1, 1, 1), dtype=np.complex64)
    deltaradec = np.array(
        [[600 * np.deg2rad(cell), 600 * np.deg2rad(cell)]])
    lm = radec_to_lmn(deltaradec + np.array([[0, d0]]),
                      phase_centre=np.array([0, d0]))

    vis_dft = im_to_vis(mod, uvw, lm[:, 0:2],
                        frequency).repeat(2).reshape(nrow, 1, 2)
    chanmap = np.array([0])

    detaper = kernels.compute_detaper_dft_seperable(
        int(npix * fftpad), kernels.unpack_kernel(kern, W, OS), W, OS)
    vis_grid_nofacet = gridder.gridder(
        uvw,
        vis_dft,
        wavelength,
        chanmap,
        int(npix * fftpad),
        cell * 3600.0,
        (0, d0),
        (0, d0),
        kern,
        W,
        OS,
        "None",  # no faceting
        "None",  # no faceting
        "I_FROM_XXYY",
        "conv_1d_axisymmetric_packed_scatter",
        do_normalize=True)
    ftvis = (np.fft.fftshift(
        np.fft.ifft2(np.fft.ifftshift(vis_grid_nofacet[0, :, :]))).reshape(
            (1, int(npix * fftpad), int(
                npix * fftpad)))).real / detaper * int(npix * fftpad)**2
    ftvis = ftvis[:,
                  int(npix * fftpad) // 2 -
                  npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix,
                  int(npix * fftpad) // 2 -
                  npix // 2:int(npix * fftpad) // 2 - npix // 2 + npix]

    detaper_facet = kernels.compute_detaper_dft_seperable(
        int(npixfacet * fftpad), kernels.unpack_kernel(kern, W, OS), W, OS)
    vis_grid_facet = gridder.gridder(
        uvw,
        vis_dft,
        wavelength,
        chanmap,
        int(npixfacet * fftpad),
        cell * 3600.0, (deltaradec + np.array([[0, d0]]))[0, :], (0, d0),
        kern,
        W,
        OS,
        "rotate",
        "phase_rotate",
        "I_FROM_XXYY",
        "conv_1d_axisymmetric_packed_scatter",
        do_normalize=True)
    ftvisfacet = (np.fft.fftshift(
        np.fft.ifft2(np.fft.ifftshift(vis_grid_facet[0, :, :]))).reshape(
            (1, int(npixfacet * fftpad), int(
                npixfacet * fftpad)))).real / detaper_facet * int(
                    npixfacet * fftpad)**2
    ftvisfacet = ftvisfacet[:,
                            int(npixfacet * fftpad) // 2 -
                            npixfacet // 2:int(npixfacet * fftpad) // 2 -
                            npixfacet // 2 + npixfacet,
                            int(npixfacet * fftpad) // 2 -
                            npixfacet // 2:int(npixfacet * fftpad) // 2 -
                            npixfacet // 2 + npixfacet]

    try:
        import matplotlib
    except ImportError:
        pass
    else:
        matplotlib.use("agg")
        from matplotlib import pyplot as plt
        plot_dir = tmp_path_factory.mktemp("wcorrection_backward")

        plt.figure()
        plt.subplot(121)
        plt.imshow(ftvis[0, 1624 - 50:1624 + 50, 1447 - 50:1447 + 50])
        plt.colorbar()
        plt.title("Offset FFT (peak={0:.1f})".format(np.max(ftvis)))
        plt.subplot(122)
        plt.imshow(ftvisfacet[0, :, :])
        plt.colorbar()
        plt.title("Faceted FFT (peak={0:.1f})".format(np.max(ftvisfacet)))
        plt.savefig(plot_dir / "facet_imaging.png")

    assert (np.abs(np.max(ftvisfacet[0, :, :]) - 1.0) < 1.0e-6)


def test_wcorrection_faceting_forward(tmp_path_factory):
    # construct kernel
    W = 5
    OS = 9
    kern = kernels.pack_kernel(kernels.kbsinc(W, oversample=OS), W, OS)
    nrow = 5000
    np.random.seed(0)
    # simulate some ficticious baselines rotated by an hour angle
    uvw = np.zeros((nrow, 3), dtype=np.float64)
    blpos = np.random.uniform(26, 10000, size=(25, 3))
    ntime = int(nrow / 25.0)
    d0 = np.pi / 4.0
    for n in range(25):
        for ih0, h0 in enumerate(
                np.linspace(np.deg2rad(-20), np.deg2rad(20), ntime)):
            s = np.sin
            c = np.cos
            R = np.array([[s(h0), c(h0), 0],
                          [-s(d0) * c(h0),
                           s(d0) * s(h0),
                           c(d0)], [c(d0) * c(h0), -c(d0) * s(h0),
                                    s(d0)]])
            uvw[n * ntime + ih0, :] = np.dot(R, blpos[n, :].T)

    pxacrossbeam = 5
    frequency = np.array([1.4e9])
    wavelength = lightspeed / frequency

    cell = np.rad2deg(
        wavelength[0] /
        (max(np.max(np.abs(uvw[:, 0])), np.max(np.abs(uvw[:, 1]))) *
            pxacrossbeam))
    npixfacet = 100
    mod = np.ones((1, 1, 1), dtype=np.complex64)
    deltaradec = np.array([[20 * np.deg2rad(cell), 20 * np.deg2rad(cell)]])
    lm = radec_to_lmn(deltaradec + np.array([[0, d0]]),
                      phase_centre=np.array([0, d0]))

    vis_dft = im_to_vis(mod, uvw, lm[:, 0:2],
                        frequency).repeat(2).reshape(nrow, 1, 2)
    chanmap = np.array([0])
    ftmod = np.ones((1, npixfacet, npixfacet),
                    dtype=np.complex64)  # point source at centre of facet
    vis_degrid = degridder.degridder(
        uvw,
        ftmod,
        wavelength,
        chanmap,
        cell * 3600.0,
        (deltaradec + np.array([[0, d0]]))[0, :],
        (0, d0),
        kern,
        W,
        OS,
        "rotate",  # no faceting
        "phase_rotate",  # no faceting
        "XXYY_FROM_I",
        "conv_1d_axisymmetric_packed_gather")

    try:
        import matplotlib
    except ImportError:
        pass
    else:
        matplotlib.use("agg")
        from matplotlib import pyplot as plt
        plot_dir = tmp_path_factory.mktemp("wcorrection_forward")

        plt.figure()
        plt.plot(vis_degrid[:, 0, 0].real,
                 label=r"$\Re(\mathtt{degrid facet})$")
        plt.plot(vis_dft[:, 0, 0].real, label=r"$\Re(\mathtt{dft})$")
        plt.plot(np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real),
                 label="Error")
        plt.legend()
        plt.xlabel("sample")
        plt.ylabel("Real of predicted")
        plt.savefig(plot_dir / "facet_degrid_vs_dft_re_packed.png")
        plt.figure()
        plt.plot(vis_degrid[:, 0, 0].imag,
                 label=r"$\Im(\mathtt{degrid facet})$")
        plt.plot(vis_dft[:, 0, 0].imag, label=r"$\Im(\mathtt{dft})$")
        plt.plot(np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag),
                 label="Error")
        plt.legend()
        plt.xlabel("sample")
        plt.ylabel("Imag of predicted")
        plt.savefig(plot_dir / "facet_degrid_vs_dft_im_packed.png")

    assert np.percentile(
        np.abs(vis_dft[:, 0, 0].real - vis_degrid[:, 0, 0].real),
        99.0) < 0.05
    assert np.percentile(
        np.abs(vis_dft[:, 0, 0].imag - vis_degrid[:, 0, 0].imag),
        99.0) < 0.05
