{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}
module Main (main) where

import qualified Codec.Picture.Png as PNG
import qualified Codec.Picture.Types as Pix
import qualified Data.ByteString.Lazy as BSL
import Data.List (intercalate)
import qualified Data.Map.Strict as Map
import Data.Map.Strict (Map)
import qualified Data.Serialize as Ser
import Data.Word
import qualified Data.Vector.Storable as VS
import System.Environment (getArgs)
import System.Exit (die, exitFailure)

import MandelHSlib
import Inferno


type Colorscheme = Double -> (Word8, Word8, Word8)

colorscheme_yellow :: Colorscheme
colorscheme_yellow fraction =
  let x = max 0 fraction ** 0.3
      bg = 0.2 * (1 - curve x (-0.1) 0.2)
  in (tow8 $ bg + curve x 0 0.4 - 0.8 * curve x 0.6 1.0
     ,tow8 $ bg + 0.8 * curve x 0.3 0.7
     ,tow8 $ bg + curve x 0.6 1.0)

colorscheme_purple :: Colorscheme
colorscheme_purple fraction =
  let x = max 0 fraction ** 0.4
  in (tow8 $ 0.05 + 0.2 * curve x 0.0 0.4 + 0.65 * curve x 0.3 0.8 - 0.6 * curve x 0.8 1.0
     ,0
     ,tow8 $ 0.1 + 0.7 * curve x 0.0 0.4 - 0.4 * curve x 0.7 1.0 - 0.2 * curve x 0.9 1.0)

tow8 :: Double -> Word8
tow8 x = round (max 0 (min 255 (x * 255)))

curve :: Double -> Double -> Double -> Double
curve x start end
  | x <= start = 0
  | x >= end = 1
  | otherwise = sin (pi/(end-start) * (x - start) - pi/2) / 2 + 0.5

colorschemes :: Map String Colorscheme
colorschemes = Map.fromList
  [("yellow", colorscheme_yellow)
  ,("purple", colorscheme_purple)
  ,("infernoinv", colorscheme_infernoinv)]

renderFractal :: Colorscheme -> Fractal -> Pix.Image Pix.PixelRGB8
renderFractal colors (Fractal (w, h) _ arr) =
  Pix.Image w h $ VS.fromList
    [val
    | fraction <- VS.toList arr
    , let (r, g, b) = colors fraction
    , val <- [r, g, b]]

usage :: String
usage =
  "Usage: mandelhs-render <infile.data> <outfile.png> <colorscheme>\n"

main :: IO ()
main = do
  -- let (w, h) = (4 * 2000, 4 * 1500)
  --     maxiter = 1024

  args <- getArgs
  case args of
    [infile, outfile, csname] -> do
      colors <- case Map.lookup csname colorschemes of
                  Nothing -> die $ "Available colorschemes: " ++ intercalate ", " (Map.keys colorschemes)
                  Just f -> return f
      datafile <- BSL.readFile infile
      fractal <- case Ser.runGetLazy Ser.get datafile of
        Left err -> die err
        Right res -> return res
      BSL.writeFile outfile $
        PNG.encodePng $
          renderFractal colors fractal

    _ -> putStr usage >> exitFailure