from psyclone.parse.algorithm import parse
_, inv_info = parse("shallow_alg.f90", api="gocean1.0")

from psyclone.psyGen import PSyFactory, TransInfo
psy = PSyFactory("gocean1.0", distributed_memory=False).create(inv_info)

from psyclone.psyir.backend.fortran import FortranWriter
fwriter = FortranWriter()
print(fwriter(psy.container))

print(psy.invokes.names)
sched = psy.invokes.get('invoke_0').schedule
print(sched.view())

trans_info = TransInfo()
print(trans_info.list)
fuse_trans = trans_info.get_trans_name('LoopFuseTrans')
omp_trans = trans_info.get_trans_name('GOceanOMPParallelLoopTrans')

# invoke0
# fuse all outer loops
fuse_trans.apply(sched[0], sched[1])
fuse_trans.apply(sched[0], sched[1])
fuse_trans.apply(sched[0], sched[1])
print(sched.view())

# fuse all inner loops
fuse_trans.apply(sched[0].loop_body[0], sched[0].loop_body[1])
fuse_trans.apply(sched[0].loop_body[0], sched[0].loop_body[1])
fuse_trans.apply(sched[0].loop_body[0], sched[0].loop_body[1])
print(sched.view())

omp_trans.apply(sched[0])
print(sched.view())

print(fwriter(psy.container))