-- YUVDiff

import Char
import qualified Data.ByteString as B
import Data.Word
import Foreign ( unsafePerformIO )
import Text.Regex
import List
import Maybe
import System


showCmdUsage = unlines [ "To locate / show difference pixel between YUV seq files.",
                         "",
                         "Usage:\tYUVDiff.exe -opt rec-yuv dec-yuv [F-MY-MX]",
                         "",
                         "opt:\tspWxH\t\tW*H pixels.",
                         "\tsmWxH\t\tW*H macroblocks.",
                         "\tmM\t\tof M mode, where M may be:",
                         "\t\t\t  QCIF, CIF, D1 (case-insensitive)",
                         "",
                         "The last optional argument is to specify which MB to show.",
                         "If this is not specified, I will show you the first diff-MB.",
                         "Where F, MX and MY are all numbers.",
                         "\tF\t\tframe",
                         "\tMX\t\tX coordinate of MB",
                         "\tMY\t\tY coordinate of MB",
                         "",
                         "Made by Levi.G. 2007-06. All rights reserved.",
                         "Email: leaveye <dot> guo <at> gmail <dot> com" ]

main = do
    args <- getArgs
    if length args < 3
      then giveUsage
      else case readOption (args!!0) of
             Just sz -> if length args > 3
                          then findDiffYUV (YUV420,sz) (readPossition (args!!3)) (args!!1) (args!!2)
                          else findDiffYUV (YUV420,sz) Nothing (args!!1) (args!!2)
             Nothing -> giveUsage
  where
    giveUsage = putStr showCmdUsage
    match = matchRegex . mkRegex
    readOption ('-':'s':'p':sz) = case match "([0-9]+)x([0-9]+)" sz of
            Just [ws,hs] -> let w = read ws
                                h = read hs
                            in  if all ((==0).(`rem` 16)) [w,h]
                                  then Just (read ws `div` 16, read hs `div` 16)
                                  else Nothing
            _            -> Nothing
    readOption ('-':'s':'m':sz) = case match "([0-9]+)x([0-9]+)" sz of
            Just [ws,hs] -> Just (read ws, read hs)
            _            -> Nothing
    readOption ('-':'m':m:_) = case toUpper m of
            'Q' ->  readOption "-sp176x144"
            'C' ->  readOption "-sp352x288"
            'D' ->  readOption "-sp720x576"
            _   ->  Nothing
    readOption _ = Nothing
    readPossition opt = case match "([0-9]+)-([0-9]+)-([0-9]+)" opt of
            Just [f,my,mx] -> Just [read f, read my, read mx]
            _              -> Nothing


type Pixel = (Word8,Word8,Word8)

type BlockData = [Pixel]

type Image = B.ByteString

type Size = (Int,Int)

type Matrix = [[Pixel]]

type Picture = Matrix

type Config = (ImageMode, Size)

data ImageMode = YUV420 | YUV422 | YUV444
  deriving (Enum, Eq, Show, Read)


dup      :: Int -> [a] -> [a]
dup _ [] = []
dup n ls = let (t,xs) = splitAt n ls
           in  t ++ t ++ (dup n xs)
chop      :: Int -> [a] -> [[a]]
chop _ [] = []
chop n ls = take n ls : chop n (drop n ls)

mkMBs :: Picture -> [BlockData]
mkMBs = unpack . transpose . pack
  where pack = chop 16 . map (chop 16)
        unpack = map concat . concat

mkBlocks :: [BlockData] -> [BlockData]
mkBlocks = unpack . transpose . pack
  where pack = chop 4 . map (chop 4)
        unpack = map concat . concat

mkPicture :: Config -> [Word8] -> Picture
mkPicture c@(m,(w,h)) img = pixArr
  where pixArr = chop (w*16) pixels
        pixels = pad $ zip3 ys (uvFix us) (uvFix vs)
        pad = take lumaSz . (++ repeat (0,128,128))
        ys = take lumaSz img
        us = take chromaSz . drop lumaSz $ img
        vs = take chromaSz . drop (lumaSz+chromaSz) $ img
        uvFix = case m of
            YUV420 -> dup (w*16) . dup 1
            YUV422 -> dup 1
            YUV444 -> id
        lumaSz = mb88Sz * 4
        chromaSz = case m of
            YUV420 -> mb88Sz
            YUV422 -> mb88Sz * 2
            YUV444 -> mb88Sz * 4
        mb88Sz = (w*h)*(8*8)

showPix :: Pixel -> String
showPix (y,u,v) = showHex y ++ showHex u ++ showHex v
  where showHex x = [toHex (x `div` 16), toHex (x `rem` 16)]
        toHex = fromJust . flip lookup (zip [0..15] "0123456789ABCDEF")

showBlockPair :: [(Pixel,Pixel)] -> String
showBlockPair = unlines . map unwords . concatMap transpose . chop 4 . map diffPairShow
  where diffPairShow (a,b) = if a /= b
                               then [showPix a, showPix b]
                               else [showPix a, "      "]

readImageFile :: Config -> FilePath -> IO [Picture]
readImageFile c fn = do
    img <- B.readFile fn
    return . map (mkPicture c) . chop oneSize . B.unpack $ img
  where oneSize = (*) mbSize . uncurry (*) . snd $ c
        mbSize = case fst c of
                   YUV420 -> 16*16+8*8*2
                   YUV422 -> 16*16+8*16*2
                   YUV444 -> 16*16*3

