"""
Program: Figure 5 - Jet Direction Volatility
Date: 22 June 2019
Author: Lebo Molefe

Predict how much jet direction changes for a tiny step in position.
"""

import numpy as np
import matplotlib.pyplot as plt

# Set up subplots
fig = plt.figure(figsize=(11, 9))

ax1 = fig.add_subplot(2, 2, 1)
ax2 = fig.add_subplot(2, 2, 2)
ax3 = fig.add_subplot(2, 2, 3)
ax4 = fig.add_subplot(2, 2, 4)

fontsize = 20
ticksize = 16


### a - Square ###
dimensions = 1
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]
W_sq = np.loadtxt('square_volatility.csv', delimiter=',', skiprows=1, unpack=False)
# Colors
map = ax1.contourf(X_sq, Y_sq, W_sq, levels=np.linspace(-1, 4, 101, endpoint=True), cmap='inferno')
ax1.set_xticks([])
ax1.set_yticks([])
x_text = -0.2
y_text = 1.1
ax1.text(x=x_text, y=y_text, s='a', fontsize=24, weight='bold', transform=ax1.transAxes)


### b - Equilateral triangle ###
dimensions = 1
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]
W_eq_tri = np.loadtxt('equilateral_triangle_volatility.csv', delimiter=',', skiprows=1, unpack=False)
# Colors
map = ax2.contourf(X_eq_tri, Y_eq_tri, W_eq_tri, levels=np.linspace(-1, 4, 101, endpoint=True), cmap='inferno')
# Fill in areas outside triangle.
x = np.linspace(-l / 2, l / 2, 100, endpoint=True)
y1 = np.sqrt(3) * x
y2 = - np.sqrt(3) * x
ax2.plot(x, y1,
        x, y2, color='k', linewidth=1)
ax2.axhline(y=np.sqrt(3) / 2 * l, color='k', linewidth=1)
ax2.fill_between(x, y1, color='w')
ax2.fill_between(x, y2, color='w')
# Set axis properties.
epsilon_x = 0.0005 * h
epsilon_y = 0.0005 * l / np.sqrt(3)
ax2.axis([-l / 2, l / 2, 0, np.sqrt(3) / 2 * l + epsilon_y])
ax2.spines['bottom'].set_visible(False)
ax2.spines['top'].set_visible(False)
ax2.spines['right'].set_visible(False)
ax2.spines['left'].set_visible(False)
ax2.set_xticks([])
ax2.set_yticks([])
x_text = -0.2
y_text = 1.1
ax2.text(x=x_text, y=y_text, s='b', fontsize=24, weight='bold', transform=ax2.transAxes)


### c - Isosceles right triangle ###
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]
W_isc_right = np.loadtxt('isosceles_right_triangle_volatility.csv', delimiter=',', skiprows=1, unpack=False)

map = ax3.contourf(X_isc_right, Y_isc_right, W_isc_right, levels=np.linspace(-1, 4, 101, endpoint=True), cmap='inferno')
ax3.plot([0, l, 0, 0], [l, 0, 0, l], linestyle='-', color='k', linewidth=1)
ax3.axis([-0.01, l, -0.01, l])
ax3.spines['top'].set_visible(False)
ax3.spines['bottom'].set_visible(False)
ax3.spines['right'].set_visible(False)
ax3.spines['left'].set_visible(False)
ax3.set_xticks([])
ax3.set_yticks([])
x_text = -0.2
y_text = 1.1
ax3.text(x=x_text, y=y_text, s='c', fontsize=24, weight='bold', transform=ax3.transAxes)


### d - 306090 right triangle ###
dimensions = 1
l = dimensions
start = 0.999
step = (1 - start) * l
X_306090, Y_306090 = np.mgrid[0:(l * start):step, 0:(l * start):step]
W_306090 = np.loadtxt('306090_triangle_volatility.csv', delimiter=',', skiprows=1, unpack=False)

map = ax4.contourf(X_306090, Y_306090, W_306090, levels=np.linspace(-1, 4, 101, endpoint=True), cmap='inferno')
ax4.plot([0, l/2, 0, 0], [np.sqrt(3)/2 * l, 0, 0, np.sqrt(3)/2 * l], linestyle='-', color='k', linewidth=1)
ax4.axis([-0.01, l, -0.01, np.sqrt(3)/2 * l])
ax4.spines['top'].set_visible(False)
ax4.spines['bottom'].set_visible(False)
ax4.spines['right'].set_visible(False)
ax4.spines['left'].set_visible(False)
ax4.set_xticks([])
ax4.set_yticks([])

x_text = -0.2
y_text = 1.1
ax4.text(x=x_text, y=y_text, s='d', fontsize=24, weight='bold', transform=ax4.transAxes)

# Colorbar
fig.subplots_adjust(bottom=0.1, top=0.9, left=0.1, right=0.8, wspace=0.5, hspace=0.4)
cbar_axes = fig.add_axes([0.83, 0.1, 0.03, 0.8])
cbar = fig.colorbar(map, cax=cbar_axes)
cbar.set_ticks(np.arange(-1, 4.01, 0.5))
cbar.ax.tick_params(labelsize=ticksize)
cbar.set_label('log$_{10}\\left(\\sqrt{\\left(\\frac{\partial \\theta}{\partial x}\\right)^2 '
               '+ \\left(\\frac{\partial \\theta}{\partial y}\\right)^2}\\right)$',fontsize=fontsize)

# Save
plt.savefig('figure05.png', dpi=600)
plt.show()