{- Jason Ganzhorn
   Haskell Project Question #4, Au08

   This program generates a 2D infinite list of modified fractal values. To
   generate the fractal table, call generateFractal and pass it a 'seed' complex
   value.
   Example: let array1 = generateFractal (Complex (1.3, 0.8))
   Simply calling generateFractal with its argument is a bad idea unless you like
   hitting Ctrl-C.
   
   Since it is rather hard to look at the infinite array passed back from
   generateFractal, I have defined a function called 'takeSquare', which, when passed
   an integer n, returns the top left n x n box of the infinite array. (Top left is
   defined as starting from (0, 0) with increasing coordinates moving right and down,
   respectively.) Please, only use 'takeSquare' on infinite arrays, it isn't designed
   to be used on arrays of finite size (due to its lack of error checking).
   Complete Example:
   let array1 = generateFractal (Complex (1.3, 0.8))
   takeSquare 10 array1
   
   And the output should be the 10 x 10 top left corner of array1.
   
   To run the unit tests, just open this file up in GHCi and type 'runTests'.
-}

import HUnit

{- define a new data type for complex numbers -}
data Complex = Complex (Double, Double)

{- threshold over which to stop calculating terms in the orbit - the orbit is the list
   of terms generated by the fractal equation as it iterates toward the threshold-}
fractalThreshold = 1000.0

{- define a show function for complex numbers so they look pretty when printed; 'a+bi' form -}
instance Show Complex where
    show (Complex (real, complex)) = 
        let complexPart = 
                if complex < 0.0 
                    then '-':(show (abs complex) ++ "i")
                    else '+':(show complex ++ "i")
        in show real ++ complexPart

{- define a function to multiply two complex numbers, following the mathematical rules -}
multiplyComplex :: Complex -> Complex -> Complex
multiplyComplex (Complex (real1, complex1)) (Complex (real2, complex2)) = Complex ((real1*real2 - complex1*complex2), (real1*complex2 + real2*complex1))

{- define a function to add two complex numbers -}
addComplex :: Complex -> Complex -> Complex
addComplex (Complex (real1, complex1)) (Complex (real2, complex2)) = Complex ((real1 + real2), (complex1 + complex2))

{- create a two-dimensional infinite array over which the fractal generator will iterate -}
createBaseMap :: [[Complex]]
createBaseMap = [map (\k -> (Complex (k, 0.0))) element | element <- infArray] where
                    infArray = [[k, k+1 ..] | k <- [1,2 ..]]

{- define a helpful function for users of this program to take a n x n square out of an array
   (note that this function is only guaranteed to work on infinite lists, as I don't do any
   checking for sizes smaller than n -}
takeSquare :: Int -> [[a]] -> [[a]]
takeSquare n array = [take n elem | elem <- (take n array)]

{- run the Mandelbrot fractal equation z = z^2 + c using the given inputs. I chose to use the difference
   between the threshold-breaking value of the equation and the threshold as the returned value.
   I found this number more varied and interesting than what is more commonly returned, the
   number of times the fractal equation runs -}
iterateFractal :: Complex -> Complex -> Double
iterateFractal z c = 
    if magnitude z > fractalThreshold
        then (magnitude z) - 100.0
        else iterateFractal (addComplex (multiplyComplex z z) c) c

{- calculate the magnitude of a complex number -}
magnitude :: Complex -> Double
magnitude (Complex (real, complex)) = sqrt (real^2 + complex^2)

{- applies the fractal equation to every element in the infinite map.
   The parameter to this function gives a starting z for the Mandelbrot equation
   at every index in the array. -}
generateFractal :: Complex -> [[Integer]]
generateFractal z = [map (\c -> ((round (iterateFractal z c)))) elem | elem <- createBaseMap]

{- This is not an easy program to test, as I really don't know what the fractal table SHOULD
   look like. However, I can test the various component functions. -}

c1 = Complex (3.0, (-2.0))
c2 = Complex (5.0, 0.0)

{- test multiplication and addition -}
test1 = TestCase (assertEqual "complex multiplication test" "15.0-10.0i" (show (multiplyComplex c1 c2)))
test2 = TestCase (assertEqual "complex addition test" "8.0-2.0i" (show (addComplex c1 c2)))

{- test the takeSquare and createBaseMap methods (mostly takeSquare) -}
array1 = takeSquare 10 createBaseMap
test3 = TestCase (assertEqual "takeSquare and createBaseMap test" 10 (length array1))
test4 = TestCase (assertEqual "takeSquare and createBaseMap test" 10 (length (head array1)) )

{- test the magnitude function -}
test5 = TestCase (assertEqual "magnitude test" 5.0 (magnitude c2))
test6 = TestCase (assertEqual "magnitude test" (sqrt 13) (magnitude c1))

tests = TestList [TestLabel "Complex Multiplication Test 1" test1, TestLabel "Complex Addition Test 1" test2, TestLabel "takeSquare and createBaseMap test 1" test3, TestLabel "takeSquare and createBaseMap test 2" test4, TestLabel "Magnitude test 1" test5, TestLabel "Magnitude test 2" test6]

runTests = runTestTT tests