"""volume.py: Functions to calculate volumes of hexahedral cells; assumes consistent length units."""

version = '31st August 2021'

import numpy as np


def _pyr(ra, ab, ac, ad, db):  # returns 6 * volume of one quad pyramid, defined by specific vectors
    return np.dot(ra, np.cross(db, ac)) + 0.5 * np.dot(ac, np.cross(ad, ab))


def pyramid_volume(apex, a, b, c, d, crs_is_right_handed = False):
    """Returns volume of a quadrilateral pyramid.

    arguments:
       apex (triple float): location of the apex of the pyramid
       a, b, c, d (each triple float): locations of corners of base of pyramid; clockwise viewed from apex
          for a left handed crs
       crs_is_right_handed (boolean, default False): set True if xyz axes of crs are right handed

    returns:
       float, being the volume of the pyramid; units are implied by crs units in use by the vertices
    """

    apex = np.array(apex)
    a = np.array(a)
    b = np.array(b)
    c = np.array(c)
    d = np.array(d)
    v = _pyr(a - apex, b - a, c - a, d - a, b - d) / 6.0
    if not crs_is_right_handed:
        v = -v
    return v


def tetrahedron_volume(a, b, c, d):
    """Returns volume of a tetrahedron.

    arguments:
       a, b, c, d (each triple float): locations of corners of tetrahedron
       crs_is_right_handed (boolean, default False): set True if xyz axes of crs are right handed

    returns:
       float, being the volume of the tetrahedron; units are implied by crs units in use by the vertices
    """

    a = np.array(a)
    b = np.array(b)
    c = np.array(c)
    d = np.array(d)

    return np.abs(np.dot(a - d, np.cross(b - d, c - d))) / 6.0


def tetra_cell_volume(cp, centre = None, off_hand = False):
    """Returns volume of single hexahedral cell with corner points cp of shape (2, 2, 2, 3); assumes bilinear faces."""

    if centre is None:
        centre = np.mean(cp.reshape((8, 3)), axis = 0)
    r0 = cp[0, 0, 0] - centre
    r1 = cp[1, 1, 1] - centre
    v = 0.0
    v += _pyr(r0, cp[0, 0, 1] - cp[0, 0, 0], cp[0, 1, 1] - cp[0, 0, 0], cp[0, 1, 0] - cp[0, 0, 0],
              cp[0, 1, 0] - cp[0, 0, 1])
    v += _pyr(r0, cp[1, 0, 0] - cp[0, 0, 0], cp[1, 0, 1] - cp[0, 0, 0], cp[0, 0, 1] - cp[0, 0, 0],
              cp[0, 0, 1] - cp[1, 0, 0])
    v += _pyr(r0, cp[0, 1, 0] - cp[0, 0, 0], cp[1, 1, 0] - cp[0, 0, 0], cp[1, 0, 0] - cp[0, 0, 0],
              cp[1, 0, 0] - cp[0, 1, 0])
    v += _pyr(r1, cp[1, 0, 1] - cp[1, 1, 1], cp[1, 0, 0] - cp[1, 1, 1], cp[1, 1, 0] - cp[1, 1, 1],
              cp[1, 1, 0] - cp[1, 0, 1])
    v += _pyr(r1, cp[1, 1, 0] - cp[1, 1, 1], cp[0, 1, 0] - cp[1, 1, 1], cp[0, 1, 1] - cp[1, 1, 1],
              cp[0, 1, 1] - cp[1, 1, 0])
    v += _pyr(r1, cp[0, 1, 1] - cp[1, 1, 1], cp[0, 0, 1] - cp[1, 1, 1], cp[1, 0, 1] - cp[1, 1, 1],
              cp[1, 0, 1] - cp[0, 1, 1])

    if off_hand:
        v = -v

    return v / 6.0


