import numpy as nm


class DGTerm:

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

    def evaluate(self, mode="weak", diff_var="u",
                 standalone=True, ret_status=False, **kwargs):
        raise NotImplemented

    @staticmethod
    def assemble_to(asm_obj, val, iels, mode="vector"):
        if (asm_obj is not None) and (iels is not None):
            if mode == "vector":
                if (len(iels) == 2) and (nm.shape(val)[0] == len(iels[0])):
                    for ii in iels[0]:
                        asm_obj[ii][iels[1]] = (asm_obj[ii][iels[1]].T + val[ii]).T
                else:
                    asm_obj[iels] = asm_obj[iels] + val

            elif mode == "matrix":
                if (len(iels) == 3) and (nm.shape(val)[0] == len(iels[0])):
                    for ii in iels[0]:
                        asm_obj[ii][iels[1], iels[2]] = asm_obj[ii][iels[1], iels[2]] + val[ii]
                else:
                    asm_obj[iels] = asm_obj[iels] + val
            else:
                raise ValueError("Unknown assebmly mode '%s'" % mode)


class AdvIntDGTerm(DGTerm):

    def __init__(self, mesh):
        DGTerm.__init__(self, mesh)
        self.vvar = "v"
        self.diff_var = "u"

    def evaluate(self, mode="weak", diff_var="u",
                 standalone=True, ret_status=False, **kwargs):
        if diff_var == self.diff_var:

            val = nm.vstack(((self.mesh.coors[1:] - self.mesh.coors[:-1]).T,
                             (self.mesh.coors[1:] - self.mesh.coors[:-1]).T/3))

            iels = ([0, 1], nm.arange(len(self.mesh.coors) - 1),
                            nm.arange(len(self.mesh.coors) - 1))
            # values go on to the diagonal, in sfepy this is assured
            # by mesh connectivity induced by basis
        else:
            val = None
            iels = None

        return val, iels


class AdvFluxDGTerm(DGTerm):

    def __init__(self, mesh, a):
        DGTerm.__init__(self, mesh)
        self.a = a
        self.vvar = None
        self.diff_var = None

    def evaluate(self, mode="weak", diff_var=None,
                 standalone=True, ret_status=False, **kwargs):

        u = kwargs.pop('u', None)
        if diff_var == self.diff_var:
            # TODO check exact integral!
            intg = self.a * u[0, :-2].T
            # intg = (u[0, :-2]  * (self.mesh.coors[1:] - self.mesh.coors[:-1]) +
            #         u[1, :-2]/2 * (self.mesh.coors[1:]**2 - self.mesh.coors[:-1]**2)).T

            # TODO is flux right?
            fl = self.a * (u[0, :-2] + u[1, :-2]).T if self.a > 0 else self.a * u[0, 1:-1].T
            fp = self.a * (u[0, 1:-1] + u[1, 1:-1]).T if self.a > 0 else self.a * u[0, 2:].T

            val = nm.vstack((fl - fp, - fl - fp + intg))

            # placement is simple, bud getting the values requires looping over neighbours
            iels = ([0, 1], nm.arange(len(self.mesh.coors) - 1))  # just fill the vector
        else:
            val = None
            iels = None

        return val, iels

@@ -1062,10 +1062,10 @@
 asse
-b
 m
+b
 ly m
@@ -2302,16 +2302,19 @@
  intg = 
+0 *
 self.a *
@@ -2547,16 +2547,17 @@
       fl
+0
  = self.
@@ -2579,47 +2579,138 @@
 + u%5B
-1, :-2%5D).T if self.a %3E 0 else 
+0, 1:-1%5D).T + abs(self.a)/2 * (u%5B0, :-2%5D - u%5B0, 1:-1%5D).T%0A            fp0 = self.a * (u%5B0, 1:-1%5D + u%5B0,  2:%5D).T + abs(
 self.a
+)/2
  * 
+(
 u%5B0,
@@ -2715,18 +2715,34 @@
 0, 1:-1%5D
-.T
+ - u%5B0,  2:%5D).T%0A%0A%0A
 %0A       
@@ -2751,11 +2751,16 @@
    f
-p
+l1
  =
+ 0 *
  sel
@@ -2772,69 +2772,185 @@
  (u%5B
-0, 1:-1%5D + u%5B1, 1:-1%5D).T if self.a %3E 0 else 
+1, :-2%5D + u%5B1, 1:-1%5D).T + 0 * abs(self.a)/2 * (u%5B1, :-2%5D - u%5B1, 1:-1%5D).T%0A            fp1 = 0 * self.a * (u%5B1, 1:-1%5D + u%5B1,  2:%5D).T + 0 * abs(
 self.a
+)/2
  * 
+(
 u%5B
-0,
+1, 1:-1%5D - u%5B1, 
  2:%5D
+)
 .T%0A%0A
@@ -2984,24 +2984,28 @@
 ((fl
+0
  - fp
+0
 , - fl
+1
  - fp
+1
  + i

Rewrite flux term, try zeroing out first order