from plt_util import plt_config, plt_export
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch
import matplotlib.ticker as ticker
import numpy as np

plt_config(plt)

## Initiate test plot data
x = np.linspace(0, 10, 14)

# -- bounds
upper_bound = np.ones_like(x) * 4.5
lower_bound = upper_bound * 0.4

# -- lines
method_1_1 = upper_bound - (upper_bound - lower_bound) * np.exp(-0.7 * x)
method_1_2 = upper_bound - (upper_bound - lower_bound) * np.exp(-0.5 * x)
method_1_3 = upper_bound - (upper_bound - lower_bound) * np.exp(-0.3 * x)
method_2 = upper_bound - (upper_bound - lower_bound) * np.exp(-0.5 * x) * np.cos(x)


## Plot
x_min, x_max = 0, 10
y_min, y_max = 2, 6
plt.plot(x, upper_bound, 'k-', label='Optimal')
plt.plot(x, method_1_1, '^b-', label='Method 1')
plt.plot(x, method_1_2, '^b--')
plt.plot(x, method_1_3, '^b-.')
plt.plot(x, method_2, '*r-', label='Method 2')
plt.xlim(x_min, x_max)
plt.ylim(y_min, y_max)
plt.grid(True, linestyle='--', color='gray', linewidth=0.5)

# -- Add arrow
arrow = FancyArrowPatch((2.5, 4.2), # start point
                        (3.8, 3),   # end point
                        arrowstyle='-|>',  # arrow style
                        mutation_scale=10, # arrow head size
                        color='black',  # arrow color
                        linewidth=0.5)  # line width
plt.gca().add_patch(arrow)
plt.text(4, 3, '$\gamma = 1, 2, 3$')

# -- Add legends and labels
plt.legend()
plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(1))
plt.gca().yaxis.set_major_locator(ticker.MultipleLocator(1))
plt.ylabel(r"$\mathbb{E}\left[S\right]$")
plt.xlabel(r'Rounds of Iteration')

# -- export and show
plt_export(filename='example_python.pdf', 
           folder='../../images', 
           format='pdf')
plt.show()