"""
Program: Sink Positions
Date: 29 March 2019 (Updated: 25 June 2019)
Author: Lebo Molefe

Calculate and plot image sink positions to produce the boundary conditions for various solid boundary shapes.
"""

# Import packages.
import numpy as np
import math
import itertools
import matplotlib.pyplot as plt


# Sink positions functions.
def sink_positions_corner(n, bubble_position, print_result):
    """
    Calculate mirror sink positions to give the boundary conditions of a bubble collapsing in a corner.
    :param n: a natural number, n = 1, 2, 3, ...; specify n to calculate sink positions for a bubble within a corner of
    angle pi/n.
    :param bubble_position: bubble position, where bubble corner; a vector
    :param print_result: True or False, prints how many sinks were calculated
    :return: list of sink_positions as arrays
    """

    s = []

    x_bubble = bubble_position[0]
    y_bubble = bubble_position[1]

    if math.atan2(y_bubble, x_bubble) <= 0 \
            or math.atan2(y_bubble, x_bubble) >= np.pi / n:
        raise ValueError('Bubble must be within corner. '
                         'The bubble angle from horizontal must be between 0 and pi/{0}.'.format(n))

    # Calculate the radially symmetric pattern for a bubble centered in a corner.
    bubble_position = np.array(bubble_position)
    # Make a copy of the array that only exists in the function so any original bubble_position array is not changed.
    bubble_position = bubble_position.copy()

    r = np.linalg.norm(bubble_position)
    t = []
    for i in range(2 * n):
        theta = np.pi / (2 * n) + (np.pi / n) * i
        t.append(theta)

    for m in range(len(t)):
        angle = t[m]
        p_x = r * np.cos(angle)
        p_y = r * np.sin(angle)
        p = (p_x, p_y)
        p = np.array(p)
        s.append(p)

    # For a bubble that is not centered in the corner, calculate offset vectors to alter the radially symmetric pattern.
    change_vector_0 = bubble_position - s[0]
    r_offset = np.linalg.norm(change_vector_0)

    theta = [0] * (2 * n)

    # TODO: Include a picture of how to get this formula.
    a = math.atan2(change_vector_0[1], change_vector_0[0])
    theta[0] = a
    theta[2 * n - 1] = -a

    for k in range(1, n):
        for j in (0, 1):
            theta[2 * k - 1 + j] = (2 * np.pi / n) * k + (-1) ** (j + 1) * a

    change_vector = [0] * (2 * n)
    change_vector[0] = change_vector_0
    for n in range(1, 2 * n):
        change_vector[n] = (r_offset * np.cos(theta[n]), r_offset * np.sin(theta[n]))
    # Turn the vectors into type numpy.ndarray so they can be added to other vectors.
    # The function enumerate gives a list of tuples which contain an index and a vector.
    for m, vector in enumerate(change_vector):
        change_vector[m] = np.array(vector)
    # Make the whole list change_vector itself a numpy.ndarray.
    change_vector = np.array(change_vector)

    # Check that the number of offset vectors is the same as the number of positions.
    if len(s) != len(change_vector):
        raise ValueError('The number of offset vectors is different than the number of sink positions.')

    # Move the radially symmetric pattern of sinks.
    for p, vector in enumerate(s):
        if p == 0:
            # Avoiding floating point precision problems.
            s[p] = bubble_position
        else:
            s[p] = s[p] + change_vector[p]
        # Make the point into a tuple instead of an array.
        s[p] = (s[p][0], s[p][1])

    number_sinks = len(s)
    if print_result:
        print('Calculated {0} sink positions.'.format(number_sinks))
    return s


