{-
    fractal : A fractal generator in hasell

    Author : Mriganka Basu Roy Chowdhury

    Install : Preferably (if you don't want cabal hell),
              form a cabal sandbox and install JuicyPixels
              and optparse-applicative. Then just
              cabal exec -- ghc fractal.hs -O3
              Please use -O3, as otherwise, the generation
              may take a lot of time.

    Modifications :
             Look into the source below, and change
             stuff if you need to. Most of the code is
             segmented, and documented, so that part
             should be pretty obvious.
-}


import Codec.Picture                            -- needs JuicyPixels
import Data.Complex                             -- in base, so chill.
import Data.List                                -- again in base.
import Options.Applicative                      -- need optparse-applicative
import Data.Semigroup       hiding  (option)    -- in base
import Text.Read                    (readMaybe) -- in base as well.
import Data.Maybe                               -- in base.
import Control.Monad                            -- in base.


type R = Float                          -- The real used, change to Double for precision
type C = Complex R
type IterF = C -> C                     -- Iteration function

type Pix        = PixelRGB8             -- The pixel type used.
type View       = (C, C)                -- Upper left, lower right complex numbers
type Dimensions = (Int, Int)            -- Width, Height

inf :: R
inf = 1e20

-------------------- For general uses, just modify this part -------

baseFunction :: C -> IterF              -- Change this function to your wishes, but
                                        -- keep the signature intact. :)
                                        -- This function currently generates mandelbrot sets.
baseFunction c z = z**2 + c

-------------- ************************************ ----------------

characteristic  :: (C -> IterF)         -- Iteration function
                -> R                    -- Escape Radius
                -> Int                  -- Iterations
                -> C                    -- Parameter
                -> Pix                  -- A double indicating color.
characteristic f r n c =
    let g = f c
        l = map snd $
                takeWhile (\(i, z) -> i <= n && realPart (abs z) < r) $
                     zip [1..] $ iterate g (0.0 :+ 0.0)
    in colorFunc n (length l) c (last l)

-------------------------------------------------------------------

data FractalOptions
    = FractalOptions
    { iterations   :: Int
    , escapeRad    :: R
    , upperLeftX   :: R
    , upperLeftY   :: R
    , lowerRightX  :: R
    , lowerRightY  :: R
    , width        :: Int
    , height       :: Int
    , output       :: FilePath
    }


main :: IO ()
main = do
    results <- execParser $ info (parser <**> helper)
                            ( header "--- fractal : A fractal generator ---"
                            <> progDesc "Generates a PNG image of a fractal generated by iteration"
                            <> fullDesc
                            )
    let (FractalOptions n r ulx uly lrx lry w h out) = results
        image =    produceFractal baseFunction n r
                        ((ulx :+ uly), (lrx :+ lry))
                        (w, h)

    putStrLn "Generating and writing the image... Hang tight."

    writePng out image

    putStrLn "Done. Enjoy!"

---------------- The parser for optparse ------------

parser :: Parser FractalOptions
parser = FractalOptions <$>
    option auto
    (   long "iterations"
    <>  short 'n'
    <>  help "The number of times the iteration happens."
    <>  metavar "INTEGER"
    <>  value 30
    )
    <*>
    option auto
    (  long "escape"
    <> short 'r'
    <> help "Escape radius"
    <> metavar "REAL"
    <> value 10.0
    )
    <*>
    option auto
    (   long "upper-left-x"
    <>  short 'X'
    <>  help "The upper left x in the frame of the image generated."
    <>  metavar "REAL"
    <>  value (-1.8)
    )
    <*>
    option auto
    (   long "upper-left-y"
    <>  short 'Y'
    <>  help "The upper left y in the frame of the image generated."
    <>  metavar "REAL"
    <>  value (1.3)
    )
    <*>
    option auto
    (   long "lower-right-x"
    <>  short 'x'
    <>  help "The upper left x in the frame of the image generated."
    <>  metavar "REAL"
    <>  value (0.8)
    )
    <*>
    option auto
    (   long "lower-right-y"
    <>  short 'y'
    <>  help "The upper left x in the frame of the image generated."
    <>  metavar "REAL"
    <>  value (-1.3)
    )
    <*>
    option auto
    (   long "width"
    <>  short 'w'
    <>  help "Width of the image generated."
    <>  metavar "INTEGER"
    <>  value 600
    )
    <*>
    option auto
    (   long "height"
    <>  short 'h'
    <>  help "Height of the image generated."
    <>  metavar "INTEGER"
    <>  value 600
    )
    <*>
    strOption
    (   long "output"
    <>  short 'o'
    <>  help "Output file name of the png image"
    <>  metavar "FILENAME"
    <>  value "output.png"
    )


------------- A few math functions ---------------------

clamp :: R -> R -> R -> R
clamp l r x = min r (max x l)

--

normal          :: R               -- center
                -> R               -- spread
                -> R               -- value
                -> R
normal c spread x = exp (- (abs (x - c))**2 / spread)

--

sigmoidCut :: R -> R -> R -> R
sigmoidCut c s x = if x < c then 0.0 else -0.5 + 1.0 / (1 + exp (-(x - c)/s))

--------------------------------------------------------


-- The function that generates the color.
-- The current configuration was obtained after many trials.
-- So perfecting it may require trial and error / theory.

colorFunc :: Int                   -- Number of iterations max
          -> Int                   -- Number of iterations taken
          -> C                     -- Current complex number
          -> C                     -- Final complex number after iteration
          -> Pix                   -- Generated color
colorFunc n m z z' =
    let w = 1 + (realPart $ abs z)
        w'= 1 + (realPart $ abs z')
        x = fromIntegral (n - m) / (fromIntegral n)
        r = normal (0.8)  0.05 x
        g = normal (0.8)  0.1 x
        b = sigmoidCut (0.8) 0.05 x

        r'= floor $ r * 255
        g'= floor $ g * 255
        b'= floor $ b * 255
    in  PixelRGB8 r' g' b'


--------------- The function that produces the fractal ------

produceFractal  :: (C -> IterF)         -- Iteration function
                -> Int                  -- Iterations
                -> R                    -- Escape radius
                -> View                 -- Region being viewed for the fractal
                -> Dimensions           -- Image dimensions
                -> Image Pix            -- Output image

produceFractal f n r (upperLeft, lowerRight) (w, h) =
    let viewWidth = realPart (lowerRight - upperLeft)
        viewHeight= imagPart (upperLeft - lowerRight)

        charProducer = characteristic f r n

        pixelSpace = [(x, y) | x <- [0..w], y <- [0..h]]

        genFunction :: Int -> Int -> Pix
        genFunction x y =
            let dx =  (fromIntegral x * viewWidth)  / (fromIntegral w)
                dy = -(fromIntegral y * viewHeight) / (fromIntegral h)
                z  = upperLeft + (dx :+ dy)
            in charProducer z

    in generateImage genFunction w h

-------------------------------------------------------------