{-# LANGUAGE BangPatterns #-}
module Main (main) where

import Control.Concurrent (killThread)
import Control.Monad (when)
import Control.Monad.Trans (liftIO)
import Data.Array.IO (IOUArray, newArray, readArray, writeArray, inRange)
import Data.IORef (newIORef, readIORef, writeIORef)
import Data.List (isPrefixOf)
import Data.PriorityQueue (PriorityQueue, newPriorityQueue, enqueue, enqueueBatch, dequeue)
import Foreign (mallocBytes, nullPtr, plusPtr, pokeArray, pokeByteOff, Word8)
import Foreign.C (CDouble)
import GHC.Conc (forkOnIO, numCapabilities)
import Graphics.UI.Gtk
import Graphics.UI.Gtk.OpenGL
import qualified Graphics.Rendering.OpenGL as GL
import Graphics.Rendering.OpenGL (($=), GLfloat)
import Numeric.QD.DoubleDouble (DoubleDouble(DoubleDouble))
import Numeric.QD.QuadDouble (QuadDouble(QuadDouble))
import qualified Numeric.QD.DoubleDouble as DD
import qualified Numeric.QD.QuadDouble as QD
import Numeric.QD.FPU.Raw (fpu_fix_start)
import Unsafe.Coerce (unsafeCoerce)

type B = Word8
type N = Int
type R = Double

convert :: (Real a, Fractional b) => a -> b
convert = realToFrac
convertDouble2CDouble :: Double -> CDouble
convertDouble2CDouble !x = unsafeCoerce x
convertCDouble2Double :: CDouble -> Double
convertCDouble2Double !x = unsafeCoerce x
convertDouble2DoubleDouble :: Double -> DoubleDouble
convertDouble2DoubleDouble !x = convertCDouble2DoubleDouble . convertDouble2CDouble $ x
convertCDouble2DoubleDouble :: CDouble -> DoubleDouble
convertCDouble2DoubleDouble !x = DoubleDouble x 0
convertDoubleDouble2Double :: DoubleDouble -> Double
convertDoubleDouble2Double !(DoubleDouble x _) = convertCDouble2Double x
convertDoubleDouble2CDouble :: DoubleDouble -> CDouble
convertDoubleDouble2CDouble !(DoubleDouble x _) = x
{-# RULES "convert/Double2CDouble" convert = convertDouble2CDouble #-}
{-# RULES "convert/CDouble2Double" convert = convertCDouble2Double #-}
{-# RULES "convert/Double2DoubleDouble" convert = convertDouble2DoubleDouble #-}
{-# RULES "convert/CDouble2DoubleDouble" convert = convertCDouble2DoubleDouble #-}
{-# RULES "convert/DoubleDouble2Double" convert = convertDoubleDouble2Double #-}
{-# RULES "convert/DoubleDouble2CDouble" convert = convertDoubleDouble2CDouble #-}

data Complex c = {-# UNPACK #-} !c :+ {-# UNPACK #-} !c deriving (Read, Show, Eq)

instance Num c => Num (Complex c) where
  {-# SPECIALIZE instance Num (Complex Float) #-}
  {-# SPECIALIZE instance Num (Complex Double) #-}
  {-# SPECIALIZE instance Num (Complex DoubleDouble) #-}
  {-# SPECIALIZE instance Num (Complex QuadDouble) #-}
  (!(a :+ b)) + (!(c :+ d)) = {-# SCC "C+" #-} ((a + c) :+ (b + d))
  (!(a :+ b)) - (!(c :+ d)) = {-# SCC "C-" #-} ((a - c) :+ (b - d))
  (!(a :+ b)) * (!(c :+ d)) = {-# SCC "C*" #-} ((a * c - b * d) :+ (a * d + b * c))
  negate !(a :+ b) = (-a) :+ (-b)
  abs x = error $ "Cx.abs: " ++ show x
  signum x = error $ "Cx.signum: " ++ show x
  fromInteger !x = fromInteger x :+ 0

class Num c => Turbo c where
  sqr :: c -> c
  sqr !x = x * x
  twice :: c -> c
  twice !x = x + x

instance Turbo Float where
instance Turbo Double where
instance Turbo CDouble where

instance Turbo c => Turbo (Complex c) where
  {-# SPECIALIZE instance Turbo (Complex Float) #-}
  {-# SPECIALIZE instance Turbo (Complex Double) #-}
  {-# SPECIALIZE instance Turbo (Complex DoubleDouble) #-}
  {-# SPECIALIZE instance Turbo (Complex QuadDouble) #-}
  sqr !(r :+ i) = (sqr r - sqr i) :+ (twice (r * i))
  twice !(r :+ i) = (twice r) :+ (twice i)

instance Turbo DoubleDouble where
  sqr !x = DD.sqr x
  twice !(DoubleDouble a b) = DoubleDouble (twice a) (twice b)

instance Turbo QuadDouble where
  sqr !x = QD.sqr x
  twice !(QuadDouble a b c d) = QuadDouble (twice a) (twice b) (twice c) (twice d)

hsv2rgb :: R -> R -> R -> (R, R, R)
hsv2rgb !h !s !v
  | s == 0 = (v, v, v)
  | h == 1 = hsv2rgb 0 s v
  | otherwise =
      let !i = floor (h * 6) `mod` 6 :: N
          !f = (h * 6) - fromIntegral i
          !p = v * (1 - s)
          !q = v * (1 - s * f)
          !t = v * (1 - s * (1 - f))
      in  case i of
            0 -> (v, t, p)
            1 -> (q, v, p)
            2 -> (p, v, t)
            3 -> (p, q, v)
            4 -> (t, p, v)
            5 -> (v, p, q)
            _ -> (0, 0, 0)

colour :: Complex Double -> Complex Double -> N -> (B, B, B)
colour !(zr:+zi) !(dzr:+dzi) !n =
  let !il2 = 1 / log 2
      !zd2 = sqr zr + sqr zi
      !dzd2 = sqr dzr + sqr dzi
      !d = (fromIntegral n :: R) - log (log zd2 / log escapeR2) * il2
      !dwell = fromIntegral (floor d :: N)
      !finala = atan2 zi zr
      !de = (log zd2 * il2) * sqrt zd2 / sqrt dzd2
      !dscale = log de * il2 + 32
      !hue = log d * il2 / 3
      !saturation = 0 `max` (log d * il2 / 8) `min` 1
      !value = 0 `max` (1 - dscale / 48) `min` 1
      !h = hue - fromIntegral (floor hue :: N)
      !k = dwell / 2
      !satf = if k - fromIntegral (floor k :: N) >= (0.5 :: R) then 0.9 else 1
      !valf = if finala < 0 then 0.9 else 1
      (!r, !g, !b) = hsv2rgb h (satf * saturation) (valf * value)
      !rr = floor $ 0 `max` 255 * r `min` 255
      !gg = floor $ 0 `max` 255 * g `min` 255
      !bb = floor $ 0 `max` 255 * b `min` 255
  in  (rr, gg, bb)

data Job c = Job !N !N !(Complex c) !(Complex c) !(Complex c) !N

priority :: Job c -> N
priority !(Job _ _ _ _ _ n) = n

addJob :: RealFloat c => N -> N -> Complex c -> N -> PriorityQueue IO (Job c) -> IOUArray (N,N) Bool -> N -> N -> IO ()
{-# SPECIALIZE addJob :: N -> N -> Complex Float -> N -> PriorityQueue IO (Job Float) -> IOUArray (N,N) Bool -> N -> N -> IO () #-}
{-# SPECIALIZE addJob :: N -> N -> Complex Double -> N -> PriorityQueue IO (Job Double) -> IOUArray (N,N) Bool -> N -> N -> IO () #-}
{-# SPECIALIZE addJob :: N -> N -> Complex DoubleDouble -> N -> PriorityQueue IO (Job DoubleDouble) -> IOUArray (N,N) Bool -> N -> N -> IO () #-}
{-# SPECIALIZE addJob :: N -> N -> Complex QuadDouble -> N -> PriorityQueue IO (Job QuadDouble) -> IOUArray (N,N) Bool -> N -> N -> IO () #-}
addJob !w !h !c !zoom todo sync !i !j = do
  already <- readArray sync (j, i)
  when (not already) $ do
    writeArray sync (j, i) True
    enqueue todo $! Job i j (coords w h c zoom i j) 0 0 0

renderer :: (Turbo c, RealFloat c) => ((N,N),(N,N)) -> (N -> N -> B -> B -> B -> IO ()) -> Complex c -> N -> IO (IO ())
{-# SPECIALIZE renderer :: ((N,N),(N,N)) -> (N -> N -> B -> B -> B -> IO ()) -> Complex Float -> N -> IO (IO ()) #-}
{-# SPECIALIZE renderer :: ((N,N),(N,N)) -> (N -> N -> B -> B -> B -> IO ()) -> Complex Double -> N -> IO (IO ()) #-}
{-# SPECIALIZE renderer :: ((N,N),(N,N)) -> (N -> N -> B -> B -> B -> IO ()) -> Complex DoubleDouble -> N -> IO (IO ()) #-}
{-# SPECIALIZE renderer :: ((N,N),(N,N)) -> (N -> N -> B -> B -> B -> IO ()) -> Complex QuadDouble -> N -> IO (IO ()) #-}
renderer rng output !c !zoom = do
  workerts <- mapM (\w -> forkOnIO w $ worker rng c zoom output w) [ 0 .. workers - 1 ]
  return $ do
    mapM_ killThread workerts

coords :: RealFloat c => N -> N -> Complex c -> N -> N -> N -> Complex c
{-# SPECIALIZE coords :: N -> N -> Complex Float -> N -> N -> N -> Complex Float #-}
{-# SPECIALIZE coords :: N -> N -> Complex Double -> N -> N -> N -> Complex Double #-}
{-# SPECIALIZE coords :: N -> N -> Complex DoubleDouble -> N -> N -> N -> Complex DoubleDouble #-}
{-# SPECIALIZE coords :: N -> N -> Complex QuadDouble -> N -> N -> N -> Complex QuadDouble #-}
coords !w !h !c !zoom !i !j = c + ( (fromIntegral (i - w`div`2) * k)
                                  :+(fromIntegral (h`div`2 - j) * k))
  where !k = convert (1/2^^zoom :: Double)

border :: N -> N -> [(N, N)]
border !w !h = concat $
  [ [ (j, i) | i <- [ 0 .. w - 1 ], j <- [ 0 .. workers - 1 ] ]
  , [ (j, i) | j <- [ 0 .. h - 1 ], i <- [ 0 .. workers - 1 ] ]
  , [ (j, i) | j <- [ 0 .. h - 1 ], i <- [ w - workers .. w - 1 ] ]
  , [ (j, i) | i <- [ 0 .. w - 1 ], j <- [ h - workers .. h - 1 ] ]
  ]

worker :: (Turbo c, RealFloat c) => ((N,N),(N,N)) -> Complex c -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO ()
{-# SPECIALIZE worker :: ((N,N),(N,N)) -> Complex Float -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO () #-}
{-# SPECIALIZE worker :: ((N,N),(N,N)) -> Complex Double -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO () #-}
{-# SPECIALIZE worker :: ((N,N),(N,N)) -> Complex DoubleDouble -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO () #-}
{-# SPECIALIZE worker :: ((N,N),(N,N)) -> Complex QuadDouble -> N -> (N -> N -> B -> B -> B -> IO ()) -> N -> IO () #-}
worker rng@((y0,x0),(y1,x1)) !c !zoom output !me = do
  sync <- newArray rng False
  queue <- newPriorityQueue priority
  let addJ = addJob w h c zoom queue sync
      js = filter mine (border w h)
      w = x1 - x0 + 1
      h = y1 - y0 + 1
  mapM_ (flip (writeArray sync) True) js
  enqueueBatch queue (map (\(j,i) -> Job i j (coords w h c zoom i j) 0 0 0) js)
  compute rng addJ output queue
  where
    mine (j, _) = j `mod` workers == me

compute :: (Turbo c, RealFloat c) => ((N,N),(N,N)) -> (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job c) -> IO ()
{-# SPECIALIZE compute :: ((N,N),(N,N)) -> (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job Float) -> IO () #-}
{-# SPECIALIZE compute :: ((N,N),(N,N)) -> (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job Double) -> IO () #-}
{-# SPECIALIZE compute :: ((N,N),(N,N)) -> (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job DoubleDouble) -> IO () #-}
{-# SPECIALIZE compute :: ((N,N),(N,N)) -> (N -> N -> IO ()) -> (N -> N -> B -> B -> B -> IO ()) -> PriorityQueue IO (Job QuadDouble) -> IO () #-}
compute rng addJ output queue = do
  mjob <- dequeue queue
  case mjob of
    Just (Job i j c z dz n) -> do
      let done' !(zr:+zi) !(dzr:+dzi) !n' = do
            let (r, g, b) = colour (convert zr :+ convert zi) (convert dzr :+ convert dzi) n'
            output i j r g b
            sequence_
              [ addJ x y
              | u <- spreadX
              , v <- spreadY
              , let x = i + u
              , let y = j + v
              , inRange rng (y, x)
              ]
          todo' z' dz' n' = enqueue queue $! Job i j c z' dz' n'
      calculate c limit z dz n done' todo'
      compute rng addJ output queue
    Nothing -> return ()

calculate :: (Turbo c, RealFloat c) => Complex c -> N -> Complex c -> Complex c -> N -> (Complex c -> Complex c -> N -> IO ()) -> (Complex c -> Complex c -> N -> IO ()) -> IO ()
{-# SPECIALIZE calculate :: Complex Float -> N -> Complex Float -> Complex Float -> N -> (Complex Float -> Complex Float -> N -> IO ()) -> (Complex Float -> Complex Float -> N -> IO ()) -> IO () #-}
{-# SPECIALIZE calculate :: Complex Double -> N -> Complex Double -> Complex Double -> N -> (Complex Double -> Complex Double -> N -> IO ()) -> (Complex Double -> Complex Double -> N -> IO ()) -> IO () #-}
{-# SPECIALIZE calculate :: Complex DoubleDouble -> N -> Complex DoubleDouble -> Complex DoubleDouble -> N -> (Complex DoubleDouble -> Complex DoubleDouble -> N -> IO ()) -> (Complex DoubleDouble -> Complex DoubleDouble -> N -> IO ()) -> IO () #-}
{-# SPECIALIZE calculate :: Complex QuadDouble -> N -> Complex QuadDouble -> Complex QuadDouble -> N -> (Complex QuadDouble -> Complex QuadDouble -> N -> IO ()) -> (Complex QuadDouble -> Complex QuadDouble -> N -> IO ()) -> IO () #-}
calculate !c !m0 !z0 !dz0 !n0 done todo = go m0 z0 dz0 n0
  where
    go !m !z@(zr:+zi) !dz !n
      | not (sqr zr + sqr zi < er2) = done z dz n
      | m <= 0 = todo z dz n
      | otherwise = go (m - 1) (sqr z + c) (let !zdz = z * dz in twice zdz + 1) (n + 1)
    !er2 = convert escapeR2

renderer' :: Real c => ((N,N),(N,N)) -> (N -> N -> B -> B -> B -> IO ()) -> Complex c -> N -> IO (IO ())
renderer' rng output !c !zoom
  | zoom < 20  = renderer rng output (f c :: Complex Float       ) zoom
  | zoom < 50  = renderer rng output (f c :: Complex Double      ) zoom
  | zoom < 100 = renderer rng output (f c :: Complex DoubleDouble) zoom
  | otherwise  = renderer rng output (f c :: Complex QuadDouble  ) zoom
  where f !(cr :+ ci) = convert cr :+ convert ci

data Args = Args{ aWidth :: N, aHeight :: N }

defaultArgs :: Args
defaultArgs = Args{ aWidth = 788, aHeight = 576 }

combineArgs :: Args -> String -> Args
combineArgs a0 s
  | "--width="  `isPrefixOf` s = a0{ aWidth  = read $ "--width="  `dropPrefix` s }
  | "--height=" `isPrefixOf` s = a0{ aHeight = read $ "--height=" `dropPrefix` s }
  | "-w=" `isPrefixOf` s = a0{ aWidth  = read $ "-w=" `dropPrefix` s }
  | "-h=" `isPrefixOf` s = a0{ aHeight = read $ "-h=" `dropPrefix` s }
  | otherwise = a0

dropPrefix :: String -> String -> String
dropPrefix p s = drop (length p) s

roundUp2 :: N -> N
roundUp2 n = head . dropWhile (< n) . iterate (2*) $ 1

main :: IO ()
main = do
  args <- foldl combineArgs defaultArgs `fmap` unsafeInitGUIForThreadedRTS
  let width = aWidth args
      height = aHeight args
      size = roundUp2 (width `max` height)
      rng = ((0, 0), (height - 1, width - 1))
  _ <- initGL
  glconfig <- glConfigNew [ GLModeRGBA, GLModeDouble ]
  canvas <- glDrawingAreaNew glconfig
  widgetSetSizeRequest canvas width height
  imgdata <- mallocBytes $ width * height * 3
  pokeArray imgdata (replicate (height * width * 3) (255 :: B))
  let output x y r g b = do
        let p = imgdata `plusPtr` ((y * width + x) * 3)
        pokeByteOff p 0 r
        pokeByteOff p 1 g
        pokeByteOff p 2 b
  window <- windowNew
  eventb <- eventBoxNew
  set window [ containerBorderWidth := 0, containerChild := eventb,windowResizable := False ]
  set eventb [ containerBorderWidth := 0, containerChild := canvas ]
  mapM_ (flip forkOnIO $ fpu_fix_start nullPtr) [ 0 .. numCapabilities - 1 ]
  stop0 <- renderer' rng output c0 zoom0
  sR <- newIORef (c0, zoom0, stop0)
  _ <- eventb `on` buttonPressEvent $ {-# SCC "cb/event" #-} tryEvent $ do
    LeftButton <- eventButton
    (x, y) <- eventCoordinates
    liftIO $ do
      (c, zoom, stop) <- readIORef sR
      stop
      pokeArray imgdata (replicate (height * width * 3) (255 :: B))
      let c' = c + ((convert x :+ convert (-y)) - (fromIntegral width :+ fromIntegral (-height)) * (0.5 :+ 0)) * ((1/2^^zoom) :+ 0)
          zoom' = zoom + 1
      stop' <- renderer' rng output c' zoom'
      writeIORef sR (c', zoom', stop')
      print (c', zoom') -- FIXME replace with GUI widgets
  _ <- onRealize canvas $ {-# SCC "cb/realize" #-}withGLDrawingArea canvas $ \_ -> do
    GL.clearColor $= (GL.Color4 0.0 0.0 0.0 0.0)
    GL.matrixMode $= GL.Projection
    GL.loadIdentity
    GL.ortho 0.0 1.0 0.0 1.0 (-1.0) 1.0
    GL.drawBuffer $= GL.BackBuffers
    [tex] <- GL.genObjectNames 1
    GL.texture GL.Texture2D $= GL.Enabled
    GL.textureBinding GL.Texture2D $= Just tex
    GL.texImage2D Nothing GL.NoProxy 0 GL.RGB' (GL.TextureSize2D (fromIntegral size) (fromIntegral size)) 0 (GL.PixelData GL.RGB GL.UnsignedByte nullPtr)
    GL.textureFilter GL.Texture2D $= ((GL.Nearest, Nothing), GL.Nearest)
    GL.textureWrapMode GL.Texture2D GL.S $= (GL.Repeated, GL.ClampToEdge)
    GL.textureWrapMode GL.Texture2D GL.T $= (GL.Repeated, GL.ClampToEdge)
  _ <- onExpose canvas $ {-# SCC "cb/expose" #-} \_ -> do
    withGLDrawingArea canvas $ \glwindow -> do
      let v :: GLfloat -> GLfloat -> GLfloat -> GLfloat -> IO ()
          v tx ty vx vy = GL.texCoord (GL.TexCoord2 tx ty) >> GL.vertex (GL.Vertex2 vx vy)
          w = fromIntegral width
          h = fromIntegral height
          sx = fromIntegral width  / fromIntegral size
          sy = fromIntegral height / fromIntegral size
      GL.clear [ GL.ColorBuffer ]
      GL.texSubImage2D Nothing 0 (GL.TexturePosition2D 0 0) (GL.TextureSize2D w h) (GL.PixelData GL.RGB GL.UnsignedByte imgdata)
      GL.renderPrimitive GL.Quads $ do
        v 0 sy 0 0 >> v 0 0 0 1 >> v sx 0 1 1 >> v sx sy 1 0
      glDrawableSwapBuffers glwindow
    return True
  _ <- onDestroy window mainQuit
  _ <- timeoutAdd (widgetQueueDraw canvas >> return True) 200
  widgetShowAll window
  mainGUI

spreadX, spreadY :: [ N ]
spreadX = [ -workers, 0, workers ]
spreadY = [ -1, 0, 1 ]

workers :: N
workers = numCapabilities

limit :: N
limit = (2^(11::N)-1)

c0 :: Complex QuadDouble
c0 = 0

zoom0 :: N
zoom0 = 6

escapeR, escapeR2 :: R
escapeR = 65536
escapeR2 = escapeR * escapeR