#!/usr/bin/env python3
# SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
# SPDX-License-Identifier: GPL-3.0-only
"""Tabulate the TIMING lines of a rugnux <prefix>_log.txt.

rugnux logs one debug line per timed section as it ends (TimingMark in common/Logger.h):

    [2026-10-07 15:29:52.114] [Rugnux] [debug] TIMING merge P1 4.210

The section name is everything between TIMING and the last field, which is the seconds. Sections
nest: an "rsm ..." line is a part of the "merge <role>" line logged after it, and every section of
a pass ends before its "pass (...)" line.

The tail is the stretch from the end of the image pass of the last output pass to the last line of
the log. Its sections are listed in the order they ended, then summed by name.

    tools/rugnux_timing.py p_log.txt          # the tail
    tools/rugnux_timing.py --all p_log.txt    # every timed section of the run
"""

import argparse
import datetime
import re
import sys

LINE = re.compile(r'^\[(\d{4}-\d\d-\d\d \d\d:\d\d:\d\d\.\d+)\] \[[^\]]*\] \[[^\]]*\] (.*)$')


def read(path):
    marks = []   # (end time, section, seconds)
    first = last = None
    with open(path, errors='replace') as f:
        for line in f:
            m = LINE.match(line.rstrip('\n'))
            if not m:
                continue
            t = datetime.datetime.strptime(m.group(1), '%Y-%m-%d %H:%M:%S.%f')
            first = first or t
            last = t
            text = m.group(2)
            if text.startswith('TIMING '):
                section, seconds = text[len('TIMING '):].rsplit(' ', 1)
                marks.append((t, section, float(seconds)))
    return first, last, marks


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('log')
    ap.add_argument('--all', action='store_true', help='every timed section, not only the tail')
    args = ap.parse_args()

    first, last, marks = read(args.log)
    if not marks:
        sys.exit(f'{args.log}: no TIMING lines')

    start = 0
    if not args.all:
        output = [i for i, m in enumerate(marks) if m[1] == 'pass (output)']
        if output:
            image = [i for i in range(output[-1]) if marks[i][1] == 'image pass']
            start = image[-1] + 1 if image else 0
    tail = marks[start:]
    t0 = marks[start - 1][0] if start > 0 else first
    wall = (last - t0).total_seconds()

    print(f'{"ends at [s]":>11} {"seconds":>9}  section')
    for t, section, seconds in tail:
        indent = '    ' if section.startswith('rsm ') else ''
        print(f'{(t - t0).total_seconds():11.2f} {seconds:9.3f}  {indent}{section}')

    print()
    # Shares are of the wall clock above. Sections nest (a pass holds its merges, a merge its rsm
    # parts), so they do not add up to 100 %.
    print(f'{"count":>5} {"seconds":>9} {"share":>6}  section (summed by name; nested, so shares overlap)')
    totals = {}
    for _, section, seconds in tail:
        n, s = totals.get(section, (0, 0.0))
        totals[section] = (n + 1, s + seconds)
    for section, (n, s) in sorted(totals.items(), key=lambda kv: -kv[1][1]):
        share = f'{100.0 * s / wall:5.1f}%' if wall > 0 else '     -'
        print(f'{n:5d} {s:9.3f} {share}  {section}')
    print()
    origin = "the end of the last output pass's image pass" if start > 0 else 'the first log line'
    print(f'{"run" if args.all else "tail"} wall clock: {wall:.2f} s, from {origin} to the last log line')


if __name__ == '__main__':
    main()