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


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

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


class FieldBase (object):
    def __init__ (self, name, fdict):
        self.name = name
        self.array = name
        self.dep = set([])
        self.fluc = False
        self.fdict = fdict

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

    def is_fluctuation (self):
        return False

    def __repr__ (self):
        return 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)

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

    def is_fluctuation (self):
        return self.fluc

    def code (self, alloc):
        transform(self.exp)
        return (a in self.dep)


class PrimaryField (FieldBase):

    def __init__ (self, name, fdict):
        super(PrimaryField,self).__init__(name, fdict)
        self.derivs = set([])

    def code (self):
        return ""


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):
        return "call {0} ( {2}, {1} )".format(self.op, self.v, self.array)


class AveragedField (FieldBase):

    def __init__ (self, w, tgt, fdict):
        if w:
            name = "{}_avg_{}".format(w, tgt)
        else:
            name = "avg_{}".format(tgt)
        super(AveragedField,self).__init__(name, fdict)
        self.tgt = fdict[tgt]
        self.dep.add(tgt)
        if w:
            self.w = fdict[w]
            self.dep.add(tgt)

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 Stage1():
    ''' conversion from tree to python data '''

    def __init__ (self, raw_tree):
        self.primary = set([])
        self.derived = {}
        self.averaged = {}
        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 '''

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

        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

        for w, tgts in src.averaged.iteritems():
            for t in tgts:
                AveragedField(w, t, self.derived)

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



class Stage3():
    ''' analyze fluctuation '''

    def __init__ (self, src):
        self.src = src
        self.derived = src.derived
        self.derivative = src.derivative

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

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

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




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

class Stage5():
    ''' analyze liveness and allocate array '''

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