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

infixr pattern 6 +
infixr pattern 7 *
infix pattern 7 /
infix pattern 8 ^

def mathExpr :=
  matcher
    | #$val as () with
      | $tgt -> if val = tgt then [()] else []
    | $ as (mathExpr') with
      | $tgt -> [fromMathExpr tgt]

def mathExpr' :=
  matcher
    | $ / $ as (mathExpr, mathExpr) with
      | Div $p1 $p2 -> [(toMathExpr' p1, toMathExpr' p2)]
      | _ -> []
    | poly $ as (multiset mathExpr) with
      | Div (Plus $ts) (Plus [Term 1 []]) -> [map toMathExpr' ts]
      | _ -> []
    | plus $ as (multiset mathExpr) with
      | Div (Plus $ts) (Plus [Term 1 []]) ->
          map (\t -> toMathExpr' (Div (Plus [t]) (Plus [Term 1 []]))) ts
      | _ -> []
    | $ + $ as (mathExpr, mathExpr) with
      | Div (Plus $ts) (Plus [Term 1 []]) ->
          matchAll (map toMathExpr' 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, map (\(x, n) -> (toMathExpr' x, n)) xs)]
      | _ -> []
    | mult $ $ as (integer, multExpr) with
      | Div (Plus [Term $n $xs]) (Plus [Term 1 []]) ->
        [(n, product' (map (\(x, n) -> toMathExpr' x ^' n) xs))]
      | _ -> []
    | $ * $ as (integer, multExpr) with
      | Div (Plus [Term $n $xs]) (Plus [Term 1 []]) ->
        [(n, product' (map (\(x, n) -> toMathExpr' x ^' n) xs))]
      | _ -> []
    | symbol $ $ as (eq, list indexExpr) with
      | Div (Plus [Term 1 [(Symbol $v $js, 1)]]) (Plus [Term 1 []]) ->
        [(v, js)]
      | _ -> []
    | apply $ $ as (eq, list mathExpr) with
      | Div (Plus [Term 1 [(Apply $v $mexprs, 1)]]) (Plus [Term 1 []]) ->
        [(v, map toMathExpr' mexprs)]
      | _ -> []
    | quote $ as (mathExpr) with
      | Div (Plus [Term 1 [(Quote $mexpr, 1)]]) (Plus [Term 1 []]) ->
        [toMathExpr' mexpr]
      | _ -> []
    | func $ $ $ $ as
        (mathExpr, list mathExpr, list mathExpr, list indexExpr) with
      | Div
          (Plus [Term 1 [(Function $name $argnames $args $js, 1)]])
          (Plus [Term 1 []]) ->
        [(name, argnames, args, js)]
      | _ -> []
    | $ as something with
      | $tgt -> [toMathExpr' tgt]

def indexExpr :=
  algebraicDataMatcher
    | sub mathExpr
    | sup mathExpr
    | user mathExpr

def polyExpr := mathExpr

def termExpr := mathExpr

def symbolExpr := mathExpr

def multExpr :=
  matcher
    | [] as () with
      | $tgt ->
        match tgt as mathExpr with
          | #0 -> [()]
          | _ -> []
    | $ ^ #$k * $ as (mathExpr, multExpr) with
      | $tgt ->
        matchAll tgt as mathExpr with
          | term _ (ncons $x #k $xs) -> (x, product' (map (uncurry (^')) xs))
    | $ ^ $ * $ as (mathExpr, integer, multExpr) with
      | $tgt ->
        matchAll tgt as mathExpr with
          | term _ (ncons $x $n $xs) -> (x, n, product' (map (uncurry (^')) xs))
    | $ ^ $ as (mathExpr, integer) with
      | $tgt ->
        match tgt as mathExpr with
          | term _ (ncons $x $n []) -> [(x, n)]
          | _ -> []
    | $ * $ as (mathExpr, multExpr) with
      | $tgt ->
        matchAll tgt as mathExpr with
          | term _ ($x :: $rs) -> (x, product' (map (uncurry (^')) rs))
    | $ as mathExpr with
      | $tgt -> [tgt]

def isSymbol %mexpr :=
  match mexpr as mathExpr with
    | symbol _ _ -> True
    | _ -> False

def isApply %mexpr :=
  match mexpr as mathExpr with
    | apply _ _ -> True
    | _ -> False

def isSimpleTerm mexpr := isSymbol mexpr || isApply mexpr

def isTerm %mexpr :=
  match mexpr as mathExpr with
    | term _ _ -> True
    | #0 -> True
    | _ -> False

def isPolynomial %mexpr :=
  match mexpr as mathExpr with
    | poly _ -> True
    | #0 -> True
    | _ -> False

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

--
-- Accessor
--
def fromMonomial $mexpr :=
  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 $mexpr :=
  match mexpr as mathExpr with
    | $p1 / $p2 -> fn p1 /' fn p2

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

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

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

def mapSymbols $fn $mexpr :=
  mapTerms
    (\match as termExpr with
      | term $a $xs ->
        a *' foldl
               (*')
               1
               (map
                  (\(x, n) -> match x as symbolExpr with
                    | symbol _ _ -> fn x ^' n
                    | apply $g $args ->
                      let args' := map (mapSymbols fn) args
                       in if args = args'
                            then x ^' n
                            else fn (capply g args') ^' n)
                  xs))
    mexpr

def scanAllTerms $mexpr $f :=
  match mexpr as mathExpr with
    | poly $ts1 / poly $ts2 -> any f (ts1 ++ ts2)

def containSymbol $x $mexpr :=
  scanAllTerms mexpr
    (\match as termExpr with
      | term _ $xs ->
        any
          (\(y, _) -> match y as symbolExpr with
            | #x -> True
            | apply _ $args -> any (containSymbol x) args
            | _ -> False)
          xs)

def containFunction $f $mexpr :=
  scanAllTerms mexpr
    (\match as termExpr with
      | term _ $xs ->
        any
          (\(y, _) -> match y as symbolExpr with
            | apply #f _     -> True
            | apply $g $args -> any (containFunction f) args
            | _ -> False)
          xs)

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

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

def rewriteSymbol $x $a $sexpr :=
  match sexpr as symbolExpr with
    | #x -> a
    | _ -> sexpr

def V.substitute %xs %ys $mexpr :=
  substitute (zip (tensorToList xs) (tensorToList ys)) mexpr

def expandAll $mexpr :=
  match mexpr as mathExpr with
    | ?isSymbol -> mexpr
    -- function application
    | apply $g $args -> capply g (map expandAll args)
    -- quote
    | quote $g -> g
    -- term (multiplication)
    | term $a $ps -> a * product (map (\(x, n) -> expandAll x ^ expandAll n) ps)
    -- polynomial
    | poly $ts -> sum (map expandAll ts)
    -- quotient
    | $p1 / $p2 -> expandAll p1 / expandAll p2

def expandAll' $mexpr :=
  match mexpr as mathExpr with
    | ?isSymbol -> mexpr
    -- function application
    | apply $g $args -> capply g (map expandAll' args)
    -- quote
    | quote $g -> g
    -- term (multiplication)
    | term $a $ps -> a *' product' (map (\(x, n) -> expandAll' x ^' expandAll' n) ps)
    -- polynomial
    | poly $ts -> sum' (map expandAll' ts)
    -- quotient
    | $p1 / $p2 -> expandAll' p1 / expandAll' p2

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

def coefficient $f $x $m :=
  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 $x $m :=
  sum
    (matchAll f as mathExpr with
      | poly (term $a (ncons #x #m (!(#x :: _) & $ts)) :: _) / _ ->
        foldl (*') a (map (uncurry (^')) ts)) / denominator f

def coefficient2 $f $x $y :=
  sum
    (matchAll f as mathExpr with
      | poly (term $a (#x :: #y :: $ts) :: _) / _ ->
        foldl (*') a (map (uncurry (^')) ts)) / denominator f