#!/usr/bin/python3

import sys
import os
import glob
import xml.etree.ElementTree as ET
from datetime import datetime

def median(elements, key):
  sorted_elem = sorted(elements, key=key)
  return sorted_elem[int(len(sorted_elem)/2)]

def get_runs(xml):
  return xml.getroot().findall('run')

def get_name(run):
  return run.attrib.get('name')

def is_successful(data):
  expected = data['run'].attrib.get('expectedVerdict')
  status = data['run'].find('column[@title="status"]').attrib.get('value')
  if expected is not None:
    return status.startswith(expected)
  return status.startswith('true') or status.startswith('false')

def get_cputime(data):
  return float(data['run'].find('column[@title="cputime"]').attrib.get('value').rstrip('s'))

def get_logfile(source, data):
  resultFile = source['path']
  resultElem = source['xml'].getroot()
  sourcefile = data['run']
  
  # remainder of the method copied from benchexec/tablegenerator
  log_folder = resultFile[0 : resultFile.rfind(".results.")] + ".logfiles/"

  # append begin of filename
  runSetName = resultElem.get("name")
  if runSetName is not None:
      blockname = resultElem.get("block")
      if blockname is None:
          log_folder += runSetName + "."
      elif blockname == runSetName:
          pass  # real runSetName is empty
      else:
          assert runSetName.endswith("." + blockname)
          runSetName = runSetName[: -(1 + len(blockname))]  # remove last chars
          log_folder += runSetName + "."
  return f"{log_folder}{os.path.basename(sourcefile.get('name'))}.log"

###############################################################################
##                                  CONFIG                                   ##
###############################################################################

# One of max, min, median
aggr = min
out = 'IA-Random-MIN.Weaver-Boogie'
filter_failures = True

# Function to retrieve key value from dictionary { source: "...(path to xml)...", run: <run> element }
attr = get_cputime
###############################################################################
##                                   MAIN                                    ##
###############################################################################

xmlfiles = sys.argv[1:]

# from all files, get all /result/run
sources = {}
for file in xmlfiles:
  print(file)
  tree = ET.parse(file)
  sources[file] = { 'path': file, 'xml': tree, 'name': tree.getroot().attrib['name'], 'runs': get_runs(tree) }

# group by @name
benchmarks = {}
for source in sources.values():
  for run in source['runs']:
    name = get_name(run)
    if name not in benchmarks:
      benchmarks[name] = []
    benchmarks[name].append({ 'source': source['path'], 'sourcename': source['name'], 'run': run })

chosen = {}
for name, runs in benchmarks.items():
  # keep only successful, unless all failed
  success = any(is_successful(run) for run in runs)
  if success and filter_failures:
    relevant = [ run for run in runs if is_successful(run) ]
  else:
    relevant = runs
  chosen[name] = aggr(relevant, key=attr)
  chosen[name]['run'].attrib['logfile'] = get_logfile(sources[chosen[name]['source']], chosen[name])
  sourceTag = ET.Element('column')
  sourceTag.attrib['title'] = 'configuration'
  sourceTag.attrib['value'] = chosen[name]['sourcename']
  chosen[name]['run'].append(sourceTag)
  

# Prepare empty template
templatefile = xmlfiles[0]
print('Using ' + templatefile + ' as XML template')
template = ET.parse(templatefile)
templateroot = template.getroot()
for run in templateroot.findall('run'):
  templateroot.remove(run)
templateroot.attrib.pop('block')
templateroot.attrib['date'] = datetime.now().isoformat()
templateroot.attrib['generator'] = 'aggregate.py'
templateroot.attrib['name'] = out
templateroot.attrib['options'] = ''

# Add chosen entries to XML
for name, run in chosen.items():
  templateroot.append(run['run'])

template.write('aggregated-' + out + '.xml')