from lark import Lark, Visitor, Transformer, v_args, Token


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):
        lval, rval = tree.children

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

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

        try:
            self.averaged[w.children[0].value] = set([x.value for x in targets])
        except IndexError:
            self.averaged[None] = 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 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 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 : "dabs"
    rxn_rate = lambda self : "rxn_rate"
    udf = lambda self, a  : a.value




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 __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*8, allocatable, dimension({1}) :: {0}"
        return real_array_decl.format(self.name, self.dim)

    def code_alloc (self):
        real_array_alloc = "allocate({0}({1}), stat=ierr) ; {0} = 0."
        return real_array_alloc.format(self.name, self.shape)

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


class Field (FieldBase):

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

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

    def code (self, alloc=None):
        real_array_loop  = """
! {1}
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
"""
        self.array = alloc[self.name] if alloc else self.name

        rhs = ExpToCode(self.fdict).transform(self.exp)
        return real_array_loop.format((self.array, rhs), self.comment)


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)


    def code (self, alloc=None):
        real_array_loop  = """
! {1}
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
"""
        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 real_array_loop.format((self.array, rhs), self.comment)

    @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

    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])

    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()

        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)

    def code (self, alloc=None):

        avg_array_sum  = """
do k = 1, nzp
do j = 1, nyp
do i = 1, nxp
{0}(i) = {0}(i) + {1}
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 avg_array_sum.format(self.name, arrname)


    def code_avg (self):

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

{0} = {0} {1} / denum
"""
        meanw = "/ avg_{}"
        dWeight = (meanw.format(self.weighted) if self.weighted else "")

        return avg_array_divide.format(self.name, 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.iteritems():
            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.iteritems():
            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.iteritems():
            dgraph[k] = v.dep

        for k,v in self.averaged.iteritems():
            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))
        map (pass1calc.update, map(AveragedField.depClosure, self.avg1))
        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))
        map (pass2calc.update, map(AveragedField.depClosure, self.avg2))
        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.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)

        narr = mask.astype(np.int).sum(axis=0).max()

        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 = {}

        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 array_codes (self):

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

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

        decl  = "\n".join([real_array_decl.format(v) for v in array_names])
        alloc = "\n".join([real_array_alloc.format(v) for v in array_names])
        free  = "\n".join([real_array_free.format(v) for v in array_names])

        return decl, alloc, free



    def print_program (self):

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

        declarr, allocarr, freearr = self.array_codes()

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

        calc1 = "\n".join(allvar[v].code(self.alloc1) for v in self.pass1)
        calc2 = "\n".join(allvar[v].code(self.alloc2) 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, declavg, declarr))
        md["module_init"] = "\n".join((allocavg, allocarr))
        md["module_finalize"] = "\n".join((freeavg, freearr))
        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





    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 >>irf, "g = ", json.dumps(dg, indent=4)
            print >>irf, "l1 = ", json.dumps(self.pass1, indent=4)
            print >>irf, "l2 = ", json.dumps(self.pass2, indent=4)
            print >>irf, "avg1 = ", json.dumps(map(repr,self.avg1), indent=4)
            print >>irf, "avg2 = ", json.dumps(map(repr,self.avg2), indent=4)

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

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