{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE OverloadedStrings     #-}

{-
    Copyright  : Copyright (C) 2014-2021 Synchrotron SOLEIL
                                         L'Orme des Merisiers Saint-Aubin
                                         BP 48 91192 GIF-sur-YVETTE CEDEX
    License    : GPL3+

    Maintainer : Picca Frédéric-Emmanuel <picca@synchrotron-soleil.fr>
    Stability  : Experimental
    Portability: GHC only (not tested)
-}

module Hkl.Binoculars.Pipes
  ( Chunk(..)
  , DataFrameSpace(..)
  , LenP(..)
  , mkJobs
  , mkInputHkl
  , mkInputQxQyQz
  , processHkl
  , processQxQyQz
  ) where

import           Bindings.HDF5.Core                (Location)
import           Control.Concurrent.Async          (mapConcurrently)
import           Control.Exception                 (throwIO)
import           Control.Monad                     (forM_, forever)
import           Control.Monad.Catch               (MonadThrow, tryJust)
import           Control.Monad.IO.Class            (MonadIO (liftIO))
import           Control.Monad.Logger              (MonadLogger)
import           Control.Monad.Trans.Cont          (cont, runCont)
import           Data.Array.Repa                   (Shape, size)
import           Data.Array.Repa.Index             (DIM1, DIM2)
import           Data.IORef                        (IORef, readIORef)
import           Data.Maybe                        (fromMaybe)
import           Data.Vector.Storable              (fromList)
import           Data.Word                         (Word16)
import           Foreign.ForeignPtr                (ForeignPtr)
import           GHC.Base                          (returnIO)
import           GHC.Conc                          (getNumCapabilities)
import           GHC.Float                         (float2Double)
import           Numeric.Units.Dimensional.NonSI   (angstrom)
import           Numeric.Units.Dimensional.Prelude (Quantity, Unit, degree,
                                                    (*~))
import           Pipes                             (Consumer, Pipe, Proxy,
                                                    await, each, runEffect,
                                                    yield, (>->))
import           Pipes.Prelude                     (tee, toListM)
import           Pipes.Safe                        (MonadSafe, SafeT,
                                                    SomeException, bracket,
                                                    catchP, displayException,
                                                    runSafeP, runSafeT)
import           System.ProgressBar                (Progress (..), ProgressBar,
                                                    Style (..), defStyle,
                                                    elapsedTime, incProgress,
                                                    newProgressBar,
                                                    renderDuration,
                                                    updateProgress)

import           Prelude                           hiding (filter)

import           Hkl.Binoculars.Common
import           Hkl.Binoculars.Config
import           Hkl.Binoculars.Projections
import           Hkl.C.Binoculars
import           Hkl.C.Geometry
import           Hkl.Detector
import           Hkl.H5                            hiding (File)
import           Hkl.Pipes
import           Hkl.Types


-- LenP

class LenP a where
  lenP :: a -> Pipe FilePath Int (SafeT IO) ()


-- Jobs

mkJobs :: LenP a => InputFn -> a -> IO ([[Chunk Int FilePath]], ProgressBar ())
mkJobs fn h5d = do
  let fns = concatMap (replicate 1) (toList fn)
  ns <- runSafeT $ toListM $ each fns >-> lenP h5d
  c' <- getNumCapabilities
  let ntot = sum ns
      c = if c' >= 2 then c' - 1 else c'
  pb <- newProgressBar defStyle{ stylePostfix=elapsedTime renderDuration }
                10 (Progress 0 ntot ())
  return $ (mkJobs' (quot ntot c) fns ns, pb)


-- Project

withSpace :: (MonadSafe m, Shape sh)
            => Detector a DIM2 -> Int -> (Space sh -> m r) -> m r
withSpace d n = bracket (liftIO $ newSpace d n) pure

project :: (MonadSafe m, Shape sh)
          => Detector a DIM2
        -> Int
        -> (Space sh -> b -> IO (DataFrameSpace sh))
        -> Pipe b (DataFrameSpace sh) m ()
project d n f = withSpace d n $ \s -> forever $ do
               df <- await
               yield =<< liftIO (f s df)

skipMalformed :: MonadSafe m =>
                Proxy a' a b' b m r
              -> Proxy a' a b' b m r
skipMalformed p = loop
  where
    loop = catchP p $ \e -> do
        liftIO $ Prelude.print $ displayException (e :: SomeException)
        loop

-- QxQyQz

class LenP a => FramesQxQyQzP a where
  framesQxQyQzP :: a -> Detector b DIM2 -> Pipe (Chunk Int FilePath) DataFrameQxQyQz (SafeT IO) ()

mkJobsQxQyQz :: LenP a => InputQxQyQz a -> IO ([[Chunk Int FilePath]], ProgressBar ())
mkJobsQxQyQz (InputQxQyQz _ fn h5d _ _ _ _ _ _) = mkJobs fn h5d

mkInputQxQyQz :: (MonadIO m, MonadLogger m, MonadThrow m, FramesQxQyQzP a)
              => BinocularsConfig
              -> (BinocularsConfig -> m a)
              -> m (InputQxQyQz a)
mkInputQxQyQz c f = do
  fs <- files c
  let d = fromMaybe defaultDetector (_binocularsInputDetector c)
  mask' <- getMask c d
  res <- getResolution c 3
  h5dpath' <- f c
  pure $ InputQxQyQz { detector = d
                     , filename = InputList fs
                     , h5dpath = h5dpath'
                     , output = case _binocularsInputInputRange c of
                                  Just r  -> destination' r (_binocularsDispatcherDestination c)
                                  Nothing -> destination' (ConfigRange []) (_binocularsDispatcherDestination c)
                     , resolutions = res
                     , centralPixel = _binocularsInputCentralpixel c
                     , sdd' = _binocularsInputSdd c
                     , detrot' = fromMaybe (0 *~ degree) ( _binocularsInputDetrot c)
                     , mask = mask'
                     }

processQxQyQz :: FramesQxQyQzP a => InputQxQyQz a -> IO ()
processQxQyQz input@(InputQxQyQz det _ h5d o res cen d r mask') = do
  pixels <- getPixelsCoordinates det cen d r
  (jobs, pb) <- mkJobsQxQyQz input
  r' <- mapConcurrently (\job -> withCubeAccumulator $ \c ->
                           runSafeT $ runEffect $
                           each job
                           >-> framesQxQyQzP h5d det
                           -- >-> filter (\(DataFrameQxQyQz _ _ _ ma) -> isJust ma)
                           >-> project det 3 (spaceQxQyQz det pixels res mask')
                           >-> tee (accumulateP c)
                           >-> progress pb
                       ) jobs
  saveCube o r'

-- Hkl

class LenP a => FramesHklP a where
  framesHklP :: a -> Detector b DIM2 -> Pipe (Chunk Int FilePath) (DataFrameHkl b) (SafeT IO) ()

mkJobsHkl :: LenP a => InputHkl a -> IO ([[Chunk Int FilePath]], ProgressBar ())
mkJobsHkl (InputHkl _ fn h5d _ _ _ _ _ _ _) = mkJobs fn h5d

mkInputHkl :: (FramesHklP a, MonadIO m, MonadThrow m)
           => BinocularsConfig
           -> (BinocularsConfig -> m a)
           -> m (InputHkl a)
mkInputHkl c f = do
  fs <- files c
  let d = fromMaybe defaultDetector (_binocularsInputDetector c)
  mask' <- getMask c d
  res <- getResolution c 3
  h5dpath' <- f c
  pure $ InputHkl
    { detector = d
    , filename = InputList fs
    , h5dpath = h5dpath'
    , output = destination'
               (fromMaybe (ConfigRange []) (_binocularsInputInputRange c))
               (_binocularsDispatcherDestination c)
    , resolutions = res
    , centralPixel = _binocularsInputCentralpixel c
    , sdd' = _binocularsInputSdd c
    , detrot' = fromMaybe (0 *~ degree) (_binocularsInputDetrot c)
    , config = c
    , mask = mask'
    }

processHkl :: FramesHklP a => InputHkl a -> IO ()
processHkl input@(InputHkl det _ h5d o res cen d r config' mask') = do
  pixels <- getPixelsCoordinates det cen d r
  (jobs, pb) <- mkJobsHkl input
  r' <- mapConcurrently (\job -> withCubeAccumulator $ \c ->
                           runEffect $ runSafeP $
                           each job
                           -- >-> tee Pipes.Prelude.print
                           >-> framesHklP h5d det
                           -- >-> filter (\(DataFrameHkl (DataFrameQxQyQz _ _ _ ma) _) -> isJust ma)
                           >-> project det 3 (spaceHkl config' det pixels res mask')
                           >-> tee (accumulateP c)
                           >-> progress pb
                       ) jobs
  saveCube o r'

  updateProgress pb $ \p@(Progress _ t _) -> p{progressDone=t}

--  Create the Cube

accumulateP :: (MonadIO m, Shape sh) => IORef (Cube' sh) -> Consumer (DataFrameSpace sh) m ()
accumulateP ref =
    forever $ do s <- await
                 liftIO $ addSpace s =<< readIORef ref

progress :: (MonadIO m, Shape sh) => ProgressBar s -> Consumer (DataFrameSpace sh) m ()
progress p = forever $ do
  _ <- await
  liftIO $ p `incProgress` 1

-- Instances

withDetectorPathP :: (MonadSafe m, Location l) => l -> Detector a DIM2 -> DetectorPath -> ((Int -> IO (ForeignPtr Word16)) -> m r) -> m r
withDetectorPathP f det (DetectorPath p) g = do
  let n = (size . shape $ det) * 2  -- hardcoded size
  withBytes n $ \buf ->
      withHdf5PathP f p $ \p' -> g (getArrayInBuffer buf det p')

nest :: [(r -> a) -> a] -> ([r] -> a) -> a
nest xs = runCont (Prelude.mapM cont xs)

withAxesPathP :: (MonadSafe m, Location l) => l -> [Hdf5Path DIM1 Double] -> ([Dataset] -> m a) -> m a
withAxesPathP f dpaths = nest (map (withHdf5PathP f) dpaths)

withGeometryPathP :: (MonadSafe m, Location l) => l -> GeometryPath -> ((Int -> IO Geometry) -> m r) -> m r
withGeometryPathP f (GeometryPathCristalK6C w m ko ka kp g d) gg =
    withHdf5PathP f w $ \w' ->
    withHdf5PathP f m $ \mu' ->
    withHdf5PathP f ko $ \komega' ->
    withHdf5PathP f ka $ \kappa' ->
    withHdf5PathP f kp $ \kphi' ->
    withHdf5PathP f g $ \gamma' ->
    withHdf5PathP f d $ \delta' -> do
      wavelength <- liftIO $ getValueWithUnit w' 0 angstrom
      mu <- liftIO $ get_position mu' 0
      komega <- liftIO $ get_position komega' 0
      kappa <- liftIO $ get_position kappa' 0
      gamma <- liftIO $ get_position gamma' 0
      delta <- liftIO $ get_position delta' 0
      gg (\j -> do
            kphi <- get_position kphi' j
            return (Geometry
                    K6c
                    (Source wavelength)
                    (fromList [mu, komega, kappa, kphi, gamma, delta])
                    Nothing))
withGeometryPathP f (GeometryPathFix w) gg =
  withHdf5PathP f w $ \w' ->
                        gg (const $
                             Geometry Fixe
                             <$> (Source <$> getValueWithUnit w' 0 angstrom)
                             <*> (fromList <$> pure [])
                             <*> pure Nothing)
withGeometryPathP f (GeometryPathMedH w as) gg =
    withHdf5PathP f w $ \w' ->
    withAxesPathP f as $ \as' ->
        gg (\j -> Geometry MedH
                 <$> (Source <$> getValueWithUnit w' 0 angstrom)
                 <*> (fromList <$> do
                         vs <- Prelude.mapM (`get_position` j) as'
                         return (0.0 : vs))
                 <*> pure Nothing)
withGeometryPathP f (GeometryPathMedV w as) gg =
    withHdf5PathP f w $ \w' ->
    withAxesPathP f as $ \as' ->
        gg (\j -> Geometry MedV
                 <$> (Source <$> getValueWithUnit w' 0 angstrom)
                 <*> (fromList <$> do
                         vs <- Prelude.mapM (`get_position` j) as'
                         return (0.0 : vs))
                 <*> pure Nothing)
withGeometryPathP _f (GeometryPathMedVEiger _w _as _eix _eiz) _gg = undefined
withGeometryPathP f (GeometryPathUhv w as) gg =
    withHdf5PathP f w $ \w' ->
    withAxesPathP f as $ \as' ->
        gg (\j -> Geometry Uhv
                 <$> (Source <$> getValueWithUnit w' 0 angstrom)
                 <*> (fromList <$> Prelude.mapM (`get_position` j) as')
                 <*> pure Nothing)

withAttenuationPathP :: (MonadSafe m, Location l) =>
                       l
                     -> AttenuationPath
                     -> ((Int -> IO Double) -> m r)
                     -> m r
withAttenuationPathP f matt g =
    case matt of
      NoAttenuation -> g (const $ returnIO 1)
      (AttenuationPath p offset coef) ->
          withHdf5PathP f p $ \p' -> g (\j -> do
                                          v <-  get_position p' (j + offset)
                                          if v == badAttenuation
                                          then throwIO (WrongAttenuation "file" (j + offset) (float2Double v))
                                          else return  (coef ** float2Double v))

withQxQyQzPath :: (MonadSafe m, Location l) =>
                 l
               -> Detector a DIM2
               -> QxQyQzPath
               -> ((Int -> IO DataFrameQxQyQz) -> m r)
               -> m r
withQxQyQzPath f det (QxQyQzPath att d dif) g =
  withAttenuationPathP f att $ \getAttenuation ->
  withDetectorPathP f det d $ \getImage ->
  withGeometryPathP f dif $ \getDiffractometer ->
  g (\j -> DataFrameQxQyQz j
          <$> getAttenuation j
          <*> getDiffractometer j
          <*> getImage j
    )

--  FramesQxQyQzP

instance LenP QxQyQzPath where
    lenP (QxQyQzPath ma (DetectorPath i) _) =
        skipMalformed $ forever $ do
                           fp <- await
                           withFileP (openH5 fp) $ \f ->
                               withHdf5PathP f i $ \i' -> do
                                    (_, ss) <- liftIO $ datasetShape i'
                                    case head ss of
                                      (Just n) -> yield $ fromIntegral n - case ma of
                                                                            NoAttenuation           -> 0
                                                                            (AttenuationPath _ off _) -> off
                                      Nothing  -> error "can not extract length"

tryYield :: IO r -> Proxy x' x () r (SafeT IO) ()
tryYield io = do
  edf <- liftIO $ tryJust selectHklBinocularsException io
  case edf of
    Left _   -> return ()
    Right df -> yield df
  where
    selectHklBinocularsException :: HklBinocularsException -> Maybe HklBinocularsException
    selectHklBinocularsException e = Just e

instance FramesQxQyQzP QxQyQzPath where
    framesQxQyQzP p det =
        skipMalformed $ forever $ do
          (Chunk fp from to) <- await
          withFileP (openH5 fp) $ \f ->
            withQxQyQzPath f det p $ \getDataFrameQxQyQz ->
            forM_ [from..to-1] (\j -> tryYield (getDataFrameQxQyQz j))

-- FramesHklP

getValueWithUnit :: Dataset -> Int -> Unit m d Double -> IO (Quantity d Double)
getValueWithUnit d j u = do
  v <- get_position d j
  return $ v *~ u

withSamplePathP :: (MonadSafe m, Location l) => l -> SamplePath -> (IO (Sample Triclinic) -> m r) -> m r
withSamplePathP f (SamplePath a b c alpha beta gamma ux uy uz) g =
    withHdf5PathP f a $ \a' ->
    withHdf5PathP f b $ \b' ->
    withHdf5PathP f c $ \c' ->
    withHdf5PathP f alpha $ \alpha' ->
    withHdf5PathP f beta $ \beta' ->
    withHdf5PathP f gamma $ \gamma' ->
    withHdf5PathP f ux $ \ux' ->
    withHdf5PathP f uy $ \uy' ->
    withHdf5PathP f uz $ \uz' ->
        g (Sample "test"
           <$> (Triclinic
                <$> getValueWithUnit a' 0 angstrom
                <*> getValueWithUnit b' 0 angstrom
                <*> getValueWithUnit c' 0 angstrom
                <*> getValueWithUnit alpha' 0 degree
                <*> getValueWithUnit beta' 0 degree
                <*> getValueWithUnit gamma' 0 degree)
           <*> (Parameter "ux"
                <$> get_position ux' 0
                <*> pure (Range 0 0))
           <*> (Parameter "uy"
                <$> get_position uy' 0
                <*> pure (Range 0 0))
           <*> (Parameter "uz"
                <$> get_position uz' 0
                <*> pure (Range 0 0)))
withSamplePathP _ (SamplePath2 s) g = g (return s)

instance LenP HklPath where
  lenP (HklPath p _)           = lenP p

instance FramesHklP HklPath where
  framesHklP (HklPath qp samp) det = skipMalformed $ forever $ do
    (Chunk fp from to) <- await
    withFileP (openH5 fp) $ \f ->
      withQxQyQzPath f det qp $ \getDataFrameQxQyQz ->
      withSamplePathP f samp $ \getSample ->
      forM_ [from..to-1] (\j -> tryYield ( DataFrameHkl
                                          <$> getDataFrameQxQyQz j
                                          <*> getSample
                                        ))
