#
# Changes
#
# 2022-09, Arjo Segers
#   Support computation of grid cell areas for 2D arrays.
#
# 2022-09, Arjo Segers
#   Support computation of distances between selected location
#   and arrays of other locations.
#
# 2023, Lewis Blake, Arjo Segers
#   Formatted using "black".
#
# 2023-08, Arjo Segers
#   Remove rounding errors when locations are very close.
#

########################################################################
###
### help
###
########################################################################


"""
**********************
``cso_mapping`` module
**********************

Tools for mapping between footprints and grid cells.


Methods
=======

"""

# modules:
import logging


########################################################################
###
### distance
###
########################################################################


def LonLatDistance(lon1, lat1, lon2, lat2):

    """
    Calculates the distance between two points given their (lat, lon) coordinates.

    Arguments (scalar or arrays of same length):

    * ``lon1``, ``lat1`` : first location in degrees east/north;
    * ``lon2``, ``lat2`` : second location in degrees east/north.

    """

    #
    # Uses the `Spherical Law Of Cosines <http://en.wikipedia.org/wiki/Spherical_law_of_cosines>`_:
    #
    # .. math::
    #     \cos(c) = \cos(a) * \cos(b) + \sin(a) * \sin(b) * \cos(C)
    #
    # where:
    #
    # * ``a = lat1`` in radians;
    # * ``b = lat2`` in radians;
    # * ``C = lon2 - lon1`` in radians
    #
    # Because the latitude range is :math:`[-\pi/2,\pi/2]` instead of :math:`[0,\pi]`,
    # and the longitude range is :math:`[-\pi, \pi]` instead of :math:`[0, 2\pi]`,
    # this transforms into:
    #
    # .. math::
    #     x = cos(c) = sin(a) * sin(b) + cos(a) * cos(b) * cos(C)
    #
    # Finally the distance is:
    #
    # .. math::
    #     d = \cos^{-1}(x)
    #
    # Original source (not available anymore?):
    #   `<http://cyberpython.wordpress.com/2010/03/31/python-calculate-the-distance-between-2-points-given-their-coordinates>`_
    #

    # modules:
    import numpy

    # tools:
    from cso_constants import EarthRadius

    # convert:
    delta = lon2 - lon1
    a = numpy.radians(lat1)
    b = numpy.radians(lat2)
    C = numpy.radians(delta)
    x = numpy.sin(a) * numpy.sin(b) + numpy.cos(a) * numpy.cos(b) * numpy.cos(C)

    # distance in radians, clip for small errors:
    phi = numpy.arccos(x.clip(-1.0, 1.0))  # radians

    # in m over the globe:
    return phi * EarthRadius  # m


# enddef  # LonLatDistance


########################################################################
###
### triangle area
###
########################################################################


def TrianglesArea(xx, yy):

    """
    Return areas of triangles defined by:

    * ``xx``, ``yy`` : corner locations as ``(n,3)`` arrays

    Uses `Heron's formula <https://en.wikipedia.org/wiki/Heron%27s_formula>`_ to compute area of triangle.
    """

    # modules:
    import numpy

    # check ..
    if (xx.shape[1] != 3) or (yy.shape[1] != 3) or (xx.shape[0] != yy.shape[0]):
        logging.error("arguments xx %s and yy %s should have shape (n,3)" % (xx.shape, yy.shape))
        raise Exception
    # endif

    # side lengths:
    a = numpy.sqrt((xx[:, 1] - xx[:, 0]) ** 2 + (yy[:, 1] - yy[:, 0]) ** 2)
    b = numpy.sqrt((xx[:, 2] - xx[:, 1]) ** 2 + (yy[:, 2] - yy[:, 1]) ** 2)
    c = numpy.sqrt((xx[:, 0] - xx[:, 2]) ** 2 + (yy[:, 0] - yy[:, 2]) ** 2)

    # semi-perimeter:
    s = 0.5 * (a + b + c)

    # squared area:
    area2 = s * (s - a) * (s - b) * (s - c)

    # area:
    return numpy.sqrt(area2.clip(0.0))


# enddef TrianglesArea

# *


