{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}
module Examples.Mandel (afun) where

import Prelude ()
import qualified Prelude as P
import Data.Array.Accelerate


type Dims = (Int, Int)
type Pos = (Double, Double)
type Viewport = (Dims     -- image size
                ,Pos      -- midpoint
                ,Double)  -- complex width of the viewport

type RGB = (Word8, Word8, Word8)

-- Arguments: viewport and maxiter
-- Result: image in row-major order
afun :: Acc (Scalar (Viewport, Int)) -> Acc (Matrix RGB)
afun (the -> T2 viewport maxiter) = mandel viewport maxiter (clrBasic maxiter)

clrBasic :: Exp Int -> Exp Int -> Exp RGB
clrBasic maxiterI nI =
    let maxiter = log (toFloating maxiterI)
        n = log (toFloating nI)
    in cond (n == maxiter)
            (T3 0 0 0)
            (let r = slope (Just (maxiter / 4)) Nothing n maxiter
                 g = slope Nothing Nothing n maxiter
                 b = slope Nothing (Just (maxiter * 3 / 4)) n maxiter
             in T3 r g b)
  where
    slope :: Maybe (Exp Double) -> Maybe (Exp Double) -> Exp Double -> Exp Double -> Exp Word8
    slope mlo mhi x m =
        (P.maybe ($ 0) (\lo' -> max 0 . ($ lo')) mlo) $ \lo ->
        (P.maybe ($ m) (\hi' -> min 255 . ($ hi')) mhi) $ \hi ->
            fromIntegral @Int @Word8 $ round @Double @Int $
                max 0 $ min 255 $ (x - lo) * 255 / (hi - lo)

mandel :: Exp Viewport -> Exp Int -> (Exp Int -> Exp RGB) -> Acc (Matrix RGB)
mandel (T3 (T2 w h) (T2 cx cy) cw) maxiter clrscheme =
    generate (I2 h w) $ \(I2 yi xi) ->
        let minx = cx - cw / 2
            ch = toFloating h / toFloating w * cw
            maxy = cy + ch / 2
            x = minx + toFloating xi / (toFloating w - 1) * cw
            y = maxy - toFloating yi / (toFloating h - 1) * ch
        in clrscheme (mandeliter x y maxiter)

mandeliter :: Exp Double -> Exp Double -> Exp Int -> Exp Int
mandeliter x y maxiter =
    let T5 _ _ _ _ n =
            while (\(T5 _ _ a2 b2 i) -> a2 + b2 < 4 && i < maxiter)
                  (\(T5 a b a2 b2 i) ->
                      let a' = a2 - b2 + x
                          b' = 2 * a * b + y
                      in T5 a' b' (a'*a') (b'*b') (i + 1))
                  (T5 x y (x*x) (y*y) 0)
    in n