"""
Program: Figure 6 - Square Figure
Date: 12 May 2019
Author: Lebo Molefe

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

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

### LOAD DATA ###
first_movie = 624
last_movie = 658

mm_per_pixel = 0.0295015341484
square_side_length = np.mean([522.784001945, 521.108087531, 515.79686456, 519.661153342])

print('square side length:', square_side_length * mm_per_pixel, 'mm')

tip1 = (0,  0)
tip2 = (-5.37681652, 522.75635102)
tip3 = (515.79686456, 514.48791942)
tip4 = (515.79686456, -5.17323393)

# 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_sq(line_start, line_stop, p, max_range, mod_2pi=True):
    """
    Perform the data and model calculations necessary for a plot comparing square data to experiment.
    :param line_start: first movie number of square data, an integer
    :param line_stop: last movie number of square data, an integer
    :param p: number of points at which to calculate square 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 square 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 CALCULATIONS
    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)

    # Shift the nondimensional data values by 0.5 since the prediction square has the origin point (0, 0) defined as the
    # middle, but the data has the top (becomes bottom since image y axis is backward) left corner defined as [0, 0].

    ########################################
    # Make lists into arrays.
    nums = np.array(nums)
    trials = np.array(trials)
    xs = np.array(xs)
    xs_nodim = xs / square_side_length - 0.5
    x_means = np.array(x_means)
    x_errs = np.array(x_errs)
    ys = np.array(ys)
    ys_nodim = ys / square_side_length - 0.5
    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 square.
    l = 1
    # Find the mean x position and the outer margins of the 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.
    e = 0.01 * l
    y = np.linspace(-l / 2 + e, l / 2 - e, p)
    # Calculate the mean line and the outer margin lines within the square.
    line = []
    line_lower = []
    line_upper = []
    for n in range(len(y)):
        line.append(np.array((x_mean, y[n])))
        line_lower.append(np.array((lower_bound, y[n])))
        line_upper.append(np.array((upper_bound, y[n])))
    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='square', dimensions=l,
                                              max_range=max_range,
                                              sink_type='3D', mod_2pi=mod_2pi)
    #####################################################################################################

    # 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])

    # Find the mean x position and the outer margins of the region where data was collected.
    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])

        # Calculate expected jet angle at the average bubble position.
        bubble_pos = (x_mean, y_mean)
        sink_pos = sink_positions_square(l=1, bubble_position=bubble_pos, w=7, print_result=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)
        ids.append(same_position_set)

    #############################
    # 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, ang_points, ang_errs, ang_expectations, x_points, y_points, ids


