"""
Program: Figure 4 - Representative Model Plots
Date: 12 March 2019
Author: Lebo Molefe

Plot model for a few representative locations within the 4 shapes.
"""

import numpy as np
import matplotlib.pyplot as plt
from jet_angle import jet_direction_vector_field, bubble_angle_along_lines

dimensions = 1
max_range = 20
sink_type = '3D'
vector_color = 'indigo'
colors = ['darkorange', 'firebrick', 'navy', 'g', 'purple']
linestyle = ['-', '--', '-.', ':', '-']

# Positions for plot labels 'a', 'b', 'c', and 'd'
x_text = -0.3
y_text = 1
fontsize = 24
labelsize = 20
ticksize = 14
legendsize = 14
linewidth1 = 2.5
linewidth2 = 3

test_positions_sq = [0.4, 0.2, 0.05]
test_positions_tri = [0.8, 0.6, 0.55, 0.3, 0.1]

test_positions_isc_right = [0.7, 0.5, 0.3, 0.1]
test_positions_306090_tri = [0.7, 0.5, 0.3, 0.2, 0.1]

### PLOT ###

# Figure and subplots.
fig = plt.figure(figsize=(10, 20))  # Or figsize = (9, 9) if only half the plots
ax1 = fig.add_subplot(4, 2, 1)
ax2 = fig.add_subplot(4, 2, 2)
ax3 = fig.add_subplot(4, 2, 3)
ax4 = fig.add_subplot(4, 2, 4)
ax5 = fig.add_subplot(4, 2, 5)
ax6 = fig.add_subplot(4, 2, 6)
ax7 = fig.add_subplot(4, 2, 7)
ax8 = fig.add_subplot(4, 2, 8)

########################################################################################################################

### FOR SUBPLOT 1 - SQUARE ###
X_sq, Y_sq, U_sq, V_sq = jet_direction_vector_field(shape='square', dimensions=1, max_range=20, sink_type=sink_type, plot=False)


### FOR SUBPLOT 2 - SQUARE ###
# Calculate lines.
l = dimensions
# Number of points at which to calculated predicted jet angle.
p = 80

y_values = []
x_values = []
lines = []

# Use -l / 2 + 0.01 * l / 2 for whole square and 0.01 * l / 2 for half a square.
for y_value in test_positions_sq:
    # Record the y value.
    y_values.append(y_value)
    # Make an empty line.
    line = []
    # Find equally spaced x-values at the selected y-position.
    e = 0.001 * l
    x_lim = l / 2
    x = np.linspace(-x_lim + e, x_lim - e, p)
    x_values.append(x)
    for n in range(len(x)):
        # A horizontal line.
        line.append((x[n], y_value))
    lines.append(line)

jet_angle_data = bubble_angle_along_lines(lines=lines, number_points=p, shape='square', dimensions=l,
                                          max_range=max_range,
                                          sink_type=sink_type,
                                          mod_2pi=True)

########################################################################################################################

### SUBPLOT a ###
l = dimensions
# Quiver plot.
ax1.quiver(X_sq, Y_sq, U_sq, V_sq, color=vector_color, pivot='mid')
for k in range(len(test_positions_sq)):
    ax1.plot(x_values[k], [y_values[k]] * p, color=colors[k % len(colors)], linestyle=linestyle[k % len(linestyle)],
             linewidth=linewidth1)
ax1.set_xlabel('$x$', fontsize=labelsize)
ax1.set_ylabel('$y$', fontsize=labelsize)
ax1.axis([-l / 2, l / 2, -l / 2, l / 2])
ax1.tick_params(axis='both', labelsize=ticksize)
# Subplot a
ax1.text(x=x_text, y=y_text, s='a',
         fontsize=fontsize, weight='bold', transform=ax1.transAxes)

### SUBPLOT b ###
# Prediction.
for k in range(len(test_positions_sq)):
    ax2.plot(x_values[k], jet_angle_data[k] * 180 / np.pi, color=colors[k % len(colors)],
             linestyle=linestyle[k % len(linestyle)], linewidth=linewidth2,
             label='y = {0}'.format(round(y_values[k], 2)))
ax2.set_xlabel('$x$', fontsize=labelsize)
ax2.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
ax2.axis([-l / 2, l / 2, 0, 360])
yticks = np.arange(0, 360 + 1, 30)
ax2.set_yticks(yticks)
ax2.tick_params(axis='both', labelsize=ticksize)

# Subplot b
ax2.text(x=x_text, y=y_text, s='b',
         fontsize=fontsize, weight='bold', transform=ax2.transAxes)

ax2.legend(loc='upper right', fontsize=legendsize)
# plt.grid()

########################################################################################################################