def sink_positions_parallel_walls(distance_between_walls, bubble_position, n, print_result=False):
    """
    Calculate mirror sink positions to give the boundary conditions of a bubble collapsing between two parallel walls.
    :param distance_between_walls: distance between parallel walls; an integer or float
    :param bubble_position: bubble position within vertical walls centered around the origin
    :param n: number of sink positions to calculate, since the number of sinks should be infinite, but function can
    calculate a finite number
    :param print_result: True or False, prints how many sinks were calculated
    :return:  list of sink_positions as arrays
    """

    s = []
    offset = distance_between_walls - bubble_position[0]

    i = 0
    while len(s) < n:
        s.append((bubble_position[0] + 2 * distance_between_walls * i, bubble_position[1]))
        s.append((offset + 2 * distance_between_walls * i, bubble_position[1]))
        s.append((bubble_position[0] - 2 * distance_between_walls * (i + 1), bubble_position[1]))
        s.append((offset - 2 * distance_between_walls * (i + 1), bubble_position[1]))
        i += 1

    number_sinks = len(s)
    if print_result:
        print('Calculated {0} sink positions.'.format(number_sinks))
    return s


def sink_positions_square(l, bubble_position, w=None, x_range=None, y_range=None, print_result=True):
    """
    Calculate mirror sink positions to give the boundary conditions of a bubble collapsing inside a square.
    :param l: square side length
    :param bubble_position: bubble position, where bubble is within the square centered at the origin; a vector
    :param w: range of symmetric reflections in the x-direction and y-direction to calculate; setting this value will
    calculate a square range of symmetric reflections centered at the origin. To set custom values, specify x_range and
    y_range instead.
    :param x_range: range of symmetric reflections in the x-direction to calculate; a vector of the form (x_1, x_2)
    :param y_range: range of symmetric reflections in the y-direction to calculate; a vector of the form (y_1, y_2)
    :param print_result: True or False, if True, prints how many sink positions were calculated
    :return: list of sink positions as tuples
    """
    s_x = []
    s_y = []
    s = []
    x_bubble = bubble_position[0]
    y_bubble = bubble_position[1]

    if x_bubble >= l / 2 \
            or x_bubble <= -l / 2 \
            or y_bubble >= l / 2 \
            or y_bubble <= -l / 2:
        print(x_bubble, y_bubble)
        raise ValueError('Bubble position must be within square at the origin (coordinates within range ({0}, {1})).'
                         .format(-l / 2, l / 2))

    if x_range is None and y_range is None and w is None:
        raise ValueError('Define the area of reflections to calculate. \n Specify w or x_range and y_range.')
    elif x_range is None and y_range is None and w is not None:
        x_range = (-w, w + 1)
        y_range = (-w, w + 1)
    elif x_range is not None and y_range is not None and w is None:
        x_range = (-x_range, x_range + 1)
        y_range = (-y_range, y_range + 1)
    elif (x_range is None and y_range is not None) \
            or (x_range is not None and y_range is None) \
            or (x_range is not None and y_range is not None and w is not None):
        raise ValueError('Specify BOTH x_range and y_range, or specify w.')

    for i, j in itertools.product(range(x_range[0], x_range[1]), range(y_range[0], y_range[1])):
        s_x.append(x_bubble * (-1) ** i + l * i)
        s_y.append(y_bubble * (-1) ** j + l * j)
        # The -1 below retrieves the last added element to s_x and s_y.
        s.append((s_x[-1], s_y[-1]))
    number_sinks = len(s)
    if print_result:
        print('Calculated {0} sink positions.'.format(number_sinks))
    return s


