Knytt Stories level map renderer

KnyttStoriesMapper.hs

-- Pass this program a world name, and run it from the Knytt Stories root directory:
-- 
--     cd "Knytt Stories"
--     ./KnyttStoriesMapper "Nifflas - The Machine"
--
-- It will create a file called "Nifflas - The Machine.png" containing a map.

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}

import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import Control.Applicative
import qualified Data.Array as A
import Data.Array (Array)
import Data.Char
import Data.Either.Unwrap (fromRight)
import Data.List (intercalate)
import Data.Map (Map)
import qualified Data.Map as M

import Data.Word (Word8)
import qualified Codec.Compression.GZip as GZip
import System.IO.Error
import qualified Data.Attoparsec.ByteString as P
import Data.Attoparsec.ByteString (Parser)

import Codec.Picture.Types
import Codec.Picture.Png (decodePng, writePng)
import qualified Vision.Primitive as Im
import qualified Vision.Image.Type as Im
import qualified Vision.Image.Class as Im
import Vision.Primitive (Z(..), (:.)(..))
import Vision.Image.Transform (crop)
import Vision.Image.RGBA.Type (RGBA, RGBAPixel(..))
import Vision.Image.JuicyPixels (toFridayRGBA, toJuicyRGBA)

import System.Environment
import qualified Data.Vector.Storable as V

---------------------------------------------------------------------------
-- Types
---------------------------------------------------------------------------

data ResourceType = Gradient | Tileset
type WorldName = String

type GradientArray = Array Word8 RGBA
type TilesetArray = Array Word8 RGBA
type SpriteMap = Map (Word8, Word8) RGBA
data Resources = Resources GradientArray TilesetArray SpriteMap

data LayerData
    = TileLayer (Array (Int, Int) Word8)
    | SpriteLayer (Array (Int, Int) (Word8, Word8)) deriving (Show)

-- Layers, tileset A, tileset B, background.
data MapData =
    MapData [LayerData] Word8 Word8 Word8 deriving (Show)

type LevelData = Map (Int, Int) MapData

---------------------------------------------------------------------------
-- Constants
---------------------------------------------------------------------------

mapWidth, mapHeight :: Int
mapWidth = 25
mapHeight = 10

tileWidth, tileHeight :: Int
tileWidth = 24
tileHeight = 24

tilesPerTilesetRow :: Int
tilesPerTilesetRow = 16

tileSize :: Im.Size
tileSize = Im.ix2 tileHeight tileWidth

layerWidth :: Int
layerWidth = mapWidth * tileWidth

layerHeight :: Int
layerHeight = mapHeight * tileHeight

layerSize :: Im.Size
layerSize = Im.ix2 layerHeight layerWidth

---------------------------------------------------------------------------
-- File I/O
---------------------------------------------------------------------------

joinPath :: [String] -> String
joinPath = intercalate "\\"

resourceDirectoryName :: ResourceType -> String
resourceDirectoryName Gradient = "Gradients"
resourceDirectoryName Tileset = "Tilesets"

resourcePrefix :: ResourceType -> String
resourcePrefix Gradient = "Gradient"
resourcePrefix Tileset = "Tileset"

resourceName :: ResourceType -> Int -> String
resourceName r num =
    resourcePrefix r ++ show num ++ ".png"

globalResourcePath :: ResourceType -> Int -> String
globalResourcePath r num =
    joinPath ["Data", resourceDirectoryName r, resourceName r num]

localResourcePath :: WorldName -> ResourceType -> Int -> String
localResourcePath worldName r num =
    joinPath ["Worlds", worldName, resourceDirectoryName r, resourceName r num]

mapBinPath :: WorldName -> String
mapBinPath worldName =
    joinPath ["Worlds", worldName, "Map.bin"]

catchNotFound :: IO a -> IO a -> IO a
catchNotFound x y =
    catchIOError x (\e -> if isDoesNotExistError e then y else ioError e)

readPng :: B.ByteString -> RGBA
readPng bs = case decodePng bs of
    Right (ImageRGBA8 im) -> toFridayRGBA im
    Right (ImageRGB8 im) -> toFridayRGBA (promoteImage im)
    Right _ -> error "invalid PNG format"
    Left _ -> error "invalid resource"

