#!/usr/bin/env python
# coding: utf-8

# In this tutorial we will show how to access and navigate the Iteration/Expression Tree (IET) rooted in an `Operator`.
# 
# 
# # Part I - Top Down
# 
# Let's start with a fairly trivial example. First of all, we disable all performance-related optimizations, to maximize the simplicity of the created IET as well as the readability of the generated code.

# In[1]:


from devito import configuration
configuration['opt'] = 'noop'
configuration['language'] = 'C'


# Then, we create a `TimeFunction` with 3 points in each of the space `Dimension`s _x_ and _y_.

# In[2]:


from devito import Grid, TimeFunction

grid = Grid(shape=(3, 3))
u = TimeFunction(name='u', grid=grid)


# We now create an `Operator` that increments by 1 all points in the computational domain.

# In[3]:


from devito import Eq, Operator

eq = Eq(u.forward, u+1)
op = Operator(eq)


# An `Operator` is an IET node that can generate, JIT-compile, and run low-level code (e.g., C). Just like all other types of IET nodes, it's got a number of metadata attached. For example, we can query an `Operator` to retrieve the input/output `Function`s.

# In[4]:


op.input


# In[5]:


op.writes


# If we print `op`, we can see how the generated code looks like.

# In[6]:


print(op)


# An `Operator` is the root of an IET that typically consists of several nested `Iteration`s and `Expression`s – two other fundamental IET node types. The user-provided SymPy equations are wrapped within `Expressions`. Loop nest embedding such expressions are constructed by suitably nesting `Iterations`.
# 
# The Devito compiler constructs the IET from a collection of `Cluster`s, which represent a higher-level intermediate representation (not covered in this tutorial).
# 
# The Devito compiler also attaches to the IET key computational properties, such as _sequential_, _parallel_, and _affine_, which are derived through data dependence analysis.
# 
# We can print the IET structure of an `Operator`, as well as the attached computational properties, using the utility function `pprint`.

# In[7]:


from devito.tools import pprint
pprint(op)


# In this example, `op` is represented as a `<Callable Kernel>`. Attached to it are metadata, such as `_headers` and `_includes`, as well as the `body`, which includes the children IET nodes. Here, the body is the concatenation of an `PointerCast` and a `List` object.
# 

# In[8]:


op._headers


# In[9]:


op._includes


# In[10]:


op.body


# We can explicitly traverse the `body` until we locate the user-provided `SymPy` equations.

# In[11]:


print(op.body.casts[0])  # Printing the PointerCast


# In[12]:


print(op.body.body[0])  # Printing the actual body


# Below we access the `Iteration` representing the time loop.

# In[13]:


t_iter = op.body.body[0].body[0]
t_iter


# We can for example inspect the `Iteration` to discover what its iteration bounds are.

# In[14]:


t_iter.limits


# And as we keep going down through the IET, we can eventually reach the `Expression` wrapping the user-provided SymPy equation.

# In[15]:


expr = t_iter.nodes[0].body[0].body[0].nodes[0].nodes[0].body[0]
expr.view


# Of course, there are mechanisms in place to, for example, find all `Expression`s in a given IET. The Devito compiler has a number of IET visitors, among which `FindNodes`, usable to retrieve all nodes of a particular type. So we easily 
# can get all `Expression`s within `op` as follows

# In[16]:


from devito.ir.iet import Expression, FindNodes
exprs = FindNodes(Expression).visit(op)
exprs[0].view