from lark import Lark, Visitor, Transformer, v_args, Token
import warnings
from jinja2 import Template


class CollectDefinitions(Visitor):

    def __init__ (self, primary, derived, averaged):
        self.primary = primary
        self.derived = derived
        self.averaged = averaged

    def varlist(self, tree):
        for v in tree.children:
            self.primary.add(v.value)
            self.derived[v.value] = PrimaryField(v.value, self.derived)

    def assign_var (self, tree):
        if len(tree.children) > 2:
            lval, lattr, rval = tree.children
        else:
            lval, rval = tree.children
            lattr = None

        attr_dict = {}

        if lattr is not None:
            for t in lattr.children:
                k, v = t.children
                attr_dict[k.value] = v.value

        if lval.value in self.derived:
            raise ValueError("duplicate definition of " + lval)
        self.derived[lval.value] = Field(lval.value, attr_dict, rval, self.derived)

    def assign_avg_var (self, tree):
        w = tree.children[0]
        targets = tree.children[1:]

        if (not w.children) or (w.children[0] is None):
            self.averaged[None] = set([x.value for x in targets])
        else:
            self.averaged[w.children[0].value] = set([x.value for x in targets])


class ExpInspector(Visitor):
    def __init__(self):
        self.fluctuation = False
        self.dep = set([])
        self.deriv = set([])

    @classmethod
    def inspect(cls, tree):
        self = cls()
        return self(tree)

    def __call__(self, tree):
        self.visit(tree)
        return self.fluctuation, self.dep, self.deriv

    def fluc(self, tree):
        self.fluctuation = True
        self.dep.add(tree.children[0].value)

    def var(self, tree):
        self.dep.add(tree.children[0].value)

    def dnx (self, tree):
        op, v = tree.children
        deriv = "{}_{}".format(op.data, v.value)
        self.dep.add(deriv)
        self.deriv.add((op.data, v.value))


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

    def __init__(self, fdict):
        self.fdict = fdict

    def arithmatic_rooted(self, name):
        try:
            exproot = self.fdict[name].exp.data
        except AttributeError:
            exproot = "something_11fasq2afa3rfzsaerqw23"

        return ((exproot == "add") or (exproot == "sub") or
                (exproot == "mul") or (exproot == "div"))

    def parenthise(self, name):
        try:
            latex = self.fdict[name].latex
            latex_given = self.fdict[name].latex_given
        except KeyError:
            warnings.warn(name + " is not found")
            latex = r"\mathrm{{{}}}".format(name)
            latex_given = None

        if self.arithmatic_rooted(name) and (latex_given is None):
            latex = "(" + latex + ")"

        return latex

    def number(self, numeral):
        return numeral

    def env(self, name):
        return r"\mathrm{{{}}}".format(name.value)

    def paren(self, name):
        return "({})".format(str(name))

    def var(self, name):
        return self.parenthise(name.value)

    def fluc(self, name):
        return self.parenthise(name.value)+ "''"

    def dnx (self, partial, b):
        fmt = r"\partial_{{{}}}"
        coord = partial.data[-1]
        op = fmt.format(coord + coord if len(partial.data) > 3 else coord)

        signature = "{}_{}".format(partial.data, b.value)

        try:
            eq = self.fdict[signature].latex
        except KeyError:
            eq = op + self.parenthise(b.value)
            warnings.warn(signature + " is not found: " + eq)

        return eq

    def icall (self, a, b):
        if a.data == "sqr":
            fcode = "({0})^2".format(b)
        elif a.data == "pow3":
            fcode = "({0})^3".format(b)
        else:
            fcode = "({0})".format(b)
        return fcode

    def fcall (self, *args):
        a = args[0]
        b = ", ".join(args[1:])
        if a == 'sqrt':
            fcode = r"\sqrt{{{}}}".format(b)
            return fcode
        elif a == 'abs':
            fcode = r"\left| {} \right|".format(b)
            return fcode
        elif a.startswith("\\"):
            fcode = r"{}{{({})}}".format(a, b)
            return fcode
        else:
            fcode = r"\mathrm{{{}}}({})".format(a, b)
            return fcode

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

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

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

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

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

    log = lambda self : "\log"
    exp = lambda self : "\exp"
    sqrt = lambda self : "sqrt"
    abs = lambda self : "abs"
    rxn_rate = lambda self : "\omega"
    udf = lambda self, a  : a.value




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

    def __init__(self, fdict):
        self.fdict = fdict

    def number(self, numeral):
        return str(float(numeral))

    def env(self, name):
        return name.value

    def paren(self, name):
        return "({})".format(str(name))

    def var(self, name):
        try:
            arrname = self.fdict[name.value].array
        except KeyError:
            arrname = name.value

        return arrname + "(i,j,k)"

    def fluc(self, name):
        try:
            arrname = self.fdict[name.value].array
        except KeyError:
            arrname = name.value

        fmt = "({0}(i,j,k) - {{0}}avg_{1}(i))"
        return fmt.format(arrname, name.value)

    def dnx (self, partial, b):
        signature = "{}_{}".format(partial.data, b.value)
        try:
            arrname = self.fdict[signature].array
        except KeyError:
            arrname = signature

        return arrname + "(i,j,k)"

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

    def fcall (self, *args):
        a = args[0]
        b = ", ".join(args[1:])
        fcode = "( {} ( {} ) )".format(a, b)
        return fcode

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

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

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

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

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

    log = lambda self : "log"
    exp = lambda self : "exp"
    sqrt = lambda self : "sqrt"
    abs = lambda self : "abs"
    rxn_rate = lambda self : "rxn_rate"
    udf = lambda self, a  : a.value




