import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# Load your uploaded file
df = pd.read_csv('/content/fio_results.csv')  # Update this path if needed

# Create combined "iodepth/numjobs" label
df['param'] = df['iodepth'].astype(str) + '/' + df['numjobs'].astype(str)

# Create pivot tables for read and write
read_df = df[df['test_type'] == 'read'].pivot_table(index='device', columns='param', values='iops')
write_df = df[df['test_type'] == 'write'].pivot_table(index='device', columns='param', values='iops')

# Define sorting function
def sort_by_numjobs_then_iodepth(param):
    iodepth, numjobs = param.split('/')
    return (int(numjobs), int(iodepth))

# Sort the columns in the desired order
sorted_columns = sorted(read_df.columns, key=sort_by_numjobs_then_iodepth)
read_df = read_df.reindex(sorted_columns, axis=1)
write_df = write_df.reindex(sorted_columns, axis=1)

# Plot side-by-side heatmaps
fig, axes = plt.subplots(1, 2, figsize=(24, 10), sharey=True)

sns.heatmap(read_df, ax=axes[0], cmap="YlGnBu", linewidths=0.5, linecolor='gray')
axes[0].set_title("Randread IOPS Heatmap")
axes[0].set_xlabel("IOdepth / Numjobs")
axes[0].set_ylabel("Device")

sns.heatmap(write_df, ax=axes[1], cmap="YlOrRd", linewidths=0.5, linecolor='gray')
axes[1].set_title("Randwrite IOPS Heatmap")
axes[1].set_xlabel("IOdepth / Numjobs")
axes[1].set_ylabel("")

plt.tight_layout()
plt.show()