"""
Program: Jet Angle
Date: 13 August 2018
Author: LM

Calculate jet angle at the bubble position due to mirror sinks.
Make a plot describing convergence of the jet angle as layers of mirror sinks are added around the bubble position, for
cases with infinite numbers of mirror sinks.
Plot a vector field of the cavitation bubble jet angle inside a triangle and a square.
"""

import numpy as np
import itertools
import matplotlib.pyplot as plt
import math
from fluid_velocity import U_2d, U_3d
from sink_positions import sink_positions_corner, sink_positions_square, sink_positions_equilateral_triangle, \
    sink_positions_isosceles_right, sink_positions_306090


# Angle at bubble position functions.
def angle_bubble_position(bubble_position, sink_positions, sink_type='2D', print_angle=True):
    """
    Calculate the angle of the fluid velocity vector at the cavitation bubble position.
    :param bubble_position: position of cavitation bubble
    :param sink_positions: array of sink positions
    :param sink_type: 2D or 3D, velocities will be calculated for a 2-dimensional or 3-dimensional sink
    :param print_angle: True or False, if True will print the angle calculated
    :return: angle in radians of fluid velocity vector at the cavitation bubble position, defined as the angle
    counterclockwise from a rightward horizontal position
    """
    # Convert to float in case bubble_position vector contains integers.
    x_bubble = float(bubble_position[0])
    y_bubble = float(bubble_position[1])

    # 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()
    # Remove sink at cavitation bubble position.
    sink_positions.remove(bubble_position)

    # Calculate the x-component and y-component of velocity. Calculate for either 2D or 3D sinks.
    if sink_type == '2D':
        U_x = U_2d(x_bubble, y_bubble, sink_positions)[0]
        U_y = U_2d(x_bubble, y_bubble, sink_positions)[1]
    elif sink_type == '3D':
        U_x = U_3d(x_bubble, y_bubble, sink_positions)[0]
        U_y = U_3d(x_bubble, y_bubble, sink_positions)[1]

    # Calculate the angle.
    angle = math.atan2(U_y, U_x)
    if print_angle:
        print('Velocity vector angle is {0} radians (or {1} degrees).'.format(angle, angle / (2 * math.pi) * 360))
    return angle


