"""
Program: Fluid Velocity Functions
Date: 26 July 2018
Author: LM

Calculate the fluid velocity at a position (x, y) in a fluid due to a 2-dimensional sink in a plane, or a 3-dimensional
sink in space.
"""

# Import packages.
import numpy as np
import itertools


# Fluid velocity functions (2D).
def u_x_2d(x, y, r, q=0.0004):
    """
    X-velocity for 2D sink - For a 2-dimensional sink at position r, calculate the x-component of fluid velocity at
    position (x, y).
    :param x: x-position in fluid
    :param y: y-position in fluid
    :param r: sink position in fluid
    :param q: sink strength
    :return: x-velocity of fluid at the position (x,y) produced by a single sink at position r
    """
    return -q / (2 * np.pi) * (x - r[0]) / ((x - r[0]) ** 2 + (y - r[1]) ** 2)


def u_y_2d(x, y, r, q=0.0004):
    """
    Y-velocity for 2D sink - For a 2-dimensional sink at position r, calculate the y-component of fluid velocity at
    position (x, y).
    :param x: x-position in fluid
    :param y: y-position in fluid
    :param r: sink position in fluid
    :param q: sink strength
    :return: y-velocity of fluid at the position (x,y) produced by a single sink at position r
    """
    return -q / (2 * np.pi) * (y - r[1]) / ((x - r[0]) ** 2 + (y - r[1]) ** 2)


def U_2d(x, y, sink_positions, sink_strengths=0.0004):
    """
    Total velocity due to a group of sinks - For several 2-dimensional sinks at positions listed in sink_positions,
    calculate the total fluid velocity vector at position (x, y) or at a grid of (x, y) positions.
    :param x: x-position in fluid, or array of x-positions from a meshgrid
    :param y: y-position in fluid, or array of y-positions from a meshgrid
    :param sink_positions: array of sink positions
    :param sink_strengths: array of sink strengths that is the same length as sink_positions, or a float if all sink
    strengths are the same
    :return: velocity vector of fluid at the position (x,y) produced by a group of sinks at the positions listed in
    sink_positions
    """
    # Check that sink_positions and sink_strengths have the same length, if sink_strengths is an array.
    if type(sink_strengths) is not float and type(sink_strengths) is not np.ndarray:
        raise ValueError('Sink strength must be given as a float or an array.')

    if type(sink_strengths) is np.ndarray:
        if len(sink_positions) != len(sink_strengths):
            raise ValueError('The number of sink positions and sink strengths given must be the same.')

    # Calculation if (x, y) is a single position.
    if not type(x) is np.ndarray and not type(y) is np.ndarray:
        U_x = 0
        U_y = 0
        # The array of sink strengths is in the same order and has the same length as the array of sink positions, so
        # use `enumerate` on sink_positions to index the sink strength.
        for i, sink in enumerate(sink_positions):
            if type(sink_strengths) is float:
                q = sink_strengths
            elif type(sink_strengths) is np.ndarray:
                q = sink_strengths[i]
            U_x += u_x_2d(x, y, r=sink, q=q)
            U_y += u_y_2d(x, y, r=sink, q=q)
            U = (U_x, U_y)
    # Calculation if x and y are arrays.
    if type(x) is np.ndarray and type(y) is np.ndarray:
        # The shape of x and y is the same if they are from the same meshgrid; either shape may be used.
        shape = x.shape
        U_x = np.zeros(shape=shape)
        U_y = np.zeros(shape=shape)
        U = np.empty(shape, dtype=object)
        for i, j in itertools.product(range(shape[0]), range(shape[1])):
            for n, sink in enumerate(sink_positions):
                if type(sink_strengths) is float:
                    q = sink_strengths
                elif type(sink_strengths) is np.ndarray:
                    q = sink_strengths[n]
                U_x[i][j] += u_x_2d(x[i][j], y[i][j], r=sink, q=q)
                U_y[i][j] += u_y_2d(x[i][j], y[i][j], r=sink, q=q)
                U[i][j] = (U_x[i][j], U_y[i][j])
    # Error if x and y are different object types.
    if not type(x) is type(y):
        print('Either x and y must both be floats, or they must both be arrays.')
    return U_x, U_y, U


