#!/usr/bin/env python3
import subprocess
import csv
import re
from argparse import ArgumentParser, Action
from itertools import product
from datetime import datetime

MODEL_TYPES = ["LI", "HI", "HO", "HOmod"]
DS_MODEL_TYPES = ["LI2HI"]
INT_CYCLES = 0
EXT_CYCLES = 0

PyPDEVS = "./pythonpdevs/main.py"
MinimalPyPDEVS = "./pythonpdevs-minimal/main.py"
OUTPUT_FILE = "./Results-PyPDEVS/DEVSTONE"

# Regex to capture times from script output
TIME_PATTERN = re.compile(
    r"Model creation time: ([\d.eE+-]+).*"
    r"Engine setup time: ([\d.eE+-]+).*"
    r"Simulation time: ([\d.eE+-]+)",
    re.S
)

def run_devstone(model_type, depth, width, minimal=False, classic=False, dynamic=False):
    """Runs the devstone_driver.py script and returns timing results as floats."""
    SCRIPT_PATH = PyPDEVS
    if minimal:
        SCRIPT_PATH = MinimalPyPDEVS

    cmd = [
        "python3", SCRIPT_PATH,
        "-m", model_type,
        "-d", str(depth),
        "-w", str(width),
        "-i", str(INT_CYCLES),
        "-e", str(EXT_CYCLES),
    ]
    if dynamic and not minimal:
        cmd.extend(["-D"])
    if classic and not minimal:
        cmd.extend(["-C"])
    result = subprocess.run(cmd, capture_output=True, text=True)
    output = result.stdout.strip()
    errors = result.stderr.strip()
    if errors:
        print(errors)
    match = TIME_PATTERN.search(output)

    if match:
        creation, setup, sim = map(float, match.groups())
        return creation, setup, sim, output
    else:
        print("Failed to parse output for", cmd)
        print(output)
        return None, None, None, output

class EnsureDS(Action):
    def __call__(self, parser, namespace, values, option_string=None):
        if not getattr(namespace, "dynamic", False):
            parser.error("--dynamic-models (-dm) can only be used with --dynamic (-D)")
        setattr(namespace, self.dest, values)

def parse_args():
    parser = ArgumentParser()
    parser.add_argument("-M", "--minimal", action="store_true", default=False, help="Uses the minimal DEVS kernel. (Less features => less overhead)")
    parser.add_argument("-C", "--classic", action="store_true", default=False, help="Selects the Classic DEVS Simulator")
    parser.add_argument("-D", "--dynamic", action="store_true", default=False, help="Enables Dynamic Structure DEVS")

    parser.add_argument("-ms", "--model-structure", nargs="+", default=MODEL_TYPES+DS_MODEL_TYPES, choices=MODEL_TYPES+DS_MODEL_TYPES, help="Select one or more model structures")

    parser.add_argument("-md", "--max-depth", default=8, type=int, help="Max depth for the generated models")
    parser.add_argument("-ds", "--depth-stepsize", default=1, type=int, help="Determines for which depths a model is generated.")

    parser.add_argument("-mw", "--max-width", default=8, type=int, help="Max width for the generated models")
    parser.add_argument("-ws", "--width-stepsize", default=1, type=int, help="Determines for which depths a model is generated.")

    parser.add_argument("-r", "--repetitions", default=1, type=int, help="The number of times a simulation is repeated")

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    print(args)

    d = (args.max_depth, args.depth_stepsize)
    w = (args.max_width, args.width_stepsize)
    if d[0] % d[1] != 0:
        raise RuntimeError("Invalid depth stepsize.")
    if w[0] % w[1] != 0:
        raise RuntimeError("Invalid width stepsize.")

    appendix = ""
    if args.minimal:
        appendix += "-minimal"
    if args.classic:
        appendix += "-classic"
    if args.dynamic:
        appendix += "-dynamic"

    with open(OUTPUT_FILE + appendix + ".csv", "w", newline="") as csvfile:
        fieldnames = ["model_type", "depth", "width", "run",
                      "creation_time", "setup_time", "simulation_time"]
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()

        for model_type, d, w in product(args.model_structure, range(d[1], d[0]+1, d[1]), range(w[1], w[0]+1, w[1])):
            for run in range(0, args.repetitions):
                creation, setup, sim, output = run_devstone(model_type, d, w, args.minimal, args.classic, args.dynamic)
                if creation is not None:
                    writer.writerow({
                        "model_type": model_type,
                        "depth": d,
                        "width": w,
                        "run": run,
                        "creation_time": creation,
                        "setup_time": setup,
                        "simulation_time": sim
                    })
                    csvfile.flush()
                else:
                    print("Skipping due to parse failure")

    print(f"\n Benchmarking complete. Results saved to {OUTPUT_FILE}")