from typing import Union, Sequence, Optional

import numpy as np
from astropy.units import Quantity

from .utilities import ExtentType, CoordinateType, unitless

# shape of a numpy array, or a single integer that is broadcast to a square shape
ShapeType = Union[int, Sequence[int]]

# a scalar quantity with optional unit attached
ScalarType = Union[float, np.ndarray, Quantity]

"""
Library of functions to create commonly used patterns

Each of the functions takes a `shape` input, which may be a scalar integer or a 2-element Sequence of integers
indicating the size (shape) in pixels of the returned field. If `shape` is a scalar, the same value is
used for both axes.

For the coordinates, the OpenGL convention is used, where the coordinates indicate the centers of the pixels.
By default, the returned pattern is assumed to cover a -1,1 x -1,1 square, 
which corresponds to a default `extent` parameter of (2.0, 2.0).
In this case, the coordinates range from -1+dx/2 to 1-dx/2, where dx=2.0/shape is the pixel size.

Exptent and shape may be specified individually to work with anisotropic pixels or rectangular patterns.
For example, a square pattern with anisotropic pixels may be described by shape=(80,100) and extent(2,2)
whereas shape=(80,100) and extent(8,10) describes square pixels that form a rectangle.

In a pupil-conjugate configuration, a disk of extent=(NA, NA) exactly covers the back pupil of the microscope objective.
The transformation matrix of the SLM should be set such that SLM coordinates correspond to normalized pupil coordinates.

The extent may have a unit of measure. In this case, other parameters (such as `radius`) may need to have
an according unit of measure.

The (0,0) coordinate is always located in the center of the pattern, which may be on a grid point (for odd shape)
or between grid points (for even shape).

The returned array has a pixel_size property attached.
"""


def coordinate_range(
    shape: ShapeType, extent: ExtentType, offset: Optional[CoordinateType] = None
) -> (Quantity, Quantity):
    """
    Returns coordinate vectors for the two coordinates (y and x).

    Given a range from `-extent/2` to `+extent/2`, uniformly divided into `shape` pixels,
    the returned coordinates correspond to the centers of these coordinates.

    Arguments:
        shape (ShapeType): size of the full grid (y, x) in pixels
        extent (ExtentType): extent of the coordinate range
        offset (Optional[CoordinateType]): offset to be added to the coordinates (optional)

    Returns:
        Tuple[Quantity, Quantity]: coordinate vectors for the two coordinates (y and x)
    """
    if np.size(shape) == 1:
        shape = np.array(shape).repeat(2)

    extent = Quantity(extent)
    if extent.size == 1:
        extent = extent.repeat(2)

    if offset is None:
        offset = 0.0 * extent  # by default, let center be (0,0) (in whatever unit)

    def c_range(res, ex, cx):
        dx = ex / res
        return np.arange(res) * dx + (0.5 * dx - 0.5 * ex + cx)

    return (
        c_range(shape[0], extent[0], offset[0]).reshape((-1, 1)),
        c_range(shape[1], extent[1], offset[1]).reshape((1, -1)),
    )


def r2_range(shape: ShapeType, extent: ExtentType):
    """Convenience function to return square distance to the origin
    Equivalent to computing cx^2 + cy^2
    """
    c0, c1 = coordinate_range(shape, extent)
    return c0**2 + c1**2


def tilt(
    shape: ShapeType,
    g: ExtentType,
    extent: ExtentType = (2.0, 2.0),
    phase_offset: float = 0.0,
):
    """Constructs a linear gradient pattern φ=2g·r

    Note, these are the Zernike tilt modes (modes 2 and 3 in the Noll index convention) with normalization
    so that :math:`\frac{\int_0^{2\pi} \int_0^1 |Z(\rho, \phi)|^2 \rho d\rho d\phi}{\int_0^{2\pi} \int_0^1 \rho d\rho d\phi} = 1`.

    Args:
        shape: see module documentation
        g(tuple of two floats): gradient vector.
          This has the unit: 1 / extent.unit.
          For the default extent of (2.0, 2.0), a value of g=(1,0)
          corresponds to having a ramp from -2 to +2 over the height of the pattern.
          With an extent of (2.0, 2.0) covering the full NA,
          this pattern causes a displacement of -2/π times the Abbe diffraction limit
          (Note: a positive x-gradient g causes the focal point to move in the _negative_ x-direction)
        extent: see module documentation
        phase_offset: optional additional phase offset to be added to the pattern
    """
    c0, c1 = coordinate_range(shape, extent * (Quantity(g) * 2.0))
    return unitless(c0 + (c1 + phase_offset))


def lens(shape: ShapeType, f: ScalarType, wavelength: ScalarType, extent: ExtentType):
    """Constructs a square texture that represents a wavefront defocus: (f-sqrt(f²+r²)) · 2π/λ

    `extent`, `wavelength` and `f` should have compatible units (typically astropy length units).

    Args:
        shape(ShapeType): see module documentation
        f(ScalarType): focal length
        wavelength(ScalarType): wavelength
        extent(ExtentType): physical extent of the SLM, same units as `f` and `wavelength`
    """
    r_sqr = r2_range(shape, extent)
    return unitless((f - np.sqrt(f**2 + r_sqr)) * (2 * np.pi / wavelength))


def propagation(
    shape: ShapeType,
    distance: ScalarType,
    wavelength: ScalarType,
    extent: ExtentType,
    refractive_index: ScalarType = 1.0,
):
    """Computes a wavefront that corresponds to digitally propagating the field in the object plane.

    k_z = sqrt(n² k_0²-k_x²-k_y²)
    φ = k_z · distance

    Args:
          shape: see module documentation
          distance (ScalarType): physical distance to propagate axially.
          refractive_index (Scalar):
          wavelength (Scalar):
            the numerical aperture, refractive index and wavelength are used
            to convert the `extent` from pupil coordinates to k-space (unit radians/meter),
          extent: extent of the returned image, in NA units. To cover the full NA with a square, use (2*NA, 2*NA)
    """
    # convert pupil coordinates to absolute k_x, k_y coordinates
    n_p2 = r2_range(shape, Quantity(extent))
    n_z = np.sqrt(np.maximum(refractive_index**2 - n_p2, 0.0))
    return unitless(2.0 * np.pi / wavelength * distance * n_z)


def disk(shape: ShapeType, radius: ScalarType = 1.0, extent: ExtentType = (2.0, 2.0)):
    """Constructs an image of a centered (ellipsoid) disk.

    (x / rx)^2 + (y / ry)^2 <= 1.0

    Args:
          shape: see module documentation
          radius (ScalarType): radius of the disk, should have the same unit as `extent`.
          extent: see module documentation
    """
    return 1.0 * (r2_range(shape, extent) < radius**2)


def gaussian(
    shape: ShapeType,
    waist: ScalarType,
    truncation_radius: ScalarType = None,
    extent: ExtentType = (2.0, 2.0),
):
    """Constructs an image of a centered Gaussian

    `waist`, `extent` and the optional `truncation_radius` should all have the same unit.

    Args:
        shape: see module documentation
        waist (ScalarType): location of the beam waist (1/e value)
            relative to half of the size of the pattern (i.e. relative to the `radius` of the square)
        truncation_radius (ScalarType): when not None, specifies the radius of a disk that is used to truncate the
            Gaussian. All values outside the disk are set to 0.
        extent: see module documentation

    """
    r_sqr = r2_range(shape, extent)
    w2inv = -1.0 / waist**2
    gauss = np.exp(unitless(r_sqr * w2inv))
    if truncation_radius is not None:
        gauss = gauss * disk(shape, truncation_radius, extent=extent)
    return unitless(gauss)