def make_allocate(name, shape, init_zero=True):
    alloc_str = f"allocate({name}({shape}), stat=ierr)\n"
    alloc_str += f"if (ierr /= 0) then\n"
    alloc_str += f"    write(0,*) 'Error: allocation of {name} failed on process', myid\n"
    alloc_str += f"    call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err)\n"
    alloc_str += f"end if"
    if init_zero:
        alloc_str += f"\n{name} = 0."
    return alloc_str


class FieldBase (object):
    def __init__ (self, name, fdict):
        self.name = name
        self.array = name
        self.dep = set([])
        self.fluc = False
        self.prime = False
        self.fdict = fdict
        self.shape = "nxp,nyp,nzp"
        self.dim = ":,:,:"

    def depends_on (self, a):
        return (a in self.dep)

    def is_primary (self):
        return self.prime

    def not_primary (self):
        return not self.prime

    def is_fluctuation (self):
        return self.fluc

    def export_on (self):
        return False

    def __repr__ (self):
        return self.name

    def checkFluctuation (self):

        fset = set([])

        for d in map(self.fdict.get, self.dep):
            fset.update(d.checkFluctuation())

        if self.is_fluctuation() or len(fset) > 0:
            fset.add(self.name)

        return fset

    def depClosure (self):

        fset = set(self.dep)

        for d in self.dep:
            fset.update(self.fdict[d].depClosure())

        return fset


    def code_decl (self):
        real_array_decl  = "real(real64), allocatable, dimension({1}) :: {0}"
        return real_array_decl.format(self.name, self.dim)

    def code_alloc (self):
        return make_allocate(self.name, self.shape)

    def code_free (self):
        real_array_free = "deallocate({})"
        return real_array_free.format(self.name)


