{-# LANGUAGE BangPatterns, DeriveDataTypeable #-}
{- |
Module      :  Fractal.RUFF.Mandelbrot.Image
Copyright   :  (c) Claude Heiland-Allen 2011
License     :  BSD3

Maintainer  :  claudiusmaximus@goto10.org
Stability   :  unstable
Portability :  portable

Generic functions to render images.

-}

module Fractal.RUFF.Mandelbrot.Image
  ( simpleImage, complexImage, imageLoop, coordinates, ascii, unicode
  ) where

import Control.Monad.ST (ST)
import Data.Array.ST (newArray, writeArray, runSTUArray)
import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef)
import Data.Array.Unboxed (UArray, (!), bounds, range)

import Fractal.RUFF.Types.Complex (Complex((:+)))
import Fractal.RUFF.Types.Tuple (Tuple2(Tuple2))
import Fractal.RUFF.Mandelbrot.Iterate (iterates, initial, Mode(Simple, DistanceEstimate), Iterate(), Output(OutSimple, OutDistanceEstimate), escapeTime, distanceEstimate, finalAngle, outUser)

-- | Render an image with the 'Simple' algorithm.  The iteration count is
--   doubled until the image is good enough, or the fixed maximum iteration
--   count is reached.
--
-- > putStr . unicode $ simpleImage 100 100 ((-1.861):+0) (0.001) 1000000000
simpleImage :: (Ord r, Floating r) => Int {- ^ width -} -> Int {- ^ height -} -> Complex r {- ^ center -} -> r {- ^ radius -} -> Int {- ^ max iterations -} -> UArray (Int, Int) Bool {- ^ image -}
{-# INLINABLE simpleImage #-}
simpleImage width height c0 r0 n0 = runSTUArray $ do
    a <- newArray bs True
    s <- newSTRef (0 :: Int)
    imageLoop s a n0 0 False 64 i0s (out s a)
  where
    (bs, cs) = coordinates width height c0 r0
    i0s = map (uncurry $ initial Simple) cs
    out s a (OutSimple{ outUser = Tuple2 j i }) = do
      writeArray a (j, i) False
      modifySTRef' s (+ 1)
    out _ _ _ = return ()
 
-- | Render an image with the 'DistanceEstimate' algorithm.  The iteration count is
--   doubled until the image is good enough, or the fixed maximum iteration
--   count is reached.  The output values are converted to 'Float'.
complexImage :: (Ord r, Real r, Floating r) => Int {- ^ width -} -> Int {- ^ height -} -> Complex r {- ^ center -} -> r {- ^ radius -} -> Int {- ^ max iterations -} -> UArray (Int, Int, Int) Float {- ^ image -}
{-# INLINABLE complexImage #-}
complexImage !width !height !c0 !r0 !n0 = runSTUArray $ do
    a <- newArray bs (-1)
    s <- newSTRef (0 :: Int)
    imageLoop s a n0 0 False 64 i0s (out s a)
  where
    bs = ((jlo,ilo,0), (jhi,ihi,2))
    (((jlo,ilo),(jhi,ihi)), cs) = coordinates width height c0 r0
    i0s = map (uncurry $ initial DistanceEstimate) cs
    out !s !a (OutDistanceEstimate{ escapeTime = et, distanceEstimate = de, finalAngle = fa, outUser = Tuple2 j i }) = {-# SCC "complexImage.out" #-} do
      writeArray a (j, i, 0) (realToFrac et)
      writeArray a (j, i, 1) (realToFrac de)
      writeArray a (j, i, 2) (realToFrac fa)
      modifySTRef' s (+ 1)
    out _ _ _ = return ()

-- | Image rendering loop.
imageLoop :: (Ord r, Floating r) => STRef s Int {- ^ escapees -} -> a {- ^ output array -} -> Int {- ^ max iterations -} -> Int {- ^ iterations -} -> Bool {- ^ prior escapees -} -> Int {- ^ iterations this phase -} -> [Iterate u r] {- ^ iterates -} -> (Output u r -> ST s ()) {- ^ output callback -} -> ST s a {- ^ output array as given -}
{-# INLINABLE imageLoop #-}
imageLoop s a !n0 !n1 !f1 !m1 is1 out = loop f1 n1 m1 is1
  where
    loop !f !n !m is = do
      writeSTRef s 0
      is' <- iterates m is out
      o <- readSTRef s
      if null is || (f && o == 0) || n > n0 then return a else loop (f || o > 0) (n + m) (m * 2) is'

-- | The parameter plane coordinates for an image, with bounds.
coordinates :: (Ord r, Floating r) => Int {- ^ width -} -> Int {- ^ height -} -> Complex r {- ^ center -} -> r {- ^ radius -} -> (((Int,Int),(Int,Int)), [(Tuple2 Int Int, Complex r)]) {- ^ (bounds, coords) -}
{-# INLINABLE coordinates #-}
coordinates !width !height !(c0r :+ c0i) !r0 = (bs, cs)
  where
    bs = ((0, 0), (height - 1, width - 1))
    cs =  [ (Tuple2 j i, c)
          | (j,i) <- range bs
          , let y = (fromIntegral j - h) / h
          , let x = (fromIntegral i - w) / w
          , let ci = c0i + r0 * y
          , let cr = c0r + r0 * x
          , let c = cr :+ ci
          ]
    w = fromIntegral $ width  `div` 2
    h = fromIntegral $ height `div` 2

-- | Convert a bit array to ascii graphics.
ascii :: UArray (Int, Int) Bool {- ^ image -} -> String {- ^ ascii -}
ascii a = unlines . map concat $ [ [ b (a ! (j, i)) | i <- [ ilo .. ihi ] ] | j <- [ jhi, jhi - 1 .. jlo ] ]
  where
    ((jlo, ilo), (jhi, ihi)) = bounds a
    b False = "  "
    b True  = "##"

-- | Convert a bit array to unicode block graphics.
unicode :: UArray (Int, Int) Bool {- ^ image -} -> String {- ^ unicode -}
unicode a = unlines [ [ b (a ! (j, i)) (a ! (j - 1, i)) | i <- [ ilo .. ihi ] ] | j <- [ jhi, jhi - 2 .. jlo ] ]
  where
    ((jlo, ilo), (jhi, ihi)) = bounds a
    b False False = ' '
    b True False = '\x2580'
    b False True = '\x2584'
    b True True = '\x2588'

-- | Strict version of 'modifySTRef'.
modifySTRef' :: STRef s a -> (a -> a) -> ST s ()
{-# INLINABLE modifySTRef' #-}
modifySTRef' s f = do
  x <- readSTRef s
  writeSTRef s $! f x