readOnePicture :: Config -> FilePath -> IO (Maybe Picture)
readOnePicture c fn = do
    hdl <- openBinaryFile fn ReadMode
    hSetBuffering hdl . BlockBuffering . Just $ oneSize
    lenr <- hGetBuf hdl 
  where oneSize = (*) mbSize . uncurry (*) . snd $ c
        mbSize = case fst c of
                   YUV420 -> 16*16+8*8*2
                   YUV422 -> 16*16+8*16*2
                   YUV444 -> 16*16*3

-- Position : (frame, ((mbY, mbX), (offY, offX), (pixY, pixX)))
type Position = (Int,(Int,Int),(Int,Int),(Int,Int))

mkPos :: Config -> Int -> (Position,Int)
mkPos (_,(w,h)) i = ((frame, (mbY, mbX), (offY, offX), (pixY, pixX)),i)
  where frame   = i `div` (w*h*256)
        offset  = i `rem` (w*h*256)
        pixY    = offset `div` (w*16)
        pixX    = offset `rem` (w*16)
        mbY     = pixY `div` 16
        offY    = pixY `rem` 16
        mbX     = pixX `div` 16
        offX    = pixX `rem` 16

backGround = [ "   C D E F|0_1_2_3_4_5_6_7_8_9_A_B_C_D_E_F|0 1 2 3  ",
               "C         |       |       |       |       |         ",
               "D         |       |       |       |       |         ",
               "E         |       |       |       |       |         ",
               "F_________|_______|_______|_______|_______|_________",
               "0 |       |       |       |       |       |       | ",
               "1 |       |       |       |       |       |       | ",
               "2 |       |       |       |       |       |       | ",
               "3 |_______|_______|_______|_______|_______|_______| ",
               "4 |       |       |       |       |       |       | ",
               "5 |       |       |       |       |       |       | ",
               "6 |       |       |       |       |       |       | ",
               "7 |_______|_______|_______|_______|_______|_______| ",
               "8 |       |       |       |       |       |       | ",
               "9 |       |       |       |       |       |       | ",
               "A |       |       |       |       |       |       | ",
               "B |_______|_______|_______|_______|_______|_______| ",
               "C |       |       |       |       |       |       | ",
               "D |       |       |       |       |       |       | ",
               "E |       |       |       |       |       |       | ",
               "F_|_______|_______|_______|_______|_______|_______|_",
               "0         |       |       |       |       |         ",
               "1         |       |       |       |       |         ",
               "2         |       |       |       |       |         ",
               "3         |_______|_______|_______|_______|         ",
               "          |                               |         " ]

placeMark :: [String] -> (Int,Int) -> ((String,String),Position) -> [String]
placeMark g (by,bx) (p,(_,(my,mx),(oy,ox),_)) = replace g (y+1)
                                              . replace (g!!(y+1)) (x+x+3)
                                              . uncurry getMark $ p
  where getMark p1 p2 = last . show
                      . sum . map (product . flip replicate 2 . (2-))
                      . findIndices (uncurry (/=))
                      $ zip (chop 2 p1) (chop 2 p2)
        x | bx == mx      = ox + 4
          | bx == mx - 1  = ox + 4 + 16
          | bx == mx + 1  = ox - 12
        y | by == my      = oy + 4
          | by == my - 1  = oy + 4 + 16
          | by == my + 1  = oy - 12
        replace [] _ _ = []
        replace (x:xs) 0 y = y : xs
        replace (x:xs) (n+1) y = x : replace xs n y

findDiffYUV :: Config -> Maybe [Int] -> FilePath -> FilePath -> IO ()
findDiffYUV c@(m,(w,h)) focusP fn1 fn2 = do
  let img1 = unsafePerformIO $ readImageFile c fn1
  let img2 = unsafePerformIO $ readImageFile c fn2
  let pixelPairs = zip ((concat . concat) img1) ((concat . concat) img2)
  let diffIdcs = findIndices (uncurry (/=)) pixelPairs
  if isNothing focusP && null diffIdcs
    then putStrLn . concat $ ["No difference between '", fn1, "' and '", fn2, "'."]
    else do
      let sameMB (f1,m1,_,_) (f2,m2,_,_) = f1 == f2 && m1 == m2
      let diffs = sort . map (mkPos c) $ diffIdcs
      let ((focusFrame,focusMB,_,_),_) = if focusP == Nothing
                          then head diffs
                          else let Just [f,my,mx] = focusP
                               in  ((f, (my, mx), (0,0), (0,0)), 0)
      let close (mby,mbx) ((my,mx),(oy,ox))
            | mby == my && mbx == mx                  = True
            | mby == my && mbx == mx - 1 && ox < 4    = True
            | mby == my && mbx == mx + 1 && ox >= 12  = True
            | mbx == mx && mby == my - 1 && oy < 4    = True
            | mbx == mx && mby == my + 1 && oy >= 12  = True
            | otherwise                               = False
      let near (frm,mb) (f,m,o,_) = f == frm && close mb (m,o)
      let (diffPoss, idcs) = unzip . filter (near (focusFrame,focusMB) . fst) $ diffs
      let diffPels = map ((\(a,b) -> (showPix a, showPix b)) . (!!) pixelPairs) idcs
      putStr . unlines $ [ "diffMarkNum = (yuv)b",
                           "center (Frame,(MB_y,MB_x)) = " ++ show ( focusFrame, focusMB ) ]
      putStr . unlines . foldl (flip placeMark focusMB) backGround $ zip diffPels diffPoss
      putStrLn "((Pixel1,Pixel2), (frame, (mbY, mbX), (offY, offX), (pixY, pixX)))"
      mapM_ print $ zip diffPels diffPoss