from __future__ import annotations

import os
import sys

from exo import proc
from exo.platforms.x86 import *
from exo.stdlib.scheduling import *

from exo import proc
from exo.platforms.x86 import *
from exo.stdlib.scheduling import *

old_split = repeat(divide_loop)
old_unroll = repeat(unroll_loop)


def old_fission_after(proc, stmt_pattern, n_lifts=1):
    def find_stmts(p):
        return [c.after() for c in p.find_all(stmt_pattern)]

    return loop_hack(autofission, find_stmts)(proc, n_lifts)



# Hide output when running through exocc.
if __name__ != "__main__" and hasattr(os, "devnull"):
    sys.stdout = open(os.devnull, "w")



@proc
def simple_math_avx2_sched(
    n: size, x: R[n] @ DRAM, y: R[n] @ DRAM
):  # pragma: no cover
    for i in seq(0, n):
        x[i] = x[i] * y[i] * y[i]
p = rename(simple_math_avx2_sched , "math_avx2_scheduled")
p = old_split(p, "i", 8, ["io", "ii"], tail="cut_and_guard")
p = stage_assn(p, "x[_] = _ #0", "xyy")
p = autolift_alloc(p, "xyy: _", keep_dims=True)
p = set_memory(p, "xyy", AVX2)
p = old_fission_after(p, "xyy[_] = _")