#!/usr/bin/env python

'''
Optimizations over an expression graph.

Optimization passes take as input an expression graph, and return a
(hopefully) simpler, equivalent graph.  This module defines the
pass infrastructure, the fusion passes and an optimization pass to
lower code to Parakeet.
'''
import operator
import math
import weakref

from collections import namedtuple

from . import local
from . import tiling
from .base import Expr, Val, AsArray, ListExpr, lazify, expr_like, ExprTrace
from .base import NotShapeable, CollectionExpr
from .filter import FilterExpr
from .local import LocalInput, LocalMapExpr, LocalMapLocationExpr, make_var
from .local import ParakeetExpr
from .map import MapExpr, Map2Expr
from .ndarray import NdArrayExpr
from .outer import OuterProductExpr
from .shuffle import ShuffleExpr
from .slice import SliceExpr
from .write_array import WriteArrayExpr
from .reduce import ReduceExpr, LocalReduceExpr
from ..dot import DotExpr

from ... import util
from ...util import Assert
from ...config import FLAGS, BoolFlag
from ...array.distarray import DistArray, LocalWrapper

try:
  import numexpr
except:
  numexpr = None

try:
  import parakeet
except:
  parakeet = None


_parakeet_blacklist = set()
_tiled_exprlist = {}
_not_idempotent_list = set()


def disable_parakeet(fn):
  "Disables parakeet optimization for this function."
  _parakeet_blacklist.add(fn)
  return fn


def not_idempotent(fn):
  "Force the result of not_idempotent ``fn`` to be evaluated at once."
  def wrapped(*args, **kw):
    result = fn(*args, **kw)
    if isinstance(result, Expr):
      result.needs_cache = True
      _not_idempotent_list.add(id(result))
    return result
  return wrapped

visited_expr = {'map_fusion': weakref.WeakValueDictionary(),
                'reduce_fusion': weakref.WeakValueDictionary(),
                'collapse_cached': weakref.WeakValueDictionary(),
                'parakeet_gen': weakref.WeakValueDictionary(),
                'rotate_slice': weakref.WeakValueDictionary(),
                'auto_tiling': weakref.WeakValueDictionary()
                }


class OptimizePass(object):
  def __init__(self):
    self.visited = visited_expr[self.name]
    #self.visited = {}

  def visit(self, op):
    if not isinstance(op, Expr):
      return op

    if op.expr_id in self.visited:
      return self.visited[op.expr_id]

    #assert not op in self.visited, 'Infinite recursion during optimization %s' % op
    #self.visited.add(op)

    util.log_debug('VISIT %s: %s', op.typename(), hash(op))
    opt_op = None
    if hasattr(self, 'visit_default'):
      opt_op = self.visit_default(op)
    elif hasattr(self, 'visit_%s' % op.typename()):
      opt_op = getattr(self, 'visit_%s' % op.typename())(op)
    else:
      opt_op = op.visit(self)

    self.visited[opt_op.expr_id] = opt_op
    return opt_op


def fusable(v):
  return isinstance(v, (MapExpr,
                        ReduceExpr,
                        ShuffleExpr,
                        NdArrayExpr,
                        SliceExpr,
                        FilterExpr,
                        Val,
                        AsArray,
                        WriteArrayExpr))


def merge_var(children, child_to_var, k, v):
  """Add a new expression with key ``k`` to the ``children`` dictionary.

  If ``k`` is already in the dictionary, than ``v`` must be equal to
  the current value stored there.
  """
  try:
    i = child_to_var.index(k)
    Assert.eq(v, children[i])
  except ValueError:
    children.append(v)
    child_to_var.append(k)


class MapMapFusion(OptimizePass):
  '''Fold sequences of Map operations together.

  map(f, map(g, map(h, x))) -> map(f . g . h, x)
  '''
  name = 'map_fusion'

  def visit_MapExpr(self, expr):
    Assert.iterable(expr.children)
    #util.log_info('VISIT %d', id(op))
    #util.log_info('VISIT %s', op.children)
    map_children = self.visit(expr.children)
    all_maps = True
    Assert.isinstance(map_children, ListExpr)
    for k, v in zip(expr.child_to_var, map_children):
      if not fusable(v):
        util.log_debug('Skipping fusion: (%s -> %s)' % (k, type(v)))
        all_maps = False
        break

    if not all_maps or isinstance(expr.op, local.ParakeetExpr) or \
       id(expr) in _not_idempotent_list:
      return expr.visit(self)

    #util.log_info('Original: %s', expr.op)
    children = []
    child_to_var = []
    combined_op = expr.op.__class__(fn=expr.op.fn, kw=expr.op.kw,
                                    pretty_fn=expr.op.pretty_fn)

    trace = ExprTrace()

    for child_expr in map_children:
      trace.fuse(child_expr.stack_trace)

      if isinstance(child_expr, MapExpr):
        for k, v in zip(child_expr.child_to_var, child_expr.children):
          merge_var(children, child_to_var, k, v)

        combined_op.add_dep(child_expr.op)
        util.log_debug('Fusion: %s <- %s', expr.expr_id, child_expr.expr_id)
      else:
        children.append(child_expr)
        key = make_var()
        combined_op.add_dep(LocalInput(idx=key))
        child_to_var.append(key)

    if isinstance(combined_op, LocalMapLocationExpr):
      combined_op.add_dep(LocalInput(idx='extent'))

    return expr_like(expr,
                     children=ListExpr(vals=children),
                     child_to_var=child_to_var,
                     op=combined_op,
                     trace=trace)


