pts to png

Main.hs

{-# LANGUAGE ScopedTypeVariables #-}

module Main where

import System.IO
import Control.Monad
import Codec.Picture
import Codec.Picture.Png
import Foreign.ForeignPtr
import Foreign.Marshal.Array
import Foreign.Ptr
import Foreign.Storable
import Data.Vector.Storable hiding ((++), map, zipWith3)
import qualified Data.Vector as V
import Data.List (zipWith3)
import GHC.Float

times :: Monad m => Int -> (Int -> m ()) -> m ()
times d a = go 0 where
  go x
    | x == d = return ()
    | otherwise = a x >> go (x+1)

alloc :: Storable a => Int -> IO (Ptr a)
alloc = callocArray

give :: Storable a => Ptr a -> Int -> a -> IO ()
give = pokeElemOff

packV :: Storable a => Ptr a -> Int -> Int -> Int -> IO (Int, Int, Vector a)
packV p offset w h = do
  fp <- newForeignPtr_ p
  return (w,h,unsafeFromForeignPtr fp offset (3*w*h))

savePNG :: FilePath -> (Int, Int, Vector Pixel8) -> IO ()
savePNG f (w,h,v) = writePng f (Image w h v :: Image PixelRGB8)

main :: IO ()
main = do
  hSetBuffering stdout NoBuffering
  putStrLn "Launched."
  count <- readLn :: IO Int
  let
    width = 1024
    height = 1024
    size = width * height
    upCount = if count `mod` size == 0
      then count
      else (count `div` size + 1) * size
  posV <- alloc (upCount*3*4) :: IO (Ptr Pixel8)
  colV <- alloc (upCount*3) :: IO (Ptr Pixel8)
  putStrLn $ "UpCount = " ++ show upCount
  times count $ \i -> do
    when (i`mod`100000==0) $ putStrLn $ "Processing " ++ show i ++ "th point..."
    [x,y,z,_,r,g,b] <- words <$> getLine
    let
      cs = [-11943.1, -37795.7, 8.19995]
      ss = [409.199, 391.48, 38.9134]
    [x,y,z] <- return $ zipWith3 (\p c s -> double2Float $ p-c) (map read [x,y,z]) cs ss
    [r,g,b] <- return $ map read [r,g,b]
    let ps = unsafeCast $ fromList [x,y,z]
    let base = (i`div`width)*2*2*width + (i`mod`width)*2
    times 2 $ \j0 -> do
      times 2 $ \j1 -> do
        times 3 $ \k -> do
          give posV (base*3+(j0+j1*2*width)*3+k) $ ps!(j0+j1*2+k*4)
    give colV (i*3+0) r
    give colV (i*3+1) g
    give colV (i*3+2) b
  times (upCount `div` size) $ \i -> do
    putStrLn $ "Generating " ++ show i ++ "th image..."
    savePNG ("pos" ++ show i ++ ".png") =<< packV posV (i*size*3*4) (width*2) (height*2)
    savePNG ("col" ++ show i ++ ".png") =<< packV colV (i*size*3) width height
  putStrLn "Done!"