{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE RebindableSyntax    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators       #-}
{-# LANGUAGE ViewPatterns        #-}

module Julia where

import Data.Array.Accelerate                              hiding ( fromInteger )
import Data.Array.Accelerate.Data.Complex

import Prelude                                            ( fromInteger )
import qualified Prelude                                  as P


julia
    :: forall a. (Num a, RealFloat a, FromIntegral Int a)
    => Int                                                -- ^ image width
    -> Int                                                -- ^ image height
    -> (Exp Float -> Exp (Complex a) -> Exp (Complex a))  -- ^ iteration function
    -> Acc (Scalar Float)                                 -- ^ current time
    -> Acc (Scalar a)                                     -- ^ centre x
    -> Acc (Scalar a)                                     -- ^ centre y
    -> Acc (Scalar a)                                     -- ^ view width
    -> Acc (Scalar Int32)                                 -- ^ iteration limit
    -> Acc (Scalar a)                                     -- ^ divergence radius
    -> Acc (Array DIM2 (Complex a, Int32))
julia screenX screenY next (the -> t) (the -> x0) (the -> y0) (the -> width) (the -> limit) (the -> radius) =
  generate
    (constant (Z :. screenY :. screenX))
    (\ix -> let z0 = complexOfPixel ix
                zn = while (\zi -> snd zi       < limit
                                && dot (fst zi) < radius)
                           step
                           (T2 z0 0)
             in zn)

  where
    complexOfPixel :: Exp DIM2 -> Exp (Complex a)
    complexOfPixel (unlift -> Z :. y :. x) =
      let
          height = P.fromIntegral screenY / P.fromIntegral screenX * width
          xmin   = x0 - width  / 2
          ymin   = y0 - height / 2
          --
          re     = xmin + (fromIntegral x * width)  / fromIntegral (constant screenX)
          im     = ymin + (fromIntegral y * height) / fromIntegral (constant screenY)
      in
      lift (re :+ im)

    -- Divergence condition
    --
    dot :: Exp (Complex a) -> Exp a
    dot (unlift -> x :+ y) = x*x + y*y

    -- Take a single step of the recurrence relation
    --
    step :: Exp (Complex a, Int32) -> Exp (Complex a, Int32)
    step (T2 z i) = T2 (next t z) (i + 1)