def angle_convergence(bubble_position, shape, dimensions, max_range, sink_type='2D', plot=True,
                      plot_color='C0', print_result=False, saveas=None):
    """
    Data describing convergence of jet angle toward a single value - Calculate the jet angle for a cavitation bubble
    inside a shape using a number of "layers" of mirror sinks around the shape up to a specified maximum range.
    :param bubble_position: position of cavitation bubble
    :param max_range: number of "layers" of mirror sinks around the shape
    TODO: Describe the layers better. Possibly make it so that the sinks can be added one at a time.
    :param shape: square or triangle, at the moment
    :param dimensions: shape dimensions; For a square, provide a scalar side length.
    For an equilateral triangle, provide a scalar side length.
    :param sink_type: 2D or 3D, velocities will be calculated for a 2-dimensional or 3-dimensional sink
    :param plot: True or False, if True, returns a plot of the convergence data
    :param plot_color: optional, color for plot points
    :param print_result: True or False, if True, prints angle calculated and array of tuples with angle for each number
    of sinks
    :param saveas: file name
    :return: Returns an array of tuples (a, b) with the number of mirror sinks as the first item a and the fluid
    velocity vector angle at the bubble position as the second item b.
    """
    angles = []
    angle_convergence = []
    differences = []

    # Calculate angle convergence for a square, equilateral triangle, isosceles right triangle, or 306090 triangle.
    for w in range(1, max_range + 1):
        if shape == 'square':
            sink_positions = sink_positions_square(dimensions, bubble_position, w, print_result=print_result)
        if shape == 'triangle':
            sink_positions = sink_positions_equilateral_triangle(dimensions, bubble_position, w,
                                                                 print_result=print_result)
        if shape == 'isosceles right':
            sink_positions = sink_positions_isosceles_right(dimensions, bubble_position, w, print_result=print_result)
        if shape == '306090':
            sink_positions = sink_positions_306090(dimensions, bubble_position, w, print_result=print_result)

        theta = angle_bubble_position(bubble_position, sink_positions, sink_type=sink_type,
                                      print_angle=print_result)
        angles.append(theta)
    # Compute the number of mirror sinks, not including the bubble sink.
        number_sinks = len(sink_positions) - 1
        # Add the last value of the angles vector and the number_sinks vector to a tuple, which is put in a list.
        angle_convergence.append((number_sinks, angles[-1]))
        # Also compute the difference between the last calculated angle and the one before.
        if len(angles) > 2:
                differences.append((number_sinks, angles[-1] - angles[-2]))

    # Plot the angle convergence points, if selected.
    if plot:
        fig = plt.figure()
        ax = fig.add_subplot(1, 1, 1)
        angle_convergence = np.array(angle_convergence)
        ax.plot(angle_convergence[:, 0], angle_convergence[:, 1] * 180/np.pi, marker='o', markersize=4,
                color=plot_color, linestyle='None')
        ax.set_title('Jet Angle Prediction for a {0} \n as Number of Sinks Increases'.format(str.capitalize(shape)),
                     fontsize=16)
        ax.set_xlabel('Number of Sinks', fontsize=12)
        ax.set_ylabel('Fluid Velocity Vector Angle (degrees) \n at Cavitation Bubble Position', fontsize=12)
        plt.tight_layout()
        if saveas is not None:
            plt.savefig(saveas)
        plt.show()

        fig = plt.figure()
        ax = fig.add_subplot(1, 1, 1)
        differences = np.array(differences)
        ax.plot(differences[:, 0], differences[:, 1] * 180/np.pi, marker='o', markersize=4, color='green',
                linestyle='None')
        ax.set_title('Difference in Jet Angle Prediction for a {0} \n Between Current and Previous Number of Sinks'
                     .format(str.capitalize(shape)), fontsize=16)
        ax.set_xlabel('Number of Sinks', fontsize=12)
        ax.set_ylabel('Difference in Fluid Velocity Vector Angle (degrees)', fontsize=12)
        plt.axhline(y=0, linewidth=1, color='k')
        plt.tight_layout()
        if saveas is not None:
            plt.savefig(saveas)
        plt.show()

    if print_result:
        print(angle_convergence)
    return angle_convergence


