"""
Program: Jet Direction Volatility Calculations
Date: 23 June 2019
Author: Lebo Molefe

Predict how much jet direction changes for a tiny step in two directions around each point. This file performs the
calculations for the plot in Figure 5.
"""

import numpy as np
import itertools
from jet_angle import angle_bubble_position, sink_positions_square, sink_positions_equilateral_triangle, \
    sink_positions_isosceles_right, sink_positions_306090


### FOR PLOTTING SQUARE - FROM JET ANGLE VECTOR FIELD FUNCTION ###
dimensions = 1
max_range = 5
sink_type = '3D'

l = dimensions
start = 0.99
step = (1 - start) * l
X_sq, Y_sq = 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_sq.shape
U = np.zeros(shape=shape)
V = np.zeros(shape=shape)
A = np.zeros(shape=shape)
W_sq = np.zeros(shape=shape)

print('square calculations')
for i, j in itertools.product(range(shape[0]), range(shape[1])):
    # Print percentage completion.
    if j == 0 and i % 4 == 0:
        print(round((i / shape[0] + j / (shape[1] * shape[0])) * 100, 1), '%')
    bubble_position = (X_sq[i][j], Y_sq[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)
    theta = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions, sink_type=sink_type,
                                  print_angle=False)
    # Subtract pi/2 because the plot's definition of angle is not the same.
    angle = theta - np.pi / 2
    U[i][j] = np.cos(angle)
    V[i][j] = np.sin(angle)
    A[i][j] = angle

# After the matrix A[i][j] is defined, find the "volatility."
for i, j in itertools.product(range(shape[0]), range(shape[1])):
    if i != 0 or j != 0 or i != (shape[0] - 1) or j != (shape[1] - 1):
        # Central difference formula
        if i > 0 and i < (shape[0] - 1):
            if 2 * np.pi - abs(A[i + 1][j] - A[i - 1][j]) < abs(A[i + 1][j] - A[i - 1][j]):
                diff_x = 2 * np.pi - abs(A[i + 1][j] - A[i - 1][j])
            else:
                diff_x = abs(A[i + 1][j] - A[i - 1][j])
        else:
            diff_x = 0
        deriv_x = diff_x / (2 * step)
        if j > 0 and j < (shape[1] - 1):
            if 2 * np.pi - abs(A[i][j + 1] - A[i][j - 1]) < abs(A[i][j + 1] - A[i][j - 1]):
                diff_y = 2 * np.pi - abs(A[i][j + 1] - A[i][j - 1])
            else:
                diff_y = abs(A[i][j + 1] - A[i][j - 1])
        else:
            diff_y = 0
        deriv_y = diff_y / (2 * step)
        # In degrees
        deriv_x = deriv_x * 180 / np.pi
        deriv_y = deriv_y * 180 / np.pi
        W_sq[i][j] = np.log10(np.sqrt(deriv_x**2 + deriv_y**2))

    else:
        W_sq[i][j] = None


### FOR PLOTTING TRIANGLE - FROM JET ANGLE VECTOR FIELD FUNCTION ###
dimensions = 1
height = dimensions * np.sqrt(3) / 2
max_range = 5
sink_type = '3D'

l = dimensions
h = np.sqrt(3) / 2 * l
start = 0.995
step = (1 - start) * l / np.sqrt(3)
h_bound = np.round(h * start, 1)
h_remainder = divmod(h_bound, step)[1]
X_eq_tri, Y_eq_tri = 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_eq_tri.shape
U = np.zeros(shape=shape)
V = np.zeros(shape=shape)
A = np.zeros(shape=shape)
W_eq_tri = np.zeros(shape=shape)

