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

Maintainer  :  claude@mathr.co.uk
Stability   :  unstable
Portability :  portable

Generic (slow) functions to render images.

-}

module Fractal.RUFF.Mandelbrot.Image
  ( simpleImage, complexImage, imageLoop, coordinates, ascii, unicode
  , Channel(..), Coordinates, border
  ) 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, amap, ixmap)
import Data.Strict.Tuple (Pair((:!:)))

import Data.Ix (Ix)
import Data.Data (Data)
import Data.Typeable (Typeable)

import Fractal.RUFF.Types.Complex (Complex((:+)), magnitude)
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 (coordinates 100 100 ((-1.861):+0) (0.001)) 1000000000
simpleImage :: (Ord r, Floating r) => Coordinates r {- ^ coordinates -} -> Int {- ^ max iterations -} -> UArray (Int, Int) Bool {- ^ image -}
{-# INLINABLE simpleImage #-}
simpleImage (bs, cs) n0 = runSTUArray $ do
    a <- newArray bs True
    s <- newSTRef (0 :: Int)
    imageLoop s a n0 0 False 64 i0s (out s a)
  where
    i0s = map (uncurry $ initial Simple) cs
    out s a (OutSimple{ outUser = 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'.
--
-- > putStr . unicode . border $ complexImage (coordinates 100 100 ((-1.861):+0) (0.001)) 1000000000
complexImage :: (Ord r, Real r, Floating r) => Coordinates r {-^ coordinates -} -> Int {- ^ max iterations -} -> UArray (Int, Int, Channel) Float {- ^ image -}
{-# INLINABLE complexImage #-}
complexImage (((jlo,ilo),(jhi,ihi)), cs) !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,minBound), (jhi,ihi,maxBound))
    (_, cx0):(_, cx1):_ = cs
    pixelSpacing = magnitude (cx1 - cx0)
    i0s = map (uncurry $ initial DistanceEstimate) cs
    out !s !a (OutDistanceEstimate{ escapeTime = et, distanceEstimate = de, finalAngle = fa, outUser = j :!: i }) = {-# SCC "complexImage.out" #-} do
      writeArray a (j, i, EscapeTime) (realToFrac et)
      writeArray a (j, i, DistanceEstimate') (realToFrac (de / pixelSpacing))
      writeArray a (j, i, FinalAngle) (realToFrac fa)
      modifySTRef' s (+ 1)
    out _ _ _ = return ()

-- | Channels in an image.
data Channel = EscapeTime {- ^ continuous dwell -} | DistanceEstimate' {- ^ normalized to pixel spacing -} | FinalAngle {- ^ in [-pi,pi] -}
  deriving (Eq, Ord, Enum, Bounded, Ix, Read, Show, Data, Typeable)

-- | 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'

-- | Image bounds and coordinates.
type Coordinates r = (((Int,Int),(Int,Int)), [(Pair Int Int, Complex r)])

-- | The parameter plane coordinates for an image, with bounds.
coordinates :: (Ord r, Floating r) => Int {- ^ width -} -> Int {- ^ height -} -> Complex r {- ^ center -} -> r {- ^ size -} -> Coordinates r
{-# INLINABLE coordinates #-}
coordinates !width !height !(c0r :+ c0i) !r0 = (bs, cs)
  where
    bs = ((0, 0), (height - 1, width - 1))
    cs =  [ (j :!: i, c)
          | (j,i) <- range bs
          , let y = (fromIntegral j - h) / h
          , let x = (fromIntegral i - w) / h
          , 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 distance estimate image to a near-boundary bit array.
--   The input image must have a DistanceEstimate' channel.
border :: UArray (Int, Int, Channel) Float {- ^ image -} -> UArray (Int, Int) Bool
border a = amap (\x -> x > 0 && x < 1) . ixmap bs (\(j, i) -> (j, i, DistanceEstimate')) $ a
  where
    ((jlo, ilo, _), (jhi, ihi, _)) = bounds a
    bs = ((jlo, ilo), (jhi, ihi))

-- | 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