check-pdata

#!/usr/bin/env python3
"""
Check pdata for consistency

Check things like PM totals in tasks, work packages, sites, etc.
"""

import copy
from pprint import pprint
import re
import sys


site_order = [
    'SRL',
    'XFEL',
    'QS',
    'CDS',
    'EGI',
    'UIO',
    'WTT',
    'SIL',
    'UPSUD',
    'INSERM',
    'EP',
]


def _mini_fsm(s):
    """mini fsm for dealing with nested {}"""
    level = 0
    characters = []
    for i, c in enumerate(s):
        if c == '{':
            level += 1
            if level == 1:
                # don't include first opening brace
                continue
        elif c == '}':
            level -= 1
            if level == 0:
                break
        characters.append(c)
    else:
        raise ValueError(f"Failed to properly parse line: {s}")
    return ''.join(characters), s[i + 1 :]


PDATA_PREFIX = '\@pdata@def'
_num_pat = re.compile(r'^\d+$')


def parse_pdata_line(line):
    """Return list of pdata fields on a given line"""
    assert line.startswith(PDATA_PREFIX)
    fields = []
    remaining = line[len(PDATA_PREFIX) :].strip()
    while remaining:
        field, remaining = _mini_fsm(remaining)
        fields.append(field)
    for i, field in enumerate(fields):
        # cast integers
        try:
            fields[i] = int(field)
        except ValueError:
            pass
    return fields


def _add_field(data, fields):
    """Add pdata fields to running dict"""
    d = data
    for f in fields[:-2]:
        d = d.setdefault(f, {})
    key, value = fields[-2:]
    # split known list entries
    if key in {'ids', 'partners'}:
        value = value.split(',')
    d[key] = value


def reformat_data(data):
    pass


def summary(data):

    invalid = False
    total_pms = 0
    site_pms = {key: 0 for key in data['site']}
    data = copy.deepcopy(data)
    for wp_id, wp in data['wp'].items():
        print(f"WP{wp['number']}: {wp['title']}")
        wp_data = data[wp_id]

        # collect pms by site assigned to the work package
        wp_pms = 0
        for site in site_order:
            rm = wp_data.get(site, {}).get('RM', 0)
            print(site, rm)
            site_pms[site] += rm
            wp_pms += rm

        # collect pms assigned to the tasks
        task_ids = wp_data['task']['ids']
        task_pms = 0
        for task_id in task_ids:
            task = data['task'][task_id]
            try:
                task_pms += task['PM']
            except (TypeError, ValueError):
                print(f"task {task['title']} has invalid pms: {task['PM']}")
        print(f"  Lead: {wp['lead']}")
        print(f"  Tasks: {len(task_ids)}")
        print(f"  Deliverables: ???")
        print(f"    WP PMs: {wp_pms}")
        print(f"  task PMs: {task_pms}")
        if wp_pms != task_pms:
            print("!!!!! task PMs don't add up !!!!")
            invalid = 1
        print()

    print()
    print(f"{'Site':12}: {'PMs':3}")
    total_pms = 0
    # for site, pms in sorted(site_pms.items(), key=itemgetter(1), reverse=True):
    assert sorted(site_pms) == sorted(site_order), set(site_pms).difference(
        set(site_order)
    )
    for site in site_order:
        pms = site_pms[site]
        print(f"{site:12}: {pms:3}")
        total_pms += pms
    print()
    print(f"{'total':12}: {total_pms:3}")
    return invalid


def main(pdata_file='draft.pdata'):
    data = {}
    with open(pdata_file) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            fields = parse_pdata_line(line)
            _add_field(data, fields)
    try:
        return summary(data)
    except Exception:
        pprint(data)
        raise


if __name__ == '__main__':
    sys.exit(main())