# Fluid velocity functions (3D).
def u_x_3d(x, y, r, q=0.0004):
    """
    X-velocity for 3D sink - For a 3-dimensional sink at position r, calculate the x-component of fluid velocity at
    position (x, y).
    :param x: x-position in fluid
    :param y: y-position in fluid
    :param r: sink position in fluid
    :param q: sink strength
    :return: x-velocity of fluid at the position (x,y) produced by a single sink at position r
    """
    return -q / (4 * np.pi) * (x - r[0]) / ((x - r[0]) ** 2 + (y - r[1]) ** 2) ** (3 / 2)


def u_y_3d(x, y, r, q=0.0004):
    """
    Y-velocity for 3D sink - For a 3-dimensional sink at position r, calculate the y-component of fluid velocity at
    position (x, y).
    :param x: x-position in fluid
    :param y: y-position in fluid
    :param r: sink position in fluid
    :param q: sink strength
    :return: y-velocity of fluid at the position (x,y) produced by a single sink at position r
    """
    return -q / (4 * np.pi) * (y - r[1]) / ((x - r[0]) ** 2 + (y - r[1]) ** 2) ** (3 / 2)


def U_3d(x, y, sink_positions, sink_strengths=0.0004):
    """
    Total velocity due to a group of sinks - For several 3-dimensional sinks at positions listed in sink_positions,
    calculate the total fluid velocity vector at position (x, y).
    :param x: x-position in fluid, or array of x-positions from a meshgrid
    :param y: y-position in fluid, or array of y-positions from a meshgrid
    :param sink_positions: array of sink positions
    :param sink_strengths: array of sink strengths that is the same length as sink_positions, or a float if all sink
    strengths are the same
    :return: velocity vector of fluid at the position (x,y) produced by a group of sinks at the positions listed in
    sink_positions
    """
    # Check that sink_positions and sink_strengths have the same length, if sink_strengths is an array.
    if type(sink_strengths) is not float and type(sink_strengths) is not np.ndarray:
        raise ValueError('Sink strength must be given as a float or an array.')

    if type(sink_strengths) is np.ndarray:
        if len(sink_positions) != len(sink_strengths):
            raise ValueError('The number of sink positions and sink strengths given must be the same.')

    # Calculation if (x, y) is a single position.
    if not type(x) is np.ndarray and not type(y) is np.ndarray:
        U_x = 0
        U_y = 0
        for i, sink in enumerate(sink_positions):
            if type(sink_strengths) is float:
                q = sink_strengths
            elif type(sink_strengths) is np.ndarray:
                q = sink_strengths[i]
            U_x += u_x_3d(x, y, r=sink, q=q)
            U_y += u_y_3d(x, y, r=sink, q=q)
            U = (U_x, U_y)
    # Calculation if x and y are arrays.
    if type(x) is np.ndarray and type(y) is np.ndarray:
        # The shape of x and y is the same if they are from the same meshgrid; either shape may be used.
        shape = x.shape
        print('Shape is' + str(shape))
        U_x = np.zeros(shape=shape)
        U_y = np.zeros(shape=shape)
        U = np.empty(shape, dtype=object)
        for i, j in itertools.product(range(shape[0]), range(shape[1])):
            for n, sink in enumerate(sink_positions):
                if type(sink_strengths) is float:
                    q = sink_strengths
                elif type(sink_strengths) is np.ndarray:
                    q = sink_strengths[n]
                U_x[i][j] += u_x_3d(x[i][j], y[i][j], r=sink, q=q)
                U_y[i][j] += u_y_3d(x[i][j], y[i][j], r=sink, q=q)
                U[i][j] = (U_x[i][j], U_y[i][j])
    # Error if x and y are different object types.
    if not type(x) is type(y):
        print('Either x and y must both be floats, or they must both be arrays.')
    return U_x, U_y, U