#!/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", "HI2LI", "LI-HI", "dLI", "dHI")
GEN_METHODS = ("start", "uniform", "end")

INT_CYCLES = 0
EXT_CYCLES = 0

PyPDEVS = "./pythonpdevs/main.py"
MinimalPyPDEVS = "./pythonpdevs-minimal/main.py"

# 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, gen_mode="uniform"):
    """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"])
        cmd.extend(["-g", gen_mode])
    if classic and not minimal:
        cmd.extend(["-C"])
    result = subprocess.run(cmd, capture_output=True, text=True)
    output = result.stdout.strip()
    print(output)
    errors = result.stderr.strip()
    match = TIME_PATTERN.search(output)

    if match:
        creation, setup, sim = map(float, match.groups())
        return creation, setup, sim, output

    print("Failed to run: ", cmd)
    print(output)
    print(errors)
    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("-p", "--path", type=str, default="./Results-PyPDEVS/", help="Path to the results directory")
    parser.add_argument("-f", '--filename', type=str, default="", help="Name of the output file")

    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("-g", "--gen-mode", choices=GEN_METHODS, default="uniform", help="Determines the spread of model creation.")

    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.")

    fieldnames = ["model_type", "depth", "width", "run",
                  "creation_time", "setup_time", "simulation_time", "generator_mode"]

    for model_type in args.model_structure:
        model_filename = f"{args.path}{model_type}"
        if args.filename:
            model_filename += f"_{args.filename}"
        else:
            model_filename += ".csv"
        with open(model_filename, "w", newline="") as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            writer.writeheader()

            for depth in range(d[1], d[0] + 1, d[1]):
                for width in range(w[1], w[0] + 1, w[1]):
                    for run in range(args.repetitions):
                        creation, setup, sim, output = run_devstone(
                            model_type, depth, width,
                            args.minimal, args.classic, args.dynamic, args.gen_mode
                        )
                        if creation is not None:
                            writer.writerow({
                                "model_type": model_type,
                                "depth": depth,
                                "width": width,
                                "run": run,
                                "creation_time": creation,
                                "setup_time": setup,
                                "simulation_time": sim,
                                "generator_mode": args.gen_mode
                            })
                            csvfile.flush()
        print(f"✅ Saved {model_type} results to {model_filename}")