-- Aluno: Salomão Rodrigues Jacinto
-- Matrícula UFSC: 15104097


import Data.Binary
import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString.Lazy as BS
import Data.ByteString.Lazy (ByteString)
import System.Environment
import System.FilePath
import System.IO


-- Header infos, offset is needed to jump to the bgr rows
data BMPHeader = BMPHeader {
    bm :: Word16,
    size :: Word32,
    u1 :: Word16,
    u2 :: Word16,
    offset :: Word32
} deriving (Show)


-- Bitmap info header, width e bpp~are needed to check the padding of the bitmap
data DIBHeader = DIBHeader {
    headerSize :: Word32,
    width :: Word32,
    height :: Word32,
    planes :: Word16,
    bpp :: Word16,
    compression :: Word32,
    dataSize :: Word32,
    ppmh :: Word32,
    ppmv :: Word32,
    palleteSize :: Word32,
    important :: Word32
} deriving (Show)


-- The headers and the data part
data Bitmap = Bitmap {
   bmp :: BMPHeader,
   dib :: DIBHeader,
   bytes :: ByteString
}

-- More infos can be found at [https://en.wikipedia.org/wiki/BMP_file_format]

-- The padding info is necessary because not every bmp image has a multiple of 4
-- row and that would crash the program. Pixel storage of the wiki above explain

-- Another interesting fact, .bmp is little-endian

-- Data type for the rgb values
data BGRPixel = BGRPixel {
   b :: Word8,
   g :: Word8,
   r :: Word8
}


-- Tuple and list of rgb values
type BGR = ([Word8], [Word8], [Word8])
type BGRList = [[Word8]]


-- The next Get's are used with runGet to get from a bytestring the values of the
-- headers

getBMPHeader :: Get BMPHeader
getBMPHeader = do
    bm <- getWord16le
    size <- getWord32le
    u1 <- getWord16le
    u2 <- getWord16le
    offset <- getWord32le
    return $ BMPHeader bm size u1 u2 offset


getDIBHeader :: Get DIBHeader
getDIBHeader = do
    headerSize <- getWord32le
    width <- getWord32le
    height <- getWord32le
    planes <- getWord16le
    bpp <- getWord16le
    compression <- getWord32le
    dataSize <- getWord32le
    ppmh <- getWord32le
    ppmv <- getWord32le
    palleteSize <- getWord32le
    important <- getWord32le
    return (DIBHeader headerSize width height planes bpp compression dataSize
            ppmh ppmv palleteSize important)


getHeaders :: Get (BMPHeader, DIBHeader)
getHeaders = do
    b <- getBMPHeader
    d <- getDIBHeader
    return (b, d)


-- The bytestring of the .bmp file creates the Bitmap, separeting the headers and
-- the data part using the offset
getBitmap :: ByteString -> Bitmap
getBitmap bytes = let
    (b, d) = runGet getHeaders bytes
    o = fromIntegral $ offset b
    in Bitmap b d (BS.drop o bytes)


-- Get the width of the rows of the bitmap
bmpWidth :: Bitmap -> Int
bmpWidth bmp = fromIntegral $ width (dib bmp)


bytesPerPixel :: Bitmap -> Int
bytesPerPixel bmp = bitsPerPixel `div` 8
    where bitsPerPixel = fromIntegral $ bpp (dib bmp)


-- The rows are multiple of four, if not a padding is added, this checks how long
-- is this padding to read the file the right way
padding :: Int -> Int -> Int
padding width bpp  = (4 - (width*bpp `mod` 4)) `mod` 4


-- Separate the rows of the data part, creating a list of rows and eliminating the
-- padding if there is any
getRows :: Int -> Int -> [Word8] -> [[Word8]]
getRows _ _ [] = []
getRows width bpp bytes = row:(getRows width bpp remainder)
    where p = padding width bpp
          size = width * bpp
          (padded, remainder) = splitAt (size+p) bytes
          row = take size padded


-- Transform the rows into type BGRPixel, splits at every 3 words8
rowToPixels :: [Word8] -> [BGRPixel]
rowToPixels [] = []
rowToPixels row = let
    ([b, g, r], remainder) = splitAt 3 row
    in (BGRPixel b g r):(rowToPixels remainder)


-- Makes a composition to go from the bitmap file to a list of BGRPixels
getPixels :: Bitmap -> [BGRPixel]
getPixels bmp = let
    w = bmpWidth bmp
    bpp = bytesPerPixel bmp
    rows = getRows w bpp . BS.unpack $ bytes bmp
    in concat $ map rowToPixels rows


-- Splits the list of BRGPixels into BGR 3-tuple, this tuple has 3 lists, one for
-- eache color, and the value are stored there
splitBGR :: [BGRPixel] -> BGR
splitBGR [] = ([], [], [])
splitBGR (BGRPixel b g r:ps) = (b:bs, g:gs, r:rs)
    where (bs, gs, rs) = splitBGR ps


-- Converts a 3-tuple into a list to use zip to write the files
toList :: (a, a, a) -> [a]
toList (b, g, r) = [b, g, r]


-- Here is used all the above functions to go from a bitmap file to a 3-tuple with
-- the list of values from each color
bmpToBGR :: Bitmap -> BGR
bmpToBGR = splitBGR . getPixels


-- Converts the bytes to float values
toFloat :: Word8 -> Float
toFloat color = fromIntegral color / 255


-- Close all the files
closeAll :: [Handle] -> IO [()]
closeAll handlers = sequence $ map hClose handlers


-- Open all the files
openAll :: [String] -> IO [Handle]
openAll filenames = sequence [openFile filename WriteMode |
                              filename <- filenames]


-- Output the content values into the filenames using specified IO function, in our
-- case it is used zip function to put each pixel bgr value into the right file
outputAll :: (Handle -> a -> IO ()) -> [Handle] -> [a] -> IO [()]
outputAll outputFunction filenames contents = do
    sequence [outputFunction filename contents |
              (filename, contents) <- zip filenames contents]


-- The suffixes that are going to name the files
suffixes :: [String]
suffixes = ["_blue.out", "_green.out", "_red.out"]


-- Concat the bitmap filename with the suffixes
concatSuffixes :: [String] -> String -> [String]
concatSuffixes suffixes prefix = map (prefix++) suffixes


-- Create the names of the output files
outnames :: String -> [FilePath]
outnames = concatSuffixes suffixes


-- Get the file name and the bitmap file and outputs to the files
splitFileText :: String -> Bitmap -> IO ()
splitFileText prefix bitmap = do
    outputs <- openAll $ outnames prefix

    outputAll hPrint outputs $ [map toFloat color |
                                color <- toList $ bmpToBGR bitmap]

    closeAll outputs

    return ()


-- Use above functions to split the bitmap rgb file into three text files of 
-- floating points
main = do
    args <- getArgs
    case args of
        [filename] -> do
            input <- BS.readFile filename

            let prefix = takeBaseName filename
            let bitmap = getBitmap input

            splitFileText prefix bitmap
        _ -> putStrLn "Usage: ./newrgb FILE.bmp"


-- This does not work if the .bmp file has a alpha channel. Just splits into rgb
-- 24bit rgb files

-- This was made by Tarcisio Eduardo Moreira Crocomo, with the help of Salomão
-- Rodrigues Jacinto and Evandro Chagas Ribeiro Rosa based on the work from
-- Gustavo Zambonim
