-- Copyright 2009, Cetin Sert

-- | ExPloRe: Experimental Plot Reconstructor

module Main where

import System.Environment
import System.Directory
import System.IO.Unsafe
import System.Exit
import Data.Word
import Foreign.Ptr
import Foreign.Storable
import Data.Array.Storable
import Control.Monad
import Codec.Image.PNG

import Help

-- | line names in legends
type Name = String

-- | file paths
type ImagePath = String

-- | legends
type Legend = [(Name, RGBA)]

-- | internal representation of an unnormalized explore invocation
data ExploreInvocation = Explore ImagePath Legend Matching ScanArea StepSize Translation
                       deriving (Show, Read, Eq)

-- | 32-bit color representation
data RGBA = RGBA !Word8 !Word8 !Word8 !Word8 -- ^ RED BLUE GREEN ALPHA
          deriving (Show, Read, Eq)

-- storable instance for retrieveal 
instance Storable RGBA where
  sizeOf _    = sizeOf (0 :: Word8) * 4
  alignment _ = 1
  peek color  = do
    let byte :: Ptr Word8 = castPtr color
    [r,g,b,a] <- mapM (byte @!) [0..3]
    return $ RGBA r g b a

-- | colour matching behaviour
data Matching = S                  -- ^ strict
              | TA Int             -- ^ shared threshold for all color components
              | TC Int Int Int Int -- ^ individual thresholds for color components
              deriving (Show, Read, Eq)

-- | scan area
--   w = width of the PNG image
--   h = height of the PNG image
data ScanArea = F                  -- ^ scan all the image area (FULL)    | 0 0 w h
              | O Int Int          -- ^ offset from LEFT TOP              | L T w h
              | M Int Int Int Int  -- ^ offset from LEFT TOP RIGHT BOTTOM | L T (w-R) (h-B)
              | B Int Int Int Int  -- ^ box         LEFT TOP RIGHT BOTTOM | L T R B
              deriving (Show, Read, Eq)

-- | vertical offset reference line
data VerticalOffsetReferenceLine = Top    -- ^ top line of plot box
                                 | Bottom -- ^ bottom line of plot box
                                 deriving (Show, Read, Eq)
-- | translation
data Translation = Identity                        -- ^ identity translation
                 | Box VerticalOffsetReferenceLine -- ^ box translation
                 | Default                         -- ^ gets normalized to the default translation
                 deriving (Show, Read, Eq)

-- | step size
data StepSize = ConstantJump  Double -- ^ only step size supported upto version 0.0.6.2
              | WidthDivision Double -- ^
              deriving (Show, Read, Eq)    

-- | normalizes scan area representations
normalizeScanArea :: Int -> Int -> ScanArea -> ScanArea
normalizeScanArea w h F           = B 0 0 w h
normalizeScanArea w h (O l t)     = B l t w h
normalizeScanArea w h (M l t r b) = B l t (w-r) (h-b)
normalizeScanArea _ _ box         = box

-- | horizontal translation function
type HTrans = Int -> Int

-- | vertical translation function
type VTrans = Int -> Int

-- | creates translation functions
makeTranslationFunctions :: Translation -> ScanArea -> (HTrans, VTrans)
makeTranslationFunctions Identity     _           = (id,      id)
makeTranslationFunctions (Box Top   ) (B l t r b) = ((− l),(− t))
makeTranslationFunctions (Box Bottom) (B l t r b) = ((− l),(b −))

(−) = (-)

-- | fuzzy color matching
colorEq :: Matching -> RGBA -> RGBA -> Bool
colorEq S            = compEq 0 0 0 0
colorEq (TA v)       = compEq v v v v
colorEq (TC r g b a) = compEq r g b a

-- | fuzzy color matching
compEq :: Int -> Int -> Int -> Int -> RGBA -> RGBA -> Bool
compEq tr tg tb ta (RGBA a b c d) (RGBA x y z w) = let fe = fuzzyEq in
  fe tr a x && fe tg b y && fe tb c z && fe ta d w

-- | simple fuzzy equality
fuzzyEq :: forall b a. (Num b, Ord b, Integral a) => b -> a -> a -> Bool
fuzzyEq 0 = (==)
fuzzyEq t = \x y -> t > (abs $ fromIntegral x - fromIntegral y)

-- | gets the middle element of a list, used as a quick median function
mid :: forall a. [a] -> a
mid xs = xs !! (fromIntegral . floor . (/ 2) . fromIntegral . length) xs

(@!) :: Storable a => Ptr a -> Int -> IO a
(@!) = peekElemOff

make2DIndexer :: Storable a => Int -> Ptr a -> Int -> Int -> a
make2DIndexer w = \p j i -> unsafePerformIO $ p @! (i + j * w)

-- | ExPloRe requires PNGs with transparency channel and no interlacing 
checkAlpha :: PNGImage -> IO ()
checkAlpha (hasAlphaChannel -> True) = return ()
checkAlpha _ = putStrLn "no alpha channel going on!!!" >> exitWith (ExitFailure 1)