def sink_positions_rectangle(l_x, l_y, bubble_position, w=None, x_range=None, y_range=None, print_result=True):
    """
    Calculate mirror sink positions to give the boundary conditions of a bubble collapsing inside a rectangle.
    # TODO: Combine this with code for square. It's the same code except for one line.
    :param l_x: rectangle width
    :param l_y: rectangle height
    :param bubble_position: bubble position, where bubble is within the rectangle centered at the origin; a vector
    :param w: range of symmetric reflections in the x-direction and y-direction to calculate; setting this value will
    calculate a square range of symmetric reflections centered at the origin. To set custom values, specify x_range and
    y_range instead.
    :param x_range: range of symmetric reflections in the x-direction to calculate; a vector of the form (x_1, x_2)
    :param y_range: range of symmetric reflections in the y-direction to calculate; a vector of the form (y_1, y_2)
    :param print_result: True or False, if True, prints how many sink positions were calculated
    :return: list of sink positions as tuples
    """
    s_x = []
    s_y = []
    s = []
    x_bubble = bubble_position[0]
    y_bubble = bubble_position[1]

    if x_bubble >= l_x / 2 \
            or x_bubble <= -l_x / 2 \
            or y_bubble >= l_y / 2 \
            or y_bubble <= -l_y / 2:
        raise ValueError('Bubble position must be within square at the origin (x-coordinate within range ({0}, {1}) '
                         'and y-coordinate within range ({2}, {3})).'
                         .format(-l_x / 2, l_x / 2, -l_y / 2, l_y / 2))

    if x_range is None and y_range is None and w is None:
        raise ValueError('Define the area of reflections to calculate. \n Specify w or x_range and y_range.')
    elif x_range is None and y_range is None and w is not None:
        x_range = (-w, w + 1)
        y_range = (-w, w + 1)
    elif x_range is not None and y_range is not None and w is None:
        x_range = (-x_range, x_range + 1)
        y_range = (-y_range, y_range + 1)
    elif (x_range is None and y_range is not None) \
            or (x_range is not None and y_range is None) \
            or (x_range is not None and y_range is not None and w is not None):
        raise ValueError('Specify BOTH x_range and y_range, or specify w.')

    for i, j in itertools.product(range(x_range[0], x_range[1]), range(y_range[0], y_range[1])):
        s_x.append(x_bubble * (-1) ** i + l_x * i)
        s_y.append(y_bubble * (-1) ** j + l_y * j)
        # The -1 below retrieves the last added element to s_x and s_y.
        s.append((s_x[-1], s_y[-1]))
    number_sinks = len(s)
    if print_result:
        print('Calculated {0} sink positions.'.format(number_sinks))
    return s


def sink_positions_equilateral_triangle(l, bubble_position, w=None, x_range=None, y_range=None, print_result=True,
                                        rotate=False):
    """
    Calculate mirror sink positions to give boundary conditions of a bubble collapsing inside an equilateral triangle.
    :param l: triangle side length
    :param bubble_position: bubble position, where bubble is within a triangle with its tip at the origin and two other
    endpoints equidistant from the x-axis in the first and fourth quadrants; a vector
    :param w: range of symmetric reflections in the x-direction and y-direction to calculate; setting this value will
    calculate a square range of symmetric reflections centered at the origin. To set custom values, specify x_range and
    y_range instead.
    :param x_range: range of symmetric reflections in the x-direction to calculate; a vector of the form (x_1, x_2)
    :param y_range: range of symmetric reflections in the y-direction to calculate; a vector of the form (y_1, y_2)
    :param print_result: True or False, if True, prints how many sink positions were calculated
    :param rotate: True or False, if True, rotates the figure by 90 degrees clockwise
    :return: list of sink positions as tuples
    """

    s_x = []
    s_y = []
    s = []
    h = l * np.cos(np.pi / 6)
    x_bubble = bubble_position[0]
    y_bubble = bubble_position[1]

    # Make sure bubble is within a triangle with tip at the origin and flat side above the tip.
    if y_bubble <= np.sqrt(3) * x_bubble \
            or y_bubble <= -np.sqrt(3) * x_bubble \
            or y_bubble >= l * np.cos(np.pi / 6):
        raise ValueError('Bubble position must be within equilateral triangle with tip at the origin and flat side '
                         'above the tip (x and y such that y >= sqrt(3) * x; and y >= -sqrt(3) * x); '
                         'and y <= l * cos(pi/6) where l = {0}).'.format(l))

    if x_range is None and y_range is None and w is None:
        raise ValueError('Define the area of reflections to calculate. \n Specify w or x_range and y_range.')
    elif x_range is None and y_range is None and w is not None:
        x_range = (-w, w + 1)
        y_range = (-w, w + 1)
    elif x_range is not None and y_range is not None and w is None:
        x_range = (-x_range, x_range + 1)
        y_range = (-y_range, y_range + 1)
    elif (x_range is None and y_range is not None) \
            or (x_range is not None and y_range is None) \
            or (x_range is not None and y_range is not None and w is not None):
        raise ValueError('Specify BOTH x_range and y_range, or specify w.')

    for i, j in itertools.product(range(x_range[0], x_range[1]), range(y_range[0], y_range[1])):
        # if abs(j) < -np.sqrt(3) / 3 * abs(i) + y_range[1]:
        x_s = x_bubble + i * (3 / 2) * l
        y_s = (-1) ** j * (y_bubble - h / 2) + j * h + (-h + (-1) ** i * h) / 2 + h / 2
        # The second-to-last term of y_s is 0 (even i) or -h (odd i)
        for k in (-1, 0, 1):
            # Avoid floating point precision error for the bubble position.
            if i == 0 and j == 0 and k == 0:
                s_x.append(x_bubble)
                s_y.append(y_bubble)
            else:
                s_x.append((x_s - k ** 2 * x_bubble) * np.cos(k * np.pi / 3) - y_s * np.sin(k * np.pi / 3)
                           + k ** 2 * x_bubble * np.cos(np.pi + k * np.pi / 3))
                s_y.append((x_s - k ** 2 * x_bubble) * np.sin(k * np.pi / 3) + y_s * np.cos(k * np.pi / 3)
                           + k ** 2 * x_bubble * np.sin(np.pi + k * np.pi / 3))
            # The -1 below retrieves the last added element to s_x and s_y.
            if not rotate:
                s.append((s_x[-1], s_y[-1]))
            if rotate:
                s.append((s_y[-1], s_x[-1]))  # This rotates the figure by 90 degrees clockwise.
    number_sinks = len(s)
    if print_result:
        print('Calculated {0} sink positions.'.format(number_sinks))
    return s


