module Math.JacobiTheta
  (
    jtheta1Dash,
    jtheta1,
    jtheta2,
    jtheta3,
    jtheta4
  )
  where
import Data.Complex

type Cplx = Complex Double

i_ :: Cplx
i_ :: Cplx
i_ = Double
0.0 forall a. a -> a -> Complex a
:+ Double
1.0

machinePrecision :: Double
machinePrecision :: Double
machinePrecision = Double
2forall a. Floating a => a -> a -> a
**(-Double
52)

areClose :: Cplx -> Cplx -> Bool
areClose :: Cplx -> Cplx -> Bool
areClose Cplx
z1 Cplx
z2 = forall a. RealFloat a => Complex a -> a
magnitude (Cplx
z1 forall a. Num a => a -> a -> a
- Cplx
z2) forall a. Ord a => a -> a -> Bool
< Double
epsilon forall a. Num a => a -> a -> a
* Double
h
  where
    epsilon :: Double
epsilon = Double
2.0 forall a. Num a => a -> a -> a
* Double
machinePrecision
    magn2 :: Double
magn2 = forall a. RealFloat a => Complex a -> a
magnitude Cplx
z2
    h :: Double
h = if Double
magn2 forall a. Ord a => a -> a -> Bool
< Double
epsilon then Double
1.0 else forall a. Ord a => a -> a -> a
max (forall a. RealFloat a => Complex a -> a
magnitude Cplx
z1) Double
magn2

square :: Cplx -> Cplx
square :: Cplx -> Cplx
square Cplx
z = Cplx
z forall a. Num a => a -> a -> a
* Cplx
z

jtheta1Alt1 :: Cplx -> Cplx -> Cplx
jtheta1Alt1 :: Cplx -> Cplx -> Cplx
jtheta1Alt1 Cplx
z Cplx
q =
  Int -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
go Int
0 (Double
0.0 forall a. a -> a -> Complex a
:+ Double
0.0) Cplx
1.0 (Cplx
1.0 forall a. Fractional a => a -> a -> a
/ Cplx
qsq) Cplx
1.0
  where 
    qsq :: Cplx
qsq = Cplx
q forall a. Num a => a -> a -> a
* Cplx
q
    go :: Int -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
    go :: Int -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
go Int
n Cplx
out Cplx
alt Cplx
q_2n Cplx
q_n_np1 
      | Int
n forall a. Ord a => a -> a -> Bool
> Int
3000 = forall a. HasCallStack => [Char] -> a
error [Char]
"Reached 3000 iterations."
      | Cplx -> Cplx -> Bool
areClose Cplx
out Cplx
outnew = Cplx
2.0 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
sqrt (forall a. Floating a => a -> a
sqrt Cplx
q) forall a. Num a => a -> a -> a
* Cplx
out
      | Bool
otherwise = Int -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
go (Int
n forall a. Num a => a -> a -> a
+ Int
1) Cplx
outnew (-Cplx
alt) Cplx
q_2np1 Cplx
q_np1_np2
        where
          q_2np1 :: Cplx
q_2np1 = Cplx
q_2n forall a. Num a => a -> a -> a
* Cplx
qsq
          q_np1_np2 :: Cplx
q_np1_np2 = Cplx
q_n_np1 forall a. Num a => a -> a -> a
* Cplx
q_2np1
          n' :: Cplx
n' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n 
          k :: Cplx
k = Cplx
2.0 forall a. Num a => a -> a -> a
* Cplx
n' forall a. Num a => a -> a -> a
+ Cplx
1.0
          outnew :: Cplx
outnew = Cplx
out forall a. Num a => a -> a -> a
+ Cplx
alt forall a. Num a => a -> a -> a
* Cplx
q_np1_np2 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
sin (Cplx
k forall a. Num a => a -> a -> a
* Cplx
z) 

-- jtheta1(z, tau) = jtheta1Alt2 (z/pi) (-i_ * tau/pi)
jtheta1Alt2 :: Cplx -> Cplx -> Cplx
jtheta1Alt2 :: Cplx -> Cplx -> Cplx
jtheta1Alt2 Cplx
z' Cplx
t' = 
  let nm :: Int