# Angle at various positions vector field.
def jet_direction_vector_field(shape, dimensions, max_range, sink_type, plot=True, plot_color='r', edit_plot=False):
    """
    Plot cavitation bubble jet direction as a vector field - For a specified shape, plot the vector representing the
    jet direction of a cavitation bubble at several points within the shape.
    :param shape: square, triangle, or corner, at the moment
    :param dimensions: shape dimensions; For a square, provide a scalar side length.
    :param max_range: number of "layers" of mirror sinks around the shape
    TODO: Describe the layers better. Possibly make it so that the sinks can be added one at a time.
    :param sink_type: 2D or 3D, velocities will be calculated for a 2-dimensional or 3-dimensional sink
    :param plot: True or False, if True, plots the vector field
    :param plot_color: optional, vector color
    :param edit_plot: True or False; if True, returns the current plot axes so that the plot can be modified; if False,
    returns the arrays X, Y, U, and V, which have the position (X[i][j], Y[i][j]) and vector (U[i][j], V[i][j]) values
    :return: Returns the vector field of cavitation bubble jet directions when they are placed at points within the
    specified shape. (Plots the vector field if `plot` is True. Returns vector field values if `plot` is False.)
    """
    if shape == 'square':
        l = dimensions
        start = 0.9
        step = (1 - start) * l
        X, Y = np.mgrid[(-l / 2 * start):(l / 2 * start):step, (-l / 2 * start):(l / 2 * start):step]

        # X and Y have the same shape so either may be used below.
        shape = X.shape
        U = np.zeros(shape=shape)
        V = np.zeros(shape=shape)
        A = np.zeros(shape=shape)

        for i, j in itertools.product(range(shape[0]), range(shape[1])):
            bubble_position = (X[i][j], Y[i][j])
            # Calculate image sink positions for this bubble position.
            sink_positions = sink_positions_square(l=dimensions, bubble_position=bubble_position, w=max_range,
                                                   print_result=False)
            angle = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions, sink_type=sink_type,
                                          print_angle=False)
            U[i][j] = np.cos(angle)
            V[i][j] = np.sin(angle)
            A[i][j] = angle

        if plot or edit_plot:
            # Plot.
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(1, 1, 1)
            # Quiver plot.
            ax.quiver(X, Y, U, V, color=plot_color, pivot='mid')
            ax.set_title('Bubble Jet Direction Inside a Square')
            ax.set_xlabel('$x$')
            ax.set_ylabel('$y$')

        if edit_plot:
            if plot:
                raise ValueError('The options `plot` and `edit_plot` cannot both be selected as True.')
            return ax.quiver(X, Y, U, V, color=plot_color, pivot='mid')

        if plot:
            plt.show()
        else:
            return X, Y, U, V

    if shape == 'triangle':
        l = dimensions
        h = np.sqrt(3) / 2 * l
        start = 0.87  # 0.89
        step = (1 - start) * l / np.sqrt(3)
        h_bound = np.round(h * start, 1)
        h_remainder = divmod(h_bound, step)[1]
        X, Y = np.mgrid[(-h_bound + h_remainder):h_bound:step, 0:h:step]

        # X and Y have the same shape so either may be used below.
        shape = X.shape
        U = np.zeros(shape=shape)
        V = np.zeros(shape=shape)
        A = np.zeros(shape=shape)

        for i, j in itertools.product(range(shape[0]), range(shape[1])):
            # Print percentage completion.
            if j == 0:
                print(round((i / shape[0] + j / (shape[1] * shape[0])) * 100, 1), '%')
            # Skip calculating positions outside the triangle. Otherwise there will be an error.
            # Also skip calculating some slightly inside to avoid them sticking out when plotting, using e below.
            x_bubble = X[i][j]
            y_bubble = Y[i][j]
            e = 0.03
            if y_bubble <= (np.sqrt(3) * x_bubble) + e * l \
                    or y_bubble <= (- np.sqrt(3) * x_bubble) + e * l \
                    or y_bubble >= (np.sqrt(3) / 2 - e / 3) * l:
                U[i][j] = None
                V[i][j] = None
                A[i][j] = None
                continue

            # Calculate the jet angle at the bubble position.
            bubble_position = (X[i][j], Y[i][j])
            # Calculate image sink positions for this bubble position.
            sink_positions = sink_positions_equilateral_triangle(l=dimensions, bubble_position=bubble_position,
                                                                 w=max_range,
                                                                 print_result=False)
            angle = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions,
                                          sink_type=sink_type,
                                          print_angle=False)
            U[i][j] = np.cos(angle)
            V[i][j] = np.sin(angle)
            A[i][j] = angle

        if plot or edit_plot:
            # Plot.
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(1, 1, 1)
            # Quiver plot.
            ax.quiver(X, Y, U, V, color=plot_color, pivot='mid', units='width', width=0.0055)
            # Fill in areas outside triangle.
            x = np.linspace(-l / 2, l / 2, 100, endpoint=True)
            y = np.linspace(0, h, 100, endpoint=False)
            x_left = [-np.sqrt(3) * l] * 100
            x_right = [np.sqrt(3) * l] * 100
            y1 = np.sqrt(3) * x
            y2 = -np.sqrt(3) * x
            plt.plot(x_left, y, color='w')
            plt.plot(x_right, y, color='w')
            plt.plot(x, y1, color='k', linewidth=1)
            plt.plot(x, y2, color='k', linewidth=1)
            plt.axhline(y=h, color='k', linewidth=1)
            plt.fill_between(x, 0 * x, y1, color='w')
            plt.fill_between(x, 0 * x, y2, color='w')
            # Set axis properties.
            epsilon_x = 0.005 * h
            epsilon_y = 0.005 * l / np.sqrt(3)
            ax.axis([-l / 2 - epsilon_x, l / 2 + epsilon_x, 0, np.sqrt(3) / 2 * l + epsilon_y])
            # ax.spines['bottom'].set_position('zero')
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.set_title('Bubble Jet Direction Inside a Triangle')
            ax.set_xlabel('$x$', horizontalalignment='right', x=1.0)
            ax.set_ylabel('$y$', horizontalalignment='right', y=1.0)

        if edit_plot:
            if plot:
                raise ValueError('The options `plot` and `edit_plot` cannot both be selected as True.')
            return ax.quiver(X, Y, U, V, color=plot_color, pivot='mid', units='width', width=0.0055)

        if plot:
            plt.show()
        else:
            return X, Y, U, V

    if shape == 'isosceles right':
        l = dimensions
        start = 0.92
        step = (1 - start) * l
        X, Y = np.mgrid[0:(l * start):step, 0:(l * start):step]

        # X and Y have the same shape so either may be used below.
        shape = X.shape
        U = np.zeros(shape=shape)
        V = np.zeros(shape=shape)
        A = np.zeros(shape=shape)

        for i, j in itertools.product(range(shape[0]), range(shape[1])):
            if j == 0:
                print(round((i / shape[0] + j / (shape[1] * shape[0])) * 100, 1), '%')
            # Skip calculating positions outside the triangle. Otherwise there will be an error.
            # Also skip calculating some slightly inside to avoid them sticking out when plotting, using e below.
            x_bubble = X[i][j]
            y_bubble = Y[i][j]
            e = 0.01
            if x_bubble <= 0 \
                    or y_bubble <= 0 \
                    or y_bubble >= -x_bubble + l - e:
                U[i][j] = None
                V[i][j] = None
                A[i][j] = None
                continue
            # Calculate the jet angle at the bubble position.
            bubble_position = (X[i][j], Y[i][j])
            # Calculate image sink positions for this bubble position.
            sink_positions = sink_positions_isosceles_right(l=dimensions, bubble_position=bubble_position, w=max_range,
                                                            print_result=False)
            angle = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions,
                                          sink_type=sink_type,
                                          print_angle=False)
            U[i][j] = np.cos(angle)
            V[i][j] = np.sin(angle)
            A[i][j] = angle

        if plot or edit_plot:
            # Plot.
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(1, 1, 1)
            # Quiver plot.
            ax.quiver(X, Y, U, V, color=plot_color, pivot='mid')
            ax.plot([0, l], [l, 0], linestyle='-', color='k', linewidth=1)
            ax.axis([0, l, 0, l])
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.set_title('Bubble Jet Direction Inside an Isosceles Right Triangle')
            ax.set_xlabel('$x$')
            ax.set_ylabel('$y$')

        if edit_plot:
            if plot:
                raise ValueError('The options `plot` and `edit_plot` cannot both be selected as True.')
            return ax.quiver(X, Y, U, V, color=plot_color, pivot='mid')

        if plot:
            plt.show()
        else:
            return X, Y, U, V

    if shape == '306090':
        l = dimensions
        start = 0.92 # 0.94
        step = (1 - start) * l
        X, Y = np.mgrid[0:(l * start):step, 0:(l * start):step]

        # X and Y have the same shape so either may be used below.
        shape = X.shape
        U = np.zeros(shape=shape)
        V = np.zeros(shape=shape)
        A = np.zeros(shape=shape)

        for i, j in itertools.product(range(shape[0]), range(shape[1])):
            if j == 0:
                print(round((i / shape[0] + j / (shape[1] * shape[0])) * 100, 1), '%')
            # Skip calculating positions outside the triangle. Otherwise there will be an error.
            # Also skip calculating some slightly inside to avoid them sticking out when plotting, using e below.
            x_bubble = X[i][j]
            y_bubble = Y[i][j]
            e = 0.03
            if x_bubble <= 0 \
                    or y_bubble <= 0 \
                    or y_bubble >= -np.sqrt(3) * x_bubble + np.sqrt(3)/2 * l - e:
                U[i][j] = None
                V[i][j] = None
                A[i][j] = None
                continue
            # Calculate the jet angle at the bubble position.
            bubble_position = (X[i][j], Y[i][j])
            # Calculate image sink positions for this bubble position.
            sink_positions = sink_positions_306090(l=dimensions, bubble_position=bubble_position, w=max_range,
                                                   print_result=False)
            angle = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions,
                                          sink_type=sink_type,
                                          print_angle=False)
            U[i][j] = np.cos(angle)
            V[i][j] = np.sin(angle)
            A[i][j] = angle

        if plot or edit_plot:
            # Plot.
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(1, 1, 1)
            # Quiver plot.
            ax.quiver(X, Y, U, V, color=plot_color, pivot='mid')
            ax.plot([0, l / 2], [np.sqrt(3)/2 * l, 0], linestyle='-', color='k', linewidth=1)
            ax.axis([0, l, 0, l])
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.set_title('Bubble Jet Direction Inside a 30$^{\\circ}$-60$^{\\circ}$-90$^{\\circ}$ Triangle')
            ax.set_xlabel('$x$')
            ax.set_ylabel('$y$')

        if edit_plot:
            if plot:
                raise ValueError('The options `plot` and `edit_plot` cannot both be selected as True.')
            return ax.quiver(X, Y, U, V, color=plot_color, pivot='mid')

        if plot:
            plt.show()
        else:
            return X, Y, U, V

    if shape == 'corner':
        n = dimensions
        l = 10
        start = 0.91
        step = (1 - start) * l
        h = 10
        h_bound = np.round(h * start, 1)
        h_remainder = divmod(h_bound, step)[1]
        X, Y = np.mgrid[0:l:step, (-h_bound + h_remainder):h_bound:step]

        # X and Y have the same shape so either may be used below.
        shape = X.shape
        U = np.zeros(shape=shape)
        V = np.zeros(shape=shape)
        A = np.zeros(shape=shape)

        for i, j in itertools.product(range(shape[0]), range(shape[1])):
            if j == 0:
                print(round((i / shape[0] + j / (shape[1] * shape[0])) * 100, 1), '%')
            # Skip calculating positions outside the corner. Otherwise there will be an error.
            # Also skip calculating some slightly inside to avoid them sticking out when plotting, using e below.
            x_bubble = X[i][j]
            y_bubble = Y[i][j]
            bubble_position = (x_bubble, y_bubble)
            e = 0.03
            if x_bubble <= 0 \
                    or y_bubble >= (np.tan(np.pi / n) * x_bubble) - e * l \
                    or y_bubble <= 0 + e * l:
                U[i][j] = None
                V[i][j] = None
                A[i][j] = None
                continue

            # Calculate the angle at the bubble position.
            sink_positions = sink_positions_corner(n=dimensions,
                                                   bubble_position=bubble_position,
                                                   print_result=False)
            angle = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions,
                                          sink_type=sink_type, print_angle=False)
            U[i][j] = np.cos(angle)
            V[i][j] = np.sin(angle)
            A[i][j] = angle

        if plot or edit_plot:
            # Plot.
            fig = plt.figure(figsize=(8, 8))
            ax = fig.add_subplot(1, 1, 1)
            # Quiver plot.
            ax.quiver(X, Y, U, V, color=plot_color, pivot='mid', units='width', width=0.0055)
            # Fill in areas outside triangle.
            x = np.linspace(0, 2 * l, 100, endpoint=False)
            y_below = 0 * x
            y1 = np.tan(np.pi / n) * x
            y2 = 0 * x
            plt.plot(x, y_below, color='w')
            plt.plot(x, y1,
                     x, y2, color='k', linewidth=1)
            # Set axis properties.
            epsilon_x = 0.005 * l
            epsilon_y = 0.005 * l
            ax.axis([0, l + epsilon_x, 0, l + epsilon_y])
            # ax.spines['bottom'].set_position('zero')
            ax.spines['top'].set_visible(False)
            ax.spines['right'].set_visible(False)
            ax.set_title('Cavitation Bubble Jet Direction In a Corner of Angle $\pi$/{0}'.format(n))
            ax.set_xlabel('$x$', horizontalalignment='right', x=1.0)
            ax.set_ylabel('$y$', horizontalalignment='right', y=1.0)

        if edit_plot:
            if plot:
                raise ValueError('The options `plot` and `edit_plot` cannot both be selected as True.')
            return ax.quiver(X, Y, U, V, color=plot_color, pivot='mid', units='width', width=0.0055)

        if plot:
            plt.show()
        else:
            return X, Y, U, V


