-- ghc -O2 -threaded --make Partest
-- ./Partest +RTS -N4 > bigoutx300x300x300.vtk

import Control.Parallel 
import Data.List

-- z^n where n = 8
mandelN = 8
mandelIterations = 100

-- This is the actual equation
-- z <- z^n + c
-- mandeliter n z c 

-- Quad core map (please don't bug me about this!*)
mapP :: (a -> b) -> [a] -> [b]
mapP _ [] = []
mapP f (w:[]) = [f w]
mapP f (w:x:[]) = fx `par` [f w,fx]
 where fx = f x
mapP f (w:x:y:[]) = fx `par` fy `par` [f w,fx,fy]
 where fx = f x
       fy = f x
mapP f (w:x:y:z:zs) = fx `par` fy `par` fz `par` (f w:fx:fy:fz : (mapP f zs))
 where fx = f x 
       fy = f y
       fz = f z

-- Partition data into n chunks
splitData n l = keeptaking len l
    where
      len = length l
      chunk =  len `div` n
      keeptaking left x = if (left <= chunk)
                          then [x]
                          else
                              xs : (keeptaking (left - chunk) ys)
                              where
                                (xs,ys) =  splitAt chunk x
      
-- Partition data and mapP it 
splitDataMapP n f l = concat $ mapP fp splits
    where
      splits = splitData n l
      fp l = map f l

-- Generate a VTK header
header dim =
    [ "# vtk DataFile Version 2.0"
    ,"3D Fractal"
    ,"ASCII"
    ,"DATASET STRUCTURED_POINTS"
    ,("DIMENSIONS "++(intercalate " " dimensions))
    ,("ORIGIN "++(intercalate " " halfdimensions))
    ,"SPACING 1 1 1"
    ,"POINT_DATA "++(show dim3)
    ,"SCALARS color FLOAT"
    ,"LOOKUP_TABLE default"
    ]
    where
      repeat3 str = [str | x <- [1..3]]
      idim  = dim
      halfdim = idim `div` 2 
      dim3 = idim^3
      dimensions = repeat3 (show idim)
      halfdimensions = repeat3 (show halfdim)

-- our main, print vtk header and print the gen values out
main = do
  let dim = 100
  mapM_ putStrLn $ header dim
  mapM_ putStrLn (map show (values (fromIntegral dim)))

-- How we get our values,
--  parallel map over our range (positive quarter)
values n = mapP (itertilwrap n) (rangegen n)

-- values n = splitDataMapP 16 (itertilwrap n) (rangegen n)

-- gen the range we operate over (the structured_points)
rangegen n = [ (x,y,z) | x <- [1..n], y <- [1..n], z <- [1..n] ]

-- we should make this rangeable but whatever,
itertilwrap n (x,y,z) = itertil mandelN mandelIterations (0,0,0) (x/n,y/n,z/n)

-- This iterations cnt times over the mandelbrot equation
--  til it hits max iters or it is done
itertil n cnt (x,y,z) (xo,yo,zo) = 
    if (cnt == 0) 
    then cnt
    else if (x^2 + y^2 + z^2 > 2.0) 
         then cnt
         else itertil n (cnt - 1) (mandeliter n (x,y,z) (xo,yo,zo)) (xo,yo,zo)

-- This is the mandelbrot equation, this is like z <- z^n + c
mandeliter n (x,y,z) (xo,yo,zo) = (nxc,nyc,nzc) -- (newx,newy,newz)
    where
      r = sqrt ( x^2 + y^2 + z^2 )
      theta = atan2 (sqrt (x^2 + y^2))  z
      phi = atan2 y x 
      rn = r**n
      thetan = theta * n
      phin = phi * n
      sinthetan = sin thetan
      newx = rn * sinthetan * (cos phin)
      newy = rn * sinthetan * (sin phin)
      newz = rn * (cos thetan)
      nxc = newx + xo
      nyc = newy + yo
      nzc = newz + zo