{-# LANGUAGE RecordWildCards, DoAndIfThenElse #-}

module Main (main)
where

import Control.Monad ( when )
import qualified Data.Array.Repa as R
import Data.Maybe ( fromJust )
import GHC.Conc ( getNumProcessors, setNumCapabilities )
import System.IO ( hPutStrLn, stderr )
import System.Exit (
  ExitCode( ExitFailure ),
  exitSuccess,
  exitWith )

import CommandLine (
  Args(Args, depth, height, input, lower_threshold, output,
       scale, slice, upper_threshold, width),
  apply_args )
import ExitCodes ( exit_arg_not_positive, exit_arg_out_of_bounds )
import Grid ( zoom )
import Volumetric (
  bracket_array,
  flip_x,
  flip_y,
  read_word16s,
  round_array,
  swap_bytes,
  write_values_to_bmp,
  write_word16s,
  z_slice )


validate_args :: Args -> IO ()
validate_args Args{..} = do
  when (scale <= 0) $ do
    hPutStrLn stderr "ERROR: scale must be greater than zero."
    exitWith (ExitFailure exit_arg_not_positive)

  when (width <= 0) $ do
    hPutStrLn stderr "ERROR: width must be greater than zero."
    exitWith (ExitFailure exit_arg_not_positive)

  when (height <= 0) $ do
    hPutStrLn stderr "ERROR: height must be greater than zero."
    exitWith (ExitFailure exit_arg_not_positive)

  when (depth <= 0) $ do
    hPutStrLn stderr "ERROR: depth must be greater than zero."
    exitWith (ExitFailure exit_arg_not_positive)

  case slice of
    Just s ->
      when (s < 0 || s > depth) $ do
        hPutStrLn stderr "ERROR: slice must be between zero and depth."
        exitWith (ExitFailure exit_arg_out_of_bounds)
    Nothing -> return ()


main :: IO ()
main = do
  args@Args{..} <- apply_args
  -- validate_args will simply exit if there's a problem.
  validate_args args

  -- The first thing we do is set the number of processors. We get the
  -- number of processors (cores) in the machine with
  -- getNumProcessors, and set it with setNumCapabilities. This is so
  -- we don't have to pass +RTS -Nfoo on the command line every time.
  num_procs <- getNumProcessors
  setNumCapabilities num_procs

  let shape = (R.Z R.:. depth R.:. height R.:. width) :: R.DIM3

  -- Determine whether we're doing 2d or 3d. If we're given a slice,
  -- assume 2d.
  let main_function = case slice of
                        Nothing -> main3d
                        Just _  -> main2d

  main_function args shape
  exitSuccess


main3d :: Args -> R.DIM3 -> IO ()
main3d Args{..} shape = do
  let zoom_factor = (scale, scale, scale)
  arr <- read_word16s input shape
  let arr_swapped = swap_bytes arr
  let arr_shaped  = R.reshape shape arr_swapped
  dbl_data <- R.computeUnboxedP $ R.map fromIntegral arr_shaped
  raw_output <- zoom dbl_data zoom_factor
  let word16_output = round_array raw_output
  -- Switch the bytes order back to what it was. This lets us use the
  -- same program to view the input/output data.
  swapped_output <- R.computeUnboxedP $ swap_bytes word16_output
  write_word16s output swapped_output


main2d :: Args -> R.DIM3 -> IO ()
main2d Args{..} shape = do
  let zoom_factor = (1 :: Int, scale, scale)
  arr <- read_word16s input shape
  arrSlice <- R.computeUnboxedP
               $ z_slice (fromJust slice)
               $ flip_x width
               $ flip_y height
               $ swap_bytes arr
  let arrSlice' = R.reshape slice3d arrSlice

  -- If zoom isn't being inlined we need to extract the slice before hand,
  -- and convert it to the require formed.
  dbl_data   <- R.computeUnboxedP $ R.map fromIntegral arrSlice'
  raw_output <- zoom dbl_data zoom_factor
  arrSlice0  <- R.computeUnboxedP $ z_slice 0 raw_output

  -- Make doubles from the thresholds which are given as Ints.
  let lt = fromIntegral lower_threshold :: Double
  let ut = fromIntegral upper_threshold :: Double

  let arr_bracketed = bracket_array lt ut arrSlice0
  values <- R.computeUnboxedP $ R.map fromIntegral arr_bracketed
  write_values_to_bmp output values

  where
    slice3d :: R.DIM3
    slice3d = (R.Z R.:. 1 R.:. height  R.:. width)