from itertools import product

import nifty_gridder as ng
import numpy as np
import pytest
from numpy.testing import assert_, assert_allclose, assert_array_almost_equal

pmp = pytest.mark.parametrize


def _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow):
    pixsize = np.pi/180/60/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=0.368*pixsize,
                            pixsize_y=pixsize)
    speedoflight, f0 = 299792458., 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
    baselines = ng.Baselines(coord=uvw, freq=freq)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
    idx = ng.getIndices(baselines, conf, flags)
    return baselines, conf, idx


def _l2error(a, b):
    return np.sqrt(np.sum(np.abs(a-b)**2)/np.maximum(np.sum(np.abs(a)**2),
                                                     np.sum(np.abs(b)**2)))


def explicit_gridder(uvw, freq, ms, wgt, nxdirty, nydirty, xpixsize, ypixsize,
                     apply_w):
    speedoflight = 299792458.
    x, y = np.meshgrid(*[-ss/2 + np.arange(ss) for ss in [nxdirty, nydirty]],
                       indexing='ij')
    x *= xpixsize
    y *= ypixsize
    res = np.zeros((nxdirty, nydirty))
    eps = x**2+y**2
    if apply_w:
        nm1 = -eps/(np.sqrt(1.-eps)+1.)
        n = nm1+1
    else:
        nm1 = 0.
        n = 1.
    for row in range(ms.shape[0]):
        for chan in range(ms.shape[1]):
            phase = (freq[chan]/speedoflight *
                     (x*uvw[row, 0] + y*uvw[row, 1] - uvw[row, 2]*nm1))
            if wgt is None:
                res += (ms[row, chan]*np.exp(2j*np.pi*phase)).real
            else:
                res += (ms[row, chan]*wgt[row, chan]*np.exp(2j*np.pi*phase)).real
    return res/n


@pmp("nxdirty", (30, 128))
@pmp("nydirty", (128, 250))
@pmp("nrow", (2, 27))
@pmp("nchan", (1, 5))
@pmp("epsilon", (1e-3, 1e-5, 1e-7, 5e-13))
@pmp("singleprec", (True, False))
@pmp("wstacking", (True, False))
@pmp("use_wgt", (True, False))
@pmp("nthreads", (1, 2))
def test_adjointness_ms2dirty(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking, use_wgt, nthreads):
    if singleprec and epsilon < 5e-6:
        return
    np.random.seed(42)
    pixsizex = np.pi/180/60/nxdirty*0.2398
    pixsizey = np.pi/180/60/nxdirty
    speedoflight, f0 = 299792458., 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsizey*f0/speedoflight)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    wgt = np.random.rand(nrow, nchan) if use_wgt else None
    dirty = np.random.rand(nxdirty, nydirty)-0.5
    if singleprec:
        ms = ms.astype("c8")
        dirty = dirty.astype("f4")
        if wgt is not None:
            wgt = wgt.astype("f4")
    dirty2 = ng.ms2dirty(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex,
                         pixsizey, epsilon, wstacking, nthreads, 0).astype("f8")
    ms2 = ng.dirty2ms(uvw, freq, dirty, wgt, pixsizex, pixsizey, epsilon,
                      wstacking, nthreads, 0).astype("c16")
    tol = 5e-5 if singleprec else 5e-13
    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty2, dirty), rtol=tol)


@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (2, 27))
@pmp("nchan", (1, 5))
@pmp("epsilon", (1e-2, 1e-3, 5e-5, 1e-7, 5e-13))
@pmp("singleprec", (True, False))
@pmp("wstacking", (True, False))
@pmp("use_wgt", (False, True))
@pmp("nthreads", (1, 2))
@pmp("fov", (1., 20.))
def test_ms2dirty_against_wdft(nxdirty, nydirty, nrow, nchan, epsilon, singleprec, wstacking, use_wgt, fov, nthreads):
    if singleprec and epsilon < 5e-5:
        return
    np.random.seed(40)
    pixsizex = fov*np.pi/180/nxdirty
    pixsizey = fov*np.pi/180/nydirty*1.1
    speedoflight, f0 = 299792458., 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsizex*f0/speedoflight)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    wgt = np.random.rand(nrow, nchan) if use_wgt else None
    if singleprec:
        ms = ms.astype("c8")
        if wgt is not None:
            wgt = wgt.astype("f4")
    dirty = ng.ms2dirty(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex,
                        pixsizey, epsilon, wstacking, nthreads, 0).astype("f8")
    ref = explicit_gridder(uvw, freq, ms, wgt, nxdirty, nydirty, pixsizex, pixsizey, wstacking)
    assert_allclose(_l2error(dirty, ref), 0, atol=epsilon)