print('equilateral triangle calculations')
for i, j in itertools.product(range(shape[0]), range(shape[1])):
    # Print percentage completion.
    if j == 0 and i % 4 == 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_eq_tri[i][j]
    y_bubble = Y_eq_tri[i][j]
    e = 0.001
    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

    # Calculate the jet angle at the bubble position.
    bubble_position = (X_eq_tri[i][j], Y_eq_tri[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)
    theta = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions,
                                  sink_type=sink_type,
                                  print_angle=False)
    # Subtract pi/2 because the plot's definition of angle is not the same.
    angle = theta - np.pi / 2
    U[i][j] = np.cos(angle)
    V[i][j] = np.sin(angle)
    A[i][j] = angle

# After the matrix A[i][j] is defined, find the "volatility."
for i, j in itertools.product(range(shape[0]), range(shape[1])):
    if i != 0 or j != 0 or i != (shape[0] - 1) or j != (shape[1] - 1):
        # Central difference formula
        if i > 0 and i < (shape[0] - 1):
            if 2 * np.pi - abs(A[i + 1][j] - A[i - 1][j]) < abs(A[i + 1][j] - A[i - 1][j]):
                diff_x = 2 * np.pi - abs(A[i + 1][j] - A[i - 1][j])
            else:
                diff_x = abs(A[i + 1][j] - A[i - 1][j])
        else:
            diff_x = 0
        deriv_x = diff_x / (2 * step)
        if j > 0 and j < (shape[1] - 1):
            if 2 * np.pi - abs(A[i][j + 1] - A[i][j - 1]) < abs(A[i][j + 1] - A[i][j - 1]):
                diff_y = 2 * np.pi - abs(A[i][j + 1] - A[i][j - 1])
            else:
                diff_y = abs(A[i][j + 1] - A[i][j - 1])
        else:
            diff_y = 0
        deriv_y = diff_y / (2 * step)
        # In degrees
        deriv_x = deriv_x * 180 / np.pi
        deriv_y = deriv_y * 180 / np.pi
        W_eq_tri[i][j] = np.log10(np.sqrt(deriv_x**2 + deriv_y**2))

    else:
        W_eq_tri[i][j] = None



### FOR PLOTTING ISOSCELES RIGHT TRIANGLE - FROM JET ANGLE VECTOR FIELD FUNCTION ###
dimensions = 1
max_range = 5
sink_type = '3D'

l = dimensions
start = 0.995
step = (1 - start) * l
X_isc_right, Y_isc_right = 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_isc_right.shape
U = np.zeros(shape=shape)
V = np.zeros(shape=shape)
A = np.zeros(shape=shape)
W_isc_right = np.zeros(shape=shape)

print('isosceles right triangle calculations')
for i, j in itertools.product(range(shape[0]), range(shape[1])):
    # Print percentage completion.
    if j == 0 and i % 4 == 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_isc_right[i][j]
    y_bubble = Y_isc_right[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_isc_right[i][j], Y_isc_right[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)
    theta = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions,
                                  sink_type=sink_type,
                                  print_angle=False)
    # Subtract pi/2 because the plot's definition of angle is not the same.
    angle = theta - np.pi / 2
    U[i][j] = np.cos(angle)
    V[i][j] = np.sin(angle)
    A[i][j] = angle

# After the matrix A[i][j] is defined, find the "volatility."
for i, j in itertools.product(range(shape[0]), range(shape[1])):
    if i != 0 or j != 0 or i != (shape[0] - 1) or j != (shape[1] - 1):
        # Central difference formula
        if i > 0 and i < (shape[0] - 1):
            if 2 * np.pi - abs(A[i + 1][j] - A[i - 1][j]) < abs(A[i + 1][j] - A[i - 1][j]):
                diff_x = 2 * np.pi - abs(A[i + 1][j] - A[i - 1][j])
            else:
                diff_x = abs(A[i + 1][j] - A[i - 1][j])
        else:
            diff_x = 0
        deriv_x = diff_x / step
        if j > 0 and j < (shape[1] - 1):
            if 2 * np.pi - abs(A[i][j + 1] - A[i][j - 1]) < abs(A[i][j + 1] - A[i][j - 1]):
                diff_y = 2 * np.pi - abs(A[i][j + 1] - A[i][j - 1])
            else:
                diff_y = abs(A[i][j + 1] - A[i][j - 1])
        else:
            diff_y = 0
        deriv_y = diff_y / step
        # In degrees
        deriv_x = deriv_x * 180 / np.pi
        deriv_y = deriv_y * 180 / np.pi
        W_isc_right[i][j] = np.log10(np.sqrt(deriv_x**2 + deriv_y**2))

    else:
        W_isc_right[i][j] = None



