#!/usr/bin/env python

"""Compile a MyriaL program into a physical plan."""
import argparse
import os
import json
import sys
import logging


from raco.catalog import FromFileCatalog
import raco.myrial.interpreter as interpreter
import raco.myrial.parser as parser
from raco.fakedb import FakeDatabase
from raco import algebra
from raco.viz import operator_to_dot
from raco.myrial.exceptions import *
from raco.backends.radish import GrappaAlgebra
from raco.backends.logical import OptLogicalAlgebra
from raco.backends.sparql import SPARQLAlgebra
from raco.backends.cpp import CCAlgebra
import raco.from_repr as from_repr
from raco.compile import compile


def print_pretty_plan(plan, indent=0):
    if isinstance(plan, algebra.DoWhile):
        children = plan.children()
        body = children[:-1]
        term = children[-1]

        spc = ' ' * indent
        print '%sDO' % spc
        for op in body:
            print_pretty_plan(op, indent + 4)
        print '%sWHILE' % spc
        print_pretty_plan(term, indent + 4)
    elif isinstance(plan, algebra.Sequence):
        print '%s%s' % (' ' * indent, plan.shortStr())
        for child in plan.children():
            print_pretty_plan(child, indent + 4)
    else:
        print '%s%s' % (' ' * indent, plan)


def parse_options(args):
    arg_parser = argparse.ArgumentParser()
    group = arg_parser.add_mutually_exclusive_group()
    group.add_argument('-p', dest='parse',
                       help="Generate AST (parse tree)", action='store_true')
    group.add_argument('-l', dest='logical',
                       help="Generate logical plan", action='store_true')
    group.add_argument('-L', dest='opt_logical',
                       help="Generate optimized logical plan",
                       action='store_true')
    group.add_argument('-d', dest='dot',
                       help="Generate dot output for logical plan",
                       action='store_true')
    group.add_argument('-D', dest='opt_dot',
                       help="Generate dot output for optimized logical plan",
                       action='store_true')
    group.add_argument('-j', dest='json',
                       help="Encode plan as JSON", action='store_true')
    group.add_argument('-f', dest='standalone', action='store_true',
                       help='Execute the program in standalone mode')
    group.add_argument('-c', dest='radish', action='store_true',
                       help='Output physical plan for Radish (Grappa) and save source program to file')
    group.add_argument('--cpp', dest='cpp', action='store_true',
                       help='Output physical plan for C++ and save source program to file')
    group.add_argument('-s', dest='sparql', action='store_true',
                       help='Output an attempt at a SPARQL translation of the input program')
    arg_parser.add_argument('-r', dest='repr',
                            help="Encode plan as Python repr", action='store_true')
    arg_parser.add_argument('-v', dest='verbose', action='store_true',
                       help='Turn on verbose DEBUG logging')
    arg_parser.add_argument('--dot-radish', dest='dot_radish', action='store_true', help='print out dot for Grappa plan')
    arg_parser.add_argument('--catalog', dest="catalog_path", default=None, help="[Optional] path to catalog file")
    arg_parser.add_argument('--plan', dest="from_repr", action='store_true', help="[Optional] input file is a plan as a python repr")
    arg_parser.add_argument('--key', action='append', help="May use this argument multiple times to specify additional arguments to compiler")
    arg_parser.add_argument('--value', action='append', help="May use this argument multiple times to specify additional arguments to compiler")
    arg_parser.add_argument('file',
                            help='File containing MyriaL source program (or a physical plan if using --plan)')

    ns = arg_parser.parse_args(args)
    return ns


def main(args):
    opt = parse_options(args)

    if opt.key is not None:
        assert opt.value is not None and len(opt.key)==len(opt.value), "Must be --key K --value V pairs"
        kwargs = dict(zip(opt.key, opt.value))
    else:
        kwargs = {}

    if opt.verbose:
        logging.basicConfig(level=logging.DEBUG)

    # Search for a catalog definition file
    if opt.catalog_path is not None:
        catalog_path = opt.catalog_path
    else:
        catalog_path = os.path.join(os.path.dirname(opt.file), 'catalog.py')

    if os.path.exists(catalog_path):
        catalog = FromFileCatalog.load_from_file(catalog_path)
    else:
        catalog = FromFileCatalog({},"")

    _parser = parser.Parser()
    processor = interpreter.StatementProcessor(catalog, True)

    statement_list = None
    plan_repr = None
    with open(opt.file) as fh:
        try:
            inp = fh.read()
            if opt.from_repr:
                plan_repr = inp
            else:
                statement_list = _parser.parse(inp)
        except MyrialCompileException as ex:
            print 'MyriaL parse error: %s' % ex
            return 1

    if opt.parse:
        if statement_list == None:
            print "No MyriaL given"
        else:
            print statement_list
    else:
        if opt.from_repr:
            pd = PhysicalPlanDispatch(from_repr=plan_repr)
        else:
            pd = PhysicalPlanDispatch(processor=processor)
            processor.evaluate(statement_list)

        if opt.logical:
            if opt.from_repr:
                raise "Options logical and --plan are incompatible"
            if opt.repr:
                print repr(processor.get_logical_plan())
            else:
                print_pretty_plan(processor.get_logical_plan())

        elif opt.dot:
            if opt.from_repr:
                raise "Options dot and --plan are incompatible"
            if opt.repr:
                raise "Options dot and -r are incompatible"
            print operator_to_dot(processor.get_logical_plan())
        elif opt.opt_dot:
            if opt.from_repr:
                raise "Options opt_dot and --plan are incompatible"
            if opt.repr:
                raise "Options opt_dot and -r are incompatible"
            print operator_to_dot(processor.get_physical_plan(
                target_alg=OptLogicalAlgebra(), **kwargs))
        elif opt.dot_radish:
            if opt.from_repr:
                raise "Options dot_radish and --plan are incompatible"
            if opt.repr:
                raise "Options dot_radish and -r are incompatible"
            print operator_to_dot(pd.get_physical_plan(target_alg=GrappaAlgebra(),**kwargs))
        elif opt.json:
            if opt.repr:
                raise "Options json and -r are incompatible"
            ppj = pd.get_json()
            print(json.dumps(ppj))
        elif opt.standalone:
            if opt.repr:
                raise "Options standalone and -r are incompatible"
            pp = pd.get_physical_plan(**kwargs)
            db = FakeDatabase()
            db.evaluate(pp)
        elif opt.opt_logical:
            if opt.from_repr:
                raise "Options opt_logical and --plan are incompatible"
            if opt.repr:
                print repr(processor.get_physical_plan(
                    target_alg=OptLogicalAlgebra(), **kwargs))
            else:
                print_pretty_plan(processor.get_physical_plan(
                    target_alg=OptLogicalAlgebra(), **kwargs))
        elif opt.radish:
            if opt.repr:
                raise "Options radish and -r are incompatible"
            # some useful kwargs
            # scan_array_repr='symmetric_array'
            pp = pd.get_physical_plan(target_alg=GrappaAlgebra(),
                                      **kwargs)
            print_pretty_plan(pp)
            c = compile(pp, **kwargs)
            fname = '{0}.cpp'.format(os.path.splitext(os.path.basename(opt.file))[0])
            with open(fname, 'w') as f:
                f.write(c)
        elif opt.cpp:
            if opt.repr:
                raise "Options cpp and -r are incompatible"
            # some useful kwargs
            # scan_array_repr='symmetric_array'
            pp = pd.get_physical_plan(target_alg=CCAlgebra(), **kwargs)
            print_pretty_plan(pp)
            c = compile(pp)
            fname = '{0}.cpp'.format(os.path.splitext(os.path.basename(opt.file))[0])
            with open(fname, 'w') as f:
                f.write(c)
        elif opt.sparql:
            if opt.repr:
                raise "Options sparql and -r are incompatible"
            pp = pd.get_physical_plan(target_alg=SPARQLAlgebra(), **kwargs)
            c = compile(pp)
            print c
        elif opt.repr:
            pp = pd.get_physical_plan(**kwargs)
            print repr(pp)
        else:
            print_pretty_plan(pd.get_physical_plan(**kwargs))

    return 0


class PhysicalPlanDispatch(object):

    def __init__(self, processor=None, from_repr=None):
        if processor:
            self.processor = processor
            self.with_repr = None
        elif from_repr:
            self.processor = None
            self.with_repr = from_repr
        else:
            raise "Requires processor or repr plan input"

    def get_physical_plan(self, target_alg=None, **kwargs):
        if self.with_repr:
            return from_repr.plan_from_repr(self.with_repr)
        else:
            if target_alg is None:
                return self.processor.get_physical_plan(**kwargs)
            else:
                return self.processor.get_physical_plan(target_alg=target_alg,
                                                        **kwargs)

    def get_json(self):
        if self.with_repr:
            return interpreter.StatementProcessor.get_json_from_physical_plan(self.get_physical_plan())
        else:
            return self.processor.get_json()



if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
