import csv
import importlib.util
import os
import re
import sys

sys.path.insert(0, 'flopcounts')

ROOT = 'measurements'
OUT = 'analysis/results_x86.csv'

size_line = re.compile(r'=== N_P=(\d+) N_ITER=(\d+) ===')


def num(pattern, text):
  m = re.search(pattern, text)
  if m:
    return float(m.group(1))
  return None


_flopcount_modules = {}


def get_flops_per_iter(variant, n_p, n_obs=2):
  variant = variant.strip()

  module_name = variant.replace('/', '_').replace('\\', '_')
  module_path = f'flopcounts/x86/{module_name}'

  if module_name in _flopcount_modules:
    module = _flopcount_modules[module_name]
  elif os.path.exists(f'{module_path}.py'):
    try:
      spec = importlib.util.spec_from_file_location(module_name,
                                                    f'{module_path}.py')
      assert spec is not None
      assert spec.loader is not None
      module = importlib.util.module_from_spec(spec)
      spec.loader.exec_module(module)
      _flopcount_modules[module_name] = module
    except Exception as e:
      print(f'Warning: Failed to load flopcount for {variant}: {e}')
      return None
  else:
    print(f'Warning: No flopcount module found for {variant}')
    return None

  if hasattr(module, 'total'):
    try:
      return module.total(n_iter=1, n_p=int(n_p), n_obs=int(n_obs))
    except Exception as e:
      print(f'Warning: Error calling total() for {variant}: {e}')
      return None
  else:
    print(f'Warning: No total() function in {module_path}')
    return None


runs = {}

for variant in sorted(os.listdir(ROOT)):
  logdir = os.path.join(ROOT, variant, 'x86')
  if not os.path.isdir(logdir):
    continue

  for fn in os.listdir(logdir):
    if not fn.endswith('.log'):
      continue
    text = open(os.path.join(logdir, fn)).read()

    marks = list(size_line.finditer(text))
    for i, m in enumerate(marks):
      n = int(m.group(1))
      n_iter = int(m.group(2))
      start = m.end()
      stop = marks[i + 1].start() if i + 1 < len(marks) else len(text)
      block = text[start:stop]

      cycles = num(r'T\(n\):\s*([\d.]+)', block)
      if cycles is None:
        continue

      key = (variant, n)
      q = num(r'Q\(n\) per iteration:\s*([\d.]+)', block)
      if key not in runs:
        runs[key] = []
      runs[key].append((cycles, n_iter, q))

with open(OUT, 'w', newline='') as f:
  w = csv.writer(f)
  w.writerow([
    'variant', 'n', 'n_iter', 'W', 'Q_per_iter', 'I', 'cycles',
    'flops_per_cycle'
  ])
  for variant, n in sorted(runs):
    entries = runs[(variant, n)]
    n_iter = entries[0][1]
    q = entries[0][2]
    for _, ni, qi in entries:
      if ni != n_iter:
        print(f'Warning: n_iter mismatch for {(variant, n)}: {ni} vs {n_iter}')
      if qi != q:
        print(f'Warning: Q_per_iter mismatch for {(variant, n)}: {qi} vs {q}')

    cycles = sum(e[0] for e in entries) / len(entries)

    w_iter = get_flops_per_iter(variant, n)

    assert w_iter is not None

    W = w_iter * n_iter
    I = w_iter / q if q else w_iter
    fc = W / cycles if cycles else 0
    w.writerow(
      [variant, n, n_iter, W, q,
       round(I, 2),
       round(cycles, 1),
       round(fc, 3)])

print('wrote', OUT, 'with', len(runs), 'rows')