### FOR PLOTTING 30-60-90 TRIANGLE - FROM JET ANGLE VECTOR FIELD FUNCTION ###
dimensions = 1
max_range = 5
sink_type = '3D'

l = dimensions
start = 0.999
step = (1 - start) * l
X_306090, Y_306090 = 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_306090.shape
U = np.zeros(shape=shape)
V = np.zeros(shape=shape)
A = np.zeros(shape=shape)
W_306090 = np.zeros(shape=shape)

print('306090 triangle calculations')
for i, j in itertools.product(range(shape[0]), range(shape[1])):
    # Print percentage completion.
    if j == 0 and i % 4 == 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_306090[i][j]
    y_bubble = Y_306090[i][j]
    e = 0.01
    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_306090[i][j], Y_306090[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)
    theta = angle_bubble_position(bubble_position=bubble_position, sink_positions=sink_positions,
                                  sink_type=sink_type,
                                  print_angle=False)
    # Subtract pi/2 because the plot's definition of angle is not the same.
    angle = theta - np.pi / 2
    U[i][j] = np.cos(angle)
    V[i][j] = np.sin(angle)
    A[i][j] = angle


# After the matrix A[i][j] is defined, find the "volatility."
for i, j in itertools.product(range(shape[0]), range(shape[1])):
    if i != 0 or j != 0 or i != (shape[0] - 1) or j != (shape[1] - 1):
        # Central difference formula
        if i > 0 and i < (shape[0] - 1):
            if 2 * np.pi - abs(A[i + 1][j] - A[i - 1][j]) < abs(A[i + 1][j] - A[i - 1][j]):
                diff_x = 2 * np.pi - abs(A[i + 1][j] - A[i - 1][j])
            else:
                diff_x = abs(A[i + 1][j] - A[i - 1][j])
        else:
            diff_x = 0
        deriv_x = diff_x / step
        if j > 0 and j < (shape[1] - 1):
            if 2 * np.pi - abs(A[i][j + 1] - A[i][j - 1]) < abs(A[i][j + 1] - A[i][j - 1]):
                diff_y = 2 * np.pi - abs(A[i][j + 1] - A[i][j - 1])
            else:
                diff_y = abs(A[i][j + 1] - A[i][j - 1])
        else:
            diff_y = 0
        deriv_y = diff_y / step
        # In degrees
        deriv_x = deriv_x * 180 / np.pi
        deriv_y = deriv_y * 180 / np.pi
        W_306090[i][j] = np.log10(np.sqrt(deriv_x**2 + deriv_y**2))

    else:
        W_306090[i][j] = None



# Convert lists to arrays.
W_sq = np.array(W_sq)
W_eq_tri = np.array(W_eq_tri)
W_isc_right = np.array(W_isc_right)
W_306090 = np.array(W_306090)

# Save data to output file.
np.savetxt('square_volatility.csv', W_sq, delimiter=',', header='W_sq')
np.savetxt('equilateral_triangle_volatility.csv', W_eq_tri, delimiter=',', header='W_eq_tri')
np.savetxt('isosceles_right_triangle_volatility.csv', W_isc_right, delimiter=',', header='W_isc_right')
np.savetxt('306090_triangle_volatility.csv', W_306090, delimiter=',', header='W_isc_right')