def sink_positions_isosceles_right(l, bubble_position, w=None, x_range=None, y_range=None, rotate=False, print_result=True):
    """
    Calculate mirror sinks to give the boundary conditions of a bubble collapsing inside an isosceles right triangle.
    :param l: isosceles right triangle side (not hypotenuse) length
    :param bubble_position: bubble position, where bubble is within an isosceles right triangle in the first quadrant,
    with one corner at the origin and with one side parallel to the x axis; a vector
    :param w: range of symmetric reflections in the x-direction and y-direction to calculate; setting this value will
    calculate a square range of symmetric reflections centered at the origin. To set custom values, specify x_range and
    y_range instead.
    :param x_range: range of symmetric reflections in the x-direction to calculate; a vector of the form (x_1, x_2)
    :param y_range: range of symmetric reflections in the y-direction to calculate; a vector of the form (y_1, y_2)
    :param rotate: True or False, if True, rotates the sink positions by 45 degrees clockwise
    :param print_result: True or False, if True, prints how many sink positions were calculated
    :return: list of sink positions as tuples
    """
    s_x = []
    s_y = []
    s = []
    x_bubble = bubble_position[0]
    y_bubble = bubble_position[1]

    # Make sure bubble is within an isosceles triangle with tip at the origin and one side flat aligned with the x-axis.
    if x_bubble <= 0 \
            or y_bubble <= 0 \
            or y_bubble >= -x_bubble + l:
        raise ValueError('Bubble position must be within isosceles right triangle with tip at the origin and flat side '
                         'aligned with the x-axis (x and y such that x >=0; y>= 0; '
                         'and y <= -x_bubble + l = -x_bubble + {0}).'.format(l))

    if x_range is None and y_range is None and w is None:
        raise ValueError('Define the area of reflections to calculate. \n Specify w or x_range and y_range.')
    elif x_range is None and y_range is None and w is not None:
        x_range = (-w, w + 1)
        y_range = (-w, w + 1)
    elif x_range is not None and y_range is not None and w is None:
        x_range = (-x_range, x_range + 1)
        y_range = (-y_range, y_range + 1)
    elif (x_range is None and y_range is not None) \
            or (x_range is not None and y_range is None) \
            or (x_range is not None and y_range is not None and w is not None):
        raise ValueError('Specify BOTH x_range and y_range, or specify w.')

    # Make the pattern of image sinks fulfilling the boundary condition of a square with side length l centered at the (l/2, l/2).
    for i, j in itertools.product(range(x_range[0], x_range[1]), range(y_range[0], y_range[1])):
        # Sink positions for a square centered at (l / 2, l / 2).
        x_s = (x_bubble - l / 2) * (-1)**i + l * i - l / 2
        y_s = (y_bubble - l / 2) * (-1)**j + l * j + l / 2
        for k in (0, 1):
            # Avoid floating point precision error for the bubble position.
            if i == 0 and j == 0 and k == 0:
                s_x.append(x_bubble)
                s_y.append(y_bubble)
            else:
                s_x.append(x_s * np.cos(k * np.pi / 2) - y_s * np.sin(k * np.pi / 2) + l)
                s_y.append(x_s * np.sin(k * np.pi / 2) + y_s * np.cos(k * np.pi / 2))

            # The -1 below retrieves the last added element to s_x and s_y.
            if not rotate:
                s.append((s_x[-1], s_y[-1]))
            if rotate:
                R = np.array([[np.cos(np.pi / 4), -np.sin(np.pi / 4)], [np.sin(np.pi / 4), np.cos(np.pi / 4)]])
                s.append(R.dot(np.array([s_x[-1], s_y[-1]])))  # This rotates the positions by 45 degrees clockwise.

    number_sinks = len(s)
    if print_result:
        print('Calculated {0} sink positions.'.format(number_sinks))
    return s