class FieldExporter (object):

    mpi_io_decl='''
! field exporter common
integer(kind=MPI_OFFSET_KIND)  :: offset
'''

    # Subarray version
    fmt_decl_subarray='''
! - file_handles and mpi_infos
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_fh
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_info
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_filetype
'''

    fmt_init_subarray='''
! init subarray datatype for {{ field_name }}
block
    integer(4) :: sizes(3), subsizes(3), starts(3)
    call MPI_INFO_CREATE({{ field_name }}_info, mpi_err)
    call MPI_FILE_OPEN(MPI_COMM_TASK,'export-{{ field_name }}.dat',MPI_MODE_WRONLY+MPI_MODE_CREATE,{{ field_name }}_info,{{ field_name }}_fh,mpi_err)
    sizes = (/ nxp, nyp, nzp /)
    subsizes = (/ {{ len_xpts }}, {{ ye }} - {{ ys }} + 1, {{ ze }} - {{ zs }} + 1 /)
    starts = (/ {{ xs }} - 1, {{ ys }} - 1, {{ zs }} - 1 /)
    call MPI_TYPE_CREATE_SUBARRAY(3, sizes, subsizes, starts, MPI_ORDER_FORTRAN, MPI_REAL8, {{ field_name }}_filetype, mpi_err)
    call MPI_TYPE_COMMIT({{ field_name }}_filetype, mpi_err)
end block
'''

    fmt_final_subarray='''
! finalize
call MPI_FILE_CLOSE({{ field_name }}_fh, mpi_err)
call MPI_INFO_FREE({{ field_name }}_info, mpi_err)
call MPI_TYPE_FREE({{ field_name }}_filetype, mpi_err)
'''

    fmt_calc_subarray='''
! write to file via MPI Subarray
count = ({{ len_xpts }}) * ({{ ye }} - {{ ys }} + 1) * ({{ ze }} - {{ zs }} + 1)
offset = export_offset(fidx) * count * 8
call MPI_FILE_WRITE_AT({{ field_name }}_fh, offset, {{ work_array }}, 1, {{ field_name }}_filetype, mpi_status, mpi_err)
'''

    # Legacy copy version (fallback)
    fmt_decl_legacy='''
! - file_handles and mpi_infos
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_fh
integer(kind=MPI_INTEGER_KIND) :: {{ field_name }}_info

! - buffer
real(real64), allocatable, dimension(:,:,:) :: {{ field_name }}_export_array
integer, allocatable, dimension(:) :: {{ field_name }}_xpts
'''

    fmt_init_legacy='''
! init
call MPI_INFO_CREATE({{ field_name }}_info, mpi_err)
call MPI_FILE_OPEN(MPI_COMM_TASK,'export-{{ field_name }}.dat',MPI_MODE_WRONLY+MPI_MODE_CREATE,{{ field_name }}_info,{{ field_name }}_fh,mpi_err)
allocate({{ field_name }}_export_array(1:{{ len_xpts }},{{ ys }}:{{ ye }},{{ zs }}:{{ ze }}), stat=ierr)
if (ierr /= 0) then
    write(0,*) 'Error: allocation of {{ field_name }}_export_array failed on process', myid
    call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err)
end if
{{ field_name }}_export_array = 0.
allocate({{ field_name }}_xpts(1:{{ len_xpts }}), stat=ierr)
if (ierr /= 0) then
    write(0,*) 'Error: allocation of {{ field_name }}_xpts failed on process', myid
    call MPI_ABORT(MPI_COMM_TASK, 1, mpi_err)
end if
{{ xpts_init }}
'''

    fmt_final_legacy='''
! finalize
call MPI_FILE_CLOSE({{ field_name }}_fh, mpi_err)
call MPI_INFO_FREE({{ field_name }}_info, mpi_err)
deallocate({{ field_name }}_export_array)
deallocate({{ field_name }}_xpts)
'''

    fmt_calc_legacy='''
! copy to array for export
do k = {{ zs }}, {{ ze }}
do j = {{ ys }}, {{ ye }}
do i = 1, {{ len_xpts }}
{{ field_name }}_export_array(i,j,k) = {{ work_array }}({{ field_name }}_xpts(i),j,k)
end do
end do
end do

! write to file
count = ({{ len_xpts }}) * ({{ ye }} - {{ ys }} + 1) * ({{ ze }} - {{ zs }} + 1)
offset = export_offset(fidx) * count * 8
call MPI_FILE_WRITE_AT({{ field_name }}_fh, offset, {{ field_name }}_export_array, count, MPI_REAL8, mpi_status, mpi_err)
'''

    def __init__ (self, name, attr, parent):
        self.name = name
        self.attr = attr
        self.parent = parent

        self.params = dict(attr)

        self.params.setdefault("xs", 1)
        self.params.setdefault("xe", "nxp")
        self.params.setdefault("ys", 1)
        self.params.setdefault("ye", "nyp")
        self.params.setdefault("zs", 1)
        self.params.setdefault("ze", "nzp")

        self.params.setdefault("field_name", self.name)

        self.params.setdefault("len_xpts", f"({self.params['xe']} - {self.params['xs']} + 1)")
        
        fmt_xpts_init = '''
do i = {{ xs }}, {{ xe }}
{{ field_name }}_xpts(i-{{ xs }}+1) = i
end do
'''
        self.params.setdefault("xpts_init", Template(fmt_xpts_init).render(**self.params))

        try:
        # Sampling at listed x coordinates
            fmt_xpts_init_list = "{{ field_name }}_xpts = (/ {{ list_xpts }} /)"

            import sys
            raw_xpts = self.params["xpts"]
            int_xpts = list(map(int, raw_xpts.split()))
            len_xpts = len(int_xpts)
            self.params["len_xpts"] = len_xpts
            self.params["list_xpts"] = ",".join(map(str, int_xpts))
            self.params["xpts_init"] = Template(fmt_xpts_init_list).render(**self.params)
        except KeyError:
            pass

        self.use_subarray = ("xpts" not in self.params)


    def code (self):
        self.params["work_array"] = self.parent.array
        if self.use_subarray:
            return Template(FieldExporter.fmt_calc_subarray).render(**self.params)
        else:
            return Template(FieldExporter.fmt_calc_legacy).render(**self.params)

    def code_decl (self):
        if self.use_subarray:
            return Template(FieldExporter.fmt_decl_subarray).render(**self.params)
        else:
            return Template(FieldExporter.fmt_decl_legacy).render(**self.params)

    def code_alloc (self):
        if self.use_subarray:
            return Template(FieldExporter.fmt_init_subarray).render(**self.params)
        else:
            return Template(FieldExporter.fmt_init_legacy).render(**self.params)

    def code_free (self):
        if self.use_subarray:
            return Template(FieldExporter.fmt_final_subarray).render(**self.params)
        else:
            return Template(FieldExporter.fmt_final_legacy).render(**self.params)