### FOR SUBPLOT 3 - TRIANGLE ###
X_tri, Y_tri, U_tri, V_tri = jet_direction_vector_field(shape='triangle', dimensions=1, max_range=20, sink_type=sink_type, plot=False)

### FOR SUBPLOT 4 - TRIANGLE ###
l = dimensions
h = np.sqrt(3) / 2 * l
# Number of points at which to calculate predicted jet angle.
p = 170

y_values = []
x_values = []
lines = []

# Make the lines.
for y_value in test_positions_tri:
    # Record the y value.
    y_values.append(y_value)
    # Make an empty line.
    line = []
    # Find equally spaced x-values at the selected y-position.
    x_lim = np.sqrt(3) / 3 * y_value
    e = 0.001 * x_lim
    x = np.linspace(-x_lim + e, x_lim - e, p)
    x_values.append(x)
    for n in range(len(x)):
        line.append(np.array((x[n], y_value)))
    lines.append(line)

jet_angle_data = bubble_angle_along_lines(lines=lines, number_points=p, shape='triangle', dimensions=l,
                                          max_range=max_range,
                                          sink_type='3D',
                                          mod_2pi=False)

########################################################################################################################

### SUBPLOT c ###
# Quiver plot.
ax3.quiver(X_tri, Y_tri, U_tri, V_tri, color=vector_color, pivot='mid', units='width', width=0.0055)
for k in range(len(test_positions_tri)):
    ax3.plot(x_values[k], [y_values[k]] * p, color=colors[k % len(colors)], linestyle=linestyle[k % len(linestyle)],
             linewidth=linewidth1)
# Fill in areas outside triangle.
x = np.linspace(-l / np.sqrt(3), l / np.sqrt(3), 100, endpoint=False)
y1 = np.sqrt(3) * x
y2 = - np.sqrt(3) * x
ax3.plot(x, y1,
         x, y2, color='k', linewidth=1)
ax3.axhline(y=np.sqrt(3) / 2 * l, color='k', linewidth=1)
# Set axis properties.
epsilon_x = 0.0005 * h
epsilon_y = 0.0005 * l / np.sqrt(3)
ax3.axis([-l / 2, l / 2, 0, np.sqrt(3) / 2 * l + epsilon_y])
ax3.spines['bottom'].set_position('zero')
ax3.spines['top'].set_visible(False)
ax3.spines['right'].set_visible(False)
ax3.set_xlabel('$x$', horizontalalignment='center', fontsize=labelsize)
ax3.set_ylabel('$y$', horizontalalignment='center', fontsize=labelsize)
ax3.tick_params(axis='both', labelsize=ticksize)

# Subplot c
ax3.text(x=x_text, y=y_text, s='c',
         fontsize=fontsize, weight='bold', transform=ax3.transAxes)

### SUBPLOT d ###
p_div_2 = divmod(p, 2)[0]
for k in range(len(test_positions_tri)):
    angles = (jet_angle_data[k] * 180 / np.pi) % 360
    index = 0
    # For all the angles, find the first one greater than 360. Use this index position.
    for i in range(1, len(angles)):
        if np.abs(angles[i] - angles[i - 1]) > 180:
            if i > index:
                index = i

    # Prediction.
    # If the plot doesn't need to be split in two, plot the prediction.
    if index == 0:
        x_lim = np.sqrt(3) / 3 * y_values[k]
        ax4.plot(x_values[k] / (2 * x_lim), (jet_angle_data[k] * 180 / np.pi) % 360, color=colors[k % len(colors)],
                 linestyle=linestyle[k % len(linestyle)],
                 linewidth=linewidth2,
                 label='y = {0}'.format(round(y_values[k], 2)))
    # Splitting the plot in two is meant to avoid strange extra lines because of taking only angles between 0 to 360.
    elif index != 0:
        x_lim = np.sqrt(3) / 3 * y_values[k]
        ax4.plot(x_values[k][0:index] / (2 * x_lim), angles[0:index], color=colors[k % len(colors)],
                 linestyle=linestyle[k % len(linestyle)], linewidth=linewidth2,
                 label='y = {0}'.format(round(y_values[k], 2)))
        ax4.plot(x_values[k][index:p] / (2 * x_lim), angles[index:p], color=colors[k % len(colors)],
                 linestyle=linestyle[k % len(linestyle)], linewidth=linewidth2)
ax4.set_xlabel('$\\tilde{x}$', fontsize=labelsize)
ax4.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
ax4.axis([-0.5, 0.5, 0, 360])
ax4.tick_params(axis='both', labelsize=ticksize)
yticks = np.arange(0, 360 + 1, 30)
ax4.set_yticks(yticks)

# Subplot d
ax4.text(x=x_text, y=y_text, s='d',
         fontsize=fontsize, weight='bold', transform=ax4.transAxes)

