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

(define $math-normalize
  (lambda [$fn $x1 $x2]
    (if (or (integer? x1) (integer? x2))
      (fn x1 x2)

;      (fn x1 x2))))
      
;      (let* {[$tmp (fn x1 x2)]
;             [$ret ((capply compose (map 2#%1 (filter 2#(%2 x1 x2) rewrite-rules))) tmp)]}
;        (if (eq? tmp ret)
;          ret
;          (debug ((capply compose (map 2#%1 (filter 2#(%2 x1 x2) rewrite-rules))) (debug ((debug fn) x1 x2)))))))))

      (let* {[$tmp (fn  x1 x2)]
             [$ret ((capply compose (map 2#%1 (filter 2#(%2 x1 x2) rewrite-rules))) tmp)]}
        ret))))

(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 (or (contain-function? cos %1) (contain-function? sin %1)) (or (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 (*' 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]})))

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