;;;;;
;;;;;
;;;;; Term Rewriting
;;;;;
;;;;;

(define $math-normalize
  (lambda [$fn $x1 $x2]
    (if (and (integer? x1) (integer? x2))
      (fn x1 x2)
      ((capply compose (map 2#%1 (filter 2#(%2 x1 x2) rewrite-rules))) (fn x1 x2)))))

(define $rewrite-rules
  {
   [id 2##t]
   [rewrite-rule-for-i 2#(and (contain-symbol? i %1) (contain-symbol? i %2))]
   [rewrite-rule-for-w-term 2#(and (contain-symbol? w %1) (contain-symbol? w %2))]
   [rewrite-rule-for-rtu-term 2#(and (contain-function? rtu %1) (contain-function? rtu %2))]
   [rewrite-rule-for-w-poly 2#(and (contain-symbol? w %1) (contain-symbol? w %2))]
   [rewrite-rule-for-rtu-poly 2#(and (contain-function? rtu %1) (contain-function? rtu %2))]
   [rewrite-rule-for-sqrt 2#(and (contain-function? sqrt %1) (contain-function? sqrt %2))]
   [rewrite-rule-for-rt 2#(and (contain-function? rt %1) (contain-function? rt %2))]
   [rewrite-rule-for-cos-and-sin 2#(or (contain-function? cos %1) (contain-function? sin %1) (contain-function? cos %2) (contain-function? sin %2))]
;   [rewrite-rule-for-∂/∂ 2#(and (contain-function-with-index? %1) (contain-function-with-index? %2))]
   })

;;
;; i
;;

(define $rewrite-rule-for-i rewrite-rule-for-i-term)

(define $rewrite-rule-for-i-term (map-terms rewrite-rule-for-i-term' $))

(define $rewrite-rule-for-i-term'
  (lambda [$term]
    (match term math-expr
      {[(* $a ,i^(& ?even? $k) $r)
        (*' a (**' -1 (quotient k 2)) r)]
       [(* $a ,i^$k $r)
        (*' a (**' -1 (quotient k 2)) r i)]
       [_ term]})))

;;
;; w
;;

(define $rewrite-rule-for-w
  (compose rewrite-rule-for-w-term
           rewrite-rule-for-w-poly $))

(define $rewrite-rule-for-w-term (map-terms rewrite-rule-for-w-term' $))
(define $rewrite-rule-for-w-poly (map-polys rewrite-rule-for-w-poly' $))

(define $rewrite-rule-for-w-term'
  (lambda [$term]
    (match term math-expr
      {[(* $a ,w^(& ?(gte? $ 3) $k) $r)
        (*' a r (**' w (remainder k 3)))]
       [_ term]})))

(define $rewrite-rule-for-w-poly'
  (lambda [$poly]
    (match poly math-expr
      {[(+ (* $a ,w^,2 $mr)
           (* $b ,w ,mr)
           $pr)
        (rewrite-rule-for-w-poly' (+' pr
                                     (*' -1 a mr)
                                     (*' (- b a) mr w)
                                     ))]
       [_ poly]})))

;;
;; rtu (include i and w)
;;

(define $rewrite-rule-for-rtu
  (compose 
           (map-terms rewrite-rule-for-rtu-term $)
           (map-polys rewrite-rule-for-rtu-poly $)
           ))

(define $rewrite-rule-for-rtu-term (map-terms rewrite-rule-for-rtu-term' $))
(define $rewrite-rule-for-rtu-poly (map-polys rewrite-rule-for-rtu-poly' $))

(define $rewrite-rule-for-rtu-term'
  (lambda [$term]
    (match term math-expr
      {[(* $a (,rtu $n)^(& ?(gte? $ n) $k) $r)
        (*' a (**' (rtu n) (remainder k n)) r)]
       [_ term]})))

(define $rewrite-rule-for-rtu-poly'
  (lambda [$poly]
    (match poly math-expr
      {[(+ (* $a (,rtu $n)^,(- n 1) $mr)
           (loop $i [2 ,(- n 1)]
             (+ (* ,a ,(rtu n)^,(- n i) ,mr) ...)
             $pr))
        (rewrite-rule-for-rtu-poly' (+' pr (*' -1 a mr)))]
       [_ poly]})))

;;
;; sqrt
;;

(define $rewrite-rule-for-sqrt (map-terms rewrite-rule-for-sqrt-term $))

(define $rewrite-rule-for-sqrt-term
  (lambda [$term]
    (match term math-expr
      {[(* $a (,sqrt $x) (,sqrt ,x) $r)
        (rewrite-rule-for-sqrt-term (*' a x r))]
       [(* $a (,sqrt (& ?term? $x)) (,sqrt (& ?term? $y)) $r)
        (let* {[$d (gcd x y)]
               [[$a1 $x1] (from-monomial (/ x d))]
               [[$a2 $y1] (from-monomial (/ y d))]}
          (*' a d
             (sqrt (*' a1 a2)) (sqrt x1) (sqrt y1)
             r))]
       [_ term]})))

;;
;; rt (include sqrt)
;;

(define $rewrite-rule-for-rt
  (map-terms rewrite-rule-for-rt-term $))

(define $rewrite-rule-for-rt-term
  (lambda [$term]
    (match term math-expr
      {[(* $a (,rt $n $x)^(& ?(gte? $ n) $k) $r)
        (*' a (**' x (quotient k n)) (**' (rt n x) (remainder k n)) r)]
       [_ term]})))

;;
;; cos, sin
;;

(define $rewrite-rule-for-cos-and-sin (map-polys rewrite-rule-for-cos-and-sin-poly $))

(define $rewrite-rule-for-cos-and-sin-poly
  (lambda [$poly]
    (match poly math-expr
      {[(+ (* $a (,cos $θ)^,2 $mr)
           (* ,a (,sin ,θ)^,2 ,mr)
           $pr)
        (rewrite-rule-for-cos-and-sin-poly (+' pr (*' a mr)))]
       [(+ (* $a $mr)
           (* ,(*' -1 a) (,sin $θ)^,2 ,mr)
           $pr)
        (rewrite-rule-for-cos-and-sin-poly (+' (*' a (cos θ)^2 mr)))]
       [(+ (* $a $mr)
           (* ,(*' -1 a) (,cos $θ)^,2 ,mr)
           $pr)
        (rewrite-rule-for-cos-and-sin-poly (+' (*' a (sin θ)^2)))]
       [_ poly]})))

;;
;; d
;;

(define $rewrite-rule-for-d (map-terms rewrite-rule-for-d-term $))

(define $rewrite-rule-for-d-term
  (lambda [$term]
    (match term math-expr
      {[(* _ (,d _) (,d _) _)
        0]
       [_ term]})))

;;
;; ∂/∂
;;

(define $rewrite-rule-for-∂/∂ (map-polys rewrite-rule-for-∂/∂-poly $))

(define $rewrite-rule-for-∂/∂-poly
  (lambda [$poly]
    (match poly math-expr
      {[(+ (* $a <apply (& $g <symbol $f $subs>) $args>^$n $mr)
           (* $b <apply <symbol ,f ?1#(eq?/m (multiset something) subs %1)> ,args>^,n ,mr)
           $pr)
        (+ (* (+ a b) (`g args)^n mr)
           pr)]
       [_ poly]})))