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

x = [1, 5, 10]
tasks = ["Paper Towel Replacement\n(Bi-UR5e)", "Items in Drawer\n(Franka)",
         "Stack Bowls\n(UR5e)", "Tupperware in Microwave\n(Bi-ARX)"]

# Define y-values for each line type
y_values = {
    "π₀": [0.9, 0.85, 0.8],
    "π₀ (scratch)": [0.7, 0.75, 0.72],
    "DP": [0.2, 0.3, 0.4],
    "Octo": [0.5, 0.6, 0.55],
    "OpenVLA": [0.1, 0.15, 0.2],
    "ACT": [0.4, 0.5, 0.6]
}

# Define markers, line styles, colors for each line type
markers = {"π₀": 'o', "π₀ (scratch)": 'o', "DP": 'o', "Octo": 'D', "OpenVLA": '*', "ACT": 'o'}
styles = {"π₀": '-', "π₀ (scratch)": '--', "DP": '-', "Octo": '-', "OpenVLA": '', "ACT": '-'}
colors = {"π₀": '#1f78b4', "π₀ (scratch)": '#1f78b4', "DP": '#e31a1c', "Octo": '#33a02c', "OpenVLA": '#6a3d9a', "ACT": '#ff7f00'}

# Set line width for enhanced visibility

# Create subplots


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

x_values = [5, 10, 20, 30, 40]
y_values = [5.94,5.72, 5.21,5.15,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


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("# Dataset", 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.46, 0.55, 1.69, 1.5, 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.set_ylim(0, 2.1)
ax2.tick_params(axis='y', labelcolor='#006400')
ax2.annotate(f"{controllability_values[-1]:.1f}", (x_values[-1], controllability_values[-1]), textcoords="offset points", xytext=(0, 10), ha='center')

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


plt.savefig(f"output/dataset_sizes.png", dpi=300)  # Save the figure in high resolution