#!/usr/bin/env python3

import argparse
from pathlib import Path
import csv
import os
import subprocess


def parse_duration(gourd_dur_string: str) -> float:
    if gourd_dur_string.endswith("s"):
        return float(gourd_dur_string.replace("s", ""))
    elif "N/A" in gourd_dur_string:
        return None
    assert False, gourd_dur_string


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("]", "")
                    .rsplit('/', 1))
            row = None
            status = "ok"
            if not "job finished" in csv_row["slurm"]:
                status = (
                    'timeout' if "timed out" in csv_row['slurm'] else
                    'memout' if "out of memory" in csv_row['slurm'] else None
                )
            row = {
                'solver': csv_row['program'],
                'formula': formula,
                'status': status,
                'exit_code': int(csv_row['exit']) if csv_row['exit'] != 'N/A' else None,
                'time': parse_duration(csv_row['wall time']),
                'peak_memory': int(csv_row['max rss']) if csv_row['max rss'] != 'N/A' else None,
            }
            """ 
            if "job finished" in csv_row["slurm"]:
                row = {
                    'solver': csv_row['program'],
                    'formula': formula,
                    'status': 'ok',
                    'exit_code': int(csv_row['exit']) if csv_row['exit'] != 'N/A' else None,
                    'time': parse_duration(csv_row['wall time']),
                    'peak_memory': int(csv_row['max rss']) if csv_row['max rss'] != 'N/A' else None,
                }
            else:
                row = {
                    'solver': csv_row['program'],
                    'formula': formula,
                    'status': (
                        'timeout' if "timed out" in csv_row['slurm'] else
                        'memout' if "out of memory" in csv_row['slurm'] else None
                    ),
                    'exit_code': int(csv_row['exit']) if csv_row['exit'] != 'N/A' else None,
                }
            """
            # assert row['status'] is not None, row
            rows[csv_row['run id']] = row
    return rows

def main(args):
    in_rows = parse_gourd_report(args.input_file)
    with open(args.output_file, "w") as csvfile:
        writer = csv.DictWriter(
            csvfile,
            fieldnames=[
                "solver",
                "formula",
                "status",
                "exit_code",
                "time",
                "peak_memory",
                "best_proof",
                "best_bound"
            ]
        )
        writer.writeheader()
        for solver_dir in args.logdir.iterdir():
            runs = list(solver_dir.iterdir())
            for index, run_dir in enumerate(runs, 1):
                row = in_rows[run_dir.name]
                if row["solver"] == "baseline":
                    for line in (run_dir / "stdout").open():
                        if "Shortest proof found" in line:
                            _, row["best_proof"] = line.split(': ')
                            row["best_proof"] = int(row["best_proof"])
                else:
                    for line in (run_dir / "stderr").open():
                        if "Best proof has" in line:
                            *_, row["best_proof"], _ = line.split(' ')
                            row["best_proof"] = int(row["best_proof"])
                        elif "Best available proof has" in line:
                            *_, row["best_proof"] = line.split(' ')
                            row["best_proof"] = int(row["best_proof"])
                        elif "lower bound is" in line:
                            _, row["best_bound"] = line.split(' lower bound is ')
                            row["best_bound"] = int(row["best_bound"])
                            if row["best_bound"] == 0:
                                row["best_bound"] = None
                        elif "lower bound on proof length is" in line:
                            _, row["best_bound"] = line.split(' lower bound on proof length is ')
                            row["best_bound"] = int(row["best_bound"])
                            if row["best_bound"] == 0:
                                row["best_bound"] = None
                print(f'Parsed {index} out of {len(runs)}', end='\r')
                writer.writerow(row)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Produces a synthetic formulas output")
    parser.add_argument(
        "logdir",
        type=Path,
        help="Directory with proof minimization 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)