loadResource :: WorldName -> ResourceType -> Int -> IO RGBA
loadResource worldName r num = do
    -- First, try reading a local file. If that fails (the file isn't found),
    -- read a global file instead.
    let localPath = localResourcePath worldName r num
    let globalPath = globalResourcePath r num
    bs <- B.readFile localPath `catchNotFound` B.readFile globalPath
    return (readPng bs)

loadResourceArray :: WorldName -> ResourceType -> IO (Array Word8 RGBA)
loadResourceArray worldName r = do
    resources <- mapM (loadResource worldName r) [0..255]
    return (A.listArray (0, 255) resources)

spritePath :: (Word8, Word8) -> String
spritePath (b, n) =
    joinPath ["Data", "Objects", "Bank" ++ show b, "Object" ++ show n ++ ".png"]

loadSprite :: (Word8, Word8) -> IO (Maybe ((Word8, Word8), RGBA))
loadSprite (b, n) = do
    io <- tryIOError $ B.readFile (spritePath (b, n))
    case io of
        Right bs -> return $ Just ((b, n), magentaTransparent $ readPng bs)
        Left e -> if isDoesNotExistError e then return Nothing else ioError e

loadSpriteMap :: IO SpriteMap
loadSpriteMap = do
    ss <- mapM loadSprite [(b, n) | b <- [0..18], n <- [1..50]]
    return $ M.fromList [(k, v) | Just (k, v) <- ss]

loadRawLevelData :: WorldName -> IO BL.ByteString
loadRawLevelData worldName = do
    GZip.decompress <$> BL.readFile (mapBinPath worldName)

---------------------------------------------------------------------------
-- Parsing level data
---------------------------------------------------------------------------

pDigit :: Parser Word8
pDigit = P.satisfy (\w -> w >= 48 && w < 58)

pInt :: Parser Int
pInt = read . map (chr . fromEnum) <$> P.many1 pDigit

pTileLayer :: Parser LayerData
pTileLayer = do
    ws <- P.count (mapWidth * mapHeight) P.anyWord8
    return $ TileLayer $ A.listArray ((0, 0), (mapHeight-1, mapWidth-1)) ws

pSpriteLayer :: Parser LayerData
pSpriteLayer = do
    hi <- P.count (mapWidth * mapHeight) P.anyWord8
    lo <- P.count (mapWidth * mapHeight) P.anyWord8
    return $ SpriteLayer $ A.listArray ((0, 0), (mapHeight-1, mapWidth-1)) (zip lo hi)

pMapData :: Parser MapData
pMapData = do
    tileLayers <- P.count 4 pTileLayer
    spriteLayers <- P.count 4 pSpriteLayer
    tB <- P.anyWord8
    tA <- P.anyWord8
    _ <- P.count 3 P.anyWord8
    bg <- P.anyWord8
    return $ MapData (tileLayers ++ spriteLayers) tA tB bg

pMapDataPair :: Parser ((Int, Int), MapData)
pMapDataPair = do
    _ <- P.string "x"; x <- pInt
    _ <- P.string "y"; y <- pInt
    _ <- P.string "\NUL"
    _ <- P.count 4 P.anyWord8
    m <- pMapData
    return ((y, x), m)

pLevelData :: Parser LevelData
pLevelData = M.fromList <$> many pMapDataPair

loadLevelData :: WorldName -> IO LevelData
loadLevelData w = do
    d <- BL.toStrict <$> loadRawLevelData w
    return $ fromRight $ P.parseOnly pLevelData d

---------------------------------------------------------------------------
-- Image helpers
---------------------------------------------------------------------------

-- Wrap an image to a given size.
wrap :: Im.Size -> RGBA -> RGBA
wrap s im =
    let Z :. h :. w = Im.manifestSize im
    in Im.fromFunction s $ \(Z :. y :. x) -> im Im.! Im.ix2 (y `mod` h) (x `mod` w)

-- Map a Word8 in [0, 255] to a Double in [0, 1].
w2f :: Word8 -> Double
w2f = (/255) . fromIntegral . fromEnum

-- Map a Double in [0, 1] to a Word8 in [0, 255].
f2w :: Double -> Word8
f2w = toEnum . round . (*255)

-- Compose two images into one. `bottom` is wrapped to `top`'s size.
compose :: RGBA -> RGBA -> RGBA
compose bottom top =
    let newSize = Im.manifestSize top
        bottom' = wrap newSize bottom
    in Im.fromFunction newSize $ \p ->
        let RGBAPixel rB gB bB aB = bottom' Im.! p
            RGBAPixel rT gT bT aT = top Im.! p
            aB' = w2f aB; aT' = w2f aT
            ovl :: Double -> Double -> Double
            ovl cB cT = (cT * aT' + cB * aB' * (1.0 - aT')) / (aT' + aB' * (1.0 - aT'))
            (~*~) :: Word8 -> Word8 -> Word8
            cB ~*~ cT = f2w $ w2f cB `ovl` w2f cT
            aO = f2w (aT' + aB' * (1.0 - aT'))
        in RGBAPixel (rB ~*~ rT) (gB ~*~ gT) (bB ~*~ bT) aO

-- Make a collage out of an array of images, all of which are assumed to
-- have the same size.
collage :: Array (Int, Int) RGBA -> RGBA
collage arr =
    let ((yMin, xMin), (yMax, xMax)) = A.bounds arr
        Z :. height :. width = Im.manifestSize (arr A.! (yMin, xMin))
        newSize = Im.ix2 ((yMax - yMin + 1) * height) ((xMax - xMin + 1) * width)
    in Im.fromFunction newSize $ \(Z :. y :. x) ->
        let (qY, rY) = divMod y height
            (qX, rX) = divMod x width
        in (arr A.! (qY + yMin, qX + xMin)) Im.! Im.ix2 rY rX

-- A blank image of the given size.
blankImg :: Im.Size -> RGBA
blankImg s@(Z :. h :. w) = Im.Manifest s (V.replicate (h * w) (RGBAPixel 0 0 0 0))

-- Make magenta transparent in the given image.
magentaTransparent :: RGBA -> RGBA
magentaTransparent = Im.map $ \input@(RGBAPixel r g b _) ->
    if (r, g, b) == (0xFF, 0x00, 0xFF)
        then RGBAPixel 0 0 0 0
        else input

---------------------------------------------------------------------------
-- Drawing maps
---------------------------------------------------------------------------

-- A blank tile image.
blankTile :: RGBA
blankTile = blankImg tileSize

-- A blank layer image.
blankLayer :: RGBA
blankLayer = blankImg layerSize

-- Crop a tile (0-127) from a tileset.
getTile :: RGBA -> Word8 -> RGBA
getTile _ 0 = blankTile
getTile tileset i =
    let (y, x) = divMod (fromEnum i) tilesPerTilesetRow
        (w, h) = (tileWidth, tileHeight)
    in crop (Im.Rect (x * w) (y * h) w h) tileset

-- Select a tile (0-255) from one of two tilesets (A and B).
drawTile :: RGBA -> RGBA -> Word8 -> RGBA
drawTile tA tB i =
    let (q, r) = divMod i 128
    in getTile (if q == 1 then tA else tB) r

-- Draw a sprite with the given (bank, number).
drawSprite :: SpriteMap -> (Word8, Word8) -> RGBA
drawSprite _ (0, 0) = blankTile
drawSprite ss x =
    case M.lookup x ss of
        Just im -> wrap tileSize im
        Nothing -> blankTile

-- Draw a layer, given two tilesets and a sprite map.
drawLayer :: RGBA -> RGBA -> SpriteMap -> LayerData -> RGBA
drawLayer _ _ ss (SpriteLayer arr) =
    collage $ fmap (drawSprite ss) arr
drawLayer tA tB _ (TileLayer arr) =
    collage $ fmap (drawTile tA tB) arr

-- Draw a map by composing layers onto a background image.
drawMap :: Resources -> MapData -> RGBA
drawMap (Resources gs ts ss) (MapData layers tA tB bg) =
    let bgImg = wrap (Im.ix2 240 600) (gs A.! bg)
        layerImgs = map (drawLayer (ts A.! tA) (ts A.! tB) ss) layers
    in foldl compose bgImg layerImgs

-- Draw a level by making a collage of all maps in it.
drawLevel :: Resources -> LevelData -> RGBA
drawLevel res lvl =
    let (ys, xs) = unzip (M.keys lvl)
        bounds = ((minimum ys, minimum xs), (maximum ys, maximum xs))
        arr = A.listArray bounds (repeat Nothing) A.// M.assocs (Just <$> lvl)
        drawMap' Nothing = blankLayer
        drawMap' (Just m) = drawMap res m
    in collage $ fmap drawMap' arr

---------------------------------------------------------------------------
-- Main
---------------------------------------------------------------------------

logItems :: LevelData -> (Word8, Word8) -> String -> IO ()
logItems lvl spr name = do
    putStrLn (name ++ " at:")
    flip mapM_ (M.assocs lvl) $ \((y, x), MapData ls _ _ _) -> do
        if spr `elem` concat [A.elems layer | SpriteLayer layer <- ls] && y >= 1015 then
            putStrLn ("x" ++ show x ++ "y" ++ show y)
        else
            return ()
    putStrLn ""

logAllItems :: LevelData -> IO ()
logAllItems level = do
    logItems level (0,  4) "Climb Powerups"
    logItems level (0,  5) "Double Jump Powerups"
    logItems level (0,  6) "High Jump Powerups"
    logItems level (0,  9) "Umbrellas"
    logItems level (0, 10) "Holograms"
    logItems level (0, 21) "Red Keys"
    logItems level (0, 22) "Yellow Keys"
    logItems level (0, 23) "Blue Keys"
    logItems level (0, 24) "Purple Keys"

renderWorld :: WorldName -> IO ()
renderWorld worldName = do
    level <- loadLevelData worldName
    logAllItems level
    
    gs <- loadResourceArray worldName Gradient
    ts <- loadResourceArray worldName Tileset
    sm <- loadSpriteMap
    putStrLn "Drawing map..."
    let img = drawLevel (Resources gs ts sm) level

    writePng (worldName ++ ".png") (toJuicyRGBA img)

main :: IO ()
main = getArgs >>= renderWorld . head