{-# LANGUAGE RecordWildCards, DoAndIfThenElse #-}

module Main
where

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

import CommandLine (Args(..), apply_args)
import ExitCodes
import Grid (zoom)
import MRI (
  flip_x,
  flip_y,
  read_word16s,
  round_array,
  swap_bytes,
  write_values_slice_to_bitmap,
  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

  -- Determine whether we're doing 2d or 3d. If we're given a slice,
  -- assume 2d.
  let mri_shape = (R.Z R.:. depth R.:. height R.:. width) :: R.DIM3

  if (isJust slice) then
    main2d args mri_shape
  else
    main3d args mri_shape

  exitSuccess

  where



main3d :: Args -> R.DIM3 -> IO ()
main3d Args{..} mri_shape = do
  let zoom_factor = (scale, scale, scale)
  arr <- read_word16s input mri_shape
  let arr'          = swap_bytes arr
  let arrMRI        = R.reshape mri_shape arr'
  dbl_data <- R.computeUnboxedP $ R.map fromIntegral arrMRI
  raw_output <- zoom dbl_data zoom_factor
  word16_output <- R.computeUnboxedP $ round_array raw_output
  write_word16s output word16_output


main2d :: Args -> R.DIM3 -> IO ()
main2d Args{..} mri_shape = do
  let zoom_factor = (1, scale, scale)
  arr <- read_word16s input mri_shape
  arrSlice <- R.computeUnboxedP
               $ z_slice (fromJust slice)
               $ flip_x width
               $ flip_y height
               $ swap_bytes arr
  let arrSlice' = R.reshape mri_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

  write_values_slice_to_bitmap arrSlice0 output
  where
    mri_slice3d :: R.DIM3
    mri_slice3d = (R.Z R.:. 1 R.:. height  R.:. width)