#
# 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: NAME "=" sum    -> assign_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         -> env
         | "(" sum ")"
         | inlinefunc "(" sum ")" -> icall
         | mathfunc "(" sum ")" -> fcall
         | derivative "(" NAME ")" -> dnx

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

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

    ?derivative: "dx"  -> dx
               | "d2x" -> d2x
               | "dy" -> dy
               | "d2y" -> d2y
               | "dz" -> dz
               | "d2z" -> d2z

    %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."
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
"""


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

    number = float

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

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

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

    def var(self, name):
        return name + "(i,j,k)"

    def env(self, name):
        return name

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

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

        return "{}_{}".format(a.data, b) + "(i,j,k)"

    def icall (self, a, b):
        if a.data == "sqr":
            return "(({0})*({0}))".format(b)
        elif a.data == "pow3":
            return "(({0})*({0})*({0}))".format(b)

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

    def neg(self, value):
        return "( - {} )".format(value)

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

    def sub(self, a, b):
        return "( {} - {} )".format(a, b)

    def mul(self, a, b):
        return "( {} * {} )".format(a, b)

    def div(self, a, b):
        return "( {} / {} )".format(a, b)

    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"

        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"

        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"

        return f_code


    def array_ci (self):

        f_code = ""

        for tup in self.derived.iteritems():
            f_code = f_code + real_array_loop.format(tup) + "\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_ci"] = self.array_ci()

        return md


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:
        for line in inputfile:
             if len(line.strip()) > 0:
                 (calc(line))


    print mod_form.format(tf.module_dict())

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