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

(define $math-normalize1
  (lambda [$x]
    (if (integer? x)
      x
      (let {[$ret ((capply compose (map 2#%1 (filter 2#(%2 x) rewrite-rules1))) x)]} ret))))
;      (let {[$ret ((capply compose (map 2#%1 (filter 2#(%2 fn x) rewrite-rules1))) (debug x))]} (debug ret)))))

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

;;
;; 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)^,1 $mr)
           (loop $i [2 ,(- n 1)]
             (+ (* ,a ,(rtu 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 (*' 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 1#(rewrite-rule-for-cos-and-sin-expr (map-polys rewrite-rule-for-cos-and-sin-poly %1)))
(define $rewrite-rule-for-cos-and-sin 1#(map-polys rewrite-rule-for-cos-and-sin-poly %1))

(define $rewrite-rule-for-cos-and-sin-expr
  (lambda [$expr]
    (match [expr expr] [math-expr math-expr]
      {[[<div (+ (* $a (,cos $θ) $mr)
                 $pr1)
              $pr2>
         (| <div (+ (* _ (| (,cos ,(/ θ 2)) (,sin ,(/ θ 2))) _) _) _>
            <div _ (+ (* _ (| (,cos ,(/ θ 2)) (,sin ,(/ θ 2))) _) _)>)]
        (rewrite-rule-for-cos-and-sin-expr (/' (+' (*' a (-' (cos (/ θ 2))^2 (sin (/ θ 2))^2) mr) pr1) pr2))]
       [[<div (+ (* $a (,sin $θ) $mr)
                 $pr1)
              $pr2>
         (| <div (+ (* _ (| (,cos ,(/ θ 2)) (,sin ,(/ θ 2))) _) _) _>
            <div _ (+ (* _ (| (,cos ,(/ θ 2)) (,sin ,(/ θ 2))) _) _)>)]
        (rewrite-rule-for-cos-and-sin-expr (/' (+' (*' (*' a 2) (*' (cos (/ θ 2)) (sin (/ θ 2))) mr) pr1) pr2))]
       [[<div $pr2
              (+ (* $a (,cos $θ) $mr)
                 $pr1)>
         (| <div (+ (* _ (| (,cos ,(/ θ 2)) (,sin ,(/ θ 2))) _) _) _>
            <div _ (+ (* _ (| (,cos ,(/ θ 2)) (,sin ,(/ θ 2))) _) _)>)]
        (rewrite-rule-for-cos-and-sin-expr (/' pr2 (+' (*' a (-' (cos (/ θ 2))^2 (sin (/ θ 2))^2) mr) pr1)))]
       [[<div $pr2
              (+ (* $a (,sin $θ) $mr)
                 $pr1)>
         (| <div (+ (* _ (| (,cos ,(/ θ 2)) (,sin ,(/ θ 2))) _) _) _>
            <div _ (+ (* _ (| (,cos ,(/ θ 2)) (,sin ,(/ θ 2))) _) _)>)]
        (rewrite-rule-for-cos-and-sin-expr (/' pr2 (+' (*' (*' a 2) (*' (cos (/ θ 2)) (sin (/ θ 2))) mr) pr1)))]
       [_ expr]})))

(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 (,cos $θ)^,2 $mr)
                   ;           (* ,(* -1 a) (,sin ,θ)^,2 ,mr)
                   ;           $pr)
                ;        (rewrite-rule-for-cos-and-sin-poly (+' pr (*' a (cos (* 2 θ)) mr)))]
       ;       [(+ (* $a (,cos $θ) (,sin ,θ) $mr)
                   ;           $pr)
                ;        (rewrite-rule-for-cos-and-sin-poly (+' pr (*' (/ a 2) (sin (* 2 θ)) mr)))]
       [(+ (* $a $mr)
           (* ,(* -1 a) (,sin $θ)^,2 ,mr)
           $pr)
        (rewrite-rule-for-cos-and-sin-poly (+' pr (*' a (cos θ)^2 mr)))]
       [(+ (* $a $mr)
           (* ,(* -1 a) (,cos $θ)^,2 ,mr)
           $pr)
        (rewrite-rule-for-cos-and-sin-poly (+' pr (*' a (sin θ)^2 mr)))]
       ; (- 1 (cos (* 2 θ))) = (* 2 (sin θ)^2)
       ;       [(+ (* $a $mr)
                   ;           (* ,(* -1 a) (,cos $θ) ,mr)
                   ;           $pr)
                ;        (rewrite-rule-for-cos-and-sin-poly (+' pr (*' (* 2 a) (sin (/ θ 2))^2 mr)))]
       [_ poly]})))

(define $rewrite-rule-for-cos-to-sin 1#(map-terms rewrite-rule-for-cos-to-sin-term' %1))

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

;;
;; 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]})))