class ReduceMapFusion(OptimizePass):
  '''Fuse reduce(f, map(g, X)) -> reduce(f . g, X)'''
  name = 'reduce_fusion'

  def visit_ReduceExpr(self, expr):
    Assert.isinstance(expr.children, ListExpr)
    old_children = self.visit(expr.children)

    for v in old_children:
      if not isinstance(v, (MapExpr, ParakeetExpr)) or id(v) in _not_idempotent_list:
        return expr.visit(self)

    combined_op = LocalReduceExpr(fn=expr.op.fn,
                                  kw=expr.op.kw,
                                  deps=[expr.op.deps[0]])

    new_children = []
    new_child_to_var = []
    trace = ExprTrace()
    for i in range(len(old_children)):
      name = expr.child_to_var[i]
      child_expr = old_children[i]
      for j in range(len(child_expr.children)):
        k = child_expr.child_to_var[j]
        v = child_expr.children[j]
        merge_var(new_children, new_child_to_var, k, v)
      combined_op.add_dep(child_expr.op)
      trace.fuse(child_expr.stack_trace)

    return expr_like(expr,
                     children=ListExpr(vals=new_children),
                     child_to_var=new_child_to_var,
                     axis=expr.axis,
                     dtype_fn=expr.dtype_fn,
                     accumulate_fn=expr.accumulate_fn,
                     op=combined_op,
                     tile_hint=expr.tile_hint,
                     trace=trace)


class CollapsedCachedExpressions(OptimizePass):
  '''Replace expressions which have already been evaluated
  with a simple value expression.

  This results in simpler local expressions when evaluating
  iterative programs.
  '''

  name = 'collapse_cached'

  def visit_default(self, expr):
    #util.log_info('Visit: %s, %s', expr.expr_id, expr.cache)
    cache = expr.cache()
    if cache is not None:
      util.log_info('Collapsing %s %s', expr.expr_id, expr.typename())
      return lazify(cache)
    else:
      return expr.visit(self)


def _find_modules(op):
  '''Find any modules referenced by the given `LocalOp` or its dependencies'''
  modules = set()
  if isinstance(op, local.FnCallExpr):
    if hasattr(op.fn, '__module__'):
      modules.add(op.fn.__module__)
    elif hasattr(op.fn, '__class__'):
      modules.add(op.fn.__class__.__module__)

  for d in op.deps:
    modules.union(_find_modules(d))

  return modules


def _parakeet_codegen(op):
  '''Given a local operation, generate an equivalent parakeet function definition.'''

  def _codegen(op):
    if isinstance(op, local.FnCallExpr):
      if util.is_lambda(op.fn) and not op.pretty_fn:
        raise local.CodegenException('Cannot codegen through a lambda expression: %s' % op.fn)

      if op.fn in _parakeet_blacklist:
        raise local.CodegenException('Blacklisted %s', op.fn)

      name = op.fn_name()

      arg_str = ','.join([_codegen(v) for v in op.deps])
      kw_str = ','.join(['%s=%s' % (k, repr(v)) for k, v in op.kw.iteritems()])
      if arg_str:
        kw_str = ',' + kw_str

      return '%s(%s %s)' % (name, arg_str, kw_str)
    elif isinstance(op, local.LocalInput):
      return op.idx
    else:
      raise local.CodegenException('Cannot codegen for %s' % type(op))

  if isinstance(op, ParakeetExpr):
    return op.source

  op_code = _codegen(op)

  module_prelude = ['import parakeet',
                    'import spartan.expr',
                    'import numpy',
                    'from spartan.expr import mathlib',
                    'from spartan import util',
                    ]

  for mod in _find_modules(op):
    module_prelude.append('import %s' % mod)

  fn_prelude = '''
@util.synchronized
@parakeet.jit
'''

  fn = '\n'.join(module_prelude)
  fn = fn + fn_prelude
  fn = fn + 'def _jit_fn'
  fn = fn + '(%s):\n  ' % ','.join(op.input_names())
  fn = fn + 'return ' + op_code
  fn = fn + '\n\n'

  # verify we can compile before proceeding
  local.compile_parakeet_source(fn)
  return fn


class ParakeetGeneration(OptimizePass):
  '''
  Replace local map/reduce operations with an equivalent
  parakeet function definition.
  '''

  name = 'parakeet_gen'
  after = [MapMapFusion, ReduceMapFusion]

  def _should_run_parakeet(self, expr):
    has_fncallexpr = False
    for dep in expr.op.deps:
      if isinstance(dep, local.FnCallExpr):
        has_fncallexpr = True
        break

    # If all deps are not FnCallExprs and len(deps) is a small number(<=2 for now)
    # it is not worth to do code generation.
    if not has_fncallexpr and len(expr.op.deps) <= 2:
      return False
    else:
      return True

  def visit_MapExpr(self, expr):
    # if we've already converted this to parakeet, stop now
    if isinstance(expr.op, local.ParakeetExpr):
      return expr.visit(self)

    if not self._should_run_parakeet(expr):
      return expr.visit(self)

    try:
      source = _parakeet_codegen(expr.op)

      # Recursively check if any children nodes can be optimized with Parakeet
      new_children = []
      for child in expr.children:
        new_children.append(self.visit(child))

      parakeet_expr = expr_like(expr,
                                op=local.ParakeetExpr(source=source, deps=expr.op.deps),
                                children=ListExpr(vals=new_children),
                                child_to_var=expr.child_to_var)

      if id(expr) in _not_idempotent_list: _not_idempotent_list.add(id(parakeet_expr))

      return parakeet_expr
    except local.CodegenException:
      util.log_info('Failed to convert to parakeet.')
      return expr.visit(self)