def LonLatTrianglesArea(xx, yy):

    """
    Return areas in m2 of triangles defined by:

    * ``xx``, ``yy`` : corner locations in degrees east/north as ``(n,3)`` arrays

    Using the radius :math:`R` of Earth (m) the area is:

    .. math::
        A ~=~ R^2 \int\limits_{x,y} \cos(y)\ dx\ dy

   Approximate this by first computing the area in degrees2
   and use the average latitude:

    .. math::
        A ~=~ R^2 \int\limits_{x,y} dx\ dy\ \cos(y_{aver})

    """

    # modules:
    import numpy

    # tools:
    from cso_constants import EarthRadius

    # areas:
    aa = TrianglesArea(numpy.radians(xx), numpy.radians(yy))  # rad2

    # average latitudes:
    yy_aver = numpy.radians(yy.sum(axis=1) / 3.0)  # rad

    # areas:
    #       rad2                                m2
    return aa * numpy.cos(yy_aver) * EarthRadius ** 2  # m2


# enddef LonLatTrianglesArea

# *


def LonLatRectanglesArea(xx, yy):

    """
    Return areas in m2 of rectangles defined by:

    * ``xx``, ``yy`` : corner locations in degrees east/north as ``(n,4)`` or ``(m,n,4)`` arrays

    Approximate area of rectangle in ``(lon,lat)`` in degrees using 2 triangles::

        3 o---o 2
          | / |
        0 o---o 1

    This routine is used to ensure that area's of footprints and grid cells
    are approximated in the same way.
    """

    # modules:
    import numpy

    # switch:
    if (len(xx.shape) == 2) and (len(yy.shape) == 2):
        # area's of first triangles:
        A1 = LonLatTrianglesArea(xx[:, 0:3], yy[:, 0:3])
        # area's of second triangles:
        A2 = LonLatTrianglesArea(
            numpy.hstack((xx[:, 2:4], xx[:, 0:1])), numpy.hstack((yy[:, 2:4], yy[:, 0:1]))
        )
        # combine:
        A = A1 + A2

    elif (len(xx.shape) == 3) and (len(yy.shape) == 3):
        # shape:
        m, n, nc = xx.shape
        # init output:
        A = numpy.zeros((m, n), float)
        # loop over rows:
        for j in range(m):
            # recursive call:
            A[j, :] = LonLatRectanglesArea(xx[j, :], yy[j, :])
        # endfor

    else:
        logging.error(
            "unsupported input shapes xx(%s) and yy(%s)" % (str(xx.shape), str(yy.shape))
        )
        raise Exception
    # endif

    # ok:
    return A


# enddef LonLatRectanglesArea


########################################################################
###
### centroids
###
########################################################################


