"""
Program: Figure 7 - Triangle Figure
Date: 3 July 2019
Author: Lebo Molefe

Plots comparing experimental data for triangle to model prediction.
"""

import os
import itertools
import numpy as np
import matplotlib.pyplot as plt
from sink_positions import sink_positions_equilateral_triangle
from jet_angle import bubble_angle_along_lines, angle_bubble_position

### LOAD DATA ###
first_movie = 518
last_movie = 610

mm_per_pixel = 0.0301456848058
triangle_height = 408.22917778
triangle_side = 2 * triangle_height * np.tan(np.pi / 6)

print('triangle side length:', triangle_side * mm_per_pixel, 'mm')

tip1 = (0, 0)
tip2 = (408.22917778, 233.8999499)
tip3 = (408.22917778, -239.24568837)

# Load data.
# os.chdir('')

P = np.loadtxt('movies.csv',
               delimiter=',',
               skiprows=1, usecols=(0, 7, 8, 9),
               unpack=False)

P_movie_numbers =     P[:, 0]
translation_stage_x = P[:, 1]
translation_stage_y = P[:, 2]
translation_stage_z = P[:, 3]

R = np.loadtxt('calibrateddata.csv',
               delimiter=',',
               skiprows=1, usecols=(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10),
               unpack=False)

movie_numbers = R[:, 0]
trial_numbers = R[:, 1]
eccentricity = R[:, 3]
roundness = R[:, 4]
bubble_x0_calibrated = R[:, 5]
bubble_y0_calibrated = R[:, 6]
area = R[:, 7]
bubble_x1_calibrated = R[:, 8]
bubble_y1_calibrated = R[:, 9]
jet_angle = R[:, 10]

# Create a threshold for identifying whether a bubble is in the trial.
bubble_appears = roundness > 0.7

### CALCULATIONS ###
def calculations_tri(line_start, line_stop, p, max_range, mod_2pi=True):
    """
    Perform the data and model calculations necessary for a plot comparing triangle data to experiment.
    :param line_start: first movie number of triangle data, an integer
    :param line_stop: last movie number of triangle data, an integer
    :param p: number of points at which to calculate triangle prediciton curve, an integer
    :param max_range: number of "layers" of mirror sinks around the shape, an integer
    :param mod_2pi: True or False, if True makes all angles between 0 and 2pi
    :return: several lists of data about the triangle experiments; They are nums, trials, xs, xs_nodim, x_means, x_errs,
    ys, ys_nodim, y_means, areas, angs, angs_degrees, jet_angle_data, ang_points, ang_errs, x_points, y_points (16 total
    lists in the output)
    """
    # DATA CLACULATIONS
    nums = []
    trials = []
    xs = []
    x_errs = []
    x_means = []
    ys = []
    y_errs = []
    y_means = []
    areas = []
    angs = []

    start = line_start
    stop = line_stop

    for num in range(start, stop + 1):
        x_variation = []
        y_variation = []
        no_bubble = 0
        for k in range(len(bubble_x0_calibrated[movie_numbers == num])):
            if bubble_appears[movie_numbers == num][k] == 0:
                no_bubble += 1
            if bubble_appears[movie_numbers == num][k] == 1:
                nums.append(num)
                trials.append(trial_numbers[movie_numbers == num][k])
                x = bubble_x0_calibrated[movie_numbers == num][k]
                y = bubble_y0_calibrated[movie_numbers == num][k]
                x_variation.append(x)
                y_variation.append(y)
                xs.append(x)
                ys.append(y)
                areas.append(area[movie_numbers == num][k])
                if mod_2pi:
                    angs.append(jet_angle[movie_numbers == num][k] % (2 * np.pi))
                else:
                    angs.append(jet_angle[movie_numbers == num][k])

        # Movies such as calibration pictures with just a couple frames and no bubble measurements will not have any
        # measured x positions. Don't include them.
        if len(x_variation) != 0:
            x_variation = np.array(x_variation)
            x_mean = np.mean(x_variation)
            y_mean = np.mean(y_variation)
            # Calculate the standard error using data type np.float64 for higher accuracy.
            x_err = np.std(x_variation, dtype=np.float64) / np.sqrt(len(x_variation))
            y_err = np.std(y_variation, dtype=np.float64) / np.sqrt(len(y_variation))
        x_errs.append(x_err)
        y_errs.append(y_err)
        x_means.append(x_mean)
        y_means.append(y_mean)
        # print('Bubble percentage is', (1 - no_bubble / len(bubble_x0_calibrated[movie_numbers == num])) * 100)

    #################################
    # Make lists into arrays.
    nums = np.array(nums)
    trials = np.array(trials)
    xs = np.array(xs)
    xs_nodim = xs / triangle_side
    x_means = np.array(x_means)
    x_errs = np.array(x_errs)
    ys = np.array(ys)
    ys_nodim = ys / triangle_side
    y_means = np.array(y_means)
    areas = np.array(areas)
    angs = np.array(angs)
    angs_degrees = angs * 180 / np.pi
    #################################

    # MODEL CALCULATIONS
    # Set the length of the triangle.
    l = 1
    # Find the mean x position and the outer margins of region where data was collected.
    x_mean = np.mean(xs_nodim)
    lower_bound = x_mean - 2 * np.std(xs_nodim)
    upper_bound = x_mean + 2 * np.std(xs_nodim)
    # Find equally spaced y-values at the selected x-position.
    y_lim = np.sqrt(3) / 3 * x_mean
    e = 0.03 * y_lim
    y = np.linspace(-y_lim + e, y_lim - e, p)
    # Find equally spaced y-values at the margin x-positions.
    # Lower bound.
    y_lim_lower = np.sqrt(3) / 3 * lower_bound
    e = 0.03 * y_lim_lower
    y_lower = np.linspace(-y_lim_lower + e, y_lim_lower - e, p)
    # Upper bound.
    y_lim_upper = np.sqrt(3) / 3 * upper_bound
    e = 0.03 * y_lim_upper
    y_upper = np.linspace(-y_lim_upper + e, y_lim_upper - e, p)
    # Calculate the mean line and the outer margin lines within the triangle.
    line = []
    line_lower = []
    line_upper = []
    for n in range(len(y)):
        # Switch x and y.
        line.append(np.array((y[n], x_mean)))
        line_lower.append(np.array((y_lower[n], lower_bound)))
        line_upper.append(np.array((y_upper[n], upper_bound)))
    line = np.array(line)
    line_lower = np.array(line_lower)
    line_upper = np.array(line_upper)
    lines = [line, line_lower, line_upper]

    #######################################################################################################
    # Calculate predicted jet angle along the 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=mod_2pi)
    # Also return y, y_lower, y_upper.
    #######################################################################################################


    # TRANSLATION STAGE POSITIONS - Collect data at the same translation stage position.
    positions = []

    # For all the movies in the group being analyzed, record the translation stage position.
    for k in range(start, stop + 1):
        # Append the position (x, y, z) for movie k. The `[0]` removes the list brackets.
        positions.append((translation_stage_x[P_movie_numbers == k][0],
                          translation_stage_y[P_movie_numbers == k][0],
                          translation_stage_z[P_movie_numbers == k][0]))

    # Group movies with the same translation stage position together.
    same_position_sets = []
    for pt in set(positions):
        s = []
        index = (translation_stage_x == pt[0]) * (translation_stage_y == pt[1]) * (translation_stage_z == pt[2])
        # Make sure only the movie numbers within the desired range are selected.
        index *= (P_movie_numbers >= start) * (P_movie_numbers <= stop)
        s.append(P_movie_numbers[index])
        # The `[0]` removes the list brackets.
        same_position_sets.append(s[0])

    # DATA
    x_mean = np.mean(xs)
    lower_bound = x_mean - 2 * np.std(xs)
    upper_bound = x_mean + 2 * np.std(xs)

    ang_sets = []
    ang_points = []
    ang_errs = []
    ang_expectations = []
    y_points = []
    x_points = []
    ids = []

    for same_position_set in same_position_sets:
        index = False * xs
        for n in same_position_set:
            index_n = (xs > lower_bound) * (xs < upper_bound) * (nums == n)
            # Change index by '+1' if index_n is 1, or add 0.
            index += index_n
            # Make `index` into a boolean array.
            index = index == 1
        ang_set = angs_degrees[index]
        ang_mean = np.mean(angs_degrees[index])
        ang_stderr = np.std(angs_degrees[index]) / np.sqrt(len(angs_degrees[index]))
        x_mean = np.mean(xs_nodim[index])
        y_mean = np.mean(ys_nodim[index])
        ids.append(same_position_set[0])

        # Calculate expected jet angle at the average bubble position.
        # Switch x and y because of how position is defined, differently than the orientation movies were taken.
        bubble_pos = (-y_mean, x_mean)
        sink_pos = sink_positions_equilateral_triangle(l=1, bubble_position=bubble_pos, w=7, print_result=False,
                                                       rotate=False)
        a = angle_bubble_position(bubble_position=bubble_pos, sink_positions=sink_pos, sink_type='3D',
                                  print_angle=False)
        if mod_2pi:
            ang_hat = a % (2 * np.pi) * 180 / np.pi  # Put it in degrees.
        else:
            ang_hat = a * 180 / np.pi  # Put it in degrees.

        ang_sets.append(ang_set)
        ang_points.append(ang_mean)
        ang_errs.append(ang_stderr)
        ang_expectations.append(ang_hat)
        x_points.append(x_mean)  # For calculating exact jet angle.
        y_points.append(y_mean)

    #################################
    # Make lists into arrays.
    ang_points = np.array(ang_points)
    ang_errs = np.array(ang_errs)
    ang_expectations = np.array(ang_expectations)
    x_points = np.array(x_points)
    y_points = np.array(y_points)
    #################################

    return nums, trials, xs, xs_nodim, x_means, x_errs, ys, ys_nodim, y_means, areas, angs, angs_degrees, \
           jet_angle_data, y, y_lower, y_upper, ang_points, ang_errs, ang_expectations, x_points, y_points

def triangle_vector_field(dimensions, max_range, sink_type, plot_color, edit_plot=False, print_result=False):
    """
    Plot cavitation bubble jet direction as a vector field - For an equilateral triangle, plot the vector representing
    the jet direction of a cavitation bubble at several points within the triangle.
    :param dimensions: shape dimensions; For a triangle, 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_color: optional, vector color
    :param edit_plot: True or False, if True, returns the current plot axes so that the plot can be modified
    :param print_result: True or False, if True, prints percentage completion of vector field calculation
    :return: Returns the vector field of cavitation bubble jet directions when they are placed at points within the
    equilateral triangle.
    """

    l = dimensions
    l = l / np.sqrt(3)
    h = np.sqrt(3) / 2 * l
    start = 0.89
    step = (1 - start) * l
    l_remainder = divmod(3 / 2 * l, step)[1]
    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, l_remainder:(3 / 2 * l):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 and print_result:
            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 x_bubble >= (np.sqrt(3) / 2 - e / 3) * l:
            U[i][j] = None
            V[i][j] = None
            A[i][j] = None
            continue

        sink_positions = sink_positions_equilateral_triangle(dimensions, bubble_position=(X[i][j], Y[i][j]),
                                                             w=max_range, print_result=False, rotate=False)
        angle = angle_bubble_position(bubble_position=(X[i][j], Y[i][j]), 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

    # Plot.
    ax = plt.gca()
    # 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 * np.sqrt(3), l / 2 * np.sqrt(3), 100, endpoint=True)
    y1 = np.sqrt(3) * x
    y2 = - np.sqrt(3) * x
    plt.plot(x, y1,
             x, y2, color='k', linewidth=1)
    plt.axhline(y=3 / 2 * l, color='k', linewidth=1)
    # Set axis properties.
    epsilon_x = 0.005 * l
    epsilon_y = 0.005 * h
    ax.axis([-h - epsilon_y, h + epsilon_y, 0, 3 / 2 * l + epsilon_x])
    # ax.spines['bottom'].set_position('zero')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.set_title('')
    ax.set_xlabel('', horizontalalignment='right', x=0.5)
    ax.set_ylabel('', horizontalalignment='right', y=0.5)

    if edit_plot:
        return ax.quiver(X, Y, U, V, color=plot_color, pivot='mid', units='width', width=0.0055)
    else:
        plt.show()


plt.figure(figsize=(8.5, 10))



### PLOT VARIABLES ###
max_range = 5  # Range of reflections
p = 200         # Points along a line

# Colors for vector field
vector_color = 'purple'
line_color = 'k'
# Colors for data plot
alpha_fill = 0.3
# Labels for figure subsections
x_text = -0.7
y_text = 1.15
fontsize = 18
labelsize = 16
ticksize = 14



### LINE 1 ###
start = 594     # 519, 537, 594, 604
stop = 603      # 536, 576, 603, 610

# Turn the plot by pi/2.
turn = 90

nums, trials, xs, xs_nodim, x_means, x_errs, ys, ys_nodim, y_means, areas, angs, angs_degrees, \
    jet_angle_data, y, y_lower, y_upper, ang_points, ang_errs, ang_expectations, x_points, y_points = \
    calculations_tri(line_start=start, line_stop=stop, p=p, max_range=max_range, mod_2pi=False)

# Print some data:
print('EXPERIMENT 1 INFORMATION')
# Number of trials
print('Total Number of Trials: ', len(trials))
# Standard error of jet angle
print('Average Errorbar Length: ', np.mean(ang_errs))
# Average bubble size
diameter = 2 * np.sqrt(np.mean(areas) / np.pi)
diameter_std = np.std(2 * np.sqrt(areas / np.pi))
print('Average Bubble Diameter: ', diameter * mm_per_pixel, 'mm')
print('Standard Deviation Bubble Diameter:', diameter_std * mm_per_pixel, 'mm')
for i in range(len(ang_points)):
    if turn+ang_points[i] - ang_expectations[i] < -180 or turn+ang_points[i] - ang_expectations[i] > 180:
        ang_expectations[i] = ang_expectations[i] % 360
print('RMSD:', np.sqrt(np.mean((turn+ang_points - ang_expectations)**2)))

for j in range(0, len(jet_angle_data)):
    jet_angle_data[j][0:int(p/2)] = jet_angle_data[j][0:int(p/2)] % (2 * np.pi)

# Subplot 1 - First vector field.
ax1 = plt.subplot2grid(shape=(4, 3), loc=(0, 0), rowspan=1, colspan=1)
l = 1

x_mean = np.mean(xs_nodim)
x_lower = x_mean - 2 * np.std(xs_nodim)
x_upper = x_mean + 2 * np.std(xs_nodim)

print('x_mean:', x_mean, 'lower_bound:', x_lower, 'upper_bound:', x_upper, '\n \n')

# The limits for the upper and lower limits for the y lines are based on the x-position.
y_mean = np.linspace(-np.sqrt(3) / 3 * x_mean, np.sqrt(3) / 3 * x_mean, p)
y_lower = np.linspace(-np.sqrt(3) / 3 * x_lower, np.sqrt(3) / 3 * x_lower, p)
y_upper = np.linspace(-np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_upper, p)
# Triangle vector field
triangle_vector_field(dimensions=l, max_range=max_range, sink_type='3D', plot_color=vector_color, edit_plot=True)
# Plot mean line and outer limits.
ax1.plot(y_mean, [x_mean] * p, color=line_color)
ax1.plot(y_lower, [x_lower] * p, color=line_color)
ax1.plot(y_upper, [x_upper] * p, color=line_color)
poly_x = [-np.sqrt(3) / 3 * x_lower, -np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_lower]
poly_y = [x_lower, x_upper, x_upper, x_lower]
ax1.fill(poly_x, poly_y, color=line_color, alpha=alpha_fill)
# Label
ax1.set_ylabel('$y$', fontsize=labelsize)
ax1.axis([-l / 2, l / 2, 0, l * np.cos(np.pi/6)])
# ax1.set_xticks
ax1.tick_params(labelsize=ticksize)
ax1.text(x=x_text, y=y_text, s='a   ($y$ = {0})'.format(np.round(x_mean, 2)),
         fontsize=fontsize, weight='bold', transform=ax1.transAxes)

# Suplot 2 - First experiment data.
ax2 = plt.subplot2grid(shape=(4, 3), loc=(0, 1), rowspan=1, colspan=2)
# DATA
x_mean = np.mean(xs_nodim)
x_lower = x_mean - 2 * np.std(xs_nodim)
x_upper = x_mean + 2 * np.std(xs_nodim)
# PREDICTION
ax2.plot(y, jet_angle_data[0] * 180 / np.pi,
         color='k', linewidth=1.5,
         label='Prediction along $y$ = {0}'.format(np.round(x_mean, 4)))
for i in range(1, 3):
    ax2.plot(y, jet_angle_data[i] * 180 / np.pi, color='k', linestyle='-', linewidth=1, alpha=0.7)
ax2.fill_between(y, jet_angle_data[1] * 180 / np.pi, jet_angle_data[2] * 180 / np.pi,
                 color='gray', alpha=0.4)
# Data errorbars.
ax2.errorbar(x=-y_points, y=turn + ang_points, yerr=ang_errs,
             fmt='o', markersize=4, color='C0', capsize=5,
             label='Data Mean and Standard Error')
ax2.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
ax2.set_xlim(-l / 2, l / 2)
yticks = np.arange(-30, 210 + 1, 60)
ax2.set_yticks(yticks)
ax2.tick_params(labelsize=ticksize)
ax2.grid()


### LINE 2 ###
start = 537         # 519, 537, 594, 604
stop = 576          # 536, 576, 603, 610

nums, trials, xs, xs_nodim, x_means, x_errs, ys, ys_nodim, y_means, areas, angs, angs_degrees, \
    jet_angle_data, y, y_lower, y_upper, ang_points, ang_errs, ang_expectations, x_points, y_points = \
    calculations_tri(line_start=start, line_stop=stop, p=p, max_range=max_range, mod_2pi=True)

# Print some data:
print('EXPERIMENT 2 INFORMATION')
# Number of trials
print('Total Number of Trials: ', len(trials))
# Standard error of jet angle
print('Average Errorbar Length: ', np.mean(ang_errs))
# Average bubble size
diameter = 2 * np.sqrt(np.mean(areas) / np.pi)
diameter_std = np.std(2 * np.sqrt(areas / np.pi))
print('Average Bubble Diameter: ', diameter * mm_per_pixel, 'mm')
print('Standard Deviation Bubble Diameter:', diameter_std * mm_per_pixel, 'mm')
print('RMSD:', np.sqrt(np.mean((turn+ang_points - ang_expectations)**2)))

# Subplot 3 - Second vector field.
ax3 = plt.subplot2grid(shape=(4, 3), loc=(1, 0), rowspan=1, colspan=1)

x_mean = np.mean(xs_nodim)
x_lower = x_mean - 2 * np.std(xs_nodim)
x_upper = x_mean + 2 * np.std(xs_nodim)

print('x_mean:', x_mean, 'lower_bound:', x_lower, 'upper_bound:', x_upper, '\n \n')

# The limits for the upper and lower limits for the y lines are based on the x-position.
y_mean = np.linspace(-np.sqrt(3) / 3 * x_mean, np.sqrt(3) / 3 * x_mean, p)
y_lower = np.linspace(-np.sqrt(3) / 3 * x_lower, np.sqrt(3) / 3 * x_lower, p)
y_upper = np.linspace(-np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_upper, p)
# Triangle vector field
triangle_vector_field(dimensions=l, max_range=max_range, sink_type='3D', plot_color=vector_color, edit_plot=True)
# Plot mean line and outer limits.
ax3.plot(y_mean, [x_mean] * p, color=line_color)
ax3.plot(y_lower, [x_lower] * p, color=line_color)
ax3.plot(y_upper, [x_upper] * p, color=line_color)
poly_x = [-np.sqrt(3) / 3 * x_lower, -np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_lower]
poly_y = [x_lower, x_upper, x_upper, x_lower]
ax3.fill(poly_x, poly_y, color=line_color, alpha=alpha_fill)
# Label
ax3.set_ylabel('$y$', fontsize=labelsize)
ax3.axis([-l / 2, l / 2, 0, l * np.cos(np.pi/6)])
# ax3.set_xticks
ax3.tick_params(labelsize=ticksize)
ax3.text(x=x_text, y=y_text, s='b   ($y$ = {0})'.format(np.round(x_mean, 2)),
         fontsize=fontsize, weight='bold', transform=ax3.transAxes)

# Suplot 4 - Second experiment data.
ax4 = plt.subplot2grid(shape=(4, 3), loc=(1, 1), rowspan=1, colspan=2)
# DATA
x_mean = np.mean(xs_nodim)
x_lower = x_mean - 2 * np.std(xs_nodim)
x_upper = x_mean + 2 * np.std(xs_nodim)
# PREDICTION
ax4.plot(y, jet_angle_data[0] * 180 / np.pi,
         color='k', linewidth=1.5,
         label='Prediction along $y$ = {0}'.format(np.round(x_mean, 4)))
for i in range(1, 3):
    ax4.plot(y, jet_angle_data[i] * 180 / np.pi, color='k', linestyle='-', linewidth=1, alpha=0.7)
ax4.fill_between(y, jet_angle_data[1] * 180 / np.pi, jet_angle_data[2] * 180 / np.pi, color='gray', alpha=0.4)
# Data errorbars.
ax4.errorbar(x=-y_points, y=turn + ang_points, yerr=ang_errs,
             fmt='o', markersize=4, color='C0', capsize=5,
             label='Data Mean and Standard Error')
ax4.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
ax4.set_xlim(-l / 2, l / 2)
yticks = np.arange(180, 360 + 1, 60)
ax4.set_yticks(yticks)
ax4.tick_params(labelsize=ticksize)
ax4.grid()




### LINE 3 ###
start = 519         # 519, 537, 594, 604
stop = 536          # 536, 576, 603, 610

nums, trials, xs, xs_nodim, x_means, x_errs, ys, ys_nodim, y_means, areas, angs, angs_degrees, \
    jet_angle_data, y, y_lower, y_upper, ang_points, ang_errs, ang_expectations, x_points, y_points = \
    calculations_tri(line_start=start, line_stop=stop, p=p, max_range=max_range, mod_2pi=True)

# Print some data:
print('EXPERIMENT 3 INFORMATION')
# Number of trials
print('Total Number of Trials: ', len(trials))
# Standard error of jet angle
print('Average Errorbar Length: ', np.mean(ang_errs))
# Average bubble size
diameter = 2 * np.sqrt(np.mean(areas) / np.pi)
diameter_std = np.std(2 * np.sqrt(areas / np.pi))
print('Average Bubble Diameter: ', diameter * mm_per_pixel, 'mm')
print('Standard Deviation Bubble Diameter:', diameter_std * mm_per_pixel, 'mm')
print('RMSD:', np.sqrt(np.mean((turn+ang_points - ang_expectations)**2)))

# Subplot 5 - Third vector field.
ax5 = plt.subplot2grid(shape=(4, 3), loc=(2, 0), rowspan=1, colspan=1)

x_mean = np.mean(xs_nodim)
x_lower = x_mean - 2 * np.std(xs_nodim)
x_upper = x_mean + 2 * np.std(xs_nodim)

print('x_mean:', x_mean, 'lower_bound:', x_lower, 'upper_bound:', x_upper, '\n \n')

# The limits for the upper and lower limits for the y lines are based on the x-position.
y_mean = np.linspace(-np.sqrt(3) / 3 * x_mean, np.sqrt(3) / 3 * x_mean, p)
y_lower = np.linspace(-np.sqrt(3) / 3 * x_lower, np.sqrt(3) / 3 * x_lower, p)
y_upper = np.linspace(-np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_upper, p)
# Triangle vector field
triangle_vector_field(dimensions=l, max_range=max_range, sink_type='3D', plot_color=vector_color, edit_plot=True)
# Plot mean line and outer limits.
ax5.plot(y_mean, [x_mean] * p, color=line_color)
ax5.plot(y_lower, [x_lower] * p, color=line_color)
ax5.plot(y_upper, [x_upper] * p, color=line_color)
poly_x = [-np.sqrt(3) / 3 * x_lower, -np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_lower]
poly_y = [x_lower, x_upper, x_upper, x_lower]
ax5.fill(poly_x, poly_y, color=line_color, alpha=alpha_fill)
# Label
ax5.set_ylabel('$y$', fontsize=labelsize)
ax5.axis([-l / 2, l / 2, 0, l * np.cos(np.pi/6)])
# ax5.set_xticks
ax5.tick_params(labelsize=ticksize)
ax5.text(x=x_text, y=y_text, s='c   ($y$ = {0})'.format(np.round(x_mean, 2)),
         fontsize=fontsize, weight='bold', transform=ax5.transAxes)

# Suplot 6 - Third experiment data.
ax6 = plt.subplot2grid(shape=(4, 3), loc=(2, 1), rowspan=1, colspan=2)
# DATA
x_mean = np.mean(xs_nodim)
x_lower = x_mean - 2 * np.std(xs_nodim)
x_upper = x_mean + 2 * np.std(xs_nodim)
# PREDICTION
ax6.plot(y, jet_angle_data[0] * 180 / np.pi,
         color='k', linewidth=1.5,
         label='Prediction along $y$ = {0}'.format(np.round(x_mean, 4)))
for i in range(1, 3):
    ax6.plot(y, jet_angle_data[i] * 180 / np.pi, color='k', linestyle='-', linewidth=1, alpha=0.7)
ax6.fill_between(y, jet_angle_data[1] * 180 / np.pi, jet_angle_data[2] * 180 / np.pi, color='gray', alpha=0.4)
# Data errorbars.
ax6.errorbar(x=-y_points, y=turn + ang_points, yerr=ang_errs,
             fmt='o', markersize=4, color='C0', capsize=5,
             label='Data Mean and Standard Error')
# Labels
ax6.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
ax6.set_xlim(-l / 2, l / 2)
yticks = np.arange(180, 360 + 1, 60)
ax6.set_yticks(yticks)
ax6.tick_params(labelsize=ticksize)
ax6.grid()



### LINE 4 ###
start = 605         # 519, 537, 594, 604
stop = 610          # 536, 576, 603, 610

nums, trials, xs, xs_nodim, x_means, x_errs, ys, ys_nodim, y_means, areas, angs, angs_degrees, \
    jet_angle_data, y, y_lower, y_upper, ang_points, ang_errs, ang_expectations, x_points, y_points = \
    calculations_tri(line_start=start, line_stop=stop, p=p, max_range=max_range, mod_2pi=True)


# Print some data:
print('EXPERIMENT 4 INFORMATION')
# Number of trials
print('Total Number of Trials: ', len(trials))
# Standard error of jet angle
print('Average Errorbar Length: ', np.mean(ang_errs))
# Average bubble size
diameter = 2 * np.sqrt(np.mean(areas) / np.pi)
diameter_std = np.std(2 * np.sqrt(areas / np.pi))
print('Average Bubble Diameter: ', diameter * mm_per_pixel, 'mm')
print('Standard Deviation Bubble Diameter:', diameter_std * mm_per_pixel, 'mm')
print('RMSD:', np.sqrt(np.mean((turn+ang_points - ang_expectations)**2)))

# Subplot 7 - Fourth vector field.
ax7 = plt.subplot2grid(shape=(4, 3), loc=(3, 0), rowspan=1, colspan=1)

x_mean = np.mean(xs_nodim)
x_lower = x_mean - 2 * np.std(xs_nodim)
x_upper = x_mean + 2 * np.std(xs_nodim)

print('x_mean:', x_mean, 'lower_bound:', x_lower, 'upper_bound:', x_upper, '\n \n')

# The limits for the upper and lower limits for the y lines are based on the x-position.
y_mean = np.linspace(-np.sqrt(3) / 3 * x_mean, np.sqrt(3) / 3 * x_mean, p)
y_lower = np.linspace(-np.sqrt(3) / 3 * x_lower, np.sqrt(3) / 3 * x_lower, p)
y_upper = np.linspace(-np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_upper, p)
# Triangle vector field
triangle_vector_field(dimensions=l, max_range=max_range, sink_type='3D', plot_color=vector_color, edit_plot=True)
# Plot mean line and outer limits.
ax7.plot(y_mean, [x_mean] * p, color=line_color)
ax7.plot(y_lower, [x_lower] * p, color=line_color)
ax7.plot(y_upper, [x_upper] * p, color=line_color)
poly_x = [-np.sqrt(3) / 3 * x_lower, -np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_upper, np.sqrt(3) / 3 * x_lower]
poly_y = [x_lower, x_upper, x_upper, x_lower]
ax7.fill(poly_x, poly_y, color=line_color, alpha=alpha_fill)
# Label
ax7.set_xlabel('$x$', fontsize=labelsize)
ax7.set_ylabel('$y$', fontsize=labelsize)
ax7.axis([-l / 2, l / 2, 0, l * np.cos(np.pi/6)])
# ax7.set_xticks
ax7.tick_params(labelsize=ticksize)
ax7.text(x=x_text, y=y_text, s='d   ($y$ = {0})'.format(np.round(x_mean, 2)),
         fontsize=fontsize, weight='bold', transform=ax7.transAxes)

# Suplot 8 - Fourth experiment data.
ax8 = plt.subplot2grid(shape=(4, 3), loc=(3, 1), rowspan=1, colspan=2)
# DATA
x_mean = np.mean(xs_nodim)
x_lower = x_mean - 2 * np.std(xs_nodim)
x_upper = x_mean + 2 * np.std(xs_nodim)
# PREDICTION
ax8.plot(y, jet_angle_data[0] * 180 / np.pi,
         color='k', linewidth=1.5,
         label='Prediction along $y$ = {0}'.format(np.round(x_mean, 4)))
for i in range(1, 3):
    ax8.plot(y, jet_angle_data[i] * 180 / np.pi, color='k', linestyle='-', linewidth=1, alpha=0.7)
ax8.fill_between(y, jet_angle_data[1] * 180 / np.pi, jet_angle_data[2] * 180 / np.pi, color='gray', alpha=0.4)
# Data errorbars.
ax8.errorbar(x=-y_points, y=turn + ang_points, yerr=ang_errs,
             fmt='o', markersize=4, color='C0', capsize=5,
             label='Data Mean and Standard Error')
# Label
ax8.set_xlabel('$x$', fontsize=labelsize)
ax8.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
ax8.set_xlim(-l / 2, l / 2)
yticks = np.arange(180, 360 + 1, 60)
ax8.set_yticks(yticks)
ax8.tick_params(labelsize=ticksize)

# for n, pt in enumerate(zip(-y_points, turn + ang_points)):
#     ax8.annotate(s='{0}'.format(ids[n]), xy=pt, xycoords='data')
ax8.grid()

#%%
plt.tight_layout(pad=1.5, rect=[0, 0.03, 1, 0.95])
plt.savefig('figure07.pdf', bbox_inches='tight', pad_inches=0)
plt.show()