#   Parakeet doesn't support taking the current Extent object as
#   an argument, which prevents enabling it for reductions.
#
#   def visit_ReduceExpr(self, expr):
#     # if we've already converted this to parakeet, stop now
#     if isinstance(expr.op, local.ParakeetExpr):
#       return expr.visit(self)
#
#     try:
#       source = codegen(expr.op)
#       return expr_like(expr,
#                        children=expr.children,
#                        axis=expr.axis,
#                        dtype_fn=expr.dtype_fn,
#                        accumulate_fn=expr.accumulate_fn,
#                        op=local.ParakeetExpr(source=source,  deps=expr.op.deps))
#     except local.CodegenException:
#       util.log_info('Failed to convert to parakeet.')
#       return expr.visit(self)


class RotateSlice(OptimizePass):
  '''
  This pass rotates slice operations to the bottom of the expression graph.

  By moving slices down, we open up opportunities for performing more
  fusion operations.

  The actual operation performed is:

  (a + b)[slice] -> broadcast(a, shape)[slice] + broadcast(b, shape)[slice]
  '''
  name = 'rotate_slice'
  rotated = {}

  def visit_SliceExpr(self, slice_expr):
    'Rotate this slice with a child map expression.'
    map_expr = slice_expr.src

    if not isinstance(map_expr, MapExpr):
      return slice_expr.visit(self)

    try:
      map_shape = map_expr.compute_shape()
    except NotShapeable:
      return slice_expr.visit(self)

    Assert.iterable(map_expr.children)
    map_children = self.visit(map_expr.children)

    children = []
    for child_expr in map_children:
      child = SliceExpr(src=child_expr,
                        idx=slice_expr.idx,
                        broadcast_to=map_shape)
      if isinstance(child_expr, MapExpr):
        # If the child is a Map, RotateSlice should traverse it again. Some
        # nodes may be traversed several times. For current implementation,
        # traversing a node several times dones't cause any problem for
        # RotationSlice.
        if self.visited.get(child_expr, None) is not None:
          del self.visited[child_expr]
        child = self.visit(child)

      children.append(child)

    if self.rotated.get(map_expr, None):
      # If this Map has been rotated, we should create a new MapExpr.
      # An example is :
      #    c = a + b
      #    d = c[1:50]
      #    e = c[2:51]
      # In this example, RotateSlice first rotate c[2:51] to a and b.
      # When RotateSlice rotates c[1:50] to a and b, it should create
      # a new Map expr(+) for a[1:50] and b[1:50]
      return MapExpr(children=ListExpr(vals=children),
                     op=map_expr.op,
                     child_to_var=map_expr.child_to_var)
    else:
      self.rotated[map_expr] = True
      return expr_like(map_expr,
                       op=map_expr.op,
                       children=ListExpr(vals=children),
                       child_to_var=map_expr.child_to_var,
                       trace=map_expr.stack_trace)


