#!/usr/bin/env python3
import csv
import matplotlib.pyplot as plt
import pandas as pd
import io

# Read the CSV file
with open('iperf_all_nodes.csv', 'r') as file:
    lines = file.readlines()

# Skip the first line and read the rest into a DataFrame using io.StringIO
data = pd.read_csv(io.StringIO(''.join(lines[1:])))

# Strip any leading or trailing spaces from column names
data.columns = data.columns.str.strip()

# Print the column names to debug
print("Column names in the DataFrame:", data.columns.tolist())

# Ensure column names are correct
expected_columns = ['Server', 'Client', 'mesh0 Min TP (Mbps)', 'mesh1 Min TP (Mbps)', 
                    'mesh0 Max TP (Mbps)', 'mesh1 Max TP (Mbps)', 'mesh0 Avg TP (Mbps)', 
                    'mesh1 Avg TP (Mbps)']
if not all(column in data.columns for column in expected_columns):
    raise ValueError("The CSV file does not contain the expected columns. Please check the column names.")

# Identify unique pairs
pairs = {}

for _, row in data.iterrows():
    server_client = (row['Server'].strip(), row['Client'].strip())
    client_server = (row['Client'].strip(), row['Server'].strip())

    if client_server in pairs:
        pairs[client_server].append(row)
    else:
        pairs[server_client] = [row]

# Prepare data for plotting
avg_data = []
# min_data = []
# max_data = []
     
for pair, measurements in pairs.items():
    if len(measurements) == 1:
        measurement = measurements[0]
        avg_data.append((pair[0], pair[1], 
                         float(measurement['mesh0 Avg TP (Mbps)']), 
                         float(measurement['mesh1 Avg TP (Mbps)'])
                        ))
    else:
        first_measurement, second_measurement = measurements
        avg_data.append((pair[0], pair[1], 
                         float(first_measurement['mesh0 Avg TP (Mbps)']), 
                         float(first_measurement['mesh1 Avg TP (Mbps)']),
                         float(second_measurement['mesh0 Avg TP (Mbps)']), 
                         float(second_measurement['mesh1 Avg TP (Mbps)'])
                        ))

# print(avg_data)

# Sort the average data
sorted_avg_data = sorted(avg_data, key=lambda x: (x[0][0], x[0][1]))

print(len(pairs))

# Determine the number of subplots needed
num_pairs = len(pairs)
cols = 8  # Number of columns for subplots
rows = (num_pairs + cols - 1) // cols  # Calculate number of rows needed
# rows = 12

# Create subplots for Average Throughput
fig_avg, axs_avg = plt.subplots(rows, cols, figsize=(25, 40))
axs_avg = axs_avg.flatten()

# print(sorted_avg_data)

# Colors for mesh0 and mesh1 bars
mesh0_color = 'skyblue'
mesh1_color = 'salmon'

# Plot Average Throughput
subplot_idx = 0
for server, client, mesh0_avg_s2c, mesh1_avg_s2c, mesh0_avg_c2s, mesh1_avg_c2s in sorted_avg_data:
    if subplot_idx < len(axs_avg):
        labels = [f'{server[-2:]} to {client[-2:]} (i0)', f'{server[-2:]} to {client[-2:]} (i1)', f'{client[-2:]} to {server[-2:]} (i0)', f'{client[-2:]} to {server[-2:]} (i1)']
        values = [mesh0_avg_s2c, mesh1_avg_s2c, mesh0_avg_c2s, mesh1_avg_c2s]
        
         # Normalize values
        max_value = max(values)
        if max_value != 0:
            values = [v / max_value * 100 for v in values]
            
        # Set colors
        colors = [mesh0_color, mesh1_color, mesh0_color, mesh1_color]
            
        axs_avg[subplot_idx].bar(labels, values, color=colors, label=f'{server} to {client}')
        axs_avg[subplot_idx].set_title(f'{server} to {client} and reverse')
        axs_avg[subplot_idx].set_ylabel('Throughput (Mbps)')
        axs_avg[subplot_idx].tick_params('x', labelrotation=90)
        axs_avg[subplot_idx].legend().set_visible(False)
        subplot_idx += 1


# Save the plots to PNG files
fig_avg.tight_layout()
fig_avg.savefig('iperf3_average_throughput.png')

# fig_min.tight_layout()
# fig_min.savefig('iperf3_minimum_throughput.png')

# fig_max.tight_layout()
# fig_max.savefig('iperf3_maximum_throughput.png')

# Close the plots to free up memory
plt.close(fig_avg)
# plt.close(fig_min)
# plt.close(fig_max)