##  Copyright (c) Eduardo Julian. All rights reserved.
##  This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
##  If a copy of the MPL was not distributed with this file,
##  You can obtain one at http://mozilla.org/MPL/2.0/.

(;module: {#;doc "Complex arithmetic."}
  lux
  (lux [math]
       (control eq
                [ord]
                number
                codec
                monad)
       (data [number "r/" Number<Real> Codec<Text,Real>]
             [text "Text/" Monoid<Text>]
             text/format
             error
             maybe
             (coll [list "List/" Monad<List>]))
       [compiler]
       (macro [ast]
              ["s" syntax #+ syntax: Syntax])))

## Based on org.apache.commons.math4.complex.Complex
## https://github.com/apache/commons-math/blob/master/src/main/java/org/apache/commons/math4/complex/Complex.java

(type: #export Complex
  {#real Real
   #imaginary Real})

(syntax: #export (complex real [?imaginary (s;opt s;any)])
  {#;doc (doc "Complex literals."
              (complex real imaginary)
              "The imaginary part can be omitted if it's 0."
              (complex real))}
  (wrap (list (` {#;;real (~ real)
                  #;;imaginary (~ (default (' 0.0)
                                    ?imaginary))}))))

(def: #export i Complex (complex 0.0 1.0))

(def: #export one Complex (complex 1.0 0.0))

(def: #export zero Complex (complex 0.0 0.0))

(def: #export (nan? complex)
  (or (number;nan? (get@ #real complex))
      (number;nan? (get@ #imaginary complex))))

(def: #export (c.= param input)
  (-> Complex Complex Bool)
  (and (r.= (get@ #real param)
            (get@ #real input))
       (r.= (get@ #imaginary param)
            (get@ #imaginary input))))

(do-template [<name> <op>]
  [(def: #export (<name> param input)
     (-> Complex Complex Complex)
     {#real (<op> (get@ #real param)
                  (get@ #real input))
      #imaginary (<op> (get@ #imaginary param)
                       (get@ #imaginary input))})]

  [c.+ r.+]
  [c.- r.-]
  )

(struct: #export _ (Eq Complex)
  (def: = c.=))

(def: #export c.negate
  (-> Complex Complex)
  (|>. (update@ #real r/negate)
       (update@ #imaginary r/negate)))

(def: #export c.signum
  (-> Complex Complex)
  (|>. (update@ #real r/signum)
       (update@ #imaginary r/signum)))

(def: #export conjugate
  (-> Complex Complex)
  (update@ #imaginary r/negate))

(def: #export (c.*' param input)
  (-> Real Complex Complex)
  {#real (r.* param
              (get@ #real input))
   #imaginary (r.* param
                   (get@ #imaginary input))})

(def: #export (c.* param input)
  (-> Complex Complex Complex)
  {#real (r.- (r.* (get@ #imaginary param)
                   (get@ #imaginary input))
              (r.* (get@ #real param)
                   (get@ #real input)))
   #imaginary (r.+ (r.* (get@ #real param)
                        (get@ #imaginary input))
                   (r.* (get@ #imaginary param)
                        (get@ #real input)))})

(def: #export (c./ param input)
  (-> Complex Complex Complex)
  (let [(^slots [#real #imaginary]) param]
    (if (r.< (r/abs imaginary)
             (r/abs real))
      (let [quot (r./ imaginary real)
            denom (|> real (r.* quot) (r.+ imaginary))]
        {#real (|> (get@ #real input) (r.* quot) (r.+ (get@ #imaginary input)) (r./ denom))
         #imaginary (|> (get@ #imaginary input) (r.* quot) (r.- (get@ #real input)) (r./ denom))})
      (let [quot (r./ real imaginary)
            denom (|> imaginary (r.* quot) (r.+ real))]
        {#real (|> (get@ #imaginary input) (r.* quot) (r.+ (get@ #real input)) (r./ denom))
         #imaginary (|> (get@ #imaginary input) (r.- (r.* quot (get@ #real input))) (r./ denom))}))))

(def: #export (c./' param subject)
  (-> Real Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (r./ param real)
     #imaginary (r./ param imaginary)}))

(def: #export (cos subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (r.* (math;cosh imaginary)
                (math;cos real))
     #imaginary (r.* (math;sinh imaginary)
                     (r/negate (math;sin real)))}))

(def: #export (cosh subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (r.* (math;cos imaginary)
                (math;cosh real))
     #imaginary (r.* (math;sin imaginary)
                     (math;sinh real))}))

(def: #export (sin subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (r.* (math;cosh imaginary)
                (math;sin real))
     #imaginary (r.* (math;sinh imaginary)
                     (math;cos real))}))

(def: #export (sinh subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (r.* (math;cos imaginary)
                (math;sinh real))
     #imaginary (r.* (math;sin imaginary)
                     (math;cosh real))}))

(def: #export (tan subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject
        r2 (r.* 2.0 real)
        i2 (r.* 2.0 imaginary)
        d (r.+ (math;cos r2) (math;cosh i2))]
    {#real (r./ d (math;sin r2))
     #imaginary (r./ d (math;sinh i2))}))

(def: #export (tanh subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject
        r2 (r.* 2.0 real)
        i2 (r.* 2.0 imaginary)
        d (r.+ (math;cosh r2) (math;cos i2))]
    {#real (r./ d (math;sinh r2))
     #imaginary (r./ d (math;sin i2))}))

(def: #export (c.abs subject)
  (-> Complex Real)
  (let [(^slots [#real #imaginary]) subject]
    (if (r.< (r/abs imaginary)
             (r/abs real))
      (if (r.= 0.0 imaginary)
        (r/abs real)
        (let [q (r./ imaginary real)]
          (r.* (math;sqrt (r.+ 1.0 (r.* q q)))
               (r/abs imaginary))))
      (if (r.= 0.0 real)
        (r/abs imaginary)
        (let [q (r./ real imaginary)]
          (r.* (math;sqrt (r.+ 1.0 (r.* q q)))
               (r/abs real))))
      )))

(def: #export (exp subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject
        r-exp (math;exp real)]
    {#real (r.* r-exp (math;cos imaginary))
     #imaginary (r.* r-exp (math;sin imaginary))}))

(def: #export (log subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (math;log (c.abs subject))
     #imaginary (math;atan2 real imaginary)}))

(do-template [<name> <type> <op>]
  [(def: #export (<name> param input)
     (-> <type> Complex Complex)
     (|> input log (<op> param) exp))]

  [pow  Complex c.*]
  [pow' Real    c.*']
  )

(def: (copy-sign sign magnitude)
  (-> Real Real Real)
  (r.* (r/signum sign) magnitude))

(def: #export (sqrt (^@ input (^slots [#real #imaginary])))
  (-> Complex Complex)
  (let [t (|> input c.abs (r.+ (r/abs real)) (r./ 2.0) math;sqrt)]
    (if (r.>= 0.0 real)
      {#real t
       #imaginary (r./ (r.* 2.0 t)
                       imaginary)}
      {#real (r./ (r.* 2.0 t)
                  (r/abs imaginary))
       #imaginary (r.* t (copy-sign imaginary 1.0))})))

(def: #export (sqrt-1z input)
  (-> Complex Complex)
  (|> (complex 1.0) (c.- (c.* input input)) sqrt))

(def: #export (reciprocal (^slots [#real #imaginary]))
  (-> Complex Complex)
  (if (r.< (r/abs imaginary)
           (r/abs real))
    (let [q (r./ imaginary real)
          scale (r./ (|> real (r.* q) (r.+ imaginary))
                     1.0)]
      {#real (r.* q scale)
       #imaginary (r/negate scale)})
    (let [q (r./ real imaginary)
          scale (r./ (|> imaginary (r.* q) (r.+ real))
                     1.0)]
      {#real scale
       #imaginary (|> scale r/negate (r.* q))})))

(def: #export (acos input)
  (-> Complex Complex)
  (|> input
      (c.+ (|> input sqrt-1z (c.* i)))
      log
      (c.* (c.negate i))))

(def: #export (asin input)
  (-> Complex Complex)
  (|> input
      sqrt-1z
      (c.+ (c.* i input))
      log
      (c.* (c.negate i))))

(def: #export (atan input)
  (-> Complex Complex)
  (|> input
      (c.+ i)
      (c./ (c.- input i))
      log
      (c.* (c./ (complex 2.0) i))))

(def: #export (argument (^slots [#real #imaginary]))
  (-> Complex Real)
  (math;atan2 real imaginary))

(def: #export (nth-roots nth input)
  (-> Nat Complex (List Complex))
  (if (n.= +0 nth)
    (list)
    (let [r-nth (|> nth nat-to-int int-to-real)
          nth-root-of-abs (math;pow (r./ r-nth 1.0)
                                    (c.abs input))
          nth-phi (|> input argument (r./ r-nth))
          slice (|> math;pi (r.* 2.0) (r./ r-nth))]
      (|> (list;n.range +0 (n.dec nth))
          (List/map (lambda [nth']
                      (let [inner (|> nth' nat-to-int int-to-real
                                      (r.* slice)
                                      (r.+ nth-phi))
                            real (r.* nth-root-of-abs
                                      (math;cos inner))
                            imaginary (r.* nth-root-of-abs
                                           (math;sin inner))]
                        {#real real
                         #imaginary imaginary})))))))

(struct: #export _ (Codec Text Complex)
  (def: (encode (^slots [#real #imaginary]))
    ($_ Text/append "(" (r/encode real) ", " (r/encode imaginary) ")"))

  (def: (decode input)
    (case (do Monad<Maybe>
            [input' (text;sub +1 (n.- +1 (text;size input)) input)]
            (text;split-with "," input'))
      #;None
      (#;Left (Text/append "Wrong syntax for complex numbers: " input))

      (#;Some [r' i'])
      (do Monad<Error>
        [r (r/decode (text;trim r'))
         i (r/decode (text;trim i'))]
        (wrap {#real r
               #imaginary i}))
      )))