{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
module Grenade.Layers.FullyConnected (
    FullyConnected (..)
  , FullyConnected' (..)
  , randomFullyConnected
  ) where

import           Control.Monad.Random hiding (fromList)

import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits

import qualified Numeric.LinearAlgebra as LA
import           Numeric.LinearAlgebra.Static

import           Grenade.Core

import           Grenade.Layers.Internal.Update

-- | A basic fully connected (or inner product) neural network layer.
data FullyConnected i o = FullyConnected
                        !(FullyConnected' i o)   -- Neuron weights
                        !(FullyConnected' i o)   -- Neuron momentum

data FullyConnected' i o = FullyConnected'
                         !(R o)   -- Bias
                         !(L o i) -- Activations

instance Show (FullyConnected i o) where
  show FullyConnected {} = "FullyConnected"

instance (KnownNat i, KnownNat o) => UpdateLayer (FullyConnected i o) where
  type Gradient (FullyConnected i o) = (FullyConnected' i o)

  runUpdate lp (FullyConnected (FullyConnected' oldBias oldActivations) (FullyConnected' oldBiasMomentum oldMomentum)) (FullyConnected' biasGradient activationGradient) =
    let (newBias, newBiasMomentum)    = descendVector (learningRate lp) (learningMomentum lp) (learningRegulariser lp) oldBias biasGradient oldBiasMomentum
        (newActivations, newMomentum) = descendMatrix (learningRate lp) (learningMomentum lp) (learningRegulariser lp) oldActivations activationGradient oldMomentum
    in FullyConnected (FullyConnected' newBias newActivations) (FullyConnected' newBiasMomentum newMomentum)

  createRandom = randomFullyConnected

instance (KnownNat i, KnownNat o) => Layer (FullyConnected i o) ('D1 i) ('D1 o) where
  type Tape (FullyConnected i o) ('D1 i) ('D1 o) = R i
  -- Do a matrix vector multiplication and return the result.
  runForwards (FullyConnected (FullyConnected' wB wN) _) (S1D v) = (v, S1D (wB + wN #> v))

  -- Run a backpropogation step for a full connected layer.
  runBackwards (FullyConnected (FullyConnected' _ wN) _) x (S1D dEdy) =
          let wB'  = dEdy
              mm'  = dEdy `outer` x
              -- calcluate derivatives for next step
              dWs  = tr wN #> dEdy
          in  (FullyConnected' wB' mm', S1D dWs)

instance (KnownNat i, KnownNat o) => Serialize (FullyConnected i o) where
  put (FullyConnected (FullyConnected' b w) _) = do
    putListOf put . LA.toList . extract $ b
    putListOf put . LA.toList . LA.flatten . extract $ w

  get = do
      let f  = fromIntegral $ natVal (Proxy :: Proxy i)
      b     <- maybe (fail "Vector of incorrect size") return . create . LA.fromList =<< getListOf get
      k     <- maybe (fail "Vector of incorrect size") return . create . LA.reshape f . LA.fromList =<< getListOf get
      let bm = konst 0
      let mm = konst 0
      return $ FullyConnected (FullyConnected' b k) (FullyConnected' bm mm)

randomFullyConnected :: (MonadRandom m, KnownNat i, KnownNat o)
                     => m (FullyConnected i o)
randomFullyConnected = do
    s1    <- getRandom
    s2    <- getRandom
    let wB = randomVector  s1 Uniform * 2 - 1
        wN = uniformSample s2 (-1) 1
        bm = konst 0
        mm = konst 0
    return $ FullyConnected (FullyConnected' wB wN) (FullyConnected' bm mm)


{-# LANGUAGE MagicHash, UnboxedTuples #-}

module Render (render, parRender) where

import Control.Parallel.Strategies

import Numeric.Vector
import Numeric.Scalar (fromScalar, scalar)
import Types
import Control.Monad.State
import Camera (rayFrom)
import Raycast (raycast)
import Shader (reflectance, sampleBRDF, probability)
import Scene (makeContext)
import VecUtil (unpack3, vec3Of)
import System.Random
import SceneRandom
import Lights

map2 :: (a -> b) -> [[a]] -> [[b]]
map2 f = map (map f)

vec23 :: Vec2d -> (Double, Double, Double)
vec23 v = case unpackV2# v of
            (# u, v #) -> (u, v, 1)

possibly :: (a -> b) -> b -> Maybe a -> b
possibly _ def Nothing = def
possibly f _ (Just x) = f x

shade :: Int -> State SceneContext Vec3d
shade bounceCount = get >>= (\ctx -> if bounceCount > s_getBounces (ss_getScene ctx) then return (vec3Of 0) else do
  let hit = ss_getHit ctx
  let brdf = rh_getShader $ hit
  col <- reflectance brdf
  let scene = ss_getScene ctx
  let lights = s_getLights scene
  let norm = rh_getNormal hit

  mReflectRay <- sampleBRDF brdf
  (reflectColor, reflectProbability) <- case mReflectRay of
    Nothing -> return (0, 0)
    Just reflectRay -> do
      prob <- probability brdf (r_getDir reflectRay)
      case raycast scene reflectRay of
        Nothing -> return $ (s_getSkyColor scene, )prob
        Just hit -> do
          let (v, s) = runState (shade (bounceCount + 1)) $ makeContext scene (ss_getGen ctx) hit
          put s
          return (v, prob) -- 1 / (2 * PI)

  let newRayPos = (rh_getPos hit) + (norm * (vec3Of 0.01))

  lighting <- sequence $ map (sampleLight newRayPos norm) lights
  let lighting' = filter ((>0) . fst) lighting
  lightProbabilities <- sequence $ map (probability brdf . snd) lighting'
  let lightColors = map fst lighting'

  let probabilitySum = reflectProbability + sum lightProbabilities

  let adjustedReflectColor = reflectColor * vec3Of reflectProbability
  let adjustedIncomingLighting = sum $ zipWith (*) lightColors (map vec3Of lightProbabilities)

  let incomingLight = adjustedReflectColor + adjustedIncomingLighting

  let lit = col * incomingLight

  reflectColor <- reflectance brdf
  return lit
  )

samplePixel :: Scene -> (Int, Int) -> Int -> Vec3d
samplePixel s (x, y) sample = let w = s_getWidth s
                                  h = s_getHeight s
                                  cam = s_getCamera s
                                  skyCol = s_getSkyColor s
                                  -- Random seed is an arbitrary formula based on pixel location to form distinct results
                                  gen = mkStdGen $ (x + y * w + x * (x + y) + (s_getSeed s)) * sample
                              in  possibly (fst . runState (shade 0) . makeContext s gen) (skyCol / (2 * pi)) $
                                  raycast s $
                                  rayFrom cam $
                                  (fromIntegral x / fromIntegral w * 2.0 - 1.0, (fromIntegral h / fromIntegral w) - fromIntegral y / fromIntegral w * 2.0)

renderPixel :: Scene -> (Int, Int) -> Vec3d
renderPixel s p = let numSamples = s_getSamples s
                      samples = map (samplePixel s p)  [1..numSamples]
                  -- Multiply by 2pi for monte-carlo integration because the volume of the integration bounds is 2pi
                  -- (integration bounds form a hemisphere)
                  in  sum samples / fromIntegral numSamples * vec3Of (2 * pi)

render :: Scene -> [[(Double, Double, Double)]]
render s = let w = s_getWidth s
               h = s_getHeight s
               cam = s_getCamera s
               skyCol = s_getSkyColor s
               seed = 1
               gen = mkStdGen seed
          in  map2 unpack3 $
              map2 (renderPixel s) $
              [ [ (x, y) | x <- [1..w] ] | y <- [1..h] ]

-- Found how to evaluate the list in parallel from
-- https://stackoverflow.com/a/5606176
parRender :: Scene -> [[(Double, Double, Double)]]
parRender s = (render s) `using` parList rdeepseq


{-# LANGUAGE ForeignFunctionInterface #-}

-- $ ghc -O2 --make wrappers.hs functions.c

import Numeric.LinearAlgebra
import Data.Packed.Development
import Foreign(Ptr,unsafePerformIO)
import Foreign.C.Types(CInt)

-----------------------------------------------------

main = do
    print $ myScale 3.0 (fromList [1..10])
    print $ myDiag $ (3><5) [1..]

-----------------------------------------------------

foreign import ccall unsafe "c_scale_vector"
    cScaleVector :: Double                -- scale
                 -> CInt -> Ptr Double    -- argument
                 -> CInt -> Ptr Double    -- result
                 -> IO CInt               -- exit code

myScale s x = unsafePerformIO $ do
    y <- createVector (dim x)
    app2 (cScaleVector s) vec x vec y "cScaleVector"
    return y

-----------------------------------------------------
-- forcing row order

foreign import ccall unsafe "c_diag"
    cDiag :: CInt -> CInt -> Ptr Double  -- argument
          -> CInt -> Ptr Double          -- result1
          -> CInt -> CInt -> Ptr Double  -- result2
          -> IO CInt                     -- exit code

myDiag m = unsafePerformIO $ do
    y <- createVector (min r c)
    z <- createMatrix RowMajor r c
    app3 cDiag mat (cmat m) vec y mat z "cDiag"
    return (y,z)
  where r = rows m
        c = cols m


{-# OPTIONS_GHC -fplugin GHC.TypeLits.KnownNat.Solver #-}

{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DataKinds                 #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts          #-}
{-# LANGUAGE GADTs                     #-}
{-# LANGUAGE RankNTypes                #-}
{-# LANGUAGE ScopedTypeVariables       #-}
{-# LANGUAGE TypeFamilies              #-}

import           Control.Applicative
import           Control.Concurrent
import           Control.Monad.Random

import           Options.Applicative

import           Graphics.Image               hiding (map, on)

import qualified Numeric.LinearAlgebra        as LA
import qualified Numeric.LinearAlgebra.Static as H

import           Grenade

data SuperResOptions = SuperResOptions FilePath         -- onnx file
                                       FilePath         -- input image

pSuperRes :: Parser SuperResOptions
pSuperRes = SuperResOptions <$> argument str (metavar "onnx") <*> argument str (metavar "image")

loadSuperResImage :: FilePath -> IO (Maybe (S ('D3 224 224 1), [[RealNum]], [[RealNum]]))
loadSuperResImage path = do
  img <- readImageRGB VU path
  -- displayImage img
  return $ do
    guard $ dims img == (224, 224)
    let imgYCbCr = toImageYCbCr img
        imgY0    = map (\(PixelYCbCr y _ _ )      -> doubleToRealNum y )  . concat . toLists $ imgYCbCr
        imgCb    = map (map (\(PixelYCbCr _ cb _) -> doubleToRealNum cb)) . toLists $ imgYCbCr
        imgCr    = map (map (\(PixelYCbCr _ _ cr) -> doubleToRealNum cr)) . toLists $ imgYCbCr

    return (S3D (H.fromList imgY0), imgCb, imgCr)

displayHighResImage :: S ('D3 672 672 1) -> [[RealNum]] -> [[RealNum]] -> IO ()
displayHighResImage (S3D m) cbs crs = do
  let realNumToPixel = PixelX . realNumToDouble
      m'  = LA.toLists $ H.extract m              :: [[RealNum]]
      m'' = map (map realNumToPixel) m'           :: [[Pixel X Double]]
      img = fromLists m''                         :: Image VU X Double

      imgBs  = fromLists $ map (map realNumToPixel) cbs :: Image VU X Double
      imgRs  = fromLists $ map (map realNumToPixel) crs :: Image VU X Double

      imgBs' = resize Bilinear Edge (672, 672) imgBs :: Image VU X Double
      imgRs' = resize Bilinear Edge (672, 672) imgRs :: Image VU X Double

      finalImg = fromImagesX [(LumaYCbCr, img), (CBlueYCbCr, imgBs'), (CRedYCbCr, imgRs')] :: Image VU YCbCr Double

  displayImage finalImg
  threadDelay 10000000

main :: IO ()
main = do
    SuperResOptions netPath imgPath <- execParser (info (pSuperRes <**> helper) idm)
    res <- loadSuperResolution netPath

    inputM <- loadSuperResImage imgPath

    case (res, inputM) of
      (Right net, Just (input, cbs, crs))  -> do
        let S3D y = runNet net input
        displayHighResImage (S3D y) cbs crs

      (Left err, _) -> print err
      _             -> putStrLn "Failed to load network and file"


{-# LANGUAGE BangPatterns #-}
module STCFourierSeries where

import           Control.DeepSeq
import           Control.Monad                      as M
import qualified Data.Array.Accelerate              as A
import qualified Data.Array.Accelerate.Data.Complex as A
import           Data.Array.Accelerate.LLVM.PTX
import           Data.Array.IArray
import           Data.Array.Repa                    as R
import           Data.Binary
import           Data.Complex
import           Data.List                          as L
import           Data.Maybe
import           Data.Vector.Generic                as VG
import           FokkerPlanck.BrownianMotion
import           FokkerPlanck.FourierSeries
import           FokkerPlanck.GreensFunction
import           FokkerPlanck.Histogram
import           FokkerPlanck.MonteCarlo
import           Foreign.CUDA.Driver                as CUDA
import           Image.IO
import           STC
import           System.Directory
import           System.Environment
import           System.FilePath
import           Text.Printf
import           Utils.Array
import           Utils.Time

main = do
  args@(gpuIDStr:numPointStr:deltaXStr:numOrientationStr:numScaleStr:thetaSigmaStr:scaleSigmaStr:maxScaleStr:deltaLogStr:taoStr:numTrailStr:maxTrailStr:phiFreqsStr:rhoFreqsStr:thetaFreqsStr:scaleFreqsStr:initDistStr:initScaleStr:histFilePath:stdStr:stdGStr:locationStr:numThreadStr:_) <-
    getArgs
  let gpuID = read gpuIDStr :: [Int]
      numPoint = read numPointStr :: Int
      deltaX = read deltaXStr :: Double
      numOrientation = read numOrientationStr :: Int
      numScale = read numScaleStr :: Int
      thetaSigma = read thetaSigmaStr :: Double
      scaleSigma = read scaleSigmaStr :: Double
      tao = read taoStr :: Double
      numTrail = read numTrailStr :: Int
      maxTrail = read maxTrailStr :: Int
      phiFreq = read phiFreqsStr :: Double
      phiFreqs = [-phiFreq .. phiFreq]
      rhoFreq = read rhoFreqsStr :: Double
      rhoFreqs = [-rhoFreq .. rhoFreq]
      thetaFreq = read thetaFreqsStr :: Double
      thetaFreqs = [-thetaFreq .. thetaFreq]
      scaleFreq = read scaleFreqsStr :: Double
      scaleFreqs = [-scaleFreq .. scaleFreq]
      initDist = read initDistStr :: [(Double, Double, Double, Double)]
      initScale = read initScaleStr :: Double
      initPoints = L.map (\(x, y, t, s) -> Point x y t s) initDist
      numThread = read numThreadStr :: Int
      folderPath = "output/test/STCFourierSeries"
      maxScale = read maxScaleStr :: Double
      halfLogPeriod = log maxScale
      deltaLog = read deltaLogStr :: Double
      std = read stdStr :: Double
      stdG = read stdGStr :: Double
      location = read locationStr :: (Int, Int)
  removePathForcibly folderPath
  createDirectoryIfMissing True folderPath
  flag <- doesFileExist histFilePath
  -- hist <-
  --   if flag
  --     then do
  --       printCurrentTime $ "read data from " L.++ histFilePath
  --       decodeFile histFilePath
  --     else runMonteCarloFourierCoefficientsGPU
  --            gpuID
  --            numThread
  --            numTrail
  --            maxTrail
  --            thetaSigma
  --            scaleSigma
  --            maxScale
  --            tao
  --            phiFreqs
  --            rhoFreqs
  --            thetaFreqs
  --            scaleFreqs
  --            deltaLog
  --            initScale
  --            histFilePath
  --            (emptyHistogram
  --               [ L.length phiFreqs
  --               , L.length rhoFreqs
  --               , L.length thetaFreqs
  --               , L.length scaleFreqs
  --               ]
  --               0)
  -- hist <-
  --   sampleLogpolar
  --     histFilePath
  --     (L.head gpuID)
  --     180
  --     180
  --     (fromIntegral $ div numPoint 2)
  --     initScale
  --     thetaSigma
  --     tao
  --     phiFreqs
  --     rhoFreqs
  --     thetaFreqs
  --     scaleFreqs
  hist <-
    sampleCartesian
      folderPath
      histFilePath
      gpuID
      (fromIntegral $ div numPoint 2)
      144
      deltaLog
      initScale
      thetaSigma
      tao
      phiFreqs
      rhoFreqs
      thetaFreqs
      scaleFreqs
  let !initSource =
        computeInitialDistribution'
          numPoint
          numPoint
          phiFreqs
          rhoFreqs
          -- thetaFreqs
          halfLogPeriod
          [L.head initPoints]
      !initSink =
        computeInitialDistribution'
          numPoint
          numPoint
          phiFreqs
          rhoFreqs
          -- thetaFreqs
          halfLogPeriod
          [L.last initPoints]
      !coefficients =
        normalizeFreqArr' std phiFreqs rhoFreqs . getNormalizedHistogramArr $
        hist
       -- = getNormalizedHistogramArr $ hist :: R.Array U DIM4 (Complex Double)
      !thetaRHarmonics =
        computeThetaRHarmonics
          numOrientation
          numScale
          thetaFreqs
          scaleFreqs
          halfLogPeriod
  plan <-
    makePlan
      folderPath
      emptyPlan
      numPoint
      numPoint
      (L.length thetaFreqs)
      (L.length scaleFreqs)
  printCurrentTime "harmonicsArray"
  -- gaussian <- gaussianFilter2D plan numPoint stdG
  harmonicsArray <-
    dftHarmonicsArray
      plan
      numPoint
      deltaX
      numPoint
      deltaX
      phiFreqs
      rhoFreqs
      thetaFreqs
      scaleFreqs
      halfLogPeriod
      (fromIntegral numPoint * sqrt 2)
      -- gaussian
  -- let harmonicsArray' =
  --       computeHarmonicsArray
  --         numPoint
  --         deltaX
  --         numPoint
  --         deltaX
  --         phiFreqs
  --         rhoFreqs
  --         thetaFreqs
  --         scaleFreqs
  --         halfLogPeriod
  --         maxScale
  --         -- (fromIntegral numPoint)
  -- harmonicsArrayGPU <-
  --   dftHarmonicsArrayGPU
  --     plan
  --     numPoint
  --     deltaX
  --     numPoint
  --     deltaX
  --     phiFreqs
  --     rhoFreqs
  --     thetaFreqs
  --     scaleFreqs
  --     halfLogPeriod
  --     maxScale
  -- initialise []
  -- dev <- device . L.head $ gpuID
  -- ctx <- CUDA.create dev []
  -- ptx <- createTargetFromContext ctx
  -- registerPinnedAllocatorWith ptx
  --Source
  printCurrentTime "Source"
  sourceArr <- convolve' Source plan coefficients harmonicsArray initSource
  -- sourceArr' <- convolve'' Source plan coefficients harmonicsArray' initSource
  -- printCurrentTime "CPU"
  -- sourceArr <-
  --   convolveGPU
  --     ptx
  --     Source
  --     plan
  --     (A.use .
  --      A.fromList
  --        (A.Z A.:. (L.length scaleFreqs) A.:. (L.length thetaFreqs) A.:.
  --         (L.length rhoFreqs) A.:.
  --         (L.length phiFreqs)) .
  --      R.toList $
  --      coefficients)
  --     (A.use .
  --      A.fromList
  --        (A.Z A.:. (L.length scaleFreqs) A.:. (L.length thetaFreqs) A.:.
  --         (L.length rhoFreqs) A.:.
  --         (L.length phiFreqs)) .
  --      R.toList $
  --      coefficients)
  --     harmonicsArrayGPU
  --     initSource
  -- printCurrentTime "GPU"
  plotDFTArrayPower
    (folderPath </> (printf "SourcePower.png"))
    numPoint
    numPoint $
    sourceArr
  plotDFTArrayThetaR
    (folderPath </> "Source.png")
    numPoint
    numPoint
    thetaRHarmonics
    sourceArr
  let sourceArr' =
        fromUnboxed (Z :. numOrientation :. numPoint :. numPoint) .
        VG.convert .
        VG.concat .
        computeFourierSeriesThetaR thetaRHarmonics . getDFTArrayVector $
        sourceArr
  plotThetaDimension folderPath "FourierSeries_" location . R.map magnitude $
    sourceArr'
  -- plotDFTArrayThetaR
  --   (folderPath </> "Source_test.png")
  --   numPoint
  --   numPoint
  --   thetaRHarmonics
  --   sourceArr'
  --Sink
  printCurrentTime "Sink"
  sinkArr <- convolve' Sink plan coefficients harmonicsArray initSink
  plotDFTArrayPower (folderPath </> (printf "SinkPower.png")) numPoint numPoint $
    sinkArr
  plotDFTArrayThetaR
    (folderPath </> "Sink.png")
    numPoint
    numPoint
    thetaRHarmonics
    sinkArr
  printCurrentTime "Completion"
  completionArr <- completionField' plan sourceArr sinkArr
  plotDFTArrayThetaR
    (folderPath </> "Completion.png")
    numPoint
    numPoint
    thetaRHarmonics
    completionArr
  plotDFTArrayPower
    (folderPath </> (printf "CompletionPower.png"))
    numPoint
    numPoint
    completionArr
  let a =
        computeFourierSeriesThetaR thetaRHarmonics . getDFTArrayVector $
        sourceArr
      b =
        computeFourierSeriesThetaR thetaRHarmonics . getDFTArrayVector $ sinkArr
  plotImageRepa (folderPath </> "CompletionR2S1.png") .
    ImageRepa 8 .
    fromUnboxed (Z :. (1 :: Int) :. numPoint :. numPoint) .
    VG.map (\x -> x ^ 2) . VG.convert . L.foldl1' (VG.zipWith (+)) $
    L.zipWith (VG.zipWith (\x y -> magnitude x * magnitude y)) a b
  printCurrentTime ""
  -- M.mapM_
  --   (\s -> do
  --      let !arr =
  --            getNormalizedHistogramArr hist :: R.Array U DIM4 (Complex Double)
  --          initArr =
  --            traverse4
  --              (fromListUnboxed (Z :. L.length phiFreqs) phiFreqs)
  --              (fromListUnboxed (Z :. L.length rhoFreqs) rhoFreqs)
  --              (fromListUnboxed (Z :. L.length thetaFreqs) thetaFreqs)
  --              (fromListUnboxed (Z :. L.length scaleFreqs) scaleFreqs)
  --              (\_ _ _ _ -> extent arr) $ \fPhi fRho fTheta fScale (Z :. scale :. theta :. rho :. phi) ->
  --              cis $
  --              (-pi * log s) * fRho (Z :. rho) / halfLogPeriod -
  --              (fPhi (Z :. phi) * t0 * pi / 180)
  --      printCurrentTime ""
  --      source <- computeUnboxedP $ arr *^ initArr
  --      plotImageRepaComplex (folderPath </> (printf "Source_%.2f.png" s)) .
  --        ImageRepa 8 .
  --        fromUnboxed (Z :. (1 :: Int) :. numPoint :. numPoint) .
  --        L.foldl1' (VU.zipWith (+)) .
  --        computeFourierSeriesThetaR
  --          numOrientation
  --          numScale
  --          thetaFreqs
  --          scaleFreqs
  --          halfLogPeriod .
  --        computeFourierSeriesR2
  --          numPoint
  --          deltaX
  --          numPoint
  --          deltaX
  --          phiFreqs
  --          rhoFreqs
  --          thetaFreqs
  --          scaleFreqs
  --          halfLogPeriod
  --          maxScale
  --          harmonicsArray $
  --        source)
  --   [1,1.25 .. s0]


-------------------------------------------------------------------------------
-- |
-- Module    :  Spaces.Action
-- Copyright :  (c) Sentenai 2017
-- License   :  BSD3
-- Maintainer:  sam@sentenai.com
-- Stability :  experimental
-- Portability: non-portable
--
-- typeclass for a discrete action space, as well as helper functions
-------------------------------------------------------------------------------
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Reinforce.Spaces.Action
  ( DiscreteActionSpace(..)
  -- , oneHot
  , oneHot'
  , allActions
  , randomChoice
  ) where

import Control.Monad.IO.Class
-- import Numeric.LinearAlgebra.Static (R)
-- import qualified Numeric.LinearAlgebra.Static as LA

import Control.MonadMWCRandom
import GHC.TypeLits
import Data.Proxy
import Data.Vector (Vector)
import qualified Data.Vector as V

-- | Mostly tags around an Enum, but includes information about the size of
-- an action space and is used in helper functions.
class (Bounded a, Enum a) => DiscreteActionSpace a where
  type Size a :: Nat

  toAction :: Int -> a
  toAction = toEnum

  fromAction :: a -> Int
  fromAction = fromEnum

-- | one-hot encode a bounded enumerable. Doesn't care if minBound is < or > 0
-- oneHot :: forall a . (KnownNat (Size a), DiscreteActionSpace a) => a -> R (Size a)
-- oneHot e = LA.vector . V.toList
--   $ V.unsafeUpd (replicateZeros (Proxy :: Proxy a)) [(fromEnum e, 1)]

-- | one-hot encode a bounded enumerable
oneHot' :: forall a . (DiscreteActionSpace a) => a -> Vector Double
oneHot' e = V.unsafeUpd (replicateZeros (Proxy :: Proxy a)) [(fromEnum e, 1)]

-- | helper function to initialize a one-hot vector
replicateZeros :: forall a . (Enum a, Bounded a) => Proxy a -> Vector Double
replicateZeros _ = V.fromList $ replicate (fromEnum (maxBound :: a) + 1) 0

-- | helper function to get all actions in a discrete action space
allActions :: DiscreteActionSpace a => [a]
allActions = [minBound..maxBound]

-- | make a uniform-random selection of an Action in a discrete action space
randomChoice
  :: forall m a . (MonadIO m , MonadMWCRandom m, DiscreteActionSpace a)
  => m a
randomChoice = toEnum . fst <$> sampleFrom uniformDist
  where
    uniformDist :: [Double]
    uniformDist = fmap (\a -> convert a / total) allActions
      where
        convert :: a -> Double
        convert = fromIntegral . fromEnum

        total :: Double
        total = sum (fmap convert allActions)


{-# LANGUAGE TupleSections #-}
module Geo.Computations.Trail
       ( -- * Types
         AvgMethod(..)
       , Selected(..)
       , PointGrouping
       , TransformGrouping
         -- * Utility Functions
       , isSelected
       , isNotSelected
       , onSelected
       , selLength
         -- * Trail Functions
         -- ** Queries
       , totalDistance
       , totalTime
       , avgSpeeds
       , slidingAverageSpeed
       , closestDistance
       , convexHull
         -- ** Transformations
       , bezierCurveAt
       , bezierCurve
       , linearTime
       , filterPoints
         -- ** Grouping Methods
       , betweenSpeeds
       , restLocations
       , spansTime
       , everyNPoints
         -- ** Group Transformations
       , intersectionOf
       , invertSelection
       , firstGrouping
       , lastGrouping
       , unionOf
       , refineGrouping
       , (/\), (\/)
         -- ** Composite Operations (Higher Level)
       , smoothRests
       , smoothTrail
        -- * Misc
       , bezierPoint
         ) where

import Text.Show.Functions ()
import Geo.Computations.Basic

import Control.Arrow (first)
import Control.Monad
import Data.Fixed (mod')
import Data.Function (on)
import Data.List as L
import Data.Maybe
import Data.Ord
import Data.Time

import Statistics.Function as F
import Statistics.Sample
import qualified Data.Vector.Unboxed as V

takeWhileEnd :: (a -> Bool) -> [a] -> (Maybe a, [a],[a])
takeWhileEnd p xs = go xs Nothing
  where
  go [] e = (e, [], [])
  go (a:as) e
    | p a = let (e',xs2,zs) = go as (Just a) in (e',a:xs2,zs)
    | otherwise = (e,[], a:as)

data AvgMethod c
  = AvgMean              -- ^ Obtain the 'mean' of the considered points
  | AvgHarmonicMean      -- ^ Obtain the 'harmonicMean'
  | AvgGeometricMean     -- ^ Obtain the 'geometricMean'
  | AvgMedian            -- ^ Obtain the median of the considered points
  | AvgEndPoints         -- ^ Compute the speed considering only the given endpoints
  | AvgMinOf [AvgMethod c] -- ^ Take the minimum of the speeds from the given methods
  | AvgWith ([c] -> Speed)

-- | @avgSpeeds n points@
-- Average speed using a window of up to @n@ seconds and averaging by taking the
-- Median ('AvgMedian').
avgSpeeds :: NominalDiffTime -> Trail Point -> [(UTCTime, Speed)]
avgSpeeds = slidingAverageSpeed AvgHarmonicMean

-- | @slidingAverageSpeed m n@ Average speed using a moving window of up to @n@ seconds
-- and an 'AvgMethod' of @m@.
slidingAverageSpeed :: AvgMethod Point -> NominalDiffTime -> Trail Point -> [(UTCTime, Speed)]
slidingAverageSpeed _ _ [] = []
slidingAverageSpeed m minTime xs =
  let pts   = map unSelect (spansTime minTime xs)
      spds  = map (getAvg m) pts
      times = map getAvgTimes pts
  in concatMap maybeToList $ zipWith (\t s -> fmap (,s) t) times spds
  where
  getTimeDiff a b = on (liftM2 diffUTCTime) pntTime a b

  --  getAvg :: [] -> AvgMethod -> Speed
  getAvg _ [] = 0
  getAvg _ [_] = 0
  getAvg m2 cs =
    let ss = getSpeedsV cs
    in case m2 of
        AvgMean -> mean ss
        AvgHarmonicMean  -> harmonicMean ss
        AvgGeometricMean -> geometricMean ss
        AvgMedian ->
          let ss' = F.sort $ getSpeedsV cs
              len = V.length ss'
              mid = len `div` 2
          in if V.length ss' < 3
             then mean ss'
             else if odd len then ss' V.! mid else mean (V.slice mid 2 ss')
        AvgEndPoints -> fromMaybe 0 $ speed (head cs) (last cs)
        AvgMinOf as -> minimum $ map (flip getAvg cs) as
        AvgWith f -> f cs
  getAvgTimes [] = Nothing
  getAvgTimes [x] = pntTime x
  getAvgTimes ps = getAvgTime (head ps) (last ps)
  getAvgTime a b = liftM2 addUTCTime (getTimeDiff b a) (pntTime a)
  getSpeedsV = V.fromList . getSpeeds
  getSpeeds zs = concatMap maybeToList $ zipWith speed zs (drop 1 zs)

-- | A PointGrouping is a function that selects segments of a trail.
--
-- Grouping point _does not_ result in deleted points. It is always true that:
--
--     forall g :: PointGrouping c -->
--     concatMap unSelect (g ts) == ts
--
-- The purpose of grouping is usually for later processing.  Any desire to drop
-- points that didn't meet a particular grouping criterion can be filled with
-- a composition with 'filter' (or directly via 'filterPoints').
type PointGrouping c = Trail c -> [Selected (Trail c)]

-- | Given a selection of coordinates, transform the selected
-- coordinates in some way (while leaving the non-selected
-- coordinates unaffected).
type TransformGrouping c = [Selected (Trail c)] -> [Selected (Trail c)]

-- | When grouping points, lists of points are either marked as 'Select' or 'NotSelect'.
data Selected a = Select {unSelect :: a} | NotSelect {unSelect :: a}
  deriving (Eq, Ord, Show)

isSelected :: Selected a -> Bool
isSelected (Select _) = True
isSelected _ = False

isNotSelected :: Selected a -> Bool
isNotSelected = not . isSelected

selLength :: Selected [a] -> Int
selLength = length . unSelect

onSelected :: (a -> b) -> (a -> b) -> Selected a -> b
onSelected f _ (Select a) = f a
onSelected _ g (NotSelect a) = g a

instance Functor Selected where
  fmap f (Select x) = Select $ f x
  fmap f (NotSelect x) = NotSelect $ f x

dropExact :: Int -> [Selected [a]] -> [Selected [a]]
dropExact _ [] = []
dropExact i (x:xs) =
  case compare (selLength x) i of
    EQ -> xs
    LT -> dropExact (i - selLength x) xs
    GT -> fmap (drop i) x : xs

-- | Groups trail segments into contiguous points within the speed
-- and all others outside of the speed.  The "speed" from point p(i)
-- to p(i+1) is associated with p(i) (except for the first speed
-- value, which is associated with both the first and second point)
betweenSpeeds :: Double -> Double -> PointGrouping Point
betweenSpeeds low hi ps =
  let spds = concatMap maybeToList $ zipWith speed ps (drop 1 ps)
      psSpds = [(p,s) | p <- ps, s <- maybeToList (listToMaybe spds) ++ spds]
      inRange x = x >= low && x <= hi
      chunk [] = []
      chunk xs@(x:_) =
        let op p = if inRange (snd x) then first Select . span p else first NotSelect . break p
            (r,rest) = op (inRange . snd) xs
        in r : chunk rest
  in map (fmap (map fst)) $ chunk psSpds

-- | A "rest point" means the coordinates remain within a given distance
-- for at least a particular amount of time.
restLocations :: Distance -> NominalDiffTime -> PointGrouping Point
restLocations d s ps =
  let consToFirst x [] = [NotSelect [x]]
      consToFirst x (a:as) = (fmap (x:) a) : as
      go [] [] = []
      go [] nonRests = [NotSelect $ reverse nonRests]
      go (a:as) nonRests =
        case takeWhileEnd ((<=) d . distance a) as of
          (Just l, close, far) ->
            case (pntTime a, pntTime l) of
              (Just t1, Just t2) ->
                let diff = diffUTCTime t2 t1
                in if diff >= s then NotSelect (reverse nonRests) : Select (a:close) : go far [] else go as (a:nonRests)
              _ -> consToFirst a $ go as nonRests
          _ -> consToFirst a $ go as nonRests
  in go ps []

-- | Chunks points into groups spanning at most the given time
-- interval.
spansTime :: NominalDiffTime -> PointGrouping Point
spansTime n ps =
  let times  = mkTimePair ps
      chunk [] = []
      chunk xs@(x:_) =
        let (good,rest) = span ((<= addUTCTime n (snd x)) . snd) xs
        in if null good then [xs] else good : chunk rest
  in map (Select . map fst) $ chunk times

-- | Intersects the given groupings
intersectionOf :: [PointGrouping Point] -> PointGrouping Point
intersectionOf gs ps =
  let groupings = map ($ ps) gs
      -- chunk :: [[Selected [pnts]]] -> pnts -> [pnts]
      chunk _ [] = []
      chunk ggs xs =
        let minLen = max 1 . minimum . concatMap (take 1) $ map (map selLength) ggs   -- FIXME this is all manner of broken
            sel = if all isSelected (concatMap (take 1) ggs) then Select else NotSelect
            (c,rest) = splitAt minLen xs
        in sel c : chunk (filter (not . null) $ map (dropExact minLen) ggs) rest
  in chunk groupings ps

-- | Union all the groupings
unionOf :: [PointGrouping Point] -> PointGrouping Point
unionOf gs ps =
  let groupings = map ($ ps) gs
      chunk _ [] = []
      chunk ggs xs =
        let getSegs = concatMap (take 1)
            segs = getSegs ggs
            len =
              if any isSelected segs
                 then max 1 . maximum . getSegs . map (map selLength) . map (filter isSelected) $ ggs
                 else max 1 . minimum . getSegs . map (map selLength) $ ggs
            sel = if any isSelected segs then Select else NotSelect
            (c,rest) = splitAt len xs
        in sel c : chunk (filter (not . null) $ map (dropExact len) ggs) rest
  in chunk groupings ps

-- | Intersection binary operator
(/\) :: [Selected (Trail a)] -> TransformGrouping a
(/\) _ [] = []
(/\) [] _ = []
(/\) xsL@(Select x:_) ysL@(Select y:_) =
  let z = if length x < length y then x else y
      xs' = selListDrop (length z) xsL
      ys' = selListDrop (length z) ysL
  in Select z : (xs' /\ ys')
(/\) xs (NotSelect y:ys) = NotSelect y : (selListDrop (length y) xs /\ ys)
(/\) (NotSelect x:xs) ys = NotSelect x : (xs /\ selListDrop (length x) ys)

-- | Union binary operator
(\/) :: [Selected (Trail a)] -> TransformGrouping a
(\/) xs [] = xs
(\/) [] ys = ys
(\/) (Select x:xs) (Select y : ys) =
  let xLen = length x
      yLen = length y
  in if xLen < yLen
       then (Select y :) (selListDrop (yLen - xLen) xs \/ ys)
       else (Select x :) (xs \/ selListDrop (xLen - yLen) ys)
(\/) (Select x:_) ys = Select x : selListDrop (length x) ys
(\/) xs (Select y:_) = Select y : selListDrop (length y) xs
(\/) xsL@(NotSelect x:xs) ysL@(NotSelect y:ys) =
  let xLen = length x
      yLen = length y
  in if xLen < yLen
        then (NotSelect x:) (xs \/ selListDrop xLen ysL)
        else (NotSelect y:) (selListDrop yLen xsL \/ ys)

selListDrop :: Int -> [Selected [a]] -> [Selected [a]]
selListDrop 0 xs = xs
selListDrop _ [] = []
selListDrop n (x:xs) =
  let x' = drop n (unSelect x)
  in fmap (const x') x : selListDrop (n - (selLength x - length x')) xs

-- | Inverts the selected/nonselected segments
invertSelection :: TransformGrouping a
invertSelection = map (onSelected NotSelect Select)

-- | @firstGrouping f ps@ only the first segment remains 'Select'ed, and only
-- if it was already selected by @f@.
firstGrouping ::  TransformGrouping a
firstGrouping ps = take 1 ps ++ map (NotSelect . unSelect) (drop 1 ps)

-- | Only the last segment, if any, is selected (note: the current
-- implementation is inefficient, using 'reverse')
lastGrouping ::  TransformGrouping a
lastGrouping ps  = let ps' = reverse ps in reverse $ take 1 ps' ++ map (NotSelect . unSelect) (drop 1 ps')

-- | Chunks the trail into groups of N points
everyNPoints ::  Int -> PointGrouping a
everyNPoints n ps
  | n <= 0 = [NotSelect ps]
  | otherwise = go ps
    where
      go [] = []
      go xs = let (h,t) = splitAt n xs in Select h : go t

-- | For every selected group, refine the selection using the second
-- grouping method.  This differs from 'IntersectionOf' by restarting
-- the second grouping algorithm at the beginning each group selected
-- by the first algorithm.
refineGrouping ::  PointGrouping a -> TransformGrouping a
refineGrouping b = concatMap (onSelected b (\x -> [NotSelect x]))

-- | Remove all points that remain 'NotSelect'ed by the given grouping algorithm.
filterPoints :: PointGrouping a -> Trail a -> Trail a
filterPoints g = concatMap unSelect . filter isSelected . g

-- Extract the time from each coordinate.  If no time is available then
-- the coordinate is dropped!
mkTimePair :: Trail Point -> [(Point,UTCTime)]
mkTimePair xs =
  let timesM = map (\x-> fmap (x,) $ pntTime x) xs
  in concatMap maybeToList timesM

-- | Construct a bezier curve using the provided trail.  Construct a
-- new trail by sampling the given bezier curve at the given times.
-- The current implementation assumes the times of the input
-- coordinates are available and all equal (Ex: all points are 5
-- seconds apart), the results will be poor if this is not the case!
bezierCurveAt :: [UTCTime] -> Trail Point -> Trail Point
bezierCurveAt _ [] = []
bezierCurveAt selectedTimes xs =
  let timesDef = mkTimePair xs
      end = last timesDef
      top = head timesDef
      tTime  = diffUTCTime (snd end) (snd top)
      times = if null selectedTimes then map snd timesDef else selectedTimes
      diffTimes = [diffUTCTime t (snd top) / tTime | t <- times]
      queryTimes = map realToFrac diffTimes
  in if tTime <= 0 || any (\x -> x < 0 || x > 1) queryTimes
        then xs -- error "bezierCurveAt has a out-of-bound time!"
        else
         if null timesDef || any (\x -> x < 0 || x > 1) queryTimes
         then xs
         else let curvePoints = (map (bezierPoint xs) queryTimes)
                  newTimes = [addUTCTime t (snd top) | t <- diffTimes]
              in zipWith (\t p -> p { pntTime = Just t}) newTimes curvePoints

bezierPoint :: [Point] -> Double -> Point
bezierPoint pnts t   = go pnts
  where
  go [] = error "GPS Package: Can not create a bezier point from an empty list"
  go [p] = p
  go ps = interpolate (go (init ps)) (go (tail ps)) t

-- | Interpolate selected points onto a bezier curve.  Note this gets
-- exponentially more expensive with the length of the segment being
-- transformed - it is not advisable to perform this operation on
-- trail segements with more than ten points!
bezierCurve ::  [Selected (Trail Point)] -> Trail Point
bezierCurve = concatMap (onSelected (bezierCurveAt []) Prelude.id)

-- | Filter out any points that go backward in time (thus must not be
-- valid if this is a trail)
linearTime :: [Point] -> [Point]
linearTime [] = []
linearTime (p:ps) = go (pntTime p) ps
  where
  go _ [] = []
  go t (x:xs) = if pntTime x < t then go t xs else x : go (pntTime x) xs

-- | Return the closest distance between two trails (or Nothing if a
-- trail is empty).  Inefficient implementation:
-- O( (n * m) * log (n * m) )
closestDistance :: Trail Point -> Trail Point -> Maybe Distance
closestDistance as bs = listToMaybe $ L.sort [distance a b | a <- as, b <- bs]

-- | Find the total distance traveled
totalDistance :: [Point] -> Distance
totalDistance as = sum $ zipWith distance as (drop 1 as)

totalTime :: Trail Point -> NominalDiffTime
totalTime [] = 0
totalTime xs@(x:_) = fromMaybe 0 $ liftM2 diffUTCTime (pntTime $ last xs) (pntTime x)

-- | Uses Grahams scan to compute the convex hull of the given points.
-- This operation requires sorting of the points, so don't try it unless
-- you have notably more memory than the list of points will consume.
convexHull :: [Point] -> [Point]
convexHull lst =
        let frst = southMost lst
        in case frst of
                Nothing -> []
                Just f  ->
                     let sorted = L.sortBy (comparing (eastZeroHeading f)) (filter (/= f) lst)
                     in case sorted of
                        (a:b:cs) -> grahamScan (b:a:f:[]) cs
                        cs       -> f : cs
  where
  grahamScan [] _ = []
  grahamScan ps [] = ps
  grahamScan (x:[]) _ = [x]
  grahamScan (p2:p1:ps) (x:xs) =
        case turn p1 p2 x of
                LeftTurn  -> grahamScan (x:p2:p1:ps) xs
                Straight  -> grahamScan (x:p2:p1:ps) xs
                _         -> grahamScan (p1:ps) (x:xs)

eastZeroHeading :: Point -> Point -> Heading
eastZeroHeading s = (`mod'` (2*pi)) . (+ pi/2) . heading s

data Turn = LeftTurn | RightTurn | Straight deriving (Eq, Ord, Show, Read, Enum)

turn :: Point -> Point -> Point -> Turn
turn a b c =
        let h1 = eastZeroHeading a b
            h2 = eastZeroHeading b c
            d  = h2 - h1
        in if d >= 0 && d < pi then LeftTurn else RightTurn

-- | Find the southmost point
southMost :: [Point] -> Maybe Point
southMost []  = Nothing
southMost cs = Just . minimumBy (comparing pntLat) $ cs

---------- COMPOSIT OPERATIONS ---------------
-- These operations are simply implemented using the previously
-- defined functions. They can serve either for concise use for novice
-- users or as instructional examples.
------------------------------------------

-- | Smooth points with rest areas using a bezierCurve.
--
-- Parameters: rest for 1 minute within 30 meters get smoothed
-- in a bezier curve over every 8 points.
smoothRests :: Trail Point -> Trail Point
smoothRests = bezierCurve . refineGrouping (everyNPoints 8) . restLocations 30 60

-- | Smooth every 7 points using a bezier curve
smoothTrail :: Trail Point -> Trail Point
smoothTrail = gSmoothSome 7

-- | Smooth every n points using a bezier curve
gSmoothSome :: Int -> Trail Point -> Trail Point
gSmoothSome n = bezierCurve . everyNPoints n


{-|
Module: Numeric.Morpheus.MatrixReduce
Description: Matrix reduce by column/row operations.
Copyright: (c) Alexander Ignatyev, 2017
License: BSD-3
Stability: experimental
Portability: POSIX
-}

module Numeric.Morpheus.MatrixReduce
(
  columnPredicate
  , rowPredicate
  , columnSum
  , rowSum
  , columnMaxIndex
  , columnMinIndex
  , rowMaxIndex
  , rowMinIndex
)

where

import Numeric.Morpheus.Utils(morpheusLayout)

import Numeric.LinearAlgebra
import Numeric.LinearAlgebra.Devel
import System.IO.Unsafe(unsafePerformIO)
import Foreign
import Foreign.C.Types
import Foreign.Ptr(Ptr)

type Predicate = R -> R -> R
foreign import ccall "wrapper" mkPredicate :: Predicate -> IO (FunPtr Predicate)

{- morpheus_column_predicate -}
foreign import ccall safe "morpheus_column_predicate"
  c_morpheus_column_predicate :: FunPtr Predicate -> CInt -> CInt -> CInt
                              -> Ptr R -> Ptr R -> IO ()

call_morpheus_column_predicate :: FunPtr Predicate
                               -> CInt -> CInt -> CInt -> CInt -> Ptr R
                               -> CInt -> Ptr R
                               -> IO ()
call_morpheus_column_predicate f rows cols xRow xCol matPtr _ vecPtr = do
  let layout = morpheusLayout xCol cols
  c_morpheus_column_predicate f layout rows cols matPtr vecPtr

-- | Scan every column of the given matrix.
-- Predicate takes and accumulator and next value of the column, returns new accumulator.
-- Returns accumulator values for every column.
columnPredicate :: (R -> R -> R) -> Matrix R -> Vector R
columnPredicate f m = unsafePerformIO $ do
  v <- createVector (cols m)
  fpred <- mkPredicate f
  apply m (apply v id) (call_morpheus_column_predicate fpred)
  return v

{- morpheus_row_predicate -}
foreign import ccall safe "morpheus_row_predicate"
  c_morpheus_row_predicate :: FunPtr Predicate -> CInt -> CInt -> CInt
                           -> Ptr R -> Ptr R -> IO ()

call_morpheus_row_predicate :: FunPtr Predicate
                            -> CInt -> CInt -> CInt -> CInt -> Ptr R
                            -> CInt -> Ptr R
                            -> IO ()
call_morpheus_row_predicate f rows cols xRow xCol matPtr _ vecPtr = do
  let layout = morpheusLayout xCol cols
  c_morpheus_row_predicate f layout rows cols matPtr vecPtr

-- | Scan every row of the given matrix.
-- Predicate takes and accumulator and next value of the row, returns new accumulator.
-- Returns accumulator values for every row.
rowPredicate :: (R -> R -> R) -> Matrix R -> Vector R
rowPredicate f m = unsafePerformIO $ do
  v <- createVector (rows m)
  fpred <- mkPredicate f
  apply m (apply v id) (call_morpheus_row_predicate fpred)
  return v

{- morpheus_column_sum -}
foreign import ccall unsafe "morpheus_column_sum"
  c_morpheus_column_sum :: CInt -> CInt -> CInt -> Ptr Double -> Ptr Double -> IO ()

call_morpheus_column_sum :: CInt -> CInt -> CInt -> CInt -> Ptr Double
                         -> CInt -> Ptr Double
                         -> IO ()
call_morpheus_column_sum rows cols xRow xCol matPtr _ vecPtr = do
  let layout = morpheusLayout xCol cols
  c_morpheus_column_sum layout rows cols matPtr vecPtr

-- | Calculates sums of elements of every column of the given matrix
columnSum :: Matrix Double -> Vector Double
columnSum m = unsafePerformIO $ do
    v <- createVector (cols m)
    apply m (apply v id) call_morpheus_column_sum
    return v

{- morpheus_row_sum -}
foreign import ccall unsafe "morpheus_row_sum"
  c_morpheus_row_sum :: CInt -> CInt -> CInt -> Ptr Double -> Ptr Double -> IO ()

call_morpheus_row_sum :: CInt -> CInt -> CInt -> CInt -> Ptr Double
                      -> CInt -> Ptr Double
                      -> IO ()
call_morpheus_row_sum rows cols xRow xCol matPtr _ vecPtr = do
  let layout = morpheusLayout xCol cols
  c_morpheus_row_sum layout rows cols matPtr vecPtr

-- | Calculates sums of elements of every row of the given matrix
rowSum :: Matrix Double -> Vector Double
rowSum m = unsafePerformIO $ do
    v <- createVector (rows m)
    apply m (apply v id) call_morpheus_row_sum
    return v

{- morpheus_column_max_index -}
foreign import ccall unsafe "morpheus_column_max_index"
  c_morpheus_column_max_index :: CInt -> CInt -> CInt -> Ptr Double -> Ptr Double -> Ptr CInt -> IO ()

call_morpheus_column_max_index :: CInt -> CInt -> CInt -> CInt -> Ptr Double
                                  -> CInt -> Ptr Double
                                  -> CInt -> Ptr CInt
                                  -> IO ()
call_morpheus_column_max_index rows cols xRow xCol matPtr _ vecPtr _ idxPtr = do
  let layout = morpheusLayout xCol cols
  c_morpheus_column_max_index layout rows cols matPtr vecPtr idxPtr

-- | Finds maximum values and their indices of every column of the given matrix
columnMaxIndex :: Matrix Double -> (Vector R, Vector I)
columnMaxIndex m = unsafePerformIO $ do
    v <- createVector (cols m)
    i <- createVector (cols m)
    apply m (apply v (apply i id)) call_morpheus_column_max_index
    return (v, i)

{- morpheus_column_min_index -}
foreign import ccall unsafe "morpheus_column_min_index"
  c_morpheus_column_min_index :: CInt -> CInt -> CInt -> Ptr Double -> Ptr Double -> Ptr CInt -> IO ()

call_morpheus_column_min_index :: CInt -> CInt -> CInt -> CInt -> Ptr Double
                                  -> CInt -> Ptr Double
                                  -> CInt -> Ptr CInt
                                  -> IO ()
call_morpheus_column_min_index rows cols xRow xCol matPtr _ vecPtr _ idxPtr = do
  let layout = morpheusLayout xCol cols
  c_morpheus_column_min_index layout rows cols matPtr vecPtr idxPtr

-- | Finds minimum values and their indices of every column of the given matrix
columnMinIndex :: Matrix Double -> (Vector R, Vector I)
columnMinIndex m = unsafePerformIO $ do
    v <- createVector (cols m)
    i <- createVector (cols m)
    apply m (apply v (apply i id)) call_morpheus_column_min_index
    return (v, i)

{- morpheus_row_max_index -}
foreign import ccall unsafe "morpheus_row_max_index"
  c_morpheus_row_max_index :: CInt -> CInt -> CInt -> Ptr Double -> Ptr Double -> Ptr CInt -> IO ()

call_morpheus_row_max_index :: CInt -> CInt -> CInt -> CInt -> Ptr Double
                               -> CInt -> Ptr Double
                               -> CInt -> Ptr CInt
                               -> IO ()
call_morpheus_row_max_index rows cols xRow xCol matPtr _ vecPtr _ idxPtr = do
  let layout = morpheusLayout xCol cols
  c_morpheus_row_max_index layout rows cols matPtr vecPtr idxPtr

-- | Finds maximum values and their indices of every row of the given matrix
rowMaxIndex :: Matrix Double -> (Vector R, Vector I)
rowMaxIndex m = unsafePerformIO $ do
    v <- createVector (rows m)
    i <- createVector (rows m)
    apply m (apply v (apply i id)) call_morpheus_row_max_index
    return (v, i)

{- morpheus_row_min_index -}
foreign import ccall unsafe "morpheus_row_min_index"
  c_morpheus_row_min_index :: CInt -> CInt -> CInt -> Ptr Double -> Ptr Double -> Ptr CInt -> IO ()

call_morpheus_row_min_index :: CInt -> CInt -> CInt -> CInt -> Ptr Double
                               -> CInt -> Ptr Double
                               -> CInt -> Ptr CInt
                               -> IO ()
call_morpheus_row_min_index rows cols xRow xCol matPtr _ vecPtr _ idxPtr = do
  let layout = morpheusLayout xCol cols
  c_morpheus_row_min_index layout rows cols matPtr vecPtr idxPtr

-- | Finds minimum values and their indices of every row of the given matrix
rowMinIndex :: Matrix Double -> (Vector R, Vector I)
rowMinIndex m = unsafePerformIO $ do
    v <- createVector (rows m)
    i <- createVector (rows m)
    apply m (apply v (apply i id)) call_morpheus_row_min_index
    return (v, i)


{-# LANGUAGE ForeignFunctionInterface #-}

-- $ ghc -O2 --make wrappers.hs functions.c

import Numeric.LinearAlgebra
import Data.Packed.Development
import Foreign(Ptr,unsafePerformIO)
import Foreign.C.Types(CInt)

-----------------------------------------------------

main = do
    print $ myDiag $ (3><5) [1..]

-----------------------------------------------------
-- arbitrary data order

foreign import ccall unsafe "c_diag"
    cDiag :: CInt                        -- matrix order
          -> CInt -> CInt -> Ptr Double  -- argument
          -> CInt -> Ptr Double          -- result1
          -> CInt -> CInt -> Ptr Double  -- result2
          -> IO CInt                     -- exit code

myDiag m = unsafePerformIO $ do
    y <- createVector (min r c)
    z <- createMatrix (orderOf m) r c
    app3 (cDiag o) mat m vec y mat z "cDiag"
    return (y,z)
  where r = rows m
        c = cols m
        o = if orderOf m == RowMajor then 1 else 0


{-# LANGUAGE ForeignFunctionInterface #-}
{-|
Module      : Grenade.Layers.Internal.Add
Description : Fast addition functions that call efficient C function
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.Internal.Add (
  addPerChannel
) where

import qualified Data.Vector.Storable        as U (unsafeFromForeignPtr0,
                                                   unsafeToForeignPtr0)

import           Foreign                     (mallocForeignPtrArray, withForeignPtr)
import           Foreign.Ptr                 (Ptr)
import           Numeric.LinearAlgebra       (Matrix, Vector, flatten)
import qualified Numeric.LinearAlgebra.Devel as U
import           System.IO.Unsafe            (unsafePerformIO)

import           Grenade.Types

-- | Add the nth element to a vector to every pixel in the nth channel of a matrix.
--   It assumes that the size of the vector and the number of channels in the matrix
--   are both equal to channels.
addPerChannel :: Int -> Int -> Int -> Matrix RealNum -> Vector RealNum -> Matrix RealNum
addPerChannel channels rows cols m b
  = let outMatSize      = rows * cols * channels
        vec             = flatten m
    in unsafePerformIO $ do
      outPtr        <- mallocForeignPtrArray outMatSize
      let (inPtr, _) = U.unsafeToForeignPtr0 vec
          (bPtr, _)  = U.unsafeToForeignPtr0 b

      withForeignPtr inPtr $ \inPtr' ->
        withForeignPtr bPtr $ \bPtr' ->
          withForeignPtr outPtr $ \outPtr' ->
            add_per_channel_cpu inPtr' channels rows cols bPtr' outPtr'

      let matVec = U.unsafeFromForeignPtr0 outPtr outMatSize
      return (U.matrixFromVector U.RowMajor (rows * channels) cols matVec)
{-# INLINE addPerChannel #-}

foreign import ccall unsafe
    add_per_channel_cpu
      :: Ptr RealNum -> Int -> Int -> Int -> Ptr RealNum -> Ptr RealNum -> IO ()