#!/usr/bin/env python3

import argparse
from pathlib import Path
from copy import deepcopy
import csv
import json
import os
import subprocess


def parse_gourd_report(path):
    rows = dict()
    with path.open() as csvfile:
        reader = csv.DictReader(csvfile)
        for csv_row in reader:
            formula = (csv_row["input args"]
                    .replace('"', '')
                    .replace("[", "")
                    .replace("]", "")
                    .split('/')[-1])
            if "job finished" in csv_row["slurm"]:
                rows[csv_row['run id']] = {"formula": formula}
    return rows


def main(args):
    in_rows = parse_gourd_report(args.input_file)
    with args.output_file.open("w") as csvfile:
        writer = csv.DictWriter(
            csvfile,
            fieldnames=[
                'formula',
                'trim',
                'n_axioms',
                'n_derived',
                'n_rat',
                'duplicates'
            ],
            quotechar="'"
        )
        writer.writeheader()
        runs = [
            run_dir
            for solver_dir in args.logdir.iterdir()
            for run_dir in solver_dir.iterdir()
        ]
        for index, run_dir in enumerate(runs):
            if run_dir.name not in in_rows:
                continue
            base_row = in_rows[run_dir.name]
            stdout = (run_dir / "stdout").read_text().strip()
            if len(stdout) == 0:
                continue
            stdout = stdout.split("\n")
            assert len(stdout) == 2, stdout
            rows = [
                {
                    **base_row,
                    **{
                        "trim": ix == 1,
                        "n_axioms": int(line.split()[0]),
                        "n_derived": int(line.split()[1]),
                        "n_rat": int(line.split()[2]),
                        "duplicates": line.split()[3],
                    }
                }
                for ix, line in enumerate(stdout)
            ]
            print(f'Parsed {index} out of {len(runs)}', end='\r')
            for row in rows:
                if row["n_rat"] > 0:
                    print(row)
                writer.writerow(row)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Produces the baseline CSV")
    parser.add_argument(
        "logdir",
        type=Path,
        help="Directory with baseline runner outputs",
    )
    parser.add_argument(
        "--input-file",
        "-i",
        type=Path,
        required=True,
        help="CSV file with the Gourd report file"
    )
    parser.add_argument(
        "--output-file",
        "-o",
        type=Path,
        required=True,
        help="CSV file with the report table",
    )
    args = parser.parse_args()
    main(args)