def tetra_volumes_slow(cp, centres = None, off_hand = False):
    """Returns volume array for all hexahedral cells assuming bilinear faces, using loop over cells."""

    # NB: deprecated, superceded by much faster function below
    # todo: handle NaNs
    # Pagoda style corner point data
    assert cp.ndim == 7

    flat = cp.reshape(-1, 2, 2, 2, 3)
    cells = flat.shape[0]
    if centres is None:
        centres = np.mean(flat, axis = (1, 2, 3))
    else:
        centres = centres.reshape((-1, 3))
    volumes = np.zeros(cells)
    for cell in range(cells):
        volumes[cell] = tetra_cell_volume(flat[cell], centre = centres[cell], off_hand = off_hand)
    return volumes.reshape(cp.shape[0:3])


def tetra_volumes(cp, centres = None, off_hand = False):
    """Returns volume array for all hexahedral cells assuming bilinear faces, using numpy operations.

    arguments:
       cp (7D numpy array of floats): cell corner point data in Pagoda 7D format [nk, nj, ni, kp, jp, ip, xyz]
       centres (optional, 4D numpy array of floats): cell centre points [nk, nj, ni, xyz]; calculated if None
       off_hand (boolean, default False): if True, the handedness of IJK space is the opposite of that for xyz
          space; if this argument is not set correctly, negative volumes will be returned

    returns:
       numpy 3D array of floats being the cell volumes [nk, nj, ni]

    note:
       length units are assumed to be consistent in x, y & z; and untis of returned volumes are implicitly
       those length units cubed
    """

    def pyrs(ra, ab, ac, ad, db):  # returns 6 * volume of one quad pyramid (for each cell)
        return np.sum(ra * np.cross(db, ac), axis = -1) + 0.5 * np.sum(ac * np.cross(ad, ab), axis = -1)

    # Pagoda style corner point data
    assert cp.ndim == 7

    flat = cp.reshape(-1, 2, 2, 2, 3)
    cells = flat.shape[0]
    if centres is None:
        centres = np.mean(flat, axis = (1, 2, 3))
    else:
        centres = centres.reshape((-1, 3))
    v = np.zeros(cells)
    # todo: numpy array operation covering all cells in grid

    r0 = flat[:, 0, 0, 0] - centres
    r1 = flat[:, 1, 1, 1] - centres

    v += pyrs(r0, flat[:, 0, 0, 1] - flat[:, 0, 0, 0], flat[:, 0, 1, 1] - flat[:, 0, 0, 0],
              flat[:, 0, 1, 0] - flat[:, 0, 0, 0], flat[:, 0, 1, 0] - flat[:, 0, 0, 1])
    v += pyrs(r0, flat[:, 1, 0, 0] - flat[:, 0, 0, 0], flat[:, 1, 0, 1] - flat[:, 0, 0, 0],
              flat[:, 0, 0, 1] - flat[:, 0, 0, 0], flat[:, 0, 0, 1] - flat[:, 1, 0, 0])
    v += pyrs(r0, flat[:, 0, 1, 0] - flat[:, 0, 0, 0], flat[:, 1, 1, 0] - flat[:, 0, 0, 0],
              flat[:, 1, 0, 0] - flat[:, 0, 0, 0], flat[:, 1, 0, 0] - flat[:, 0, 1, 0])
    v += pyrs(r1, flat[:, 1, 0, 1] - flat[:, 1, 1, 1], flat[:, 1, 0, 0] - flat[:, 1, 1, 1],
              flat[:, 1, 1, 0] - flat[:, 1, 1, 1], flat[:, 1, 1, 0] - flat[:, 1, 0, 1])
    v += pyrs(r1, flat[:, 1, 1, 0] - flat[:, 1, 1, 1], flat[:, 0, 1, 0] - flat[:, 1, 1, 1],
              flat[:, 0, 1, 1] - flat[:, 1, 1, 1], flat[:, 0, 1, 1] - flat[:, 1, 1, 0])
    v += pyrs(r1, flat[:, 0, 1, 1] - flat[:, 1, 1, 1], flat[:, 0, 0, 1] - flat[:, 1, 1, 1],
              flat[:, 1, 0, 1] - flat[:, 1, 1, 1], flat[:, 1, 0, 1] - flat[:, 0, 1, 1])

    v /= 6.0

    if off_hand:
        v = np.negative(v)
    return v.reshape(cp.shape[0:3])
