module Main where

import Codec.Picture
import qualified Codec.Picture.Types as M
import Control.Arrow (second)
import Control.Concurrent.ParallelIO
import Control.Monad
import Control.Monad.ST
import Data.List (zip5)
import Data.Word (Word8)
import Options.Applicative
import System.Environment (getArgs)
import System.Random (Random, randomRIO)

{-MCCOMMENT

Required packages:

  - base
  - containers
  - random
  - parallel-io
  - optparse-applicative
  - JuicyPixels

Instructions:

  Build:
    > ghc --make -O2 -threaded -o render Main.hs

  Usage help:
    > ./render --help
    > ./render img --help
    > ./render seq --help

Examples (please use the configuration here for my submission):

  Generate a single image (~1.76s):
    > ./render +RTS -N1 -RTS img -o {out}.png -a 0.361247 -w 3840 -h 2160 -g 400000 -b 60000

  Generate an image sequence (~114s with 6 threads on 6 cores, please create the target directory first):
    > ./render +RTS -N{threads} -RTS seq -f 30 -o {out_dir}

  Generate a video from an image sequence:
    > ffmpeg -r 30 -i {out_dir}/%d.png -c:v libx264 -pix_fmt yuv420p {out}.mp4

Feel free to play with it, for example passing custom keyframes.
(You may also want to edit 'animateGalaxyStars' and 'animateBackgroundStars'.)

Cheers!

-}

{--- TYPES ---}

type File = String
type Directory = String

type Width = Int
type Height = Int

type X = Int
type Y = Int

type Radius = Int
type Angle = Double

type Pos = (X, Y)
type GalaxyStar = (Angle, Radius)
type Speed = Double

type AspectRatio = Double
type Spin = Double

type Fps = Int
type Duration = Int
type Moment = Int
type Frame = Int
type Keyframe = (Moment, Double)

type ColorTemperature = Int
type Color = (Word8, Word8, Word8)

{-- ANIMATION CONFIGURATION --}