-- | check presence of arguments, show usage otherwise
checkArgs :: [String] -> IO ()
checkArgs [] = putStrLn help >> exitSuccess
checkArgs _  = return ()

hrule :: Char -> String
hrule = replicate 80 


-- | main
main :: IO ()
main = do
  putStrLn $ "ExPloRe 0.0.7.1 : Experimental Plot Reconstructor"
  putStrLn $ hrule '='
  args <- getArgs
  checkArgs args
  putStrLn $ "invocation file"
  putStrLn $ hrule '-'
  putStrLn $ tail (filter (/= '"') $ concatMap (('\n':) . show) args) ++ "\n"
  putStrLn $ hrule '-'
  explore <- makeInvocation args
  run explore

-- | supports invocation files and provides backwards-compatibility with
--   0.0.6.2 style command line arguments
makeInvocation :: [String] -> IO ExploreInvocation

makeInvocation [invocationPath] = do
  makeInvocation =<< (liftM lines . readFile $ invocationPath)

makeInvocation args = do
  putStrLn $ "args  : " ++ show args
  let (imgPath:legend_:matching_:area_:step_:oargs) = args
  matching    <- return . read $ matching_
  area        <- return . read $ area_
  lfexists    <- doesFileExist legend_
  legend      <- case lfexists of
                   True -> liftM read . readFile $ legend_
                   _    ->     return . read     $ legend_
  stepSize    <- return $ case step_ of
                            ('/':parts) -> WidthDivision $ read parts
                            size        -> ConstantJump  $ read size
  translation <- return $ case null oargs of
                            True -> Box Bottom
                            _    -> read . head $ oargs  
  return $ Explore imgPath legend matching area stepSize translation


-- | new run function
run :: ExploreInvocation -> IO ()
run invocation@(Explore imgPath legend matching area stepSize translation) = do
  -- initialize image
  Right img <- loadPNGFile imgPath
  let bitmap  = imageData  img
  let (wu,hu) = dimensions img
  let (w,h)   = (fromIntegral wu, fromIntegral hu)
  args <- getArgs
  putStrLn $ "image : " ++ imgPath
  putStrLn $ "width : " ++ show w
  putStrLn $ "height: " ++ show h
  checkAlpha img

  -- initialize scan
  let box@(B l t r b) = normalizeScanArea w h area :: ScanArea
  let start = fromIntegral l :: Double
  let end   = fromIntegral r :: Double
  let step  = case stepSize of
                ConstantJump  size  -> size
                WidthDivision parts -> (end - start) / parts
  let trans = case translation of
                Default -> Box Bottom
                any     -> any
  let (@#)  = make2DIndexer w
  let rows  = [t..b]
  let cols  = takeWhile (< r) $ map (floor . (start +) . (step *)) [0..]
  let icols = zip [1..] cols
  let (~=)  = colorEq matching
  putStrLn $ "box   : " ++ show box
  putStrLn $ "width : " ++ show (r - l)
  putStrLn $ "height: " ++ show (b - t)
  putStrLn $ "trans.: " ++ show trans
  putStrLn $ "match : " ++ show matching
  putStrLn $ "step  : " ++ show stepSize
  putStrLn $ "step  : " ++ (show $ convertStepSize (end - start) stepSize)

  -- display legend
  putStrLn $ "#lines: " ++ (show . length $ legend)
  putStrLn $ ""
  mapM_ (putStrLn . show) legend

  -- scan bitmap
  mapM_ (scan bitmap trans box icols rows (@#) (~=)) legend

-- | converts step-size representations
convertStepSize :: Double -> StepSize -> StepSize
convertStepSize width (ConstantJump  size ) = WidthDivision $ width / size
convertStepSize width (WidthDivision parts) = ConstantJump  $ width / parts

-- | a very simple scanning algorithm that works only if no lines
--   occlude each other.
scan bitmap translation box@(B l t r b) icols rows (@#) (~=) (name,color) = do
  putStrLn $ ""
  putStrLn $ hrule '-'
  putStrLn $ name
  putStrLn $ ""
  putStrLn $ show color
  putStrLn $ ""
  putStrLn $ "box        : " ++ show box 
  putStrLn $ "width      : " ++ show (r - l)
  putStrLn $ "height     : " ++ show (b - t)
  putStrLn $ "translation: " ++ show translation
  putStrLn $ ""
  putStrLn $ "x\tpoint\tmid y\tmatching ys"
  putStrLn $ ""
  withStorableArray bitmap $ \byte -> do
    let pixel :: Ptr RGBA = castPtr byte
    let (hof,vof) = makeTranslationFunctions translation box
    forM_ icols $ \(n,j) -> do
      let ys = map vof $ flip filter rows $ \i -> (pixel @# i) j ~= color
      let my = mid ys
      let x  = hof j
      putStrLn $ case not . null $ ys of
        True  -> show x ++ "\t" ++ show n ++ "\t" ++ show my ++ "\t" ++ show ys
        False -> show x ++ "\t" ++ show n ++ "\t   \t[]"

-- Copyright 2010, Cetin Sert