import seaborn as sns
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
# Adjusting the line thickness to better match the provided example

fig, ax = plt.subplots( figsize=(5, 4))

x_values = [8287, 77664, 532150,1126876,2070965,3163485]
y_values = [9.46, 6.94, 5.81, 5.70, 5.09, 5.02]
y_values = np.exp(y_values)

# Set line width for each line plot
line_width = 1.5
x = []
# Iterate over each subplot (task) and plot the lines with specified styles, markers, and adjusted line width
# for i, task in enumerate(tasks):

# Adding a centralized legend that appears above the plot
# fig.legend(y_values, loc='upper center', bbox_to_anchor=(0.5, 1.05), ncol=3, frameon=False, markerscale=1.5)

fig, ax1 = plt.subplots(figsize=(5, 4))

# Plot Perplexity (left y-axis)
ax1.plot(x_values, y_values, marker='o', linestyle='-', color='#1f78b4', linewidth=line_width)
ax1.annotate(f"{y_values[-1]:.1f}", (x_values[-1], y_values[-1]), textcoords="offset points", xytext=(0, 10), ha='center')
ax1.set_xscale('log')
ax1.set_xlabel("# Trajectory", fontsize=14)
ax1.set_ylabel("Perplexity", fontsize=14, color='#1f78b4')
ax1.tick_params(axis='y', labelcolor='#1f78b4')

# Create a twin y-axis for controllability (right y-axis)
ax2 = ax1.twinx()
controllability_values = [0.,0.10,1.20,1.41,1.56, 1.87]  # Example values for controllability
ax2.plot(x_values, controllability_values, marker='s', linestyle='--', color='#006400', linewidth=line_width)
ax2.set_ylabel("Delta PSNR", fontsize=14, color='#006400')
ax2.annotate(f"{controllability_values[-1]:.1f}", (x_values[-1], controllability_values[-1]), textcoords="offset points", xytext=(0, 10), ha='center')

ax2.set_ylim(0, 2.1)
ax2.tick_params(axis='y', labelcolor='#006400')

# Save the figure in high resolution
plt.tight_layout()
#plt.show()

plt.savefig(f"output/traj_sizes.png", dpi=300)