{-
Joris Sebastijonas Urbaitis
s1104617

Supersampled mandelbrot set renderer

It's best to compile it before running for speed.
Otherwise just run main (no arguments).
It can take a couple of minutes to render.

I wrote it using Notepad++ on Windows, so formatting might be weird on emacs or the like.

Would prefer a book token.
-}

import Graphics.Rendering.OpenGL
import Graphics.UI.GLUT
import Data.List
import Data.Fixed

-- Not sure if this is in the standard somewhere so I wrote it
even' :: GLfloat -> Bool
even' f = f `mod'` 2 == 0

-- Type and helper functions for RGB colour
type Colour = (GLfloat, GLfloat, GLfloat)

--These two functions are for supersampling
colAdd :: Colour -> Colour -> Colour
colAdd (r1, g1, b1) (r2, g2, b2) = (r1 + r2, g1 + g2, b1 + b2)

colDiv :: Colour -> GLfloat -> Colour
colDiv (r, g, b) f = (r / f, g / f, b / f)

-- Type for pixels - first two floats are coordinates, then the colour
type Point = (GLfloat, GLfloat, Colour)

-- Complex numbers - wrote this before I became aware of Data.Complex
data Complex = Complex GLfloat GLfloat
	deriving (Eq)
	
instance Num Complex where
	Complex x y + Complex x' y' = Complex (x+x') (y+y')
	Complex x y - Complex x' y' = Complex (x-x') (y-y')
	Complex x y * Complex x' y' = Complex (x*x' - y*y') (x*y' + y*x')
	abs (Complex x y) 			= undefined
	signum (Complex x y) 		= undefined
	fromInteger i				= undefined
	
instance Show Complex where
	show (Complex x y) = show x ++ " + " ++ show y ++ "i"
	
cAbs :: Complex -> GLfloat
cAbs (Complex x y) = x^2 + y^2

cRe :: Complex -> GLfloat
cRe (Complex x y) = x

cIm :: Complex -> GLfloat
cIm (Complex x y) = y
	  
-- Main functions for OpenGL, mostly stolen from LSystem.hs (Turtle graphics file)
main :: IO ()
main = do 
	initialWindowSize $= Size 600 600
	(progname, _) <- getArgsAndInitialize
	createWindow "Mandelbrot"
	displayCallback $= drawMandelbrotGL
	mainLoop
  
drawMandelbrotGL :: IO ()
drawMandelbrotGL = do 
	clear [ColorBuffer]
	renderPrimitive Points $ mapM_ (\p -> renderPoint p) genMandelbrot
	flush
  where
    renderPoint (x, y, (r, g, b)) = do { color $ (Color3 r g b); vertex $ (Vertex2 x y) }

-- Function for generating the image
-- It first generates a 1200x1200 image, then supersamples it to remove aliasing						
genMandelbrot :: [Point]
genMandelbrot = [ cP x y | x <- [0..1199], y <- [0..1199], even' x, even' y ]
  where
    cP x y = (mx x, my y, combCol [ genColour min max fac (x+u) (y+v) 120 | u <- [0, 1], v <- [0, 1] ] )
    combCol (c1:c2:c3:c4:[]) = colDiv (c1 `colAdd` c2 `colAdd` c3 `colAdd` c4) 4.0
    mx x = (x - 600) / 600.0
    my y = (y - 600) / 600.0
    width = 0.1111111111111
    min = Complex (-1.05373672712496) (-0.356741858456567)
    max = Complex (-1.05373672712496 + width) (-0.356741858456567 + width)
    fac = Complex ((cRe max - cRe min) / 1199.0) ((cIm max - cIm min) / 1199.0)

-- Function that calculates the iterations and returns an appropriate colour
genColour :: Complex -> Complex -> Complex -> GLfloat -> GLfloat -> GLfloat -> Colour
genColour min max fac x y itMax = f itMax c
  where
    c = Complex (cRe min + x * (cRe fac)) (cIm max - y * (cIm fac))
    f it z  | it == 0 && cAbs z <= 4 = getColour itMax it
            | cAbs z > 4  		   	 = getColour itMax it
            | otherwise				 = f (it - 1) (z^2 + c)

-- Gets the colour of a pixel given a number of iterations
getColour :: GLfloat -> GLfloat -> Colour
getColour itMax it | ci <= 0.5 = (ci * 2, ci * 2, ci * 2)
                   | otherwise = (1.0 - (ci - 0.5) * 2, 1.0 - (ci - 0.5) * 2, 1.0)
  where 
    ci = it / itMax