def bubble_angle_along_lines(lines, number_points, shape, dimensions, max_range, sink_type='2D', mod_2pi=True):
    """
    Plot collapsing bubble jet angle as a function of bubble position along a line.
    :param lines: list of lines, where each line is an array of points (and each point is an array with two elements)
    :param number_points: number of points in each line
    :param shape: triangle only, at the moment
    :param dimensions: shape dimension; For an equilateral triangle, provide a scalar side length.
    :param max_range: number of "layers" of mirror sinks around the shape
    :param sink_type: 2D or 3D, velocities will be calculate for a 2-dimensional or 3-dimensional sink
    :param color: optional, line color; string or list # TODO: Add color.
    :param saveas: file name # TODO: Add saving.
    :return:
    """

    # Gather jet angle data for each line.
    n = len(lines)
    p = number_points
    jet_angle_data = np.zeros((n, p))

    for n, line in enumerate(lines):
        jet_angles = []
        # Calculate jet angle at each point on the line.
        for i in range(len(line)):
            line_position = (line[i][0], line[i][1])
            if shape == 'triangle':
                sink_positions = sink_positions_equilateral_triangle(l=dimensions,
                                                                     bubble_position=line_position,
                                                                     w=max_range,
                                                                     print_result=False)
            elif shape == 'square':
                sink_positions = sink_positions_square(l=dimensions,
                                                       bubble_position=line_position,
                                                       w=max_range,
                                                       print_result=False)
            elif shape == 'isosceles right':
                sink_positions = sink_positions_isosceles_right(l=dimensions,
                                                                bubble_position=line_position,
                                                                w=max_range,
                                                                print_result=False)
            elif shape == '306090':
                sink_positions = sink_positions_306090(l=dimensions,
                                                       bubble_position=line_position,
                                                       w=max_range,
                                                       print_result=False)
            elif shape == 'corner':
                sink_positions = sink_positions_corner(n=dimensions,
                                                       bubble_position=line_position,
                                                       print_result=False)
            angle = angle_bubble_position(bubble_position=line_position,
                                          sink_positions=sink_positions,
                                          sink_type=sink_type,
                                          print_angle=False)
            # Since atan2 returns values from -pi to pi, take the angle modulo 2pi, which will make the negative
            # values the correct positive ones.
            if mod_2pi:
                angle = angle % (2 * np.pi)
            jet_angles.append(angle)
        jet_angle_data[n] = jet_angles

    return jet_angle_data