;;;;;
;;;;;
;;;;; Mathematics Expressions
;;;;;
;;;;;

(define $math-expr
  (matcher
    {[,$val []
      {[$tgt (if (eq? val tgt)
               {[]}
               {})]}]
     [$ [math-expr']
      {[$tgt {(from-math-expr tgt)}]}]
     }))

(define $math-expr'
  (matcher
    {[<div $ $> [math-expr math-expr]
      {[<Div $p1 $p2> {[(to-math-expr' p1) (to-math-expr' p2)]}]
       [_ {}]}]
     [<poly $> [(multiset math-expr)]
      {[<Div <Plus $ts> <Plus {<Term 1 {}> @{}}>> {(map to-math-expr' ts)}]
       [_ {}]}]
     [<plus $> [plus-expr]
      {[<Div <Plus $ts> <Plus {<Term 1 {}> @{}}>> {(to-math-expr' <Div <Plus ts> <Plus {<Term 1 {}>}>>)}]
       [_ {}]}]
     [<term $ $> [integer (assoc-multiset math-expr)]
      {[<Div <Plus {<Term $n $xs> @{}}> <Plus {<Term 1 {}> @{}}>> {[n (map 2#[(to-math-expr' %1) %2] xs)]}]
       [_ {}]}]
     [<mult $ $> [integer mult-expr]
      {[<Div <Plus {<Term $n $xs> @{}}> <Plus {<Term 1 {}> @{}}>> {[n (product' (map 2#(**' (to-math-expr' %1) %2) xs))]}]
       [_ {}]}]
     [<symbol $> [eq]
      {[<Div <Plus {<Term 1 {[<Symbol $v {}> 1] @{}}> @{}}> <Plus {<Term 1 {}> @{}}>> {v}]
       [_ {}]}]
     [<symbol $ $> [eq (list index-expr)]
      {[<Div <Plus {<Term 1 {[<Symbol $v $js> 1] @{}}> @{}}> <Plus {<Term 1 {}> @{}}>> {[v js]}]
       [_ {}]}]
     [<apply $ $> [eq (list math-expr)]
      {[<Div <Plus {<Term 1 {[<Apply $v $mexprs> 1] @{}}> @{}}>
             <Plus {<Term 1 {}> @{}}>> 
        {[v (map to-math-expr' mexprs)]}]
       [_ {}]}]
     [<quote $> [math-expr]
      {[<Div <Plus {<Term 1 {[<Quote $mexpr> 1] @{}}> @{}}>
             <Plus {<Term 1 {}> @{}}>>
        {(to-math-expr' mexpr)}]
       [_ {}]}]
     [$ [something]
      {[$tgt {(to-math-expr' tgt)}]}]
     }))

(define $index-expr
  (algebraic-data-matcher
    {<sub math-expr> <sup math-expr>}))

(define $poly-expr math-expr)
(define $term-expr math-expr)
(define $symbol-expr math-expr)

(define $plus-expr
  (matcher
    {[<nil> []
      {[$tgt (if (eq? tgt 0)
               {[]}
               {})]}]
     [<cons $ $> [math-expr plus-expr]
      {[$tgt (match-all tgt math-expr
               [<poly <cons $t $ts>> [t (sum' ts)]])]}]
     [$ [math-expr]
      {[$tgt {tgt}]}]
     }))

(define $mult-expr
  (matcher
    {[<nil> []
      {[$tgt (match tgt math-expr
               {[,0 {[]}]
                [_ {}]})]}]
     [<cons $ $> [math-expr mult-expr]
      {[$tgt (match tgt math-expr
               {[<term _ $xs>
                 (match-all xs (assoc-multiset math-expr)
                   [<cons $x $rs>
                    [x (product' (map 2#(**' %1 %2) rs))]])]
                [_ {}]})]}]
     [<ncons $ ,$k $> [math-expr mult-expr]
      {[$tgt (match tgt math-expr
               {[<term _ $xs>
                 (match-all xs (list [math-expr integer])
                   [<join $hs <cons [$x (& ?(gte? $ k) $n)] $ts>>
                    [x (product' (map 2#(**' %1 %2) {@hs [x (- n k)] @ts}))]])]
                [_ {}]})]}]
     [<ncons $ $ $> [math-expr integer mult-expr]
      {[$tgt (match tgt math-expr
               {[<term _ $xs>
                 (match-all xs (list [math-expr integer])
                   [<join $hs <cons [$x $n] $ts>>
                    [x n (product' (map 2#(**' %1 %2) {@hs @ts}))]])]
                [_ {}]})]}]
     [$ [math-expr]
      {[$tgt {tgt}]}]
     }))

;;
;; Predicate
;;
(define $symbol?
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<symbol _ _> #t]
       [_ #f]})))

(define $tensor-symbol?
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<symbol _ <join _ <cons (| <sub ?symbol?> <sup ?symbol?>) _>>> #t]
       [_ #f]})))

(define $apply?
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<apply _ _> #t]
       [_ #f]})))

(define $simple-term? 1#(or (symbol? %1) (apply? %1)))

(define $term?
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<term _ _> #t]
       [,0 #t]
       [_ #f]})))

(define $polynomial?
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<poly _> #t]
       [,0 #t]
       [_ #f]})))

(define $monomial?
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<div <poly <cons <term _ _> <nil>>>
             <poly <cons <term _ _> <nil>>>>
        #t]
       [,0 #t]
       [_ #f]})))

;;
;; Accessor
;;

(define $symbol-indices
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<symbol _ $js> js]
       [_ undefined]})))

(define $from-monomial
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<div <term $a $xs>
             <term $b $ys>>
        [(/ a b)
         (/ (foldl *' 1 (map 2#(**' %1 %2) xs))
            (foldl *' 1 (map 2#(**' %1 %2) ys)))]]})))

;;
;; Map
;;
(define $map-polys
  (lambda [$fn $mexpr]
    (match mexpr math-expr
      {[<div $p1 $p2>
        (/' (fn p1) (fn p2))]})))

(define $from-poly
  (lambda [$mexpr]
    (match mexpr math-expr
      {[<div <poly $ts1> $q>
        (map (lambda [$t1] (/' t1 q))
             ts1)]})))

(define $map-poly
  (lambda [$fn $mexpr]
    (match mexpr math-expr
      {[<div <poly $ts1> $q>
        (foldl +' 0 (map (lambda [$t1] (fn (/' t1 q)))
                         ts1))]})))

(define $map-terms
  (lambda [$fn $mexpr]
    (match mexpr math-expr
      {[<div <poly $ts1> <poly $ts2>>
        (/' (foldl +' 0 (map fn ts1))
            (foldl +' 0 (map fn ts2)))]})))

(define $map-symbols
  (lambda [$fn $mexpr]
    (map-terms (lambda [$term]
                 (match term term-expr
                   {[<term $a $xs>
                     (*' a (foldl *' 1 (map 2#(match %1 symbol-expr
                                                {[<symbol _ _> (**' (fn %1) %2)]
                                                 [<apply $g $args>
                                                  (let {[$args'(map (map-symbols fn $) args)]}
                                                    (if (eq? args args')
                                                      (**' %1 %2)
                                                      (**' (fn (capply g args'))
                                                          %2)))
                                                  ]})
                                            xs)))]}))
               mexpr)))

(define $contain-symbol?
  (lambda [$x $mexpr]
    (any id (match mexpr math-expr
              {[<div <poly $ts1> <poly $ts2>>
                (map (lambda [$term]
                       (match term term-expr
                         {[<term _ $xs>
                           (any id (map 2#(match %1 symbol-expr
                                            {[,x #t]
                                             [<apply _ $args> (any id (map (contain-symbol? x $) args))]
                                             [_ #f]})
                                        xs))]}))
                     {@ts1 @ts2})]}))))

(define $contain-function?
  (lambda [$f $mexpr]
    (any id (match mexpr math-expr
              {[<div <poly $ts1> <poly $ts2>>
                (map (lambda [$term]
                       (match term term-expr
                         {[<term _ $xs>
                           (any id (map 2#(match %1 symbol-expr
                                            {[<apply $g $args>
                                              (if (eq? f g)
                                                #t
                                                (any id (map (contain-function? f $) args)))]
                                             [_ #f]})
                                        xs))]}))
                     {@ts1 @ts2})]}))))

(define $contain-function-with-order?
  (lambda [$f $n $mexpr]
    (any id (match mexpr math-expr
              {[<div <poly $ts1> <poly $ts2>>
                (map (lambda [$term]
                       (match term term-expr
                         {[<term _ $xs>
                           (any id (map 2#(match %1 symbol-expr
                                            {[<apply $g $args>
                                              (if (and (eq? f g) (gte? %2 n))
                                                #t
                                                (any id (map (contain-function-with-order? f n $) args)))]
                                             [_ #f]})
                                        xs))]}))
                     {@ts1 @ts2})]}))))

(define $contain-function-with-index?
  (lambda [$mexpr]
    (any id (match mexpr math-expr
              {[<div <poly $ts1> <poly $ts2>>
                (map (lambda [$term]
                       (match term term-expr
                         {[<term _ $xs>
                           (any id (map 2#(match %1 symbol-expr
                                            {[<apply (& ?scalar? $f) $args>
                                              (match f math-expr
                                                {[<symbol _ !<nil>> #t]
                                                 [_ (any id (map (contain-function-with-index? $) args))]})]
                                             [<apply _ $args>
                                              (any id (map (contain-function-with-index? $) args))]
                                             [_ #f]})
                                        xs))]}))
                     {@ts1 @ts2})]}))))
  
(define $find-symbols-from-poly
  (lambda [$poly]
    (match-all poly math-expr
      [<poly <cons <term _ <cons (& <symbol _> $s) _>> _>> s])))

;;;
;;; Substitute
;;;
(define $substitute
  (lambda [$ls $mexpr]
    (match ls (list [symbol-expr math-expr])
      {[<nil> mexpr]
       [<cons [$x $a] $rs>
        (substitute rs (substitute' x a mexpr))]})))

(define $substitute'
  (lambda [$x $a $mexpr]
    (map-symbols (rewrite-symbol x a $) mexpr)))

(define $rewrite-symbol 
  (lambda [$x $a $sexpr]
    (match sexpr symbol-expr
      {[,x a]
       [_ sexpr]})))

(define $V.substitute
  (lambda [%xs %ys $mexpr]
    (substitute (zip (tensor-to-list xs) (tensor-to-list ys)) mexpr)))

(define $expand-all
  (lambda [$mexpr]
    (match mexpr math-expr
      {
       [?symbol? mexpr]
       ; function application
       [<apply $g $args>
        (capply g (map expand-all args))]
       ; quote
       [<quote $g> g]
       ; term (multiplication)
       [<term $a $ps>
        (* a (product (map 2#(** (expand-all %1) (expand-all %2)) ps)))]
       ; polynomial
       [<poly $ts> (sum (map (expand-all $) ts))]
       ; quotient
       [(/ $p1 $p2)
        (let {[$p1' (expand-all p1)]
              [$p2' (expand-all p2)]}
          (/ p1' p2'))]
       })))

(define $expand-all'
  (lambda [$mexpr]
    (match mexpr math-expr
      {
       [?symbol? mexpr]
       ; function application
       [<apply $g $args>
        (capply g (map expand-all' args))]
       ; quote
       [<quote $g> g]
       ; term (multiplication)
       [<term $a $ps>
        (*' a (product' (map 2#(**' (expand-all' %1) (expand-all' %2)) ps)))]
       ; polynomial
       [<poly $ts> (sum' (map (expand-all' $) ts))]
       ; quotient
       [(/ $p1 $p2)
        (let {[$p1' (expand-all' p1)]
              [$p2' (expand-all' p2)]}
          (/' p1' p2'))]
       })))

;;;
;;; Coefficient
;;;
(define $coefficients
  (lambda [$f $x]
    (let {[$m (capply max {0 @(match-all f math-expr
                                [<div <poly <cons <term $a <ncons ,x $k $ts>> _>> _> k])})]}
      (map (coefficient f x $) (between 0 m)))))

(define $coefficient
  (lambda [$f $x $m]
    (if (eq? m 0)
      (/ (sum (match-all f math-expr
                [<div <poly <cons <term $a (& !<cons ,x _> $ts)> _>> _>
                 (foldl *' a (map 2#(**' %1 %2) ts))]))
         (denominator f))
      (coefficient' f x m))))

(define $coefficient'
  (lambda [$f $x $m]
    (/ (sum (match-all f math-expr
              [<div <poly <cons <term $a <ncons ,x $k $ts>> _>> _>
               (if (eq? m k)
                 (foldl *' a (map 2#(**' %1 %2) ts))
                 0)]))
       (denominator f))))

(define $coefficient2
  (lambda [$f $x $y]
    (/ (sum (match-all f math-expr
              [<div <poly <cons <term $a <cons ,x <cons ,y $ts>>> _>> _>
               (foldl *' a (map 2#(**' %1 %2) ts))
               ]))
       (denominator f))))