def LonLatPolygonCentroids(xx, yy, maxlevel=5, _level=0, indent=""):

    """
    Let :math:`N` polygons be defined with :math:`n_c` corners in the arguments:
    This method returns arrays with :math:`(P,N)` locations of points distributed within the polygons,
    as well as the fraction of the area corresponding to each point.

    Arguments:

    * ``xx``, ``yy`` : arrays of shape :math:`(N,n_c)` with corner locations in degrees east/north

    Return values:

    * ``xxp``, ``yyp`` : arrays of shape :math:`(P,N)` with the coordinates of the points
    * ``wwp`` : arrays of shape :math:`(P,N)` with area in m2 assigned to each point

    The polygon is recursively devided in triangles, the points are their centroids.
    First the centroid of the polygon is determined, and the first set of triangles is formed
    from with the sides of the polygon as base and the centroid as top::

          o---o
          |\ /|
          | * |
          |/ \|
          o---o

    Each triangle is devided into 2 new triangles with the middle of the longest side as their top
    and the two shortest sides as their base::

              2
              o
          0 / | \ 2
          /   |   \
        o-----*-----o
        0    1    1

    Recursion is repeated up to ``maxlevels`` deep;
    with "maxlevels==0" a single point is returned.
    For a 4-sided polygon the number of points is:

    * 1 for ``maxlevels=0``;
    * 4 for ``maxlevels=1``;
    * 8 for ``maxlevels=2``;
    * 16 for ``maxlevels=3``;
    * 32 for ``maxlevels=4``;
    * 64 for ``maxlevels=5`` (default).

    """

    # modules:
    import numpy

    # dims:
    np, nc = numpy.shape(xx)

    # centroid:
    xc = xx.mean(axis=1)
    yc = yy.mean(axis=1)

    # stop?
    if _level == maxlevel:

        # compute area; triangel or polygon?
        if nc == 3:
            # lowest level, triangles:
            A = LonLatTrianglesArea(xx, yy)
        else:
            # init sum over triangles:
            A = 0.0
            # loop over corners:
            for i in range(nc):
                # start and end point:
                x1, y1 = xx[:, i], yy[:, i]
                if i == nc - 1:
                    x2, y2 = xx[:, 0], yy[:, 0]
                else:
                    x2, y2 = xx[:, i + 1], yy[:, i + 1]
                # endif
                # add areas:
                A += LonLatTrianglesArea(numpy.array([x1, x2, xc]).T, numpy.array([y1, y2, yc]).T)
            # endfor
        # endif # triangle or polygon

        # return center:
        xxc = xc.reshape((1, np))
        yyc = yc.reshape((1, np))
        wwc = A.reshape((1, np))

    else:

        # devide in triangles?
        if nc != 3:
            # loop over corners:
            for i in range(nc):
                # start and end point:
                x1, y1 = xx[:, i], yy[:, i]
                if i == nc - 1:
                    x2, y2 = xx[:, 0], yy[:, 0]
                else:
                    x2, y2 = xx[:, i + 1], yy[:, i + 1]
                # endif
                # new triangle:
                xxt, yyt, wwt = LonLatPolygonCentroids(
                    numpy.array([x1, x2, xc]).T,
                    numpy.array([y1, y2, yc]).T,
                    _level=_level + 1,
                    maxlevel=maxlevel,
                )
                # extend:
                if i == 0:
                    xxc, yyc, wwc = xxt, yyt, wwt
                else:
                    xxc = numpy.vstack((xxc, xxt))
                    yyc = numpy.vstack((yyc, yyt))
                    wwc = numpy.vstack((wwc, wwt))
                # endif
            # endfor # corners

        else:

            #          top
            #           2
            #           o
            #       0 / | \ 2
            #       /   |   \
            #     o-----*-----o
            #     0     1     1
            #  base0   mid   base1

            # range for selecting all pixels:
            jj = numpy.arange(np)

            # search for longest edge;
            # compute squared edge lengths:
            dd = numpy.zeros((np, nc))
            for i in range(nc):
                dd[:, i] = (xx[:, i] - xx[:, i - 1]) ** 2 + (yy[:, i] - yy[:, i - 1]) ** 2
            # endfor
            # edges with longest length are the base:
            iibase = numpy.argmax(dd, axis=1)
            # corners of base:
            xxbase0 = xx[jj, iibase - 1]
            yybase0 = yy[jj, iibase - 1]
            xxbase1 = xx[jj, iibase]
            yybase1 = yy[jj, iibase]
            # midpoints:
            xxmid = (xxbase0 + xxbase1) / 2.0
            yymid = (yybase0 + yybase1) / 2.0
            # oposite corner is top:
            iitop = (iibase + 1) % nc
            xxtop = xx[jj, iitop]
            yytop = yy[jj, iitop]

            # devide triangle in 2 sub-triangles:

            # ~ triangle 1
            xxt, yyt, wwt = LonLatPolygonCentroids(
                numpy.array([xxbase0, xxmid, xxtop]).T,
                numpy.array([yybase0, yymid, yytop]).T,
                _level=_level + 1,
                maxlevel=maxlevel,
            )
            # extend:
            xxc = xxt
            yyc = yyt
            wwc = wwt

            # ~ triangle 2
            xxt, yyt, wwt = LonLatPolygonCentroids(
                numpy.array([xxmid, xxbase1, xxtop]).T,
                numpy.array([yymid, yybase1, yytop]).T,
                _level=_level + 1,
                maxlevel=maxlevel,
            )
            # extend:
            xxc = numpy.vstack((xxc, xxt))
            yyc = numpy.vstack((yyc, yyt))
            wwc = numpy.vstack((wwc, wwt))

        # endif # multi or triangle

    # endif # recursion

    # ok
    return xxc, yyc, wwc


# enddef LonLatPolygonCentroids


########################################################################
###
### end
###
########################################################################
