from dace.frontend.fortran import fortran_parser
from dace.transformation.interstate import LoopToMap


def main():
    code = """
    PROGRAM foo
        IMPLICIT NONE
        REAL INP1(NBLOCKS, KLEV)
        INTEGER, PARAMETER  :: KLEV = 137
        INTEGER, PARAMETER  :: NBLOCKS = 8


        CALL foo_test_function(NBLOCKS, KLEV, INP1)

    END PROGRAM

    SUBROUTINE foo_test_function(NBLOCKS, KLEV, INP1)
        INTEGER, PARAMETER  :: KLEV = 137
        INTEGER, PARAMETER  :: NBLOCKS = 1
        REAL INP1(NBLOCKS, KLEV)

        DO JN=1,NBLOCKS
            DO JK=1,KLEV
                INP1(JN, JK) = (JN-1) * KLEV + (JK-1)
            ENDDO
        ENDDO
    END SUBROUTINE foo_test_function
    """

    sdfg = fortran_parser.create_sdfg_from_string(code, "test_loop_map_parallel")
    sdfg.simplify()
    sdfg.apply_transformations_repeated([LoopToMap])
    sdfg.simplify()


if __name__ == '__main__':
    main()