from raco import algebra
import raco.language as language
from pipelines import Pipelined
from raco.utility import emit

import logging
LOG = logging.getLogger(__name__)

"""
Apply rules to an expression
If successful, output will
only involve operators in the
target algebra
"""


def optimize_by_rules(expr, rules):
    for rule in rules:
        def recursiverule(expr):
            newexpr = rule(expr)
            LOG.debug("apply rule %s\n--- %s => %s", rule, expr, newexpr)
            newexpr.apply(recursiverule)
            return newexpr
        expr = recursiverule(expr)
    return expr


def optimize(expr, target, **kwargs):
    """Fire the rule-based optimizer on an expression.  Fire all rules in the
    target algebra."""
    assert isinstance(expr, algebra.Operator)
    assert isinstance(target, language.Algebra), type(target)

    return optimize_by_rules(expr, target.opt_rules(**kwargs))


def compile(expr):
    """Compile physical plan to linearized form for execution"""
    # TODO: Fix this
    algebra.reset()
    exprcode = []

    # TODO, actually use Parallel[Store...]]? Right now assumes it
    if isinstance(expr, (algebra.Sequence, algebra.Parallel)):
        assert len(expr.children()) == 1, "expected single expression only"
        store_expr = expr.children()[0]
    else:
        store_expr = expr

    assert isinstance(store_expr, algebra.Store)
    assert len(store_expr.children()) == 1, "expected single expression only"

    lang = store_expr.language

    if isinstance(store_expr, Pipelined):
        body = lang.body(store_expr.compilePipeline())
    else:
        body = lang.body(expr)

    exprcode.append(emit(body))
    return emit(*exprcode)
@@ -112,16 +112,49 @@
 ort emit
+%0Aimport raco.viz as viz%0Aimport os
 %0A%0Aimport
@@ -197,16 +197,17 @@
 ame__)%0A%0A
+%0A
 %22%22%22%0AAppl
@@ -313,43 +313,637 @@
 %22%0A%0A%0A
-def optimize_by_rules(expr, rules):
+class PlanWriter():%0A%0A    def __init__(self, template=%22wip-%2502d.physical.dot%22, limit=20):%0A        self.ind = 0%0A        self.template = template%0A        self.limit = limit%0A        self.enabled = os.environ.get('RACO_OPTIMIZER_GRAPHS') in %5C%0A            %5B'true', 'True', 't', 'T', '1', 'yes', 'y'%5D%0A%0A    def write_if_enabled(self, plan, title):%0A        if self.enabled:%0A            with open(self.template %25 self.ind, 'w') as dwf:%0A                dwf.write(viz.operator_to_dot(plan, title=title))%0A%0A        self.ind += 1%0A%0A%0Adef optimize_by_rules(expr, rules):%0A    writer = PlanWriter()%0A    writer.write_if_enabled(expr, %22before rules%22)%0A
 %0A   
@@ -985,27 +985,24 @@
 ursiverule(e
-xpr
 ):%0A         
@@ -1008,19 +1008,16 @@
     newe
-xpr
  = rule(
@@ -1013,28 +1013,79 @@
 ewe = rule(e
-xpr)
+)%0A            writer.write_if_enabled(newe, str(rule))%0A
 %0A           
@@ -1163,19 +1163,16 @@
     newe
-xpr
 .apply(r
@@ -1185,16 +1185,17 @@
 verule)%0A
+%0A
         
@@ -1209,19 +1209,16 @@
 urn newe
-xpr
 %0A       
@@ -1241,24 +1241,25 @@
 erule(expr)%0A
+%0A
     return e

create dots for all optimization steps for RACO_OPTIMIZER_GRAPHS