from __future__ import annotations

from functools import cached_property
from typing import Any

import numpy as np
import numpy.fft as np_fft
import numpy.typing as npt
import torch

import mpol.constants as const
from mpol.exceptions import CellSizeError
from mpol.utils import get_maximum_cell_size



[docs]
class GridCoords:
    r"""
    The GridCoords object uses desired image dimensions (via the ``cell_size`` and
    ``npix`` arguments) to define a corresponding Fourier plane grid.

    Parameters
    ----------
    cell_size : float
        width of a single square pixel in [arcsec]
    npix : int
        number of pixels in the width of the image

    The Fourier grid is defined over the domain :math:`[-u,+u]`, :math:`[-v,+v]`, even
    though for real images, technically we could use an RFFT grid from :math:`[0,+u]` to
    :math:`[-v,+v]`. The reason we opt for a full FFT grid in this instance is
    implementation simplicity.

    Images (and their corresponding Fourier transform quantities) are represented as
    two-dimensional arrays packed as ``[y, x]`` and ``[v, u]``.  This means that an
    image with dimensions ``(npix, npix)`` will also have a corresponding FFT Fourier
    grid with shape ``(npix, npix)``. The native :class:`~mpol.gridding.GridCoords`
    representation assumes the Fourier grid (and thus image) are laid out as one might
    normally expect an image (i.e., no ``np.fft.fftshift`` has been applied).

    After the object is initialized, instance variables can be accessed, for example

    >>> myCoords = GridCoords(cell_size=0.005, npix=512)
    >>> myCoords.img_ext

    :ivar dl: image-plane cell spacing in RA direction (assumed to be positive)
        [radians]
    :ivar dm: image-plane cell spacing in DEC direction [radians]
    :ivar img_ext: The length-4 list of (left, right, bottom, top) expected by routines
        like ``matplotlib.pyplot.imshow`` in the ``extent`` parameter assuming
        ``origin='lower'``. Units of [arcsec]
    :ivar packed_x_centers_2D: 2D array of l increasing, with fftshifted applied
        [arcseconds]. Useful for directly evaluating some function to create a
        packed cube.
    :ivar packed_y_centers_2D: 2D array of m increasing, with fftshifted applied
        [arcseconds]. Useful for directly evaluating some function to create a
        packed cube.
    :ivar sky_x_centers_2D: 2D array of l arranged for evaluating a sky image
        [arcseconds]. l coordinate increases to the left (as on sky).
    :ivar sky_y_centers_2D: 2D array of m arranged for evaluating a sky image
        [arcseconds].
    :ivar du: Fourier-plane cell spacing in East-West direction
        [:math:`\lambda`]
    :ivar dv: Fourier-plane cell spacing in North-South direction
        [:math:`\lambda`]
    :ivar u_centers: 1D array of cell centers in East-West direction
        [:math:`\lambda`].
    :ivar v_centers: 1D array of cell centers in North-West direction
        [:math:`\lambda`].
    :ivar u_edges: 1D array of cell edges in East-West direction
        [:math:`\lambda`].
    :ivar v_edges: 1D array of cell edges in North-South direction
        [:math:`\lambda`].
    :ivar u_bin_min: minimum u edge [:math:`\lambda`]
    :ivar u_bin_max: maximum u edge [:math:`\lambda`]
    :ivar v_bin_min: minimum v edge [:math:`\lambda`]
    :ivar v_bin_max: maximum v edge [:math:`\lambda`]
    :ivar max_grid: maximum spatial frequency enclosed by Fourier grid
        [:math:`\lambda`]
    :ivar vis_ext: length-4 list of (left, right, bottom, top) expected by routines
        like ``matplotlib.pyplot.imshow`` in the ``extent`` parameter assuming
        ``origin='lower'``. Units of [:math:`\lambda`]
    :ivar vis_ext_Mlam: like vis_ext, but in units of [:math:`\mathrm{M}\lambda`].
    """

    def __init__(self, cell_size: float, npix: int):
        if npix <= 0 or not (npix % 2 == 0):
            raise ValueError("Image must have a positive and even number of pixels.")

        if cell_size <= 0:
            raise ValueError("cell_size must be a positive real number.")

        # Imply to users that GridCoords instance is read-only and new instance
        # is the approach if changing values
        self._cell_size = cell_size
        self._npix = npix

        # Image related
        self._image_pixel_width = cell_size * const.arcsec  # [radians]
        self._image_centers = self._image_pixel_width * (
            np.arange(npix) - npix // 2
        )  # [radians]

        # Spatial frequency related
        # These properties are identical for both u & v and defined here
        # edges and centers return fftspaced arrays (not packed, though)
        # All units in [λ]
        self._uv_pixel_width = 1 / (npix * self._image_pixel_width)
        self.uv_edges = self._uv_pixel_width * (np.arange(npix + 1) - npix // 2 - 0.5)
        self._uv_centers = self._uv_pixel_width * (np.arange(npix) - npix // 2)
        self._min_uv = float(self.uv_edges.min())
        self._max_uv = float(self.uv_edges.max())

    def __repr__(self):
        return f"GridCoords(cell_size={self.cell_size:.2e}, npix={self.npix})"
        # the output spatial frequencies of the FFT routine

    @property
    def cell_size(self) -> float:
        return self._cell_size

    @property
    def npix(self) -> int:
        return self._npix

    @property
    def dl(self) -> float:
        return self._image_pixel_width  # [radians]

    @property
    def dm(self) -> float:
        return self._image_pixel_width  # [radians]

    @property
    def l_centers(self) -> npt.NDArray[np.floating[Any]]:
        return self._image_centers

    @property
    def m_centers(self) -> npt.NDArray[np.floating[Any]]:
        return self._image_centers

    @property
    def npix_u(self) -> int:
        return self.npix

    @property
    def npix_v(self) -> int:
        return self.npix

    @property
    def du(self) -> float:
        return self._uv_pixel_width

    @property
    def dv(self) -> float:
        return self._uv_pixel_width

    @property
    def u_edges(self) -> npt.NDArray[np.floating[Any]]:
        return self.uv_edges

    @property
    def v_edges(self) -> npt.NDArray[np.floating[Any]]:
        return self.uv_edges

    @property
    def u_centers(self) -> npt.NDArray[np.floating[Any]]:
        return self._uv_centers

    @property
    def v_centers(self) -> npt.NDArray[np.floating[Any]]:
        return self._uv_centers

    @property
    def u_bin_min(self) -> float:
        return self._min_uv

    @property
    def v_bin_min(self) -> float:
        return self._min_uv

    @property
    def u_bin_max(self) -> float:
        return self._max_uv

    @property
    def v_bin_max(self) -> float:
        return self._max_uv

    @property
    def max_uv_grid_value(self) -> float:
        return self._max_uv

    @property
    def img_ext(self) -> list[float]:
        # calculate the image extent
        # say we had 10 pixels representing centers -5, -4, -3, ...
        # it should go from -5.5 to +4.5
        lmax = self.cell_size * (self.npix // 2 - 0.5)
        lmin = -self.cell_size * (self.npix // 2 + 0.5)
        return [lmax, lmin, lmin, lmax]  # arcsecs

    @property
    def vis_ext(self) -> list[float]:
        return [
            self.u_bin_min,
            self.u_bin_max,
            self.v_bin_min,
            self.v_bin_max,
        ]  # [λ]

    @property
    def vis_ext_Mlam(self) -> list[float]:
        return [1e-6 * edge for edge in self.vis_ext]

    @cached_property
    def ground_u_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # only useful for plotting a sky_vis
        # uu increasing, no fftshift
        # tile replicates the 1D u_centers array to a 2D array the size of the full
        # UV grid
        return np.tile(self.u_centers, (self.npix_u, 1))

    @cached_property
    def ground_v_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # only useful for plotting a sky_vis
        # vv increasing, no fftshift
        return np.tile(self.v_centers, (self.npix_v, 1)).T

    @cached_property
    def packed_u_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # for evaluating a packed vis
        # uu increasing, fftshifted
        return np_fft.fftshift(self.ground_u_centers_2D)

    @cached_property
    def packed_v_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # for evaluating a packed vis
        # vv increasing + fftshifted
        return np_fft.fftshift(self.ground_v_centers_2D)

    @cached_property
    def ground_q_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        return np.sqrt(
            self.ground_u_centers_2D**2 + self.ground_v_centers_2D**2
        )  # [kλ]

    @cached_property
    def sky_phi_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # https://en.wikipedia.org/wiki/Atan2
        return np.arctan2( # type: ignore
            self.ground_v_centers_2D, self.ground_u_centers_2D
        )  # (pi, pi]

    @cached_property
    def packed_q_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # for evaluating a packed vis in polar coordinates
        # q increasing, fftshifted
        return np_fft.fftshift(self.ground_q_centers_2D)

    @cached_property
    def packed_phi_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # for evaluating a packed vis in polar coordinates
        # phi increasing, fftshifted
        return np_fft.fftshift(self.sky_phi_centers_2D)

    @cached_property
    def q_max(self) -> float:
        # outer edge [lambda]
        return float(np.abs(self.packed_q_centers_2D).max() + np.sqrt(2) * self.du)

    @cached_property
    def x_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        return np.tile(self.l_centers / const.arcsec, (self.npix, 1))  # [arcsec]

    @cached_property
    def y_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        return np.tile(self.m_centers / const.arcsec, (self.npix, 1)).T  # [arcsec]

    @cached_property
    def packed_x_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        return np.fft.fftshift(self.x_centers_2D)  # [arcsec]

    @cached_property
    def packed_y_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        return np.fft.fftshift(self.y_centers_2D)  # [arcsec]

    @property
    def sky_x_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # for evaluating a sky image
        # ll mirrored, increasing, no fftshift
        return np.fliplr(self.x_centers_2D)  # [arcsec]

    @property
    def sky_y_centers_2D(self) -> npt.NDArray[np.floating[Any]]:
        # for evaluating a sky image
        # mm increasing, no fftshift
        return self.y_centers_2D  # [arcsec]


[docs]
    def check_data_fit(
        self,
        uu: torch.Tensor | npt.NDArray[np.floating[Any]],
        vv: torch.Tensor | npt.NDArray[np.floating[Any]],
    ) -> bool:
        r"""
        Test whether loose data visibilities fit within the Fourier grid defined by
        cell_size and npix.

        Parameters
        ----------
        uu : :class:`torch.Tensor`
            u spatial frequency coordinates.
            Units of [:math:`\lambda`]
        vv : :class:`torch.Tensor`
            v spatial frequency coordinates.
            Units of [:math:`\lambda`]

        Returns
        -------
        bool
            ``True`` if all visibilities fit within the Fourier grid defined by
            ``[u_bin_min, u_bin_max, v_bin_min, v_bin_max]``. Otherwise a
            :class:`mpol.exceptions.CellSizeError` is raised on the first violated
            boundary.
        """

        # we need this routine to work with both numpy.ndarray or torch.Tensor
        # because it is called for DirtyImager setup (numpy only)
        # so we'll cast to tensor as a precaution
        uu = torch.as_tensor(uu)
        vv = torch.as_tensor(vv)

        # max freq in dataset
        max_uu_vv = np.abs(np.concatenate([uu, vv])).max()
        max_uu_vv = torch.max(torch.abs(torch.concatenate([uu, vv]))).item()

        # max freq needed to support dataset
        max_cell_size = get_maximum_cell_size(max_uu_vv)

        if torch.max(torch.abs(uu)) > self.max_uv_grid_value:
            raise CellSizeError(
                "Dataset contains uu spatial frequency measurements larger "
                "than those in the proposed model image. "
                f"Decrease cell_size below {max_cell_size} arcsec."
            )

        if torch.max(torch.abs(vv)) > self.max_uv_grid_value:
            raise CellSizeError(
                "Dataset contains vv spatial frequency measurements larger "
                "than those in the proposed model image. "
                f"Decrease cell_size below {max_cell_size} arcsec."
            )

        return True


    def __eq__(self, other: Any) -> bool:
        if not isinstance(other, GridCoords):
            # don't attempt to compare against different types
            return NotImplemented

        # GridCoords objects are considered equal if they have the same cell_size and
        # npix, since all other attributes are derived from these two core properties.
        return bool(self.cell_size == other.cell_size and self.npix == other.npix)