-- Galaxy stars animation strategy.
animateGalaxyStars :: Fps -> Duration -> Speed -> Int -> [GalaxyStar] -> [[GalaxyStar]]
animateGalaxyStars fps duration s n gs =
    gss' ++ zipWith (drop . (dropsPerFrame *)) [0..] gss''
  where
    dropsPerFrame = n `div` (destroyIn * fps)
    (gss', gss'') = splitAt (fps * destroyAfter) movingGs
    movingGs = map (\i -> map (moveGalaxyStar s i) gs) [0..duration * fps]
    destroyIn = 8
    destroyAfter = duration - destroyIn

-- Background stars animation strategy.
animateBackgroundStars :: Fps -> Duration -> Int -> [Pos] -> [[Pos]]
animateBackgroundStars fps duration n rs = map rsAtFrame [0..]
  where
    rsAtFrame = flip drop rs . ceiling . (dropsPerFrame *)
    dropsPerFrame = fromIntegral n / fromIntegral (destroyUntil * fps)
    destroyUntil = duration

{--- MAIN ---}

main :: IO ()
main = do
  opts <- execParser argsParser
  let ProgOptions{modeOptions=mOpts} = opts in case mOpts of
    o@ImageOptions{} -> renderImage opts mOpts
    o@SequenceOptions{} -> renderSequence opts mOpts

{--- CLI PARAM PARSING ---}

data ProgOptions = ProgOptions
  {
    modeOptions :: ModeOptions,
    width :: Width,
    height :: Height,
    nGalaxyStars :: Int,
    nBackgroundStars :: Int
  }

data ModeOptions =
  ImageOptions
    {
      aspectRatio :: AspectRatio,
      spin :: Spin,
      out :: File
    } |
  SequenceOptions
    {
      fps :: Fps,
      starMovementSpeed :: Speed,
      aspectRatioKeyframes :: [Keyframe],
      spinKeyframes :: [Keyframe],
      out :: Directory
    }

-- Parser for command line arguments.
argsParser :: ParserInfo ProgOptions
argsParser =
  info
    (helpOption <*> versionOption <*> programOptions)
    (fullDesc <> header "URGE - The UnReal Galaxy Engine")
  where
    helpOption = abortOption ShowHelpText (long "help" <>
                                           help "Display this message")
    versionOption = infoOption "1.0" (long "version" <>
                                      short 'v' <>
                                      help "Show version")
    programOptions = ProgOptions <$>
      hsubparser (imageMode <> videoMode) <*>
      option auto (long "width" <>
                   short 'w'<>
                   metavar "INT" <>
                   value 1920 <>
                   help "Width in pixels") <*>
      option auto (long "height" <>
                   short 'h'<>
                   metavar "INT" <>
                   value 1080 <>
                   help "Height in pixels") <*>
      option auto (long "galaxyStars" <>
                   short 'g' <>
                   metavar "INT" <>
                   value 115000 <>
                   help "Number of galaxy stars") <*>
      option auto (long "backgroundStars" <>
                   short 'b'<>
                   metavar "INT" <>
                   value 2200 <>
                   help "Number of stars in the background")
    imageMode = command
      "img"
      (info imgOptions (progDesc "Generate a single image"))
    imgOptions = ImageOptions <$>
      option auto (long "eAspectRatio" <>
                   short 'a' <>
                   metavar "FLOAT" <>
                   value 0.42124721603563486 <>
                   help "Galaxy ellipse aspect ratio") <*>
      option auto (long "eSpin" <>
                   short 's' <>
                   metavar "FLOAT" <>
                   value 8.0 <>
                   help "Galaxy ellipse \"spin\"") <*>
      strOption (long "out" <>
                 short 'o' <>
                 metavar "STRING" <>
                 value "out.png" <>
                 help "Output file (.png)")
    videoMode = command
      "seq"
      (info videoOptions (progDesc "Generate an image sequence"))
    videoOptions = SequenceOptions <$>
      option auto (long "fps" <>
                   short 'f' <>
                   metavar "INT" <>
                   value 30 <>
                   help "Frames per second") <*>
      option auto (long "starMovementSpeed" <>
                   short 'm' <>
                   metavar "FLOAT" <>
                   value 0.02 <>
                   help "Galaxy star movement speed factor") <*>
      option auto (long "aspectRatioKeyframes" <>
                   short 'a' <>
                   metavar "[(INT, FLOAT)]" <>
                   value [(0, 42.0), (7, 16.0), (18, 0.03), (29, 0.03), (31, 0.0), (32, 0.0), (38, 15.0), (45, 250.0)] <>
                   help "Galaxy ellipse aspect ratio keyframes") <*>
      option auto (long "spinKeyframes" <>
                   short 's' <>
                   metavar "[(INT, FLOAT)]" <>
                   value [(0, 8.0), (18, 8.0), (29, 0.0), (45, 0.0)] <>
                   help "Galaxy ellipse spin keyframes") <*>
      strOption (long "out" <>
                 short 'o' <>
                 metavar "STRING" <>
                 value "out" <>
                 help "Output directory")

{--- IMAGE RENDERING ---}

-- Render a single image.
renderImage :: ProgOptions -> ModeOptions -> IO ()
renderImage
  ProgOptions{width=w, height=h, nGalaxyStars=nGs, nBackgroundStars=nBs}
  ImageOptions{aspectRatio=ar, spin=s, out=file} = do
    let r = galaxyRadiusFromHeight h
    rs <- randomStarPositions w h nBs
    gs <- randomGalaxyStars r nGs
    let galaxy = genGalaxy w h ar (s / fromIntegral r) gs rs
    savePngImage file (ImageRGB8 galaxy)

-- Render an image sequence.
renderSequence :: ProgOptions -> ModeOptions -> IO ()
renderSequence
  ProgOptions{width=w, height=h, nGalaxyStars=nGs, nBackgroundStars=nBs}
  SequenceOptions{
    fps=fps,
    starMovementSpeed=s,
    aspectRatioKeyframes=aks,
    spinKeyframes=sks,
    out=dir
    } = do
      putStrLn "You might want to ask your favorite subordinate MC for a cup of coffee now..."
      gs <- randomGalaxyStars (galaxyRadiusFromHeight h) nGs
      rs <- randomStarPositions w h nBs
      let duration = fst (last aks)
      let gss = animateGalaxyStars fps duration s nGs gs
      let rss = animateBackgroundStars fps duration nBs rs
      let ars = seqFromKeyframes fps aks
      let r = fromIntegral (galaxyRadiusFromHeight h)
      let sks' = map (second (/ r)) sks
      let ess = seqFromKeyframes fps sks'
      let dir' = if last dir == '/' then dir else dir ++ "/"
      let tasks = genTasks dir' w h (zip5 [0..] ars ess gss rss)
      parallel_ tasks
      putStrLn "Done."
      stopGlobalPool

-- Create a sequence of single frames from keyframes, distributing equally.
seqFromKeyframes :: Fps -> [Keyframe] -> [Double]
seqFromKeyframes fps = fst . foldl toSeq ([], (-1, 0.0))
  where
    toSeq (acc, w1) w2
      | fst w1 == -1 = (acc, w2)
      | otherwise = (acc ++ ps w1 w2, w2)
    ps (t1, p1) (t2, p2) = pointsInRange (fps * (t2 - t1)) p1 p2

-- Generate points with equal distances to each other within a given range.
pointsInRange :: Int -> Double -> Double -> [Double]
pointsInRange n a b
  | a == b = replicate n a
  | otherwise = map (\i -> a `go` dist i) [0..pred n]
  where
    dist i = (fromIntegral i / fromIntegral (pred n)) * diff
    diff = abs (b - a)
    go = if a < b then (+) else (-)

-- Generate I/O tasks which can be executed in parallel.
genTasks :: String -> Width -> Height -> [(Frame, AspectRatio, Spin, [GalaxyStar], [Pos])] -> [IO ()]
genTasks dir w h = map genTask
  where
    genTask ps@(i, ar, s, gs, rs) = do
      when (i > 0 && i `mod` 50 == 0) $
          putStrLn $ "Rendered ~" ++ show i ++ " images..."
      let galaxy = genGalaxy w h ar s gs rs
      savePngImage (dir ++ show i ++ ".png") (ImageRGB8 galaxy)

-- Generate the galaxy image.
genGalaxy :: Width -> Height -> AspectRatio -> Spin -> [GalaxyStar] -> [Pos] -> Image PixelRGB8
genGalaxy w h ar s gs rs = runST $ do
  im <- M.createMutableImage w h blackPixel
  let gps = map (galaxyStarPos w h ar s) gs
  let gps' = filter (inImageBounds w h) gps
  mapM_ (\(x, y) -> writePixel im x y whitePixel) rs
  mapM_ (\(x, y) -> writePixel im x y (colorPixelAt w h (x, y))) gps'
  M.unsafeFreezeImage im

-- Generate random star positions outside of the galaxy.
randomStarPositions :: Width -> Height -> Int -> IO [Pos]
randomStarPositions w h n = mapM (const $ randomPos w h) [1..n]

-- Generate the galaxy stars.
randomGalaxyStars :: Radius -> Int -> IO [GalaxyStar]
randomGalaxyStars r n = mapM (const $ randomGalaxyStar r) [1..n]

-- Generate a random star inside the galaxy.
randomGalaxyStar :: Radius -> IO GalaxyStar
randomGalaxyStar r = do
  a <- randomBetween 0 (2 * pi)
  r <- randomBetween 1 r
  return (a, r)

{--- GRAPHICS/2D UTILS ---}

-- Calculate the distance of a position from the center.
radius :: Width -> Height -> Pos -> Double
radius w h (x, y) = let (centerX, centerY) = center w h in
    sqrt $ fromIntegral $ ((x - centerX) ^ 2) + ((y - centerY) ^ 2)

-- Check whether a position is inside the allowed range.
inImageBounds :: Width -> Height -> Pos -> Bool
inImageBounds w h (x, y) = x > 0 && y > 0 && x < w && y < h

-- Find the center of an image.
center :: Width -> Height -> Pos
center w h = (w `div` 2, fromIntegral h `div` 2)

-- Calculate the max galaxy radius.
galaxyRadiusFromHeight :: Height -> Radius
galaxyRadiusFromHeight h = (7 * h) `div` 16

-- Move a star following its trajectory.
moveGalaxyStar :: Speed -> Int -> GalaxyStar -> GalaxyStar
moveGalaxyStar s n (a, r) = (a + fromIntegral n * s, r)

-- Calculate a galaxy star's position in the 2-dimensional space.
galaxyStarPos :: Width -> Height -> AspectRatio -> Spin -> GalaxyStar -> Pos
galaxyStarPos w h ar s (a, r) = (x, y)
  where
    x = round (centerX + x' * sin b + y' * cos b)
    y = round (centerY + x' * cos b - y' * sin b)
    x' = sin a * fromIntegral r
    y' = cos a * fromIntegral r * ar
    b = s * fromIntegral r
    centerX = fromIntegral centerX'
    centerY = fromIntegral centerY'
    (centerX', centerY') = center w h

{--- COLOR UTILS ---}

blackPixel :: PixelRGB8
blackPixel = PixelRGB8 0 0 0

whitePixel :: PixelRGB8
whitePixel = PixelRGB8 255 255 255

-- Generate a colored star pixel based on its distance from the galaxy center.
colorPixelAt :: Width -> Height -> Pos -> PixelRGB8
colorPixelAt w h p = PixelRGB8 r g b
  where
    (r, g, b) = temperatureToRgb t
    t = 1700 + round (radius w h p / radius w h (0, 0) * 6000)

-- Convert a color temperature (in K) to a color.
-- (based on http://www.tannerhelland.com/4435/convert-temperature-rgb-algorithm-code/)
temperatureToRgb :: ColorTemperature -> Color
temperatureToRgb t
  | t >= 6590 && t <= 6680 = (255, 249, 255)
  | otherwise = (to8BitRange r, to8BitRange g, to8BitRange b)
  where
    r
      | t' <= 66 = 255.0
      | otherwise = 329.698727446 * ((t' - 60.0) ** (-0.1332047592))
    g
      | t' <= 66 = 99.4708025861 * log t' - 161.1195681661
      | otherwise = 288.1221695283 * ((t' - 60.0) ** (-0.0755148492))
    b
      | t' >= 66 = 255.0
      | t' <= 19.1 = 0.0
      | otherwise = 138.5177312231 * log (t' - 10.0) - 305.0447927307
    t' = fromIntegral t / 100.0
    to8BitRange = max 0 . min 255 . round

{--- RANDOMIZATION UTILS ---}

-- Generate a random position in the 2-dimensional space.
randomPos :: Width -> Height -> IO Pos
randomPos w h = do
  x <- randomBetween 0 (pred w)
  y <- randomBetween 0 (pred h)
  return (x, y)

-- Generate a random number between given limits (both included).
randomBetween :: Random a => a -> a -> IO a
randomBetween a b = randomRIO (a, b)