@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 100))
@pmp("nchan", (1, 10))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
def test_adjointness(nxdirty, nydirty, nrow, nchan, epsilon):
    np.random.seed(42)
    bl, gconf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    grid = ng.vis2grid_c(bl, gconf, idx, vis)
    dirty = gconf.grid2dirty_c(grid)
    dirty2 = (np.random.rand(*dirty.shape)-0.5 +
              1j*(np.random.rand(*dirty.shape)-0.5))
    ms2 = bl.vis2ms(ng.grid2vis_c(bl, gconf, idx, gconf.dirty2grid_c(dirty2)),
                    idx)
    assert_allclose(np.vdot(ms, ms2).real, np.vdot(dirty, dirty2).real)


@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 127))
@pmp("nchan", (1, 17))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
def test_adjointness_wgridding(nxdirty, nydirty, nrow, nchan, epsilon):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    vis = np.random.rand(*idx.shape)-0.5 + 1j*(np.random.rand(*idx.shape)-0.5)
    dirty = np.random.rand(conf.Nxdirty(), conf.Nydirty())-0.5
    dirty2 = ng.vis2dirty(bl, conf, idx, vis, do_wstacking=True)
    vis2 = ng.dirty2vis(bl, conf, idx, dirty, do_wstacking=True)
    assert_allclose(np.vdot(vis, vis2).real, np.vdot(dirty2, dirty),
                    rtol=2e-13)


@pmp("nxdirty", (128,))
@pmp("nydirty", (128,))
@pmp("nrow", (1000,))
@pmp("nchan", (10,))
@pmp("epsilon", (2e-13,))
def test_hoisted_grid_allocation(nxdirty, nydirty, nrow, nchan, epsilon):
    np.random.seed(42)
    gconf = ng.GridderConfig(1024, 1024, 2e-13, 2.0, 2.0)
    freq = np.linspace(.856e9, 2*.856e9, nchan)
    f0 = freq[freq.shape[0]//2]

    pixsize = np.pi/180/60/nxdirty  # assume 1 arcmin FOV
    speedoflight = 299792458.
    uvw = (np.random.rand(nrow, 3)-0.5) / (pixsize*f0/speedoflight)
    baselines = ng.Baselines(coord=uvw, freq=freq)
    gconf = ng.GridderConfig(nxdirty=nxdirty, nydirty=nydirty,
                             epsilon=epsilon, pixsize_x=pixsize,
                             pixsize_y=pixsize)

    flags = np.zeros((nrow, nchan), dtype=np.bool)
    weights = np.random.rand(nrow, nchan)
    idx = ng.getIndices(baselines, gconf, flags)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = baselines.ms2vis(ms, idx)

    # ----------------------------------
    # Test complex grid case (ms2grid_c)
    # ----------------------------------
    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
    grid = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=None)
    grid2 = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=cplx_grid)

    # Almost same result
    assert_array_almost_equal(grid, grid2)
    assert grid.dtype == cplx_grid.dtype == grid2.dtype
    assert id(grid2) == id(cplx_grid)  # Same grid object

    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
    grid = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=None, wgt=weights)
    grid2 = ng.ms2grid_c(baselines, gconf, idx, ms, grid_in=cplx_grid,
                         wgt=weights)

    # Almost same result
    assert_array_almost_equal(grid, grid2)
    assert grid.dtype == cplx_grid.dtype == grid2.dtype
    assert id(grid2) == id(cplx_grid)  # Same grid object

    # ----------------------------------
    # Test complex grid case (vis2grid_c)
    # ----------------------------------
    cplx_grid = np.zeros((gconf.Nu(), gconf.Nv()), dtype=np.complex128)
    grid = ng.vis2grid_c(baselines, gconf, idx, vis, grid_in=None)
    grid2 = ng.vis2grid_c(baselines, gconf, idx, vis, grid_in=cplx_grid)

    # Almost same result
    assert_array_almost_equal(grid, grid2)
    assert grid.dtype == cplx_grid.dtype == grid2.dtype
    assert id(grid2) == id(cplx_grid)  # Same grid object


def test_pickling():
    np.random.seed(42)
    try:
        import cPickle as pickle
    except ImportError:
        import pickle

    import nifty_gridder as ng

    # Have to use cPickle and pickle protocol 2 for this to work
    # unfortunately
    pickle_protocol = 2

    gc = ng.GridderConfig(1024, 1024, 2e-13, 2.0, 2.0)
    pickled = pickle.dumps(gc, pickle_protocol)
    gc2 = pickle.loads(pickled)

    assert_(gc is not gc2)
    assert_(gc.Nxdirty() == gc2.Nxdirty())
    assert_(gc.Nydirty() == gc2.Nydirty())
    assert_(gc.Epsilon() == gc2.Epsilon())
    assert_(gc.Pixsize_x() == gc2.Pixsize_x())
    assert_(gc.Pixsize_y() == gc2.Pixsize_y())