def sink_positions_306090(l, bubble_position, w=None, x_range=None, y_range=None, rotate=False, a=None, print_result=True):
    """
    Calculate mirror sinks to give the boundary conditions of a bubble collapsing inside a 30-60-90 degree triangle.
    :param l: length of hypotenuse of the 30-60-90 degree triangle
    :param bubble_position: bubble position, where bubble is within a 30-60-90 degree right triangle in the first
    quadrant, with one corner at the origin and with one side parallel to the x axis; a vector
    :param w: range of symmetric reflections in the x-direction and y-direction to calculate; setting this value will
    calculate a square range of symmetric reflections centered at the origin. To set custom values, specify x_range and
    y_range instead.
    :param x_range: range of symmetric reflections in the x-direction to calculate; a vector of the form (x_1, x_2)
    :param y_range: range of symmetric reflections in the y-direction to calculate; a vector of the form (y_1, y_2)
    :param rotate: True or False, if True, rotates the figure by 45 degrees clockwise
    :param print_result: True or False, if True, prints how many sink positions were calculated
    :return: list of sink positions as tuples
    """

    s_x = []
    s_y = []
    s = []
    x_bubble = bubble_position[0]
    y_bubble = bubble_position[1]

    # Make sure bubble is within a 30-60-90 triangle with tip at the origin and long side flat aligned with the y-axis.
    if x_bubble <= 0 \
            or y_bubble <= 0 \
            or y_bubble >= -np.sqrt(3) * x_bubble + np.sqrt(3)/2 * l:
        raise ValueError('Bubble position must be within 30-60-90 right triangle with tip at the origin and short side '
                         'aligned with the x-axis (x and y such that x >=0; y>= 0; '
                         'and y <= -sqrt(3) * x_bubble + sqrt(3)/2 * l where l = {0}).'.format(l))

    if x_range is None and y_range is None and w is None:
        raise ValueError('Define the area of reflections to calculate. \n Specify w or x_range and y_range.')
    elif x_range is None and y_range is None and w is not None:
        x_range = (-w, w + 1)
        y_range = (-w, w + 1)
    elif x_range is not None and y_range is not None and w is None:
        x_range = (-x_range, x_range + 1)
        y_range = (-y_range, y_range + 1)
    elif (x_range is None and y_range is not None) \
            or (x_range is not None and y_range is None) \
            or (x_range is not None and y_range is not None and w is not None):
        raise ValueError('Specify BOTH x_range and y_range, or specify w.')

    # Make the pattern of image sinks fulfilling the boundary condition of a rectangle with side lengths h and 3l/2
    # centered at the point (l/2, h/2), where h is the length of the long side of the triangle.
    h = l * np.cos(np.pi / 6) # Or h = np.sqrt(l**2 - (l / 2)**2)
    for i, j in itertools.product(range(x_range[0], x_range[1]), range(y_range[0], y_range[1])):
        # Sink positions for a square centered at (l / 2, h / 2).
        x_s = (x_bubble - 3/4 * l) * (-1) ** i + 3/2 * l * i + 3/4 * l
        y_s = (y_bubble - h / 2) * (-1) ** j + h * j - h / 2
        for k in (-1, 0, 1):
            # Avoid floating point precision error for the bubble position.
            if i == 0 and j == 0 and k == 0:
                s.append((x_bubble, y_bubble))
            else:
                s_x.append(x_s * np.cos(k * np.pi / 3) - y_s * np.sin(k * np.pi / 3))
                s_y.append(x_s * np.sin(k * np.pi / 3) + y_s * np.cos(k * np.pi / 3) + h)

                # The -1 below retrieves the last added element to s_x and s_y.
                if not rotate:
                    s.append((s_x[-1], s_y[-1]))
                if rotate:
                    R = np.array([[np.cos(a), -np.sin(a)], [np.sin(a), np.cos(a)]])
                    s.append(R.dot(np.array([s_x[-1], s_y[-1]])))  # This rotates the figure by angle a degrees clockwise.

                s_x.append((x_s - 3/2 * l) * np.cos(k * np.pi / 3) - (y_s + h) * np.sin(k * np.pi / 3))
                s_y.append((x_s - 3/2 * l) * np.sin(k * np.pi / 3) + (y_s + h) * np.cos(k * np.pi / 3) + h)

                # The -1 below retrieves the last added element to s_x and s_y.
                if not rotate:
                    s.append((s_x[-1], s_y[-1]))
                if rotate:
                    R = np.array([[np.cos(a), -np.sin(a)], [np.sin(a), np.cos(a)]])
                    s.append(R.dot(np.array([s_x[-1], s_y[-1]])))  # This rotates the figure by angle a degrees clockwise.

    number_sinks = len(s)
    if print_result:
        print('Calculated {0} sink positions'.format(number_sinks))
    return s


