(;module: {#;doc "Complex arithmetic."}
  lux
  (lux [math]
       (control [eq #+ Eq]
                number
                codec
                ["M" monad #+ do Monad]
                ["p" parser])
       (data [number "f/" Number<Frac> Codec<Text,Frac>]
             [text "text/" Monoid<Text>]
             text/format
             ["E" error]
             [maybe]
             (coll [list "L/" Monad<List>]))
       [meta]
       (meta [code]
             ["s" syntax #+ syntax: Syntax])))

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

(syntax: #export (complex real [?imaginary (p;maybe 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 (~ (maybe;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 (not-a-number? complex)
  (or (number;not-a-number? (get@ #real complex))
      (number;not-a-number? (get@ #imaginary complex))))

(def: #export (c.= param input)
  (-> Complex Complex Bool)
  (and (f.= (get@ #real param)
            (get@ #real input))
       (f.= (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.+ f.+]
  [c.- f.-]
  )

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

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

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

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

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

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

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

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

(def: #export (c.% param input)
  (-> Complex Complex Complex)
  (let [scaled (c./ param input)
        quotient (|> scaled
                     (update@ #real math;floor)
                     (update@ #imaginary math;floor))]
    (c.- (c.* quotient param)
         input)))

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

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

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

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

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

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

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

(struct: #export _ (Number Complex)
  (def: + c.+)
  (def: - c.-)
  (def: * c.*)
  (def: / c./)
  (def: % c.%)
  (def: (negate x)
    (|> x
        (update@ #real f/negate)
        (update@ #imaginary f/negate)))
  (def: abs c.abs)
  (def: (signum x)
    (|> x
        (update@ #real f/signum)
        (update@ #imaginary f/signum))))

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

(def: #export (log subject)
  (-> Complex Complex)
  (let [(^slots [#real #imaginary]) subject]
    {#real (|> subject c.abs (get@ #real) math;log)
     #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' Frac    c.*']
  )

(def: (copy-sign sign magnitude)
  (-> Frac Frac Frac)
  (f.* (f/signum sign) magnitude))

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

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

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

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

(def: #export (asin input)
  (-> Complex Complex)
  (|> input
      root2-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 Frac)
  (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-frac)
          nth-root-of-abs (|> input c.abs (get@ #real) (math;pow (f./ r-nth 1.0)))
          nth-phi (|> input argument (f./ r-nth))
          slice (|> math;pi (f.* 2.0) (f./ r-nth))]
      (|> (list;n.range +0 (n.dec nth))
          (L/map (function [nth']
                   (let [inner (|> nth' nat-to-int int-to-frac
                                   (f.* slice)
                                   (f.+ nth-phi))
                         real (f.* nth-root-of-abs
                                   (math;cos inner))
                         imaginary (f.* nth-root-of-abs
                                        (math;sin inner))]
                     {#real real
                      #imaginary imaginary})))))))