class AutomaticTiling(OptimizePass):
  '''
  Automatically partition all the arrays.
  We build a min cost max flow DAG for the expr graph, where the cost is the communication cost
  through network. We estimate the cost according to the behavior of each expr. For _builtin_ expr
  and simple map and reduce expr, we can easily estimate the cost. However, for the user defined shuffle
  expr, we can only guess the cost or let users define the cost for us.

  All Exprs:
    [Val, AsArray, DistArray]: Already partitioned array
    [NdArrayExpr]: new array needs to be partitioned
    CollectionExpr([DictExpr, ListExpr, TupleExpr]): self.vals

    [MapExpr]: self.children
    [ReduceExpr]: self.children
    [ShuffleExpr]: self.array, self.fn_kw
    [DotExpr]: self.matrix_a, self.matrix_b

    [SliceExpr, FilterExpr, CheckpointExpr, TileOpExpr]: self.src or self.array
    [WriteArrayExpr]: self.array, self.data

    [TransposeExpr, ReshapeExpr]: self.array
  '''

  name = 'auto_tiling'
  node_type = namedtuple('node_type', ['expr', 'tiling', 'children', 'parents'])
  num_node_per_group = 4
  inited = False
  cost_model = {'map': {(0, 0): 0, (0, 1): 1, (0, 2): 1,
                        (0, 3): 1, (0, 4): 1, (0, 5): 2,
                        (1, 0): 1, (1, 1): 0, (1, 2): 1,
                        (1, 3): 1, (1, 4): 2, (1, 5): 1,
                        (2, 0): 1, (2, 1): 1, (2, 2): 0,
                        (2, 3): 2, (2, 4): 1, (2, 5): 1,
                        (3, 0): 0, (3, 1): 0, (3, 2): 1,
                        (3, 3): 0, (3, 4): 1, (3, 5): 1,
                        (4, 0): 0, (4, 1): 1, (4, 2): 0,
                        (4, 3): 1, (4, 4): 0, (4, 5): 1,
                        (5, 0): 1, (5, 1): 0, (5, 2): 0,
                        (5, 3): 1, (5, 4): 1, (5, 5): 0},
                'map2': {(0, 0): 0, (0, 1): 1, (0, 2): 1, (0, -1): 1,
                         (1, 0): 1, (1, 1): 0, (1, 2): 1, (1, -1): 1,
                         (2, 0): 1, (2, 1): 1, (2, 2): 0, (2, -1): 1,
                         (3, 0): 0, (3, 1): 0, (3, 2): 1, (3, -1): 1,
                         (4, 0): 0, (4, 1): 1, (4, 2): 0, (4, -1): 1,
                         (5, 0): 1, (5, 1): 0, (5, 2): 0, (5, -1): 1}
                }

  def init(self, expr):
    self.cur_node_id = 1
    self.edges = {}
    self.nodes = {0: self.node_type([], -1, [], [])}
    self.expr_to_nodes = {}
    self.split_nodes = {}
    self.groups = []
    self.init_expr = id(expr)
    self.inited = True

    self.tiled_exprlist = _tiled_exprlist
    #self.tiled_exprlist = {}

  def add_edge(self, edge_from, edge_to, edge_cost=0):
    #util.log_warn('add_edge:%d %d cost:%d', edge_from, edge_to, edge_cost)
    if (edge_from, edge_to) not in self.edges:
      self.nodes[edge_from].parents.append(edge_to)
      self.nodes[edge_to].children.append(edge_from)
    self.edges[(edge_from, edge_to)] = edge_cost

  def remove_edge(self, edge_from, edge_to):
    del self.edges[(edge_from, edge_to)]
    self.nodes[edge_from].parents.remove(edge_to)
    self.nodes[edge_to].children.remove(edge_from)

  def add_split_nodes(self, nodes):
    group_id = len(self.groups)
    for node in nodes:
      self.split_nodes[node] = group_id
    self.groups.append(tuple(nodes))
    #util.log_warn('add_split_nodes:%s', nodes)

  def visit_children(self, children, except_child=None):
    child_ids = []
    for child in children:
      if isinstance(child, (Expr, DistArray)) and id(child) != id(except_child):
        child_ids.extend(self.visit_default(child))
    return child_ids

  def visit_NdArrayExpr(self, expr):
    # new array need to be partitioned
    if len(expr.shape) > 1 and expr.shape[1] > 1:
      new_nodes = []

      self.nodes[self.cur_node_id] = self.node_type([expr], 0, [], [])
      new_nodes.append(self.cur_node_id)
      self.add_edge(0, self.cur_node_id, 0)
      self.cur_node_id += 1

      self.nodes[self.cur_node_id] = self.node_type([expr], 1, [], [])
      new_nodes.append(self.cur_node_id)
      self.add_edge(0, self.cur_node_id, 0)
      self.cur_node_id += 1

      self.nodes[self.cur_node_id] = self.node_type([expr], 2, [], [])
      new_nodes.append(self.cur_node_id)
      self.add_edge(0, self.cur_node_id, 0)
      self.cur_node_id += 1

      cost = reduce(operator.mul, expr.shape, 1)
      for i in range(3, self.num_node_per_group):
        self.nodes[self.cur_node_id] = self.node_type([expr], i, [], [])
        new_nodes.append(self.cur_node_id)
        self.add_edge(0, self.cur_node_id, cost)
        self.cur_node_id += 1

      self.add_split_nodes(new_nodes)
      return new_nodes
    else:
      self.nodes[self.cur_node_id] = self.node_type([expr], 0, [], [])
      self.add_edge(0, self.cur_node_id, 0)
      self.cur_node_id += 1
      return [self.cur_node_id - 1]

  def visit_MapExpr(self, expr):
    largest = max(expr.children.vals, key=lambda v: reduce(operator.mul, v.shape, 1))

    child_ids = self.visit_children([largest])
    other_child_ids = self.visit_children(expr.children.vals, largest)
    kw_ids = self.visit_children(expr.op.kw['fn_kw'].itervalues()) if 'fn_kw' in expr.op.kw else []

    # one input map, reuse child expr
    if len(other_child_ids) == 0:
      for child_id in child_ids: self.nodes[child_id].expr.append(expr)
      return child_ids

    if child_ids[0] in self.split_nodes:
      tiling_types = range(self.num_node_per_group)
      self.add_split_nodes(range(self.cur_node_id, self.cur_node_id + self.num_node_per_group))
      expr_node_ids = range(self.cur_node_id, self.cur_node_id + self.num_node_per_group)
    else:
      tiling_types = (self.nodes[child_ids[0]].tiling,)
      expr_node_ids = [self.cur_node_id]

    for (tiling_type, map_child_id) in zip(tiling_types, child_ids):
      self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
      self.add_edge(map_child_id, self.cur_node_id, 0)

      for child_id in other_child_ids:
        child = self.nodes[child_id]
        e_cost = self.cost_model['map'][(child.tiling, tiling_type)] * reduce(operator.mul, child.expr[0].shape, 1)
        self.add_edge(child_id, self.cur_node_id, e_cost)

      for child_id in kw_ids:
        self.add_edge(child_id, self.cur_node_id, reduce(operator.mul, self.nodes[child_id].expr[0].shape, 1))
      self.cur_node_id += 1

    return expr_node_ids

  def visit_ReduceExpr(self, expr):
    child_ids = self.visit_children(expr.children.vals)
    cost = reduce(operator.mul, expr.shape, 1)
    self.nodes[self.cur_node_id] = self.node_type([expr], 0, [], [])
    for child_id in child_ids:
      child = self.nodes[child_id]
      e_cost = 0 if expr.axis is None or child.tiling == 3 or (1-expr.axis) == (child.tiling % 4) else cost
      self.add_edge(child_id, self.cur_node_id, e_cost)
    self.cur_node_id += 1
    return [self.cur_node_id - 1]

  def visit_Map2Expr(self, expr):
    child_id_groups = []
    for array in expr.arrays:
      child_id_groups.append(self.visit_children([array]))

    copy_nodes = []
    for axis, child_ids in zip(expr.axes, child_id_groups):
      if isinstance(axis, tuple): axis = 2
      if axis is None: axis = -1
      self.nodes[self.cur_node_id] = self.node_type([expr], axis, [], [])

      cost = reduce(operator.mul, self.nodes[child_ids[0]].expr[0].shape, 1)
      for child_id in child_ids:
        child = self.nodes[child_id]
        e_cost = self.cost_model['map2'][(child.tiling, axis)] * cost
        self.add_edge(child_id, self.cur_node_id, e_cost)

      if len(copy_nodes) > 0: self.add_edge(self.cur_node_id, copy_nodes[0], 0)
      copy_nodes.append(self.cur_node_id)
      self.cur_node_id += 1

    e_cost = reduce(operator.mul, expr.shape, 1)
    inter_node_id = copy_nodes[0]
    if len(expr.shape) > 1 and expr.shape[1] > 1:
      child_ids = []
      for tiling_type in range(self.num_node_per_group):
        self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
        self.add_edge(inter_node_id, self.cur_node_id, (tiling_type / 3 + 1) * e_cost)
        child_ids.append(self.cur_node_id)
        self.cur_node_id += 1
      self.add_split_nodes(child_ids)
      return child_ids
    else:
      self.nodes[self.cur_node_id] = self.node_type([expr], 0, [], [])
      self.add_edge(inter_node_id, self.cur_node_id, e_cost)
      self.cur_node_id += 1
      return [self.cur_node_id - 1]

  def visit_OuterProductExpr(self, expr):
    child_id_groups = []
    for array in expr.arrays:
      child_id_groups.append(self.visit_children([array]))

    copy_nodes = []
    for axis, child_ids in zip(expr.axes, child_id_groups):
      self.nodes[self.cur_node_id] = self.node_type([expr], axis, [], [])

      cost = reduce(operator.mul, self.nodes[child_ids[0]].expr[0].shape, 1)
      for child_id in child_ids:
        child = self.nodes[child_id]
        e_cost = 0 if axis is not None and (axis == (child.tiling % 4) or child.tiling == 3) else cost
        self.add_edge(child_id, self.cur_node_id, e_cost)

      if len(copy_nodes) > 0: self.add_edge(self.cur_node_id, copy_nodes[0], cost)
      copy_nodes.append(self.cur_node_id)
      self.cur_node_id += 1

    e_cost = reduce(operator.mul, expr.shape, 1)
    inter_node_id = copy_nodes[0]
    if len(expr.shape) > 1 and expr.shape[1] > 1:
      child_ids = []
      for tiling_type in range(self.num_node_per_group):
        self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
        self.add_edge(inter_node_id, self.cur_node_id, (tiling_type / 3 + 1) * e_cost)
        child_ids.append(self.cur_node_id)
        self.cur_node_id += 1
      self.add_split_nodes(child_ids)
      return child_ids
    else:
      self.nodes[self.cur_node_id] = self.node_type([expr], 0, [], [])
      self.add_edge(inter_node_id, self.cur_node_id, e_cost)
      self.cur_node_id += 1
      return [self.cur_node_id - 1]

  def visit_ShuffleExpr(self, expr):
    for child in expr.fn_kw.itervalues():
      if isinstance(child, (Expr, DistArray)) and hash(child) not in expr.cost_hint:
        cost = reduce(operator.mul, child.shape, 1)
        expr.cost_hint[hash(child)] = {'00': cost, '01': cost, '10': cost, '11': cost}
    if expr.target is not None and hash(expr.target) not in expr.cost_hint:
      cost = reduce(operator.mul, expr.target.shape, 1)
      expr.cost_hint[hash(expr.target)] = {'00': cost, '01': cost, '10': cost, '11': cost}

    child_ids = self.visit_children([expr.array])
    other_child_ids = self.visit_children(expr.fn_kw.itervalues())

    # calc the copy cost
    if len(other_child_ids) == 0:
      expr_node_ids = child_ids
    else:
      if child_ids[0] in self.split_nodes:
        tiling_types = (0, 1)
        self.add_split_nodes(self.cur_node_id, self.cur_node_id + 1)
        expr_node_ids = [self.cur_node_id, self.cur_node_id + 1]
      else:
        tiling_types = (self.nodes[child_ids[0]].tiling,)
        expr_node_ids = [self.cur_node_id]

      for i in xrange(len(tiling_types)):
        tiling_type = tiling_types[i]
        self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
        self.add_edge(child_ids[i], self.cur_node_id, 0)

        for child in expr.fn_kw.itervalues():
          if isinstance(child, (Expr, DistArray)):
            for child_id in self.expr_to_nodes[hash(child)]:
              self.add_edge(child_id, self.cur_node_id, expr.cost_hint[hash(child)]['%d%d' % (self.nodes[child_id].tiling, tiling_type)])
        self.cur_node_id += 1

    # calculate update cost
    if expr.target is not None:
      other_child_ids = expr_node_ids
      child_ids = self.visit_children([expr.target])

      if child_ids[0] in self.split_nodes:
        tiling_types = (0, 1)
        self.add_split_nodes(self.cur_node_id, self.cur_node_id + 1)
        expr_node_ids = [self.cur_node_id, self.cur_node_id + 1]
      else:
        tiling_types = (self.nodes[child_ids[0]].tiling,)
        expr_node_ids = [self.cur_node_id]

      for i in xrange(len(tiling_types)):
        tiling_type = tiling_types[i]
        self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
        self.add_edge(child_ids[i], self.cur_node_id, 0)

        for child_id in other_child_ids:
          self.add_edge(child_id, self.cur_node_id, expr.cost_hint[hash(expr.target)]['%d%d' % (self.nodes[child_id].tiling, tiling_type)])
        self.cur_node_id += 1
      return expr_node_ids

    if expr.shape == expr.array.shape or len(other_child_ids) > 0:
      if len(other_child_ids) == 0:
        for child_id in expr_node_ids: self.nodes[child_id].expr.append(expr)
      return expr_node_ids

    # fix result tiling
    if len(expr.shape) <= 1 or 1 in expr.shape[:2]:
      tiling_type = 0 if len(expr.shape) <= 1 else 1 - expr.shape.index(1)
      self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
      for child_id in expr_node_ids: self.add_edge(child_id, self.cur_node_id, 0)
      self.cur_node_id += 1
      return [self.cur_node_id - 1]
    else:
      child_ids = expr_node_ids
      expr_node_ids = []
      for child_id in child_ids:
        self.nodes[self.cur_node_id] = self.node_type([expr], self.nodes[child_id].tiling, [], [])
        expr_node_ids.append(self.cur_node_id)
        self.add_edge(child_id, self.cur_node_id, 0)
        self.cur_node_id += 1
      if len(expr_node_ids) > 1: self.add_split_nodes(*expr_node_ids)
    return expr_node_ids

  def visit_DotExpr(self, expr):
    child_ids = self.visit_children([expr.matrix_a])
    other_child_ids = self.visit_children([expr.matrix_b])

    # calc copy cost
    if len(other_child_ids) != 0:
      if child_ids[0] in self.split_nodes:
        tiling_types = range(self.num_node_per_group)
        self.add_split_nodes(range(self.cur_node_id, self.cur_node_id + self.num_node_per_group))
        expr_node_ids = range(self.cur_node_id, self.cur_node_id + self.num_node_per_group)
      else:
        tiling_types = (self.nodes[child_ids[0]].tiling,)
        expr_node_ids = [self.cur_node_id]

      for i in xrange(len(tiling_types)):
        tiling_type = tiling_types[i]
        self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
        self.add_edge(child_ids[i], self.cur_node_id, 0)

        for child_id in other_child_ids:
          child = self.nodes[child_id]
          e_cost = 0 if child.tiling in (0, 3) and tiling_type in (1, 2) else reduce(operator.mul, child.expr[0].shape, 1)
          self.add_edge(child_id, self.cur_node_id, e_cost)
        self.cur_node_id += 1

      child_ids = expr_node_ids

    # calc update cost
    if len(expr.shape) == 1 or expr.shape[1] == 1:
      tiling_types = (0,)
      expr_node_ids = [self.cur_node_id]
    else:
      tiling_types = range(self.num_node_per_group)
      self.add_split_nodes(range(self.cur_node_id, self.cur_node_id + self.num_node_per_group))
      expr_node_ids = range(self.cur_node_id, self.cur_node_id + self.num_node_per_group)

    cost = reduce(operator.mul, expr.shape, 1)
    for tiling_type in tiling_types:
      self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
      for child_id in child_ids:
        e_cost = 0 if self.nodes[child_id].tiling in (0, 3) and tiling_type == 0 else cost
        self.add_edge(child_id, self.cur_node_id, e_cost)
      self.cur_node_id += 1
    return expr_node_ids

  def visit_WriteArrayExpr(self, expr):
    child_ids = self.visit_children([expr.array])
    if isinstance(expr.data, (Expr, DistArray)):
      data_child_ids = self.visit_children([expr.data])
      if child_ids[0] in self.split_nodes:
        tiling_types = range(self.num_node_per_group)
        self.add_split_nodes(range(self.cur_node_id, self.cur_node_id + self.num_node_per_group))
        expr_node_ids = range(self.cur_node_id, self.cur_node_id + self.num_node_per_group)
      else:
        tiling_types = (self.nodes[child_ids[0]].tiling,)
        expr_node_ids = [self.cur_node_id]

      for (tiling_type, map_child_id) in zip(tiling_types, child_ids):
        self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
        self.add_edge(map_child_id, self.cur_node_id, 0)

        for child_id in data_child_ids:
          child = self.nodes[child_id]
          e_cost = self.cost_model['map'][(child.tiling, tiling_type)] * reduce(operator.mul, child.expr[0].shape, 1)
          self.add_edge(child_id, self.cur_node_id, e_cost)

        self.cur_node_id += 1
      return expr_node_ids
    else:
      for child_id in child_ids: self.nodes[child_id].expr.append(expr)
      return child_ids

  def visit_aligned_nodes(self, expr, reverse_cost=False):
    array = expr.src if hasattr(expr, 'src') else expr.array
    child_ids = self.visit_children([array])
    if child_ids[0] in self.split_nodes:
      tiling_types = range(self.num_node_per_group)
      self.add_split_nodes(range(self.cur_node_id, self.cur_node_id + self.num_node_per_group))
      expr_node_ids = range(self.cur_node_id, self.cur_node_id + self.num_node_per_group)

      for tiling_type in tiling_types:
        self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
        orig_tiling = -1
        if tiling_type in (2, 3):
          orig_tiling = tiling_type
        elif tiling_type < 2:
          orig_tiling = reverse_cost ^ tiling_type
        else:
          orig_tiling = (reverse_cost ^ (tiling_type % 4)) + 4
        self.add_edge(child_ids[orig_tiling], self.cur_node_id, 0)
        self.cur_node_id += 1
    else:
      child_tiling = self.nodes[child_ids[0]].tiling
      if child_tiling in (2, 3):
        tiling_type = child_tiling
      elif child_tiling < 2:
        tiling_type = reverse_cost ^ child_tiling
      else:
        tiling_type = (reverse_cost ^ (child_tiling % 4)) + 4
      self.nodes[self.cur_node_id] = self.node_type([expr], tiling_type, [], [])
      self.add_edge(child_ids[0], self.cur_node_id, 0)
      expr_node_ids = [self.cur_node_id]
      self.cur_node_id += 1

    return expr_node_ids

  def visit_TransposeExpr(self, expr):
    return self.visit_aligned_nodes(expr, reverse_cost=True)

  def visit_ReshapeExpr(self, expr):
    return self.visit_aligned_nodes(expr, reverse_cost=True)

  def visit_SliceExpr(self, expr):
    return self.visit_aligned_nodes(expr)

  def visit_FilterExpr(self, expr):
    child_ids = self.visit_children([expr.src])
    for child_id in child_ids: self.nodes[child_id].expr.append(expr)
    return child_ids

  def visit_CheckpointExpr(self, expr):
    child_ids = self.visit_children([expr.src])
    for child_id in child_ids: self.nodes[child_id].expr.append(expr)
    return child_ids

  def visit_TileOpExpr(self, expr):
    child_ids = self.visit_children([expr.array])
    for child_id in child_ids: self.nodes[child_id].expr.append(expr)
    return child_ids

  def generate_edges(self, s=0):
    edges = []
    self.nodes[s].parents.sort(key=lambda x: reduce(operator.mul, self.nodes[x].expr[0].shape, 1))
    for parent_id in self.nodes[s].parents:
      edges.append((s, parent_id, self.edges[(s, parent_id)]))
      if parent_id not in self.visited_nodes:
        edges.extend(self.generate_edges(parent_id))
        self.visited_nodes.add(parent_id)
    return edges

  def tile_expr(self, expr, tiling):
    if isinstance(expr, (NdArrayExpr, ReduceExpr, Map2Expr, OuterProductExpr)) and len(expr.shape) > 0:
      expr.tile_hint = list(expr.shape)
      if tiling >= 3:  # duplicate tiling
        print 'dup_tiling', tiling
      elif tiling == 2 and len(expr.shape) > 1:  # block tiling
        expr.tile_hint[0] = int(math.ceil(float(expr.tile_hint[0]) / math.sqrt(FLAGS.num_workers)))
        expr.tile_hint[1] = int(math.ceil(float(expr.tile_hint[1]) / math.sqrt(FLAGS.num_workers)))
        print 'block_tiling', expr.tile_hint, expr.expr_id
      elif len(expr.shape) > tiling:
        expr.tile_hint[tiling] = int(math.ceil(float(expr.tile_hint[tiling]) / FLAGS.num_workers))

  def calc_tiling(self, expr):
    # add T node for graph
    self.nodes[self.cur_node_id] = self.node_type([expr], -1, [], [])
    self.add_edge(self.cur_node_id - 1, self.cur_node_id, 0)
    if self.cur_node_id - 1 in self.split_nodes:
      for i in range(2, self.num_node_per_group + 1):
        self.add_edge(self.cur_node_id - i, self.cur_node_id, 0)
    self.cur_node_id += 1

    # compute best tiling for all exprs
    self.visited_nodes = set()
    edges = self.generate_edges()
    util.log_debug('num of groups %d' % len(self.groups))

    nodes = []
    if FLAGS.tiling_alg == 'maxedge':
      nodes = tiling.maxedge_tiling(self.cur_node_id - 1, edges, self.groups)
      #print 'maxedge', nodes
    elif FLAGS.tiling_alg == 'mincost':
      nodes = tiling.mincost_tiling(self.cur_node_id - 1, edges, self.groups)
      #print 'mincost', nodes
    elif FLAGS.tiling_alg == 'best':
      nodes = tiling.best_tiling(self.cur_node_id - 1, edges, self.groups)
      #print 'best', nodes
    elif FLAGS.tiling_alg == 'worse':
      nodes = tiling.worse_tiling(self.cur_node_id - 1, edges, self.groups)
      #print 'worse', nodes
    else:
      assert 'Unknown tiling algorithm'

    # give expr the best tiling hint
    for node_id in nodes:
      node = self.nodes[node_id]
      if node.tiling < 0 or node.tiling is None: continue
      for cur_expr in node.expr:
        self.tiled_exprlist[hash(cur_expr)] = node.tiling
        self.tile_expr(cur_expr, node.tiling)

    self.inited = False
    return expr

  def tile_cached_expr(self, expr):
    if not isinstance(expr, Expr) or isinstance(expr, (Val, AsArray)): return

    self.tile_expr(expr, self.tiled_exprlist[hash(expr)])

    if hasattr(expr, 'array'): self.tile_cached_expr(expr.array)
    if hasattr(expr, 'src'): self.tile_cached_expr(expr.src)
    if hasattr(expr, 'target'): self.tile_cached_expr(expr.target)
    if isinstance(expr, DotExpr):
      self.tile_cached_expr(expr.matrix_a)
      self.tile_cached_expr(expr.matrix_b)

    if hasattr(expr, 'children'):
      for child in expr.children.vals:
        self.tile_cached_expr(child)
    if hasattr(expr, 'fn_kw') and expr.fn_kw is not None:
      for child in expr.fn_kw.itervalues():
        self.tile_cached_expr(child)
    if hasattr(expr, 'op') and 'fn_kw' in expr.op.kw:
      for child in expr.op.kw['fn_kw'].itervalues():
        self.tile_cached_expr(child)

  def visit_default(self, expr):
    if not self.inited: self.init(expr)
    if hash(expr) in self.tiled_exprlist:
      self.tile_cached_expr(expr)
      tiling = self.tiled_exprlist[hash(expr)]
      self.nodes[self.cur_node_id] = self.node_type([expr], tiling, [], [])
      expr_node_ids = [self.cur_node_id]
      self.add_edge(0, self.cur_node_id, 0)
      self.cur_node_id += 1

    elif hash(expr) in self.expr_to_nodes:
      # cached expr
      expr_node_ids = self.expr_to_nodes[hash(expr)]
      for node_id in expr_node_ids:
        node = self.nodes[node_id]
        node.expr.append(expr)

    elif isinstance(expr, DistArray) or (isinstance(expr, (Val, AsArray)) and isinstance(expr.val, DistArray)):
      # already partitioned array
      array = expr if isinstance(expr, DistArray) else expr.val

      if isinstance(array, LocalWrapper):
        expr_node_ids = []
      else:
        tile_shape = array.tile_shape()
        tiling = 2
        for i in range(len(tile_shape)):
          if tile_shape[i] == array.shape[i]:
            tiling = 1 - i
            break

        self.nodes[self.cur_node_id] = self.node_type([expr], tiling, [], [])
        expr_node_ids = [self.cur_node_id]
        self.add_edge(0, self.cur_node_id, 0)
        self.cur_node_id += 1

    elif isinstance(expr, CollectionExpr):
      # DictExpr, ListExpr, TupleExpr
      children = expr.itervalues() if expr.typename() == 'DictExpr' else expr.vals
      child_ids = self.visit_children(children)
      self.nodes[self.cur_node_id] = self.node_type([expr], -1, [], [])
      expr_node_ids = [self.cur_node_id]
      for child_id in child_ids:
        self.add_edge(child_id, self.cur_node_id, 0)
      self.cur_node_id += 1

    elif hasattr(self, 'visit_%s' % expr.typename()):
      expr_node_ids = getattr(self, 'visit_%s' % expr.typename())(expr)

    else:
      util.log_debug("Skip expr:%s", expr.typename())
      expr_node_ids = []

    # add T node for graph and compute the min cost flow
    if id(expr) == self.init_expr: return self.calc_tiling(expr)

    self.expr_to_nodes[hash(expr)] = expr_node_ids
    return expr_node_ids