nm = forall a b. (RealFrac a, Integral b) => a -> b
round (Double
0.5 forall a. Num a => a -> a -> a
- forall a. Complex a -> a
realPart Cplx
z') in
  let np :: Int
np = Int
nm forall a. Num a => a -> a -> a
+ Int
1 in
  Int -> Int -> Cplx -> (Cplx, Cplx) -> Cplx
go Int
nm Int
np (Double
0.0 forall a. a -> a -> Complex a
:+ Double
0.0) (if forall a. Integral a => a -> Bool
even Int
np then (-Cplx
1, Cplx
1) else (Cplx
1, -Cplx
1)) 
  where
    go :: Int -> Int -> Cplx -> (Cplx, Cplx) -> Cplx
    go :: Int -> Int -> Cplx -> (Cplx, Cplx) -> Cplx
go Int
nminus Int
nplus Cplx
series (Cplx
altm, Cplx
altp)
      | Int
nplus forall a. Num a => a -> a -> a
- Int
nminus forall a. Ord a => a -> a -> Bool
> Int
3000 = forall a. HasCallStack => [Char] -> a
error [Char]
"Reached 3000 iterations."
      | (Int
nplus forall a. Num a => a -> a -> a
- Int
nminus forall a. Ord a => a -> a -> Bool
> Int
2) Bool -> Bool -> Bool
&& Cplx -> Cplx -> Bool
areClose Cplx
series Cplx
newseries = 
          Cplx
series forall a. Fractional a => a -> a -> a
/ forall a. Floating a => a -> a
sqrt (forall a. Floating a => a
pi forall a. Num a => a -> a -> a
* Cplx
t')
      | Bool
otherwise = Int -> Int -> Cplx -> (Cplx, Cplx) -> Cplx
go (Int
nminus forall a. Num a => a -> a -> a
- Int
1) (Int
nplus forall a. Num a => a -> a -> a
+ Int
1) Cplx
newseries (-Cplx
altm, -Cplx
altp)
        where 
          nminus' :: Cplx
nminus' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nminus
          nplus' :: Cplx
nplus' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
nplus
          termm :: Cplx
termm = Cplx
altm forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
exp (- Cplx -> Cplx
square (Cplx
nminus' forall a. Num a => a -> a -> a
- Cplx
0.5 forall a. Num a => a -> a -> a
+ Cplx
z') forall a. Fractional a => a -> a -> a
/ Cplx
t')
          termp :: Cplx
termp = Cplx
altp forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
exp (- Cplx -> Cplx
square (Cplx
nplus' forall a. Num a => a -> a -> a
- Cplx
0.5 forall a. Num a => a -> a -> a
+ Cplx
z') forall a. Fractional a => a -> a -> a
/ Cplx
t')
          newseries :: Cplx
newseries = Cplx
series forall a. Num a => a -> a -> a
+ Cplx
termm forall a. Num a => a -> a -> a
+ Cplx
termp

falpha :: Cplx -> Cplx -> Cplx
falpha :: Cplx -> Cplx -> Cplx
falpha Cplx
z Cplx
tau = 
  forall a. Floating a => a -> a
sqrt (-Cplx
i_ forall a. Num a => a -> a -> a
* Cplx
tau) forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
exp (Cplx
i_ forall a. Fractional a => a -> a -> a
/ Cplx
tau forall a. Num a => a -> a -> a
* Cplx
z forall a. Num a => a -> a -> a
* Cplx
z forall a. Fractional a => a -> a -> a
/ forall a. Floating a => a
pi)

jtheta1Alt :: Cplx -> Cplx -> Cplx
jtheta1Alt :: Cplx -> Cplx -> Cplx
jtheta1Alt Cplx
z Cplx
tau = 
  if forall a. Complex a -> a
imagPart Cplx
tau forall a. Ord a => a -> a -> Bool
< Double
1.3 
    then
      let w :: Cplx
w = forall a. Floating a => a
pi forall a. Num a => a -> a -> a
* Cplx
tau in 
      Cplx
i_ forall a. Num a => a -> a -> a
* Cplx -> Cplx -> Cplx
jtheta1Alt2 (Cplx
z forall a. Fractional a => a -> a -> a
/ Cplx
w) (Cplx
i_ forall a. Fractional a => a -> a -> a
/ Cplx
w) forall a. Fractional a => a -> a -> a
/ Cplx -> Cplx -> Cplx
falpha Cplx
z Cplx
tau
    else
      Cplx
i_ forall a. Num a => a -> a -> a
* Cplx -> Cplx -> Cplx
jtheta1Alt1 (Cplx
z forall a. Fractional a => a -> a -> a
/ Cplx
tau) (forall a. Floating a => a -> a
exp (-Cplx
i_ forall a. Num a => a -> a -> a
* forall a. Floating a => a
pi forall a. Fractional a => a -> a -> a
/ Cplx
tau)) forall a. Fractional a => a -> a -> a
/ Cplx -> Cplx -> Cplx
falpha Cplx
z Cplx
tau

tauFromQ :: Cplx -> Cplx
tauFromQ :: Cplx -> Cplx
tauFromQ Cplx
q = -Cplx
i_ forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
log Cplx
q forall a. Fractional a => a -> a -> a
/ forall a. Floating a => a
pi

checkQ :: Cplx -> Cplx
checkQ :: Cplx -> Cplx
checkQ Cplx
q
  | forall a. RealFloat a => Complex a -> a
magnitude Cplx
q forall a. Ord a => a -> a -> Bool
>= Double
1 = 
    forall a. HasCallStack => [Char] -> a
error [Char]
"The modulus of the nome must be smaller than one."
  | forall a. Complex a -> a
imagPart Cplx
q forall a. Eq a => a -> a -> Bool
== Double
0 Bool -> Bool -> Bool
&& forall a. Complex a -> a
realPart Cplx
q forall a. Ord a => a -> a -> Bool
<= Double
0 = 
    forall a. HasCallStack => [Char] -> a
error [Char]
"If the nome is real, it must be positive."
  | Bool
otherwise = Cplx
q

getTauFromQ :: Cplx -> Cplx
getTauFromQ :: Cplx -> Cplx
getTauFromQ = Cplx -> Cplx
tauFromQ forall b c a. (b -> c) -> (a -> b) -> a -> c
. Cplx -> Cplx
checkQ

expM :: Cplx -> Cplx -> Cplx
expM :: Cplx -> Cplx -> Cplx
expM Cplx
z Cplx
tau = forall a. Floating a => a -> a
exp (Cplx
i_ forall a. Num a => a -> a -> a
* (Cplx
z forall a. Num a => a -> a -> a
+ Cplx
tau forall a. Num a => a -> a -> a
* forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Cplx
4))

-- | First Jacobi theta function
jtheta1 :: 
     Cplx -- ^ z
  -> Cplx -- ^ q, the nome
  -> Cplx
jtheta1 :: Cplx -> Cplx -> Cplx
jtheta1 Cplx
z Cplx
q = Cplx -> Cplx -> Cplx
jtheta1Alt Cplx
z (Cplx -> Cplx
getTauFromQ Cplx
q)

-- | Second Jacobi theta function
jtheta2 :: 
     Cplx -- ^ z
  -> Cplx -- ^ q, the nome
  -> Cplx
jtheta2 :: Cplx -> Cplx -> Cplx
jtheta2 Cplx
z = Cplx -> Cplx -> Cplx
jtheta1 (Cplx
z forall a. Num a => a -> a -> a
+ forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Cplx
2)

-- | Third Jacobi theta function
jtheta3 :: 
     Cplx -- ^ z
  -> Cplx -- ^ q, the nome
  -> Cplx
jtheta3 :: Cplx -> Cplx -> Cplx
jtheta3 Cplx
z Cplx
q = Cplx -> Cplx -> Cplx
jtheta2 (Cplx
z forall a. Num a => a -> a -> a
- forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Cplx
2 forall a. Num a => a -> a -> a
* Cplx
tau) Cplx
q forall a. Num a => a -> a -> a
* Cplx -> Cplx -> Cplx
expM (-Cplx
z) Cplx
tau
  where
    tau :: Cplx
tau = Cplx -> Cplx
tauFromQ Cplx
q

-- | Fourth Jacobi theta function
jtheta4 :: 
     Cplx -- ^ z
  -> Cplx -- ^ q, the nome
  -> Cplx
jtheta4 :: Cplx -> Cplx -> Cplx
jtheta4 Cplx
z = Cplx -> Cplx -> Cplx
jtheta3 (Cplx
z forall a. Num a => a -> a -> a
+ forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Cplx
2)

-- | Derivative of the first Jacobi theta function
jtheta1Dash :: 
     Cplx -- ^ z
  -> Cplx -- ^ q, the nome
  -> Cplx
jtheta1Dash :: Cplx -> Cplx -> Cplx
jtheta1Dash Cplx
z Cplx
q = 
  Int -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
go Int
0 (Double
0.0 forall a. a -> a -> Complex a
:+ Double
0.0) Cplx
1.0 (Cplx
1.0 forall a. Fractional a => a -> a -> a
/ Cplx
qsq) Cplx
1.0
  where 
    q' :: Cplx
q' = Cplx -> Cplx
checkQ Cplx
q
    qsq :: Cplx
qsq = Cplx
q' forall a. Num a => a -> a -> a
* Cplx
q'
    go :: Int -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
    go :: Int -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
go Int
n Cplx
out Cplx
alt Cplx
q_2n Cplx
q_n_np1 
      | Int
n forall a. Ord a => a -> a -> Bool
> Int
3000 = forall a. HasCallStack => [Char] -> a
error [Char]
"Reached 3000 iterations."
      | Cplx -> Cplx -> Bool
areClose Cplx
out Cplx
outnew = Cplx
2.0 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
sqrt (forall a. Floating a => a -> a
sqrt Cplx
q) forall a. Num a => a -> a -> a
* Cplx
out
      | Bool
otherwise = Int -> Cplx -> Cplx -> Cplx -> Cplx -> Cplx
go (Int
n forall a. Num a => a -> a -> a
+ Int
1) Cplx
outnew (-Cplx
alt) Cplx
q_2np1 Cplx
q_np1_np2
        where
          q_2np1 :: Cplx
q_2np1 = Cplx
q_2n forall a. Num a => a -> a -> a
* Cplx
qsq
          q_np1_np2 :: Cplx
q_np1_np2 = Cplx
q_n_np1 forall a. Num a => a -> a -> a
* Cplx
q_2np1
          n' :: Cplx
n' = forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n 
          k :: Cplx
k = Cplx
2.0 forall a. Num a => a -> a -> a
* Cplx
n' forall a. Num a => a -> a -> a
+ Cplx
1.0
          outnew :: Cplx
outnew = Cplx
out forall a. Num a => a -> a -> a
+ Cplx
k forall a. Num a => a -> a -> a
* Cplx
alt forall a. Num a => a -> a -> a
* Cplx
q_np1_np2 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> a
cos (Cplx
k forall a. Num a => a -> a -> a
* Cplx
z)