class Field (FieldBase):

    def __init__ (self, name, attr, exp, fdict):
        super(Field,self).__init__(name, fdict)
        self.attr = attr
        self.exp = exp
        self.fluc, self.dep, self.derivs = ExpInspector.inspect(exp)
        self.comment = self.name + " = " + ExpToCode(self.fdict).transform(self.exp)

        self.latex_given = self.attr.get("latex")
        if self.latex_given is None:
            self.latex = ExpToLatex(self.fdict).transform(self.exp)
        else:
            self.latex = self.latex_given

        self.exporter = None
        try:
            if self.attr["export"]:
                self.exporter = FieldExporter(self.name, self.attr, self)
        except KeyError:
            pass


        '''
        for a in self.dep:
            if a not in self.fdict:
                raise StandardError(a + " is not defined")
        '''

    def export_on (self):
        return self.exporter is not None

    def code (self, alloc=None):
        real_array_loop  = """
! {{ comment }}
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{{ array }}(i,j,k) = {{ rhs }}
end do
end do
end do
"""
        self.array = alloc[self.name] if alloc else self.name

        rhs = ExpToCode(self.fdict).transform(self.exp)
        calculation_code = Template(real_array_loop).render(comment=self.comment, array=self.array, rhs=rhs)

        export_code = ( self.exporter.code() if self.export_on() else "")

        return calculation_code + export_code


