--
--
-- Mathematics Expressions
--
--

inductive pattern MathExpr :=
  | div MathExpr MathExpr
  | (/) MathExpr MathExpr
  | plus [MathExpr]
  | poly [MathExpr]
  | term Integer [(MathExpr, Integer)]
  | mult Integer MathExpr
  | (+) MathExpr MathExpr
  | (*) MathExpr MathExpr
  | (^) MathExpr Integer
  | symbol String [IndexExpr]
  | apply1 (MathExpr -> MathExpr) MathExpr
  | apply2 (MathExpr -> MathExpr -> MathExpr) MathExpr MathExpr
  | apply3 (MathExpr -> MathExpr -> MathExpr -> MathExpr) MathExpr MathExpr MathExpr
  | apply4 (MathExpr -> MathExpr -> MathExpr -> MathExpr -> MathExpr) MathExpr MathExpr MathExpr MathExpr
  | quote MathExpr
  | func MathExpr [MathExpr]

inductive pattern IndexExpr :=
  | sub MathExpr
  | sup MathExpr
  | user MathExpr

def indexExpr : Matcher IndexExpr :=
  matcher
    | sub $ as (mathExpr) with
        | Sub $e -> [e]
        | _ -> []
    | sup $ as (mathExpr) with
        | Sup $e -> [e]
        | _ -> []
    | user $ as (mathExpr) with
        | User $e -> [e]
        | _ -> []
    | #$val as () with
        | $tgt -> if val = tgt then [()] else []
    | $ as something with
        | $tgt -> [tgt]

def mathExpr : Matcher MathExpr :=
  matcher
    | div $ $ as (mathExpr, mathExpr) with
        | Div $p1 $p2 -> [(p1, p2)]
        | _ -> []
    | $ / $ as (mathExpr, mathExpr) with
        | Div $p1 $p2 -> [(p1, p2)]
        | _ -> []
    | poly $ as (multiset mathExpr) with
        | Div (Plus $ts) (Plus [Term 1 []]) -> [ts]
        | _ -> []
    | plus $ as (multiset mathExpr) with
        | Div (Plus $ts) (Plus [Term 1 []]) -> [ts]
        | _ -> []
    | $ + $ as (mathExpr, mathExpr) with
        | Div (Plus $ts) (Plus [Term 1 []]) ->
            matchAll ts as multiset something with
              | $t :: $tss -> (t, sum' tss)
        | _ -> []
    | term $ $ as (integer, assocMultiset mathExpr) with
        | Div (Plus [Term $n $xs]) (Plus [Term 1 []]) ->
            [(n, xs)]
        | _ -> []
    | mult $ $ as (integer, multExpr) with
        | Div (Plus [Term $n $xs]) (Plus [Term 1 []]) ->
            [(n, product' (map (\(x, n) -> x ^' n) xs))]
        | _ -> []
    | $ * $ as (integer, multExpr) with
        | Div (Plus [Term $n $xs]) (Plus [Term 1 []]) ->
            [(n, product' (map (\(x, n) -> x ^' n) xs))]
        | _ -> []
    | symbol $ $ as (something, list indexExpr) with
        | Div (Plus [Term 1 [(Symbol $v $js, 1)]]) (Plus [Term 1 []]) ->
            [(v, js)]
        | _ -> []
    | apply1 $ $ as (something, mathExpr) with
        | Div (Plus [Term 1 [(Apply1 $v $a1, 1)]]) (Plus [Term 1 []]) ->
            [(v, a1)]
        | _ -> []
    | apply2 $ $ $ as (something, mathExpr, mathExpr) with
        | Div (Plus [Term 1 [(Apply2 $v $a1 $a2, 1)]]) (Plus [Term 1 []]) ->
            [(v, a1, a2)]
        | _ -> []
    | apply3 $ $ $ $ as (something, mathExpr, mathExpr, mathExpr) with
        | Div (Plus [Term 1 [(Apply3 $v $a1 $a2 $a3, 1)]]) (Plus [Term 1 []]) ->
            [(v, a1, a2, a3)]
        | _ -> []
    | apply4 $ $ $ $ $ as (something, mathExpr, mathExpr, mathExpr, mathExpr) with
        | Div (Plus [Term 1 [(Apply4 $v $a1 $a2 $a3 $a4, 1)]]) (Plus [Term 1 []]) ->
            [(v, a1, a2, a3, a4)]
        | _ -> []
    | quote $ as (mathExpr) with
        | Div (Plus [Term 1 [(Quote $mexpr, 1)]]) (Plus [Term 1 []]) ->
            [mexpr]
        | _ -> []
    | func $ $ as (mathExpr, list mathExpr) with
        | Div
            (Plus [Term 1 [(Function $name $args, 1)]])
            (Plus [Term 1 []]) ->
            [(name, args)]
        | _ -> []
    | #$val as () with
        | $tgt -> if val = tgt then [()] else []
    | $ as something with
        | $tgt -> [tgt]