# Plotting function.
def plot_sink_positions(sink_positions, color='C0', alpha=None, markersize=8,
                        bubble_position=None, include_bubble=True,
                        title=None, saveas=None,
                        figsize=None, tick_label_size=None, edit_plot=False):
    """
    Plot sink positions and optionally the fluid velocity vector field produced by those sinks.
    :param sink_positions: array of sink positions
    :param color: sink color, optional
    :param alpha: transparency of sink markers, where alpha = 1 is fully opaque
    :param markersize: size of sink points
    :param bubble_position: position of bubble, optional; Include this to remove the bubble
    from the plot using the include_bubble keyword argument; a vector
    :param include_bubble: True or False, if False, removes the bubble from the plot
    :param title: title for the plot; Enter a shape name, it will be preceded by the phrase 'Sink Positions for a '...
    :param saveas: file name
    :param figsize: tuple indicating size of plot
    :param tick_label_size: size of number labels for axis ticks
    :param edit_plot: True or False, if True returns the current plot axes so that the plot can be modified
    :return: plot of sink positions and, optionally, fluid velocity vector field produced by those sinks
    """
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(1, 1, 1)

    # Remove bubble if selected.
    if not include_bubble and bubble_position in sink_positions:
        # Make a copy of the array that only exists in the function so the original sink_positions array is not changed.
        sink_positions = sink_positions.copy()
        sink_positions.remove(bubble_position)

    # Plot.
    sink_positions = np.array(sink_positions)
    ax.plot(sink_positions[:, 0], sink_positions[:, 1], marker='o', markersize=markersize, color=color,
            linestyle='None', alpha=alpha)

    # Label.
    if title is not None:
        ax.set_title('Sink Positions for a ' + title)
    if tick_label_size is not None:
        ax.tick_params(labelsize=tick_label_size)

    # Save, if desired.
    if saveas is not None:
        plt.savefig(saveas)

    if edit_plot:
        return plt.gca()

    plt.show()