import pysu2
from mpi4py import MPI

def execute_su2():
    # Import mpi4py for parallel run
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    num_zones = 2
    # Initialize the corresponding driver of SU2, this includes solver preprocessing
    SU2Driver = pysu2.CMultizoneDriver(config_path, num_zones, comm) # type: ignore

    # Get all the markers defined on this rank and their associated indices.
    allMarkerIDs = SU2Driver.GetMarkerIndices()

    # Time loop is defined in Python so that we have acces to SU2 functionalities at each time step
    comm.Barrier()

    # Time iteration preprocessing
    SU2Driver.Preprocess(0)

    # Run one time-step (static: one simulation)
    SU2Driver.Run()

    # Update the solver for the next time iteration
    SU2Driver.Update()

    # Monitor the solver and output solution to file if required
    SU2Driver.Monitor(0)

    for izone in range(num_zones):
        SU2Driver.SelectZone(izone)
        # Get all the markers defined on this rank and their associated indices.
        allMarkerIDs = SU2Driver.GetMarkerIndices()
        print(allMarkerIDs)

    # Output the solution to file
    SU2Driver.Output(0)

    # Finalize the solver and exit cleanly
    SU2Driver.Finalize()