def multExpr : Matcher MathExpr :=
  matcher
    | ($ ^ $) * $ as (mathExpr, integer, multExpr) with
        | $tgt ->
            matchAll tgt as mathExpr with
              | term _ (($x, $n) :: $rs) -> (x, n, product' (map (\(x, n) -> x ^' n) rs))
    | $ ^ $ as (mathExpr, integer) with
        | $tgt ->
            match tgt as mathExpr with
              | term _ (($x, $n) :: []) -> [(x, n)]
              | _ -> []
    | $ * $ as (mathExpr, multExpr) with
        | $tgt ->
            matchAll tgt as mathExpr with
              | term _ (($x, $n) :: $rs) -> (x ^' n, product' (map (\(x, n) -> x ^' n) rs))
    | #$val as () with
        | $tgt -> if val = tgt then [()] else []
    | $ as something with
        | $tgt -> [tgt]

def termExpr : Matcher MathExpr := mathExpr

def isSymbol (mexpr: MathExpr) : Bool :=
  match mexpr as mathExpr with
    | symbol _ _ -> True
    | _ -> False

def isApply (mexpr: MathExpr) : Bool :=
  match mexpr as mathExpr with
    | apply1 _ _ -> True
    | apply2 _ _ _ -> True
    | apply3 _ _ _ _ -> True
    | apply4 _ _ _ _ _ -> True
    | _ -> False

def isSimpleTerm (mexpr: MathExpr) : Bool := isSymbol mexpr || isApply mexpr

def isTerm (mexpr: MathExpr) : Bool :=
  match mexpr as mathExpr with
    | term _ _ -> True
    | #0 -> True
    | _ -> False

def isPolynomial (mexpr: MathExpr) : Bool :=
  match mexpr as mathExpr with
    | poly _ -> True
    | #0 -> True
    | _ -> False

def isMonomial (mexpr: MathExpr) : Bool :=
  match mexpr as mathExpr with
    | poly [term _ _] / poly [term _ _] -> True
    | #0 -> True
    | _ -> False

--
-- Accessor
--
def fromMonomial (mexpr: MathExpr) : (MathExpr, MathExpr) :=
  match mexpr as mathExpr with
    | term $a $xs / term $b $ys ->
      (a / b, foldl (*') 1 (map (uncurry (^')) xs) / foldl (*') 1 (map (uncurry (^')) ys))

--
-- Map
--
def mapPolys (fn: MathExpr -> MathExpr) (mexpr: MathExpr) : MathExpr :=
  match mexpr as mathExpr with
    | $p1 / $p2 -> fn p1 /' fn p2

def fromPoly (mexpr: MathExpr) : [MathExpr] :=
  match mexpr as mathExpr with
    | poly $ts1 / $q -> map (\t1 -> t1 /' q) ts1

def mapPoly (fn: MathExpr -> MathExpr) (mexpr: MathExpr) : MathExpr :=
  match mexpr as mathExpr with
    | poly $ts1 / $q -> foldl (+') 0 (map (\t1 -> fn (t1 /' q)) ts1)