class FluctuationField (FieldBase):

    def __init__ (self, w, field, fset, fdict):

        super(FluctuationField,self).__init__(self.id(w,field), fdict)

        if w is not None:
            self.w = w + "_"
        else:
            self.w = ""

        self.field = fdict[field]
        self.dep = self.field.dep - fset
        for df in self.field.dep & fset:
            self.dep.add(self.id(w,df))

        self.comment = ExpToCode(self.fdict).transform(self.field.exp)

        if self.field.is_fluctuation():
            self.comment = self.comment.format(self.w)

        self.comment = self.name + " = " + self.comment

    def code (self, alloc=None):
        real_array_loop  = """
! {{ comment }}
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{{ array }}(i,j,k) = {{ rhs }}
end do
end do
end do
"""
        self.array = alloc[self.name] if alloc else self.name

        rhs = ExpToCode(self.fdict).transform(self.field.exp)

        if self.field.is_fluctuation():
            rhs = rhs.format(self.w)

        return Template(real_array_loop).render(comment=self.comment, array=self.array, rhs=rhs)

    @classmethod
    def id (cls, w, field):
        if w:
            name = "{}____{}_avg".format(field, w)
        else:
            name = "{}____avg".format(field)
        return name


class PrimaryField (FieldBase):

    def __init__ (self, name, fdict):
        super(PrimaryField,self).__init__(name, fdict)
        self.derivs = set([])
        self.prime = True
        self.latex = name
        self.latex_given = None

    def code (self, alloc=None):
        return "! {} is read from file".format(self.name)

    def code_decl (self):
        return "! {} is read from file".format(self.name)

    def code_alloc (self):
        return "! {} is read from file".format(self.name)

    def code_free (self):
        return "! {} is read from file".format(self.name)


class DerivedField (FieldBase):

    def __init__ (self, op, v, fdict):
        name = "{}_{}".format(op, v)
        super(DerivedField,self).__init__(name, fdict)
        self.op = op
        self.v = v
        self.dep = set([v])

        fmt = r"\partial_{{{}}}"
        coord = op[-1]
        partial = fmt.format(coord + coord if len(op) > 3 else coord)

        self.latex = partial + "(" + fdict[v].latex + ")"

    def code (self, alloc=None):
        self.array = alloc[self.name] if alloc else self.name
        varray = alloc[self.v] if alloc else self.v
        return "call {0} ( {2}, {1} )".format(self.op, varray, self.array)


class AveragedField (FieldBase):

    @classmethod
    def id (cls, w, tgt):
        if w:
            return "{}_avg_{}".format(w, tgt)
        else:
            return "avg_{}".format(tgt)


    def __init__ (self, w, tgt, fdict):

        name = self.id(w,tgt)
        super(AveragedField,self).__init__(name, fdict)
        self.shape = "nxp"
        self.dim = ":"

        tfield = fdict[tgt]
        self.fset = tfield.checkFluctuation()

        self.latex = r"\left\langle {} \right\rangle".format(tfield.latex)

        if not self.fset:
            self.tgt = tgt
            self.dep.add(tgt)
        else:
            ftgt = FluctuationField.id(w,tgt)
            self.tgt = ftgt
            self.dep.add(ftgt)

        self.weighted = w
        if w:
            self.w = fdict[w]
            self.dep.add(w)
            self.latex += ("_{{{}}}".format(w))


    def code (self, alloc=None):

        avg_array_sum  = """
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{{ name }}(i) = {{ name }}(i) + {{ arrname }}
end do
end do
end do
"""
        arrname = self.fdict[self.tgt].array + "(i,j,k)"
        if self.weighted is not None:
            arrname = arrname + " * " + self.w.array + "(i,j,k)"

        return Template(avg_array_sum).render(name=self.name, arrname=arrname)


    def code_avg (self):

        avg_array_divide  = """
call MPI_ALLREDUCE(MPI_IN_PLACE, {{ name }}, nxp, MPI_REAL8, MPI_SUM, MPI_COMM_TASK, mpi_err)

{{ name }} = {{ name }} {{ dWeight }} / denum
"""
        dWeight = (f"/ avg_{self.weighted}" if self.weighted else "")

        return Template(avg_array_divide).render(name=self.name, dWeight=dWeight)

    def isWeighted (self):
        return self.weighted is not None

    def pass1 (self):
        return not self.pass2()

    def pass2 (self):
        return len(self.fset) > 0


