import os
import sys
import json
import yaml
import matplotlib.pyplot as plt
import numpy as np
import re

# Load device weights from YAML
with open("yosys/device_count.yaml") as f:
    device_weights = yaml.safe_load(f)

type_map = {
    "_muxig.json": "muxig",
    "_muxig_abc.json": "muxig abc",
    "_bbdd.json": "bbdd",
    "_bbdd_abc.json": "bbdd abc",
    "_bdd.json": "bdd",
    "_bdd_abc.json": "bdd abc",
    "_golden.json": "golden",
    "_golden_abc.json": "golden abc"
}


data_dir = sys.argv[1]
device_counts = dict()
time_data = dict()
memory_data = dict()
print("Collecting summary from ", data_dir)
#for variant in ["muxig","muxig_abc", "bdd", "bbdd", "bbdd_abc", "golden_abc", "golden"]:
for filename in os.listdir(data_dir):
    file_path = os.path.join(data_dir, filename)
    if not os.path.isfile(file_path):
        continue
    
    for suffix, type_name in type_map.items():
        if filename.endswith(suffix):
            with open(file_path) as f:
                json_data = json.load(f)
            # Assumes single module
            module_data = list(json_data["modules"].values())[0]
            num_cells = module_data["num_cells_by_type"]
            device_sum = 0;
            for cell_type in num_cells:
                device_sum = device_sum + (num_cells[cell_type] * device_weights[cell_type]["4T"])
                device_sum = device_sum + (num_cells[cell_type] * device_weights[cell_type]["3T"])
            device_counts[type_name] = device_sum

    if filename.endswith("_time.txt"):
        time_path = os.path.join(data_dir, filename)
        #get extension name
        basename = filename.rsplit('.', 1)[0]  # e.g. "xx_bdd_time"
        parts = basename.split('_')
        opt_name = parts[-2]  # second last part is the middle part
        time_data.setdefault(opt_name, dict())
        memory_data.setdefault(opt_name, dict())
        with open(time_path, 'r') as file:
            for line in file:
                if line.startswith("Time:"):
                    time_str = line.split()[1]  # e.g. "0:00.03"
                    minutes, seconds = time_str.split(':')
                    if (filename.startswith("yosys")):
                        yosys_seconds = int(minutes) * 60 + float(seconds)
                        time_data[opt_name]["yosys"] = yosys_seconds * 1000
                    else:
                        opt_seconds = int(minutes) * 60 + float(seconds)
                        time_data[opt_name]["opt"] = opt_seconds * 1000
                match = re.search(r'(\d+)', line)
                if match:
                    if (filename.startswith("yosys")):
                        memory_data[opt_name]["yosys"] = int(match.group(1))
                    else:
                        memory_data[opt_name]["opt"] = int(match.group(1))

#print(device_counts)
title = "SUMMARY"
types = [str(v) for v in list(device_counts.keys())]
data = [str(v) for v in list(device_counts.values())]
cell_width = max(len(s) for s in types)

# Calculate width based on the data row
formatted_data = [s.center(cell_width) for s in data]
formatted_types = [s.center(cell_width) for s in types]
data_row = f"| {' | '.join(formatted_data)} |"
types_row = f"| {' | '.join(formatted_types)} |"

box_width = len(types_row)

# Print the box
print("-" * box_width)
print(f"| {title.center(box_width - 4)} |")
print("-" * box_width)
print(types_row)
print("-" * box_width)
print(data_row)
print("-" * box_width)

print("Runtime Summary")
for opt in time_data.keys():
    t_data = time_data[opt]
    m_data = memory_data[opt]

    # Time calculation
    t_num = t_data["opt"] / t_data["yosys"] * 100
    print("Optimization", opt, 
          "Time yosys:", t_data["yosys"], 
          "Time opt:", t_data["opt"], 
          "Overhead in %:", f"{t_num:.2f}")
    #print(f"{t_data['yosys']},{t_data['opt']},{t_num:.2f}")

    # Memory calculation
    m_num = m_data["opt"] / m_data["yosys"] * 100
    print("Optimization", opt, 
          "Memory yosys:", m_data["yosys"], 
          "Memory opt:", m_data["opt"], 
          "Overhead in %:", f"{m_num:.2f}")
    print(f"{t_data['yosys']},{t_data['opt']},{t_num:.2f},{m_data['yosys']},{m_data['opt']},{m_num:.2f}")