#
# This example shows how to write a basic calculator with variables.
#

from lark import Lark, Transformer, v_args


try:
    input = raw_input   # For Python2 compatibility
except NameError:
    pass



calc_grammar = """
    ?varlist: "[" [NAME ("," NAME)*] "]"

    ?start: statement*

    ?statement: NAME "=" sum -> assign_var
              | avg "{" [sum ("," sum)*] "}" -> assign_avg_var
              | varlist

    ?sum: product
        | sum "+" product   -> add
        | sum "-" product   -> sub

    ?product: atom
        | product "*" atom  -> mul
        | product "/" atom  -> div

    ?atom: NUMBER           -> number
         | "-" atom         -> neg
         | NAME             -> var
         | NAME "'"         -> fluc
         | "$" NAME         -> env
         | "(" sum ")"
         | inlinefunc "(" sum ")" -> icall
         | mathfunc "(" sum ")" -> fcall
         | derivative "(" NAME ")" -> dnx

    avg: "avg" [NAME]

    ?inlinefunc: "sqr" -> sqr
               | "pow3" -> pow3

    ?mathfunc: "log" -> log
             | "exp" -> exp
             | "sqrt" -> sqrt
             | "rxn_rate" -> rxn_rate

    ?derivative: "ddx"  -> ddx
               | "dd2x" -> dd2x
               | "ddy" -> ddy
               | "dd2y" -> dd2y
               | "ddz" -> ddz
               | "dd2z" -> dd2z

    %import common.CNAME -> NAME
    %import common.NUMBER
    %import common.WS

    %ignore WS
"""

real_array_decl  = "real*8, allocatable, dimension(:,:,:) :: {}"
real_array_alloc = "allocate({0}(nxp,nyp,nzp), stat=ierr) ; {0} = 0."

avg_array_decl  = "real*8, allocatable, dimension(:) :: {}"
avg_array_alloc = "allocate({0}(nxp), stat=ierr) ; {0} = 0."

real_array_free  = "deallocate({})"

real_array_loop  = """
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{0[0]}(i,j,k) = {0[1]}
end do
end do
end do
"""

avg_array_sum  = """
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{0}(i) = {0}(i) + {1} {2}
end do
end do
end do
"""

avg_array_divide  = "{0} = {0} / denum {1}"

real_array_diff  = "call {0[0]} ( {0[0]}_{0[1]}, {0[1]} )"

@v_args(inline=True)    # Affects the signatures of the methods
class CalculateTree(Transformer):

    def __init__(self):
        self.primary = []
        self.derived = {}
        self.averaged = {}
        self.derivatives = {}
        self.dependency = {}

    def varlist(self, *args):
        for arg in args:
            self.primary.append(arg.value)
        return self.primary

    def assign_var(self, name, (value, dep)):
        self.derived[name.value] = value
        self.dependency[name.value] = dep
        return "{} = {}".format(name, value)

    def assign_avg_var(self, weight, *args): #name, (value, dep)):

        fmt = "avg_{}"
        if weight is not None:
            wvalue, wdep = self.var(weight)
            self.averaged[fmt.format(weight)] = wvalue, None
            self.dependency[fmt.format(weight)] = wdep
            fmt = weight + "_" + fmt

        for i, (value, dep) in enumerate(args):
            self.averaged[fmt.format(i)] = value, weight
            self.dependency[fmt.format(i)] = dep + (wdep if weight is not None else [])

        return ""

    def avg(self, *args):
        try:
            return args[0]
        except IndexError:
            return None

    def number(self, numeral):
        return (numeral, [])

    def var(self, name):
        dep = []
        if name not in self.primary:
            dep.append(name.value)
        return (name + "(i,j,k)", dep)

    def fluc(self, name):
        dep = []
        if name not in self.primary:
            dep.append(name.value)
        fmt = "({0}(i,j,k) - avg_{0}(i))"
        return (fmt.format(name) , dep)

    def env(self, name):
        return (name, [])

    def dnx (self, partial, b):
        signature = "{}_{}".format(partial.data, b)

        self.derivatives[signature] = (partial.data, b.value)

        self.dependency[signature] = [b.value]

        return (signature + "(i,j,k)", [signature])

    def icall (self, a, (b, dep)):
        fcode = "({0})".format(b)

        if a.data == "sqr":
            fcode = "(({0})*({0}))".format(b)
        elif a.data == "pow3":
            fcode = "(({0})*({0})*({0}))".format(b)

        return (fcode, dep)

    def fcall (self, a, (b, dep)):
        fcode = "( {} ( {} ) )".format(a, b)
        return (fcode, dep)

    def neg(self, (b, dep)):
        fcode = "( - {} )".format(b)
        return (fcode, dep)

    def add(self, (a, adep), (b, bdep)):
        fcode = "( {} + {} )".format(a, b)
        return (fcode, adep + bdep)

    def sub(self, (a, adep), (b, bdep)):
        fcode = "( {} - {} )".format(a, b)
        return (fcode, adep + bdep)

    def mul(self, (a, adep), (b, bdep)):
        fcode = "( {} * {} )".format(a, b)
        return (fcode, adep + bdep)

    def div(self, (a, adep), (b, bdep)):
        fcode = "( {} / {} )".format(a, b)
        return (fcode, adep + bdep)

    log = lambda self : "log"
    exp = lambda self : "exp"
    sqrt = lambda self : "sqrt"
    rxn_rate = lambda self : "rxn_rate"

    def array_decl (self):

        f_code = ""

        for var in self.derived.iterkeys():
            f_code = f_code + real_array_decl.format(var) + "\n"

        for var in self.derivatives.iterkeys():
            f_code = f_code + real_array_decl.format(var) + "\n"

        for var in self.averaged.iterkeys():
            f_code = f_code + avg_array_decl.format(var) + "\n"

        return f_code


    def array_init (self):

        f_code = ""

        for var in self.derived.iterkeys():
            f_code = f_code + real_array_alloc.format(var) + "\n"

        for var in self.derivatives.iterkeys():
            f_code = f_code + real_array_alloc.format(var) + "\n"

        for var in self.averaged.iterkeys():
            f_code = f_code + avg_array_alloc.format(var) + "\n"

        return f_code


    def array_final (self):

        f_code = ""

        for var in self.derived.iterkeys():
            f_code = f_code + real_array_free.format(var) + "\n"

        for var in self.derivatives.iterkeys():
            f_code = f_code + real_array_free.format(var) + "\n"

        for var in self.averaged.iterkeys():
            f_code = f_code + real_array_free.format(var) + "\n"

        return f_code


    def array_pass1 (self):

        f_code = ""

        code_dict = {}

        for tup in self.derived.iteritems():
            code_dict[tup[0]] = real_array_loop.format(tup) + "\n"

        for tup in self.derivatives.iteritems():
            code_dict[tup[0]] = real_array_diff.format(tup[1]) + "\n"

        wfmt = "* {}(i,j,k)"
        for k, (v,w) in self.averaged.iteritems():
            code_dict[k] = avg_array_sum.format(k, v, wfmt.format(w) if w is not None else "") + "\n"

        for var in self.sort_vars():
            f_code = f_code + code_dict[var]

        return f_code


    def array_pass1_avg (self):

        f_code = ""

        meanw = "/ avg_{}"

        for k, (v,w) in self.averaged.iteritems():
            if k.startswith("avg"):
                f_code = avg_array_divide.format(k, meanw.format(w) if w is not None else "") + "\n" + f_code
            else:
                f_code = f_code + avg_array_divide.format(k, meanw.format(w) if w is not None else "") + "\n"

        return f_code


    def module_dict (self):
        md = {}
        md["module_name"] = "terms"
        md["module_data"] = self.array_decl()
        md["module_init"] = self.array_init()
        md["module_finalize"] = self.array_final()
        md["module_pass1"] = self.array_pass1()
        md["module_pass1_avg"] = self.array_pass1_avg()

        return md

    def sort_vars (self):
        order = []
        remain = set(self.derived.iterkeys()) | set(self.derivatives.iterkeys()) | set(self.averaged.iterkeys())

        while len(remain) > 0:
            for v in remain:
                if len(set(self.dependency[v]) & remain) == 0:
                    order.append(v)
                    remain.remove(v)
                    break

        return order


tf=CalculateTree()

calc_parser = Lark(calc_grammar, parser='lalr' , transformer=tf)
calc = calc_parser.parse

import sys

def main():
    while True:
        try:
            s = input('> ')
        except EOFError:
            break
        print(calc(s))


def test():

    with open("resources/m_template.f90") as template_file:
        mod_form = template_file.read()

    with open("terms.input") as inputfile:
        (calc(inputfile.read()))

    print mod_form.format(tf.module_dict())

    print "! ", tf.derived
    print "! ", tf.derivatives
    print "! ", tf.dependency

    print "! ", tf.sort_vars()

if __name__ == '__main__':
    test()
    # main()