@pmp("nxdirty", (128, 300))
@pmp("nydirty", (128, 250))
@pmp("nrow", (1, 39))
@pmp("nchan", (1, 70))
@pmp("epsilon", (1e-2, 1e-7, 2e-13))
@pmp("weight", (True, False))
def test_applyholo(nxdirty, nydirty, nrow, nchan, epsilon, weight):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    grid = ng.vis2grid_c(bl, conf, idx, vis)
    wgt = np.random.rand(*vis.shape) if weight else None
    y0 = ng.apply_holo(bl, conf, idx, grid, wgt)
    y1 = ng.vis2grid_c(bl, conf, idx, ng.grid2vis_c(bl, conf, idx, grid, wgt),
                       wgt=wgt)
    assert_allclose(y0, y1)


@pmp("nxdirty", (300,))
@pmp("nydirty", (250,))
@pmp("nrow", (10, 100))
@pmp("nchan", (1, 10))
@pmp("epsilon", (1e-2, 1e-7))
@pmp("du", (-2, 0, 1, 4))
@pmp("dv", (-1, 0, 1))
@pmp("weight", (True, False))
def test_correlations(nxdirty, nydirty, nrow, nchan, epsilon, du, dv, weight):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    if conf.Supp() <= abs(du) or conf.Supp() <= dv:
        return
    wgt = np.random.rand(*idx.shape) if weight else None
    y0 = ng.get_correlations(bl, conf, idx, du, dv, wgt)
    pts = ((nxdirty, nydirty),
           (nxdirty + 1, nydirty),
           (nxdirty, nydirty//2),
           (0, 0),
           (2*nxdirty-1, 2*nydirty-1))
    for pp in pts:
        grid = np.zeros((2*nxdirty, 2*nydirty))+0j
        grid[pp] = 1
        y1 = ng.apply_holo(bl, conf, idx, grid, wgt)
        ind = (pp[0]+du) % (2*nxdirty), (pp[1]+dv) % (2*nydirty)
        assert_allclose(y0[pp], y1[ind].real)
        assert_allclose(np.zeros_like(y1), y1.imag)


@pmp("nrow", (1, 100))
@pmp("nchan", (1, 10))
@pmp("epsilon", (1e-2, 1e-7, 2e-12))
@pmp("weight", (True, False))
def test_no_correlation(nrow, nchan, epsilon, weight):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(128, 128, epsilon, nchan, nrow)
    wgt = np.random.rand(*idx.shape) if weight else None
    w = conf.Supp()
    for uu in range(-2*w, 2*w):
        for vv in range(-2*w, 2*w):
            if abs(uu) >= w or abs(vv) >= w:
                with pytest.raises(RuntimeError):
                    ng.get_correlations(bl, conf, idx, uu, vv, wgt)
            else:
                ng.get_correlations(bl, conf, idx, uu, vv, wgt)


@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100))
@pmp("nchan", (1, 10))
@pmp("fov", (1,))
def test_wstack_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
    epsilon = 1e-7
    np.random.seed(40)
    pixsize = fov*np.pi/180/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
                            pixsize_y=0.5*pixsize)
    speedoflight, f0 = 299792458., 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
    bl = ng.Baselines(coord=uvw, freq=freq)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
    idx = ng.getIndices(bl, conf, flags)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    res0 = np.zeros((nxdirty, nydirty), dtype=np.float64)
    mi = np.min(np.abs(uvw[:, 2]))*freq[0]/speedoflight
    ma = np.max(np.abs(uvw[:, 2]))*freq[-1]/speedoflight
    nplanes = 10000
    ws = mi + np.arange(nplanes)*(ma-mi)/(nplanes-1)
    for ii in range(len(ws)-1):
        wmax = ws[ii+1] if ii != nplanes-2 else ws[ii+1] + abs(ws[ii+1])
        wmax = ws[ii+1] if ii != nplanes-2 else np.inf
        iind = ng.getIndices(bl, conf, flags, wmin=ws[ii], wmax=wmax)
        if len(iind) == 0:
            continue
        dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
        res0 += conf.apply_wscreen(dd, 0.5*(ws[ii+1]+ws[ii]), adjoint=True).real
    res1 = explicit_gridder(uvw, freq, ms, None, nxdirty, nydirty, pixsize,
                            0.5*pixsize, True)
    assert_(_l2error(res0, res1) < 1e-4)  # Very inaccurate


@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100))
@pmp("nchan", (1,))
@pmp("fov", (50,))
def test_wplane_against_wdft(nxdirty, nydirty, nchan, nrow, fov):
    epsilon = 1e-10
    np.random.seed(40)
    pixsize = fov*np.pi/180/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
                            pixsize_y=0.2*pixsize)
    speedoflight, f0 = 299792458., 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
    uvw[:, 2] = uvw[0, 2]  # use exactly one random w
    bl = ng.Baselines(coord=uvw, freq=freq)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
    idx = ng.getIndices(bl, conf, flags)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    w = np.min(uvw[:, 2])*freq[0]/speedoflight
    iind = ng.getIndices(bl, conf, flags)
    dd = conf.grid2dirty_c(ng.vis2grid_c(bl, conf, iind, bl.ms2vis(ms, iind)))
    res0 = conf.apply_wscreen(dd, w, adjoint=True).real
    res1 = explicit_gridder(uvw, freq, ms, None, nxdirty, nydirty, pixsize,
                            0.2*pixsize, True)
    assert_(_l2error(res0, res1) < epsilon)


@pmp('nxdirty', [16, 64])
@pmp('nydirty', [64])
@pmp("nrow", (10, 100))
@pmp("nchan", (1, 10))
@pmp("fov", (1, 10))
@pmp("epsilon", (1e-2, 1e-7, 2e-12))
def test_wgridder_against_wdft(nxdirty, nydirty, nchan, nrow, fov, epsilon):
    np.random.seed(40)
    pixsize = fov*np.pi/180/nxdirty
    conf = ng.GridderConfig(nxdirty=nxdirty,
                            nydirty=nydirty,
                            epsilon=epsilon,
                            pixsize_x=pixsize,
                            pixsize_y=0.5*pixsize)
    speedoflight, f0 = 299792458., 1e9
    freq = f0 + np.arange(nchan)*(f0/nchan)
    uvw = (np.random.rand(nrow, 3)-0.5)/(pixsize*f0/speedoflight)
    bl = ng.Baselines(coord=uvw, freq=freq)
    flags = np.zeros((nrow, nchan), dtype=np.bool)
    idx = ng.getIndices(bl, conf, flags)
    ms = np.random.rand(nrow, nchan)-0.5 + 1j*(np.random.rand(nrow, nchan)-0.5)
    vis = bl.ms2vis(ms, idx)
    res0 = ng.vis2dirty(bl, conf, idx, vis, do_wstacking=True)
    res1 = explicit_gridder(uvw, freq, ms, None, nxdirty, nydirty, pixsize,
                            0.5*pixsize, True)
    assert_(_l2error(res0, res1) < epsilon)


@pmp('nxdirty', (32, 16))
@pmp('nydirty', (32, 64))
@pmp("nrow", (100, 127))
@pmp("nchan", (1, 5))
@pmp('epsilon', (1e-3, 1e-6))
def test_holo_from_correlations(nxdirty, nydirty, nchan, nrow, epsilon):
    np.random.seed(42)
    bl, conf, idx = _init_gridder(nxdirty, nydirty, epsilon, nchan, nrow)
    nx, ny = 2*nxdirty, 2*nydirty
    grid = np.random.randn(nx, ny) + 1.0j*np.random.randn(nx, ny)
    W = conf.Supp()
    res0 = np.zeros_like(grid)

    # Precompute 2*W**2-2*W+1 images (naive: 4*W**2-4*W+1)
    d0, d1 = [], []
    for du, dv in product(range(-W + 1, W), range(1, W)):
        d0.append(ng.get_correlations(bl, conf, idx, du, dv))
    for du in range(1, W):
        d1.append(ng.get_correlations(bl, conf, idx, du, 0))
    d2 = ng.get_correlations(bl, conf, idx, 0, 0)

    # Apply
    for ii, (du, dv) in enumerate(product(range(-W + 1, W), range(1, W))):
        tmp = np.roll(grid, -du, axis=0)
        res0 += d0[ii]*np.roll(tmp, -dv, axis=1)
        res0 += np.roll(d0[ii]*tmp, dv, axis=1)
    for ii, du in enumerate(range(1, W)):
        res0 += d1[ii]*np.roll(grid, -du, axis=0)
        res0 += np.roll(d1[ii]*grid, du, axis=0)
    res0 += d2*grid
    res1 = ng.apply_holo(bl, conf, idx, grid)
    assert_allclose(res0, res1, rtol=1e-13, atol=1e-13)