class Stage1():
    ''' conversion from tree to python data '''

    def __init__ (self, raw_tree):
        self.primary = set([])
        self.derived = {}
        self.averaged = {}

        # Construct Field objects
        CollectDefinitions( self.primary, self.derived, self.averaged ).visit(raw_tree)

    def __repr__ (self):
        return "\n".join(map(str, [self.primary, self.derived, self.averaged]))


class Stage2():
    ''' expand derivatives and averages'''

    def __init__ (self, src):
        self.src = src
        self.primary = src.primary
        self.derived = src.derived
        self.derivative = {}
        self.averaged = {}

        # Construct Derivative Field objects

        dset = set([])
        for k, v in self.derived.items():
            dset.update(v.derivs)

        for tup in dset:
            a = DerivedField(tup[0], tup[1], self.derived)
            self.derived[a.name] = a
            self.derivative[tup] = a

        # Construct Averaged Field objects

        for w, tgts in src.averaged.items():
            for t in tgts:
                a = AveragedField(w, t, self.derived)
                self.averaged[a.name] = a
                for ff in a.fset:
                    b = FluctuationField(w, ff, a.fset, self.derived)
                    self.derived[b.name] = b

    def __repr__ (self):
        return "\n".join(map(str, [self.derived, self.derivative, self.averaged]))

    def dependency (self):
        dgraph = {}

        for k,v in self.derived.items():
            dgraph[k] = v.dep

        for k,v in self.averaged.items():
            dgraph[k] = v.dep

        return dgraph


class Stage3():
    ''' calculate execution order '''

    def __init__ (self, src):
        self.src = src
        self.primary = src.primary
        self.derived = src.derived
        self.averaged = src.averaged
        self.dependency = src.dependency()

        assert set(self.derived.keys()).isdisjoint(self.averaged.keys())


        self.avg1 = set(filter(AveragedField.pass1, self.averaged.values()))
        self.avg2 = set(filter(AveragedField.pass2, self.averaged.values()))

        pass1calc = set(map(repr, self.avg1))
        for x in self.avg1:
            pass1calc.update(x.depClosure())
        self.pass1 = self.sort_vars_new(self.dependency, pass1calc - self.primary)

        # self.sort_vars(self.dependency, pass1calc - self.primary -set(map(repr, self.avg1)))

        pass2calc = set(map(repr, self.avg2))
        for x in self.avg2:
            pass2calc.update(x.depClosure())
        self.pass2 = self.sort_vars_new(self.dependency, pass2calc - self.primary)


    def __repr__ (self):
        return "\n".join(map(str, [self.pass1, self.pass2]))


    def sort_vars (self, dependency, group):
        order = []
        remain = list(group)
        remain.sort()

        while len(remain) > 0:
            for v in remain:
                if dependency[v].isdisjoint(remain):
                    order.append(v)
                    remain.remove(v)

        return order

    def calc_size (self, ordered, remaining):
        count = 0
        dep_union = set([])

        g = self.dependency
        
        for v in remaining:
            dep_union |= set(g[v])
            
        for v in ordered:
            if v in dep_union:
                count += 1
            
        return count


    def sort_vars_new (self, dependency, group):
        order = []
        remain = list(group)
        remain.sort()

        while len(remain) > 0:
            
            candidate = []
            for v in remain:
                if set(dependency[v]).isdisjoint(remain):
                    candidate.append(v)

            impact = {}
            
            size0 = self.calc_size(set(order), set(remain))

            for v in candidate:
                impact[v] = self.calc_size(set(order) | set([v]), set(remain) - set([v])) - size0
                
            candidate.sort(key=impact.get)
            
            order.append(candidate[0])
            remain.remove(candidate[0])

        return order


    def print_program (self):

        allvar = dict(self.derived)
        allvar.update(self.averaged)

        decl  = "\n".join(allvar[v].code_decl() for v in set(self.pass1+self.pass2))
        alloc = "\n".join(allvar[v].code_alloc() for v in set(self.pass1+self.pass2))
        free  = "\n".join(allvar[v].code_free() for v in set(self.pass1+self.pass2))

        calc1 = "\n".join(allvar[v].code() for v in self.pass1)
        calc2 = "\n".join(allvar[v].code() for v in self.pass2)

        set1 = [a.name for a in filter(AveragedField.pass1, self.averaged.values())]
        set2 = [a.name for a in filter(AveragedField.pass2, self.averaged.values())]

        set1.sort()
        set2.sort()

        avg1 = "\n".join(allvar[v].code_avg() for v in set1)
        avg2 = "\n".join(allvar[v].code_avg() for v in set2)

        hfmt = 'character (len = *), parameter :: output_header="{}"'
        declh = hfmt.format(" ".join(["x"] + set1 + set2))

        avg_array_write = '''
integer :: i

open (200, file="qEdge_X.dat")

write (200,*) output_header

do i=1,nxp
 write (200,'({0}e20.10)') real(i)*hxp, {1}
end do

close (200)
'''

        avgarr = "{}(i)"
        write_avg = avg_array_write.format(
                len(self.averaged)+1,
                ", ".join(map(avgarr.format, set1+set2))
                )

        md = {}
        md["module_name"] = "terms"
        md["module_data"] = "\n".join((declh, decl))
        md["module_init"] = alloc
        md["module_finalize"] = free
        md["module_pass1"] = calc1
        md["module_pass1_avg"] = avg1
        md["module_pass2"] = calc2
        md["module_pass2_avg"] = avg2
        md["module_write_result"] = write_avg

        return md