ax4.legend(loc='center right', fontsize=legendsize)
# plt.grid()

expanded = True
if expanded:
    ########################################################################################################################

    ### FOR SUBPLOT 5 - ISOSCELES RIGHT TRIANGLE ###
    X_isc_tri, Y_isc_tri, U_isc_tri, V_isc_tri = jet_direction_vector_field(shape='isosceles right', dimensions=1,
                                                                            max_range=20, sink_type=sink_type, plot=False)

    ### FOR SUBPLOT 6 - ISOSCELES RIGHT TRIANGLE ###
    # Calculate lines.
    l = dimensions
    # Number of points at which to calculated predicted jet angle.
    p = 120

    y_values = []
    x_values = []
    lines = []

    # Use -l / 2 + 0.01 * l / 2 for whole square and 0.01 * l / 2 for half a square.
    for y_value in test_positions_isc_right:
        # Record the y value.
        y_values.append(y_value)
        # Make an empty line.
        line = []
        # Find equally spaced x-values at the selected y-position.
        e = 0.001 * l
        x_lim = l - y_value
        x = np.linspace(0 + e, x_lim - e, p)
        x_values.append(x)
        for n in range(len(x)):
            # A horizontal line.
            line.append((x[n], y_value))
        lines.append(line)

    jet_angle_data = bubble_angle_along_lines(lines=lines, number_points=p, shape='isosceles right', dimensions=l,
                                              max_range=max_range,
                                              sink_type=sink_type,
                                              mod_2pi=True)

    ########################################################################################################################

    ### SUBPLOT e ###
    l = dimensions
    # Quiver plot.
    ax5.quiver(X_isc_tri, Y_isc_tri, U_isc_tri, V_isc_tri, color=vector_color, pivot='mid')
    for k in range(len(test_positions_isc_right)):
        ax5.plot(x_values[k], [y_values[k]] * p, color=colors[k % len(colors)],
                 linestyle=linestyle[k % len(linestyle)], linewidth=linewidth1)
    ax5.plot([0, l], [l, 0], linestyle='-', color='k', linewidth=1)
    ax5.spines['top'].set_visible(False)
    ax5.spines['right'].set_visible(False)
    ax5.set_xlabel('$x$', fontsize=labelsize)
    ax5.set_ylabel('$y$', fontsize=labelsize)
    ax5.axis([0, l, 0, l])
    ax5.tick_params(axis='both', labelsize=ticksize)
    # Subplot e
    ax5.text(x=x_text, y=y_text, s='e',
             fontsize=fontsize, weight='bold', transform=ax5.transAxes)

    ### SUBPLOT f ###
    # Subplot f
    for k in range(len(test_positions_isc_right)):
        angles = (jet_angle_data[k] * 180 / np.pi) % 360
        index = 0
        # For all the angles, find the first one greater than 360. Use this index position.
        for i in range(1, len(angles)):
            if np.abs(angles[i] - angles[i - 1]) > 180:
                if i > index:
                    index = i

        # Prediction.
        # If the plot doesn't need to be split in two, plot the prediction.
        if index == 0:
            x_lim = l - y_values[k]
            ax6.plot(x_values[k] / x_lim, (jet_angle_data[k] * 180 / np.pi) % 360, color=colors[k % len(colors)],
                     linestyle=linestyle[k % len(linestyle)],
                     linewidth=linewidth2,
                     label='y = {0}'.format(round(y_values[k], 2)))
        # Splitting the plot in two is meant to avoid strange extra lines because of taking only angles between 0 to 360.
        elif index != 0:
            x_lim = l - y_values[k]
            ax6.plot(x_values[k][0:index] / x_lim, angles[0:index], color=colors[k % len(colors)],
                     linestyle=linestyle[k % len(linestyle)], linewidth=linewidth2,
                     label='y = {0}'.format(round(y_values[k], 2)))
            ax6.plot(x_values[k][index:p] / x_lim, angles[index:p], color=colors[k % len(colors)],
                     linestyle=linestyle[k % len(linestyle)], linewidth=linewidth2)
    ax6.set_xlabel('$\\tilde{x}$', fontsize=labelsize)
    ax6.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
    ax6.axis([0, l, 0, 360])
    ax6.tick_params(axis='both', labelsize=ticksize)
    yticks = np.arange(0, 360 + 1, 30)
    ax6.set_yticks(yticks)

    # Subplot f
    ax6.text(x=x_text, y=y_text, s='f',
             fontsize=fontsize, weight='bold', transform=ax6.transAxes)

    ax6.legend(loc='center right', fontsize=legendsize)
    # plt.grid()

    ########################################################################################################################

    ### FOR SUBPLOT 7 - 30-60-90 TRIANGLE ###
    X_306090_tri, Y_306090_tri, U_306090_tri, V_306090_tri = jet_direction_vector_field(shape='306090', dimensions=1,
                                                                                        max_range=20, sink_type='3D',
                                                                                        plot=False)

    ### FOR SUBPLOT 8 - 30-60-90 TRIANGLE ###
    l = dimensions
    h = np.sqrt(3) / 2 * l
    # Number of points at which to calculate predicted jet angle.
    p = 170

    y_values = []
    x_values = []
    lines = []

    # Make the lines.
    for y_value in test_positions_306090_tri:
        # Record the y value.
        y_values.append(y_value)
        # Make an empty line.
        line = []
        # Find equally spaced x-values at the selected y-position.
        x_lim = l/2 - np.sqrt(3)/3 * y_value
        e = 0.001 * x_lim
        x = np.linspace(0 + e, x_lim - e, p)
        x_values.append(x)
        for n in range(len(x)):
            line.append(np.array((x[n], y_value)))
        lines.append(line)

    jet_angle_data = bubble_angle_along_lines(lines=lines, number_points=p, shape='306090', dimensions=l,
                                              max_range=max_range,
                                              sink_type='3D',
                                              mod_2pi=False)

    ########################################################################################################################

    ### SUBPLOT g ###
    # Quiver plot.
    ax7.quiver(X_306090_tri, Y_306090_tri, U_306090_tri, V_306090_tri, color=vector_color, pivot='mid', units='width', width=0.0055)
    for k in range(len(test_positions_306090_tri)):
        ax7.plot(x_values[k], [y_values[k]] * p, color=colors[k % len(colors)],
                 linestyle=linestyle[k % len(linestyle)], linewidth=linewidth1)
    # Set axis properties.
    ax7.plot([0, l / 2], [np.sqrt(3) / 2 * l, 0], linestyle='-', color='k', linewidth=1)
    epsilon_y = 0.0005 * l / np.sqrt(3)
    ax7.axis([0, l, 0, np.sqrt(3) / 2 * l + epsilon_y])
    ax7.spines['top'].set_visible(False)
    ax7.spines['right'].set_visible(False)
    ax7.set_xlabel('$x$', fontsize=labelsize)
    ax7.set_ylabel('$y$', fontsize=labelsize)
    ax7.tick_params(axis='both', labelsize=ticksize)

    # Subplot g
    ax7.text(x=x_text, y=y_text, s='g',
             fontsize=fontsize, weight='bold', transform=ax7.transAxes)

    ### SUBPLOT h ###
    p_div_2 = divmod(p, 2)[0]
    for k in range(len(test_positions_306090_tri)):
        angles = (jet_angle_data[k] * 180 / np.pi) % 360
        index = 0
        # For all the angles, find the first one greater than 360. Use this index position.
        for i in range(1, len(angles)):
            if np.abs(angles[i] - angles[i - 1]) > 180:
                if i > index:
                    index = i

        # Prediction.
        # If the plot doesn't need to be split in two, plot the prediction.
        if index == 0:
            x_lim = l/2 - np.sqrt(3) / 3 * y_values[k]
            ax8.plot(x_values[k] / x_lim, (jet_angle_data[k] * 180 / np.pi) % 360, color=colors[k % len(colors)],
                     linestyle=linestyle[k % len(linestyle)],
                     linewidth=linewidth2,
                     label='y = {0}'.format(round(y_values[k], 2)))
        # Splitting the plot in two is meant to avoid strange extra lines because of taking only angles between 0 to 360.
        elif index != 0:
            x_lim = l/2 - np.sqrt(3) / 3 * y_values[k]
            ax8.plot(x_values[k][0:index] / x_lim, angles[0:index], color=colors[k % len(colors)],
                     linestyle=linestyle[k % len(linestyle)], linewidth=linewidth2,
                     label='y = {0}'.format(round(y_values[k], 2)))
            ax8.plot(x_values[k][index:p] / x_lim, angles[index:p], color=colors[k % len(colors)],
                     linestyle=linestyle[k % len(linestyle)], linewidth=linewidth2)
    ax8.set_xlabel('$\\tilde{x}$', fontsize=labelsize)
    ax8.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
    ax8.tick_params(axis='both', labelsize=ticksize)
    ax8.axis([0, 1, 0, 360])
    yticks = np.arange(0, 360 + 1, 30)
    ax8.set_yticks(yticks)

    # Subplot h
    ax8.text(x=x_text, y=y_text, s='h',
             fontsize=fontsize, weight='bold', transform=ax8.transAxes)

    ax8.legend(loc='center right', fontsize=legendsize)
    # plt.grid()


plt.tight_layout(pad=3, rect=[0, 0.03, 1, 0.95])


plt.savefig('figure04.jpg', format='jpg', dpi=500)
plt.show()