def mapTerms (fn: MathExpr -> MathExpr) (mexpr: MathExpr) : MathExpr :=
  match mexpr as mathExpr with
    | poly $ts1 / poly $ts2 ->
        foldl (+') 0 (map fn ts1) /' foldl (+') 0 (map fn ts2)

def mapSymbols (fn: MathExpr -> MathExpr) (mexpr: MathExpr) : MathExpr :=
  mapTerms
    (\match as mathExpr with
      | term $a $xs ->
          a *' foldl
                (*')
                1
                (map
                  (\(x, n) -> match x as mathExpr with
                      | symbol _ _ -> fn x ^' n
                      | apply1 $g $a1 ->
                          let a1' := mapSymbols fn a1
                          in if a1 = a1'
                              then x ^' n
                              else fn (g a1') ^' n
                      | apply2 $g $a1 $a2 ->
                          let a1' := mapSymbols fn a1
                              a2' := mapSymbols fn a2
                          in if a1 = a1' && a2 = a2'
                              then x ^' n
                              else fn (g a1' a2') ^' n
                      | apply3 $g $a1 $a2 $a3 ->
                          let a1' := mapSymbols fn a1
                              a2' := mapSymbols fn a2
                              a3' := mapSymbols fn a3
                          in if a1 = a1' && a2 = a2' && a3 = a3'
                              then x ^' n
                              else fn (g a1' a2' a3') ^' n
                      | apply4 $g $a1 $a2 $a3 $a4 ->
                          let a1' := mapSymbols fn a1
                              a2' := mapSymbols fn a2
                              a3' := mapSymbols fn a3
                              a4' := mapSymbols fn a4
                          in if a1 = a1' && a2 = a2' && a3 = a3' && a4 = a4'
                              then x ^' n
                              else fn (g a1' a2' a3' a4') ^' n
                      | func _ $args ->
                          let args' := map (mapSymbols fn) args
                          in if args = args'
                              then x ^' n
                              else fn (updateFunctionArgs x args') ^' n)
                  xs))
    mexpr

def scanAllTerms (mexpr: MathExpr) (f: MathExpr -> Bool) : Bool :=
  match mexpr as mathExpr with
    | poly $ts1 / poly $ts2 -> any f (ts1 ++ ts2)
    | _ -> not ((debug2 "scanAllTerms" mexpr) = mexpr) -- TODO: if tensorMap is inserted correctly, we canremove this

def containSymbol (x: MathExpr) (mexpr: MathExpr) : Bool :=
  scanAllTerms mexpr
    (\t -> match t as mathExpr with
      | term _ $xs ->
          any
            (\(y, _) -> match y as mathExpr with
              | #x -> True
              | apply1 _ $a1 -> containSymbol x a1
              | apply2 _ $a1 $a2 -> containSymbol x a1 || containSymbol x a2
              | apply3 _ $a1 $a2 $a3 -> containSymbol x a1 || containSymbol x a2 || containSymbol x a3
              | apply4 _ $a1 $a2 $a3 $a4 -> containSymbol x a1 || containSymbol x a2 || containSymbol x a3 || containSymbol x a4
              | _ -> False)
            xs)

def containFunction1 (f : MathExpr -> MathExpr) (mexpr: MathExpr) : Bool :=
  scanAllTerms mexpr
    (\t -> match t as mathExpr with
      | term _ $xs ->
          any
            (\(y, _) -> match y as mathExpr with
              | apply1 #f _ -> True
              | apply1 _ $a1 -> containFunction1 f a1
              | apply2 _ $a1 $a2 -> containFunction1 f a1 || containFunction1 f a2
              | apply3 _ $a1 $a2 $a3 -> containFunction1 f a1 || containFunction1 f a2 || containFunction1 f a3
              | apply4 _ $a1 $a2 $a3 $a4 -> containFunction1 f a1 || containFunction1 f a2 || containFunction1 f a3 || containFunction1 f a4
              | _ -> False) xs)

def containFunction2 (f : MathExpr -> MathExpr -> MathExpr) (mexpr: MathExpr) : Bool :=
  scanAllTerms mexpr
    (\t -> match t as mathExpr with
      | term _ $xs ->
          any
            (\(y, _) -> match y as mathExpr with
              | apply2 #f _ _ -> True
              | apply1 _ $a1 -> containFunction2 f a1
              | apply2 _ $a1 $a2 -> containFunction2 f a1 || containFunction2 f a2
              | apply3 _ $a1 $a2 $a3 -> containFunction2 f a1 || containFunction2 f a2 || containFunction2 f a3
              | apply4 _ $a1 $a2 $a3 $a4 -> containFunction2 f a1 || containFunction2 f a2 || containFunction2 f a3 || containFunction2 f a4
              | _ -> False)
            xs)

def containFunction3 (f : MathExpr -> MathExpr -> MathExpr -> MathExpr) (mexpr: MathExpr) : Bool :=
  scanAllTerms mexpr
    (\t -> match t as mathExpr with
      | term _ $xs ->
          any
            (\(y, _) -> match y as mathExpr with
              | apply3 #f _ _ _ -> True
              | apply1 _ $a1 -> containFunction3 f a1
              | apply2 _ $a1 $a2 -> containFunction3 f a1 || containFunction3 f a2
              | apply3 _ $a1 $a2 $a3 -> containFunction3 f a1 || containFunction3 f a2 || containFunction3 f a3
              | apply4 _ $a1 $a2 $a3 $a4 -> containFunction3 f a1 || containFunction3 f a2 || containFunction3 f a3 || containFunction3 f a4
              | _ -> False)
            xs)

def containFunction4 (f : MathExpr -> MathExpr -> MathExpr -> MathExpr -> MathExpr) (mexpr: MathExpr) : Bool :=
  scanAllTerms mexpr
    (\t -> match t as mathExpr with
      | term _ $xs ->
          any
            (\(y, _) -> match y as mathExpr with
              | apply4 #f _ _ _ _ -> True
              | apply1 _ $a1 -> containFunction4 f a1
              | apply2 _ $a1 $a2 -> containFunction4 f a1 || containFunction4 f a2
              | apply3 _ $a1 $a2 $a3 -> containFunction4 f a1 || containFunction4 f a2 || containFunction4 f a3
              | apply4 _ $a1 $a2 $a3 $a4 -> containFunction4 f a1 || containFunction4 f a2 || containFunction4 f a3 || containFunction4 f a4
              | _ -> False)
            xs)

--
-- Substitute
--
def substitute (ls: [(MathExpr, MathExpr)]) (mexpr: MathExpr) : MathExpr :=
  match ls as list (mathExpr, mathExpr) with
    | [] -> mathNormalize mexpr
    | ($x, $a) :: $rs -> substitute rs (substitute' x a mexpr)

def substitute' (x: MathExpr) (a: MathExpr) (mexpr: MathExpr) : MathExpr := 
  mapSymbols (rewriteSymbol x a) mexpr

def rewriteSymbol (x: MathExpr) (a: MathExpr) (sexpr: MathExpr) : MathExpr :=
  match sexpr as mathExpr with
    | #x -> a
    | _ -> sexpr

def V.substitute (xs: Vector MathExpr) (ys: Vector MathExpr) (mexpr: MathExpr) : MathExpr :=
  substitute (zip (tensorToList xs) (tensorToList ys)) mexpr

def expandAll (mexpr: MathExpr) : MathExpr :=
  match mexpr as mathExpr with
    | ?isInteger -> mexpr
    | ?isSymbol -> mexpr
    -- function application
    | apply1 $g $a1 -> `(g (expandAll a1))
    | apply2 $g $a1 $a2 -> `(g (expandAll a1) (expandAll a2))
    | apply3 $g $a1 $a2 $a3 -> `(g (expandAll a1) (expandAll a2) (expandAll a3))
    | apply4 $g $a1 $a2 $a3 $a4 -> `(g (expandAll a1) (expandAll a2) (expandAll a3) (expandAll a4))
    -- quote
    | quote $g -> g
    -- term (multiplication)
    | term $a $ps -> a * product (map (\(x, n) -> expandAll x ^ n) ps)
    -- polynomial
    | poly $ts -> sum (map expandAll ts)
    -- quotient
    | $p1 / $p2 -> expandAll p1 / expandAll p2

def expandAll' (mexpr: MathExpr) : MathExpr :=
  match mexpr as mathExpr with
    | ?isInteger -> mexpr
    | ?isSymbol -> mexpr
    -- function application
    | apply1 $g $a1 -> `(g (expandAll' a1))
    | apply2 $g $a1 $a2 -> `(g (expandAll' a1) (expandAll' a2))
    | apply3 $g $a1 $a2 $a3 -> `(g (expandAll' a1) (expandAll' a2) (expandAll' a3))
    | apply4 $g $a1 $a2 $a3 $a4 -> `(g (expandAll' a1) (expandAll' a2) (expandAll' a3) (expandAll' a4))
    -- quote
    | quote $g -> g
    -- term (multiplication)
    | term $a $ps -> a *' product' (map (\(x, n) -> expandAll' x ^' n) ps)
    -- polynomial
    | poly $ts -> sum' (map expandAll' ts)
    -- quotient
    | $p1 / $p2 -> expandAll' p1 / expandAll' p2

--
-- Coefficient
--
def coefficients (f: MathExpr) (x: MathExpr) : [MathExpr] :=
  let m := maximum (0 :: (matchAll f as mathExpr with
                           | poly (term $a ((#x, $k) :: $ts) :: _) / _ -> k))
  in map (coefficient f x) (between 0 m)

def coefficient (f: MathExpr) (x: MathExpr) (m: Integer) : MathExpr :=
  if m = 0
    then sum (matchAll f as mathExpr with
               | poly (term $a (!((#x, _) :: _) & $ts) :: _) / _ ->
                 foldl (*') a (map (uncurry (^')) ts)) / denominator f
    else coefficient' f x m

def coefficient' (f: MathExpr) (x: MathExpr) (m: Integer) : MathExpr :=
  sum
    (matchAll f as mathExpr with
      | poly (term $a ((#x, #m) :: (!((#x, _) :: _) & $ts)) :: _) /_ ->
        foldl (*') a (map (uncurry (^')) ts)) /' denominator f

def L./ (xs: [MathExpr]) (ys: [MathExpr]) : ([MathExpr], [MathExpr]) :=
  if length xs < length ys
    then ([], xs)
    else match (ys, xs) as (list mathExpr, list mathExpr) with
      | ($y :: $yrs, $x :: $xrs) ->
        let (zs, rs) := L./ (map2 (-) (take (length yrs) xrs) (map (* (x / y)) yrs) ++ drop (length yrs) xrs) ys
         in (x / y :: zs, rs)