import re
import matplotlib.pyplot as plt
from collections import defaultdict

# --------------------------
# Define log files
# --------------------------
corona_log = 'report.log'
spacx_logs = ['gb_comm.log', 'interposer_comm.log', 'ni_comm.log']

# --------------------------
# Regex patterns
# --------------------------
timestamp_pattern = re.compile(r"(\d+)\s+ns")
corona_packet_gen_pattern = re.compile(r"Flits generated for packet of id (\d+)")
corona_packet_forward_pattern = re.compile(r"forwarded data with ID: (\d+)")

spacx_gb_input_pattern = re.compile(r"GlobalBuffer: Received transaction is an input feature with id: (\d+)")
spacx_pe_arrival_pattern = re.compile(r"Packet with id (\d+) arrived at Parent PE")

# --------------------------
# Parse Corona log
# --------------------------
corona_packet_times = {}
corona_max_time = 0

with open(corona_log, 'r') as file:
    for line in file:
        time_match = timestamp_pattern.search(line)
        if time_match:
            time = int(time_match.group(1))
            corona_max_time = max(corona_max_time, time)

            gen_match = corona_packet_gen_pattern.search(line)
            if gen_match:
                pkt_id = int(gen_match.group(1))
                corona_packet_times[pkt_id] = {'generated': time, 'forwarded': None}

            forward_match = corona_packet_forward_pattern.search(line)
            if forward_match:
                pkt_id = int(forward_match.group(1))
                if pkt_id in corona_packet_times:
                    corona_packet_times[pkt_id]['forwarded'] = time

# --------------------------
# Parse SPACX logs
# --------------------------
spacx_packet_times = defaultdict(lambda: {'input': None, 'arrived': None})
spacx_max_time = 0

for log_file in spacx_logs:
    with open(log_file, 'r') as file:
        for line in file:
            time_match = timestamp_pattern.search(line)
            if time_match:
                time = int(time_match.group(1))
                spacx_max_time = max(spacx_max_time, time)

                gb_match = spacx_gb_input_pattern.search(line)
                if gb_match:
                    pkt_id = int(gb_match.group(1))
                    spacx_packet_times[pkt_id]['input'] = time

                pe_match = spacx_pe_arrival_pattern.search(line)
                if pe_match:
                    pkt_id = int(pe_match.group(1))
                    spacx_packet_times[pkt_id]['arrived'] = time

# --------------------------
# Calculate latencies
# --------------------------
corona_latencies = []
for pkt_id, times in corona_packet_times.items():
    if times['generated'] is not None and times['forwarded'] is not None:
        corona_latencies.append(times['forwarded'] - times['generated'])

spacx_latencies = []
for pkt_id, times in spacx_packet_times.items():
    if times['input'] is not None and times['arrived'] is not None:
        spacx_latencies.append(times['arrived'] - times['input'])

# --------------------------
# Print total simulation times
# --------------------------
print(f"Corona Total Simulation Time: {corona_max_time} ns")
print(f"SPACX Total Simulation Time: {spacx_max_time} ns")

# --------------------------
# Plot latency comparison
# --------------------------
plt.figure(figsize=(10, 6))
plt.hist(corona_latencies, bins=10, alpha=0.7, label='Corona Latency (ns)', color='blue')
plt.hist(spacx_latencies, bins=10, alpha=0.7, label='SPACX Latency (ns)', color='green')
plt.xlabel('Latency (ns)')
plt.ylabel('Number of Packets')
plt.title('Packet Latency Comparison')
plt.legend()
plt.grid(True)
plt.show()