def apply_pass(klass, dag):
  if not getattr(FLAGS, 'opt_' + klass.name):
    util.log_debug('Pass %s disabled', klass.name)
    return dag

  util.log_debug('Starting pass %s', klass.name)
  p = klass()
  result = p.visit(dag)
  util.log_debug('Finished pass %s', klass.name)
  return result


passes = []


def optimize(dag):
  if not FLAGS.optimization:
    util.log_info('Optimizations disabled')
    return dag

  util.log_info('Optimization: applying %d passes', len(passes))
  for p in passes:
    dag = apply_pass(p, dag)

  return dag


def add_optimization(klass, default):
  passes.append(klass)

  flagname = 'opt_' + klass.name
  #setattr(Flags, flagname, add_bool_flag(flagname, default=default))
  FLAGS.add(BoolFlag(flagname, default=default, help='Enable %s optimization' % klass.__name__))

  #util.log_info('Passes: %s', passes)

add_optimization(CollapsedCachedExpressions, True)
add_optimization(AutomaticTiling, True)
add_optimization(RotateSlice, False)
add_optimization(MapMapFusion, True)
if parakeet is not None:
  add_optimization(ParakeetGeneration, False)
add_optimization(ReduceMapFusion, True)

FLAGS.add(BoolFlag('optimization', default=True))