def square_vector_field(dimensions, max_range, sink_type, plot_color='r', edit_plot=False):
    """
    Plot cavitation bubble jet direction as a vector field - For a square, plot the vector representing the
    jet direction of a cavitation bubble at several points within the square.
    :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_color: optional, vector color
    :param edit_plot: True or False, if True, returns the current plot axes so that the plot can be modified
    :return: Returns the vector field of cavitation bubble jet directions when they are placed at points within the
    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])):
        sink_positions = sink_positions_square(dimensions, bubble_position=(X[i][j], Y[i][j]), w=max_range, print_result=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')
    ax.set_title('')
    ax.set_xlabel('')
    ax.set_ylabel('')

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

plt.figure(figsize=(9.5, 6))


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

### PLOT VARIABLES ###
max_range = 10  # Range of reflections
p = 200 # Points along a line
# Colors for vector field
vector_color = 'purple'
line_color = 'k'
xticks = np.arange(-0.4, 0.4 + 0.01, 0.2)
# Colors for data plot
prediction_line_color = 'k'
fill_color = 'gray'
point_color = 'C0'
point_size = 4
yticks = np.arange(0, 180 + 1, 30)
alpha_line = 0.7  # Upper and lower bound prediction line transparency
alpha_fill = 0.3  # Fill transparency
# Labels for figure subsections
x_text = -0.6
y_text = 1.15
fontsize = 16
labelsize = 14
ticksize = 11

# Line 1.
start = 625
stop = 638

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, ang_expectations, x_points, y_points, ids = \
    calculations_sq(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')
print('RMSD:', np.sqrt(np.mean((ang_points - ang_expectations)**2)))
print('MAE:', np.mean(np.abs(ang_points - ang_expectations)))

# Subplot 1 - First vector field.
ax1 = plt.subplot2grid(shape=(2, 3), loc=(0, 0), rowspan=1, colspan=1)
l = 1
y = np.linspace(-l / 2, l / 2, p)

x_mean = np.mean(xs_nodim)
lower_bound = x_mean - 2 * np.std(xs_nodim)
upper_bound = x_mean + 2 * np.std(xs_nodim)

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

square_vector_field(dimensions=l, max_range=max_range, sink_type='3D', plot_color=vector_color, edit_plot=True)
ax1.plot(y, [x_mean] * len(y), color=line_color)
ax1.plot(y, [lower_bound] * len(y), color=line_color)
ax1.plot(y, [upper_bound] * len(y), color=line_color)
ax1.fill_between(y, y1=lower_bound, y2=upper_bound, color=line_color, alpha=alpha_fill)
ax1.set_ylabel('$y$', fontsize=labelsize)
ax1.axis([-l / 2, l / 2, -l / 2, l / 2])
ax1.set_xticks(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)

# Subplot 2 - First experiment data.
ax2 = plt.subplot2grid(shape=(2, 3), loc=(0, 1), rowspan=1, colspan=2)
# DATA
x_mean = np.mean(xs_nodim)
lower_bound = x_mean - 2 * np.std(xs_nodim)
upper_bound = x_mean + 2 * np.std(xs_nodim)

# PREDICTION
l = 1
y = np.linspace(-l / 2, l / 2, p)
ax2.plot(-y, turn + jet_angle_data[0] * 180 / np.pi, color=prediction_line_color, linewidth=1.5,
         label='Prediction along $x$ = {0}'.format(np.round(x_mean, 4)))
for i in range(1, 3):
    ax2.plot(-y, turn + jet_angle_data[i] * 180 / np.pi, color='k', linestyle='-', linewidth=1, alpha=alpha_line)
ax2.fill_between(-y, turn + jet_angle_data[1] * 180 / np.pi, turn + jet_angle_data[2] * 180 / np.pi,
                 color=fill_color, alpha=alpha_fill)
# Data errorbars.
ax2.errorbar(x=-y_points, y=turn + ang_points, yerr=ang_errs,
             fmt='o', markersize=point_size, color=point_color, capsize=5,
             label='Data Mean and Standard Error')
# for n, pt in enumerate(zip(-y_points, turn + ang_points)):
#     ax2.annotate(s='{0}'.format(ids[n]), xy=pt, xycoords='data')
ax2.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
ax2.set_xlim(-l / 2, l / 2)
ax2.set_yticks(yticks)
ax2.tick_params(labelsize=ticksize)
ax2.grid()

# Line 2.
start = 644
stop = 658

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, ang_expectations, x_points, y_points, ids = \
    calculations_sq(line_start=start, line_stop=stop, p=p, max_range=max_range, mod_2pi=False)

# Print some data:
print('EXPERIMENT 2 INFORMATION')
# Number of trials
print('Total Number of Trials: ', len(nums))
# 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((ang_points - ang_expectations)**2)))
print('MAE:', np.mean(np.abs(ang_points - ang_expectations)))

# Subplot 3 - Second vector field.
ax3 = plt.subplot2grid(shape=(2, 3), loc=(1, 0), rowspan=1, colspan=1)
l = 1
y = np.linspace(-l / 2, l / 2, p)

x_mean = np.mean(xs_nodim)
lower_bound = x_mean - 2 * np.std(xs_nodim)
upper_bound = x_mean + 2 * np.std(xs_nodim)

print('x_mean:', x_mean, 'lower_bound:', lower_bound, 'upper_bound:', upper_bound)

square_vector_field(dimensions=l, max_range=max_range, sink_type='3D', plot_color=vector_color, edit_plot=True)
ax3.plot(y, [x_mean] * len(y), color=line_color)
ax3.plot(y, [lower_bound] * len(y), color=line_color, alpha=alpha_line)
ax3.plot(y, [upper_bound] * len(y), color=line_color, alpha=alpha_line)
ax3.fill_between(y, y1=lower_bound, y2=upper_bound, color=line_color, alpha=alpha_fill)
ax3.set_xlabel('$x$', fontsize=labelsize)
ax3.set_ylabel('$y$', fontsize=labelsize)
ax3.axis([-l / 2, l / 2, -l / 2, l / 2])
ax3.set_xticks(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)

# Subplot 4 - Second experiment data.
ax4 = plt.subplot2grid(shape=(2, 3), loc=(1, 1), rowspan=1, colspan=2)
# DATA
x_mean = np.mean(xs_nodim)
lower_bound = x_mean - 2 * np.std(xs_nodim)
upper_bound = x_mean + 2 * np.std(xs_nodim)

# PREDICTION
l = 1
y = np.linspace(-l / 2, l / 2, p)
ax4.plot(-y, turn + jet_angle_data[0] * 180 / np.pi, color=prediction_line_color, linewidth=1.5,
         label='Prediction along $x$ = {0}'.format(np.round(x_mean, 4)))
for i in range(1, 3):
    ax4.plot(-y, turn + jet_angle_data[i] * 180 / np.pi, color='k', linestyle='-', linewidth=1, alpha=alpha_line)
ax4.fill_between(-y, turn + jet_angle_data[1] * 180 / np.pi, turn + jet_angle_data[2] * 180 / np.pi, color=fill_color,
                 alpha=alpha_fill)
# Data errorbars.
ax4.errorbar(x=-y_points, y=turn + ang_points, yerr=ang_errs,
             fmt='o', markersize=point_size, color=point_color, capsize=5,
             label='Data Mean and Standard Error')
# for n, pt in enumerate(zip(y_points, turn - ang_points)):
#     ax4.annotate(s='{0}'.format(ids[n]), xy=pt, xycoords='data')
ax4.set_xlabel('$x$', fontsize=labelsize)
ax4.set_ylabel('$\\theta_j$ (degrees)', fontsize=labelsize)
ax4.set_xlim(-l / 2, l / 2)
ax4.set_yticks(yticks)
ax4.tick_params(labelsize=ticksize)
ax4.grid()

#%%
plt.tight_layout(pad=2)
plt.savefig('figure06.pdf', bbox_inches='tight', pad_inches=0)
plt.show()