class Stage4():
    ''' analyze liveness and allocate array '''
    def __init__ (self, src):
        self.src = src
        self.primary = src.primary
        self.derived = src.derived
        self.averaged = src.averaged
        self.dependency = src.dependency
        self.avg1 = src.avg1
        self.avg2 = src.avg2
        self.pass1 = src.pass1
        self.pass2 = src.pass2

        self.array_name = "xyzbuffer{}"

        narr1, alloc1 = (self.allocate_arr(self.pass1))
        narr2, alloc2 = (self.allocate_arr(self.pass2))

        self.narr = max(narr1, narr2)

        self.alloc1 = alloc1
        self.alloc2 = alloc2


    def liveness (self, l1, g):
        import numpy as np
        img = np.zeros((len(l1), len(l1)))
        for i, v in enumerate(l1):
            for j in range(i, len(l1)):
                img[i,i:j+1] = img[i,i:j+1] + (1 if v in g[l1[j]] else 0)
        return img > 0


    def allocate_arr (self, l):

        import numpy as np

        dg = self.dependency

        mask = self.liveness(l, dg)

        try:
            narr = mask.astype(int).sum(axis=0).max()
        except ValueError:
            narr = 0

        array_pool = set([self.array_name.format(i) for i in range(narr)])

        livesets = [set([])] + [set(np.asarray(l)[row]) for row in mask.T]

        var2arr = { p : p for p in self.primary }

        for i, (s0, s1) in enumerate(zip(livesets[:-1], livesets[1:])):

            array_pool.update(map(var2arr.get, s0 - s1))

            for new in s1 - s0:
                var2arr[new] = array_pool.pop()

        return narr, var2arr


    def work_array_codes (self):

        array_names = [self.array_name.format(i) for i in range(self.narr)]

        real_array_decl  = "real(real64), allocatable, dimension(:,:,:) :: {0}"
        decl  = "\n".join([real_array_decl.format(v) for v in array_names])
        alloc = "\n".join([make_allocate(v, "nxp,nyp,nzp") for v in array_names])
        free  = "\n".join(["deallocate({})".format(v) for v in array_names])

        return decl, alloc, free


    def write_avg_codes (self, avglist):

        avg_array_write = '''
real(real64), dimension(nxp) :: xbuffer
integer :: i

open (200, file="qEdge_X.dat")
write (200,*) output_header
do i=1,nxp
 write (200,'({{ num_args }}e20.10)') real(i)*hxp, {{ formatted_avglist }}
end do
close (200)

open (200, file="d1.dat")
{{ deriv1_lines }}
close (200)

open (200, file="d2.dat")
{{ deriv2_lines }}
close (200)
'''

        avgarr = "{}(i)"
        deriv1_avgarr = """call ddx1d ( xbuffer, {} ) ; write (200,*) xbuffer"""
        deriv2_avgarr = """call d2dx1d ( xbuffer, {} ) ; write (200,*) xbuffer"""

        num_args = len(self.averaged) + 1
        formatted_avglist = ", ".join(map(avgarr.format, avglist))
        deriv1_lines = "\n".join(map(deriv1_avgarr.format, avglist))
        deriv2_lines = "\n".join(map(deriv2_avgarr.format, avglist))

        write_avg = Template(avg_array_write).render(
            num_args=num_args,
            formatted_avglist=formatted_avglist,
            deriv1_lines=deriv1_lines,
            deriv2_lines=deriv2_lines
        )
        return write_avg


    def print_program (self):

        allvar = dict(self.derived)
        allvar.update(self.averaged)

        set1 = sorted([a.name for a in filter(AveragedField.pass1, self.averaged.values())])
        set2 = sorted([a.name for a in filter(AveragedField.pass2, self.averaged.values())])

        set_export_on = list(filter(lambda x: x.export_on(), self.derived.values()))

        ffmt = 'logical, parameter :: pass2_required={}'
        declf = ffmt.format('.true.' if len(set2) > 0 else '.false.')

        hfmt = 'character (len = *), parameter :: output_header="{}"'
        declh = hfmt.format(" ".join(["x"] + set1 + set2))

        declarr, allocarr, freearr = self.work_array_codes()

        declavg  = "\n".join(self.averaged[v].code_decl() for v in sorted(self.averaged))
        allocavg = "\n".join(self.averaged[v].code_alloc() for v in sorted(self.averaged))
        freeavg  = "\n".join(self.averaged[v].code_free() for v in sorted(self.averaged))

        decl_export  = "\n".join(v.exporter.code_decl()  for v in set_export_on)
        alloc_export = "\n".join(v.exporter.code_alloc() for v in set_export_on)
        free_export  = "\n".join(v.exporter.code_free()  for v in set_export_on)

        sub_calc1 = "\n".join(allvar[v].code(self.alloc1) for v in self.pass1)
        sub_calc2 = "\n".join(allvar[v].code(self.alloc2) for v in self.pass2)

        sub_avg1 = "\n".join(allvar[v].code_avg() for v in set1)
        sub_avg2 = "\n".join(allvar[v].code_avg() for v in set2)

        sub_write_avg = self.write_avg_codes(set1+set2)

        md = {}
        md["module_name"] = "terms"
        md["module_data"] = "\n".join((declf, declh, declavg, FieldExporter.mpi_io_decl, decl_export, declarr))
        md["module_init"] = "\n".join((allocavg, alloc_export, allocarr))
        md["module_finalize"] = "\n".join((freeavg, free_export, freearr))
        md["module_pass1"]     = sub_calc1
        md["module_pass1_avg"] = sub_avg1
        md["module_pass2"]     = sub_calc2
        md["module_pass2_avg"] = sub_avg2
        md["module_write_result"] = sub_write_avg

        return md





    def save_ir (self):

        import json

        dg = {k:list(v) for k,v in self.dependency.items()}

        with open("ir2.py", "w") as irf:

            print("g = ", json.dumps(dg, indent=4), file=irf)
            print("l1 = ", json.dumps(self.pass1, indent=4), file=irf)
            print("l2 = ", json.dumps(self.pass2, indent=4), file=irf)
            print("avg1 = ", json.dumps(list(map(repr,self.avg1)), indent=4), file=irf)
            print("avg2 = ", json.dumps(list(map(repr,self.avg2)), indent=4), file=irf)

class Stage5():
    ''' pass1 and pass2 seperation and calculation ordering '''

class Stage6():
    ''' generate fortran code '''