"""Functions for working with LOFAR single station data"""

import numpy as np
import numexpr as ne
import numba
from astropy.coordinates import SkyCoord, SkyOffsetFrame, CartesianRepresentation


__all__ = ["nearfield_imager", "sky_imager", "ground_imager", "skycoord_to_lmn"]

__version__ = "1.5.0"
SPEED_OF_LIGHT = 299792458.0


def skycoord_to_lmn(pos: SkyCoord, phasecentre: SkyCoord):
    """
    Convert astropy sky coordinates into the l,m,n coordinate system
    relative to a phase centre.

    The l,m,n is a RHS coordinate system with
    * its origin on the sky sphere
    * m,n and the celestial north on the same plane
    * l,m a tangential plane of the sky sphere

    Note that this means that l increases east-wards

    This function was taken from https://github.com/SKA-ScienceDataProcessor/algorithm-reference-library
    """

    # Determine relative sky position
    todc = pos.transform_to(SkyOffsetFrame(origin=phasecentre))
    dc = todc.represent_as(CartesianRepresentation)
    dc /= dc.norm()

    # Do coordinate transformation - astropy's relative coordinates do
    # not quite follow imaging conventions
    return dc.y.value, dc.z.value, dc.x.value - 1


@numba.jit(parallel=True)
def sky_imager(visibilities, baselines, freq, npix_l, npix_m):
    """
    Sky imager

    Args:
        visibilities: Numpy array with visibilities, shape [num_antennas x num_antennas]
        baselines: Numpy array with distances between antennas, shape [num_antennas, num_antennas, 3]
        freq: frequency
        npix_l: Number of pixels in l-direction
        npix_m: Number of pixels in m-direction

    Returns:
        np.array(float): Real valued array of shape [npix_l, npix_m]
    """
    img = np.zeros((npix_m, npix_l), dtype=np.complex128)

    for m_ix in range(npix_m):
        m = -1 + m_ix * 2 / npix_m
        for l_ix in range(npix_l):
            l = 1 - l_ix * 2 / npix_l
            img[m_ix, l_ix] = np.mean(visibilities * np.exp(-2j * np.pi * freq *
                                                            (baselines[:, :, 0] * l + baselines[:, :, 1] * m) /
                                                            SPEED_OF_LIGHT))
    return np.real(img)


def ground_imager(visibilities, freq, npix_p, npix_q, dims, station_pqr, height=1.5):
    """Do a Fourier transform for ground imaging"""
    img = np.zeros([npix_q, npix_p], dtype=np.complex128)

    for q_ix, q in enumerate(np.linspace(dims[2], dims[3], npix_q)):
        for p_ix, p in enumerate(np.linspace(dims[0], dims[1], npix_p)):
            r = height
            pqr = np.array([p, q, r], dtype=np.float32)
            antdist = np.linalg.norm(station_pqr - pqr[np.newaxis, :], axis=1)
            groundbase = antdist[:, np.newaxis] - antdist[np.newaxis, :]
            img[q_ix, p_ix] = np.mean(visibilities * np.exp(-2j * np.pi * freq * (-groundbase) / SPEED_OF_LIGHT))

    return img


def nearfield_imager(visibilities, baseline_indices, freqs, npix_p, npix_q, extent, station_pqr, height=1.5,
                     max_memory_mb=200):
    """
    Nearfield imager

    Args:
        visibilities: Numpy array with visibilities, shape [num_visibilities x num_frequencies]
        baseline_indices: List with tuples of antenna numbers in visibilities, shape [2 x num_visibilities]
        freqs: List of frequencies
        npix_p: Number of pixels in p-direction
        npix_q: Number of pixels in q-direction
        extent: Extent (in m) that the image should span
        station_pqr: PQR coordinates of stations
        height: Height of image in metre
        max_memory_mb: Maximum amount of memory to use for the biggest array. Higher may improve performance.

    Returns:
        np.array(complex): Complex valued array of shape [npix_p, npix_q]
    """
    z = height
    x = np.linspace(extent[0], extent[1], npix_p)
    y = np.linspace(extent[2], extent[3], npix_q)

    posx, posy = np.meshgrid(x, y)
    posxyz = np.transpose(np.array([posx, posy, z * np.ones_like(posx)]), [1, 2, 0])

    diff_vectors = (station_pqr[:, None, None, :] - posxyz[None, :, :, :])
    distances = np.linalg.norm(diff_vectors, axis=3)

    vis_chunksize = max_memory_mb * 1024 * 1024 // (8 * npix_p * npix_q)

    bl_diff = np.zeros((vis_chunksize, npix_q, npix_p), dtype=np.float64)
    img = np.zeros((npix_q, npix_p), dtype=np.complex128)
    for vis_chunkstart in range(0, len(baseline_indices), vis_chunksize):
        vis_chunkend = min(vis_chunkstart + vis_chunksize, baseline_indices.shape[0])
        # For the last chunk, bl_diff_chunk is a bit smaller than bl_diff
        bl_diff_chunk = bl_diff[:vis_chunkend - vis_chunkstart, :]
        np.add(distances[baseline_indices[vis_chunkstart:vis_chunkend, 0]],
               -distances[baseline_indices[vis_chunkstart:vis_chunkend, 1]], out=bl_diff_chunk)

        j2pi = 1j * 2 * np.pi
        for ifreq, freq in enumerate(freqs):
            v = visibilities[vis_chunkstart:vis_chunkend, ifreq][:, None, None]
            lamb = SPEED_OF_LIGHT / freq

            # v[:,np.newaxis,np.newaxis]*np.exp(-2j*np.pi*freq/c*groundbase_pixels[:,:,:]/c)
            # groundbase_pixels=nvis x npix x npix
            np.add(img, np.sum(ne.evaluate("v * exp(j2pi * bl_diff_chunk / lamb)"),axis=0), out=img)
    img /= len(freqs) * len(baseline_indices)

    return img