--Haskell Quil compiler
--Copyright Laurence Emms 2018

module Lib (compile) where

import Data.Complex
import Data.List
import Register
import Instruction
import ClassicalCircuit
import QuantumCircuit
import Util

testRXCircuit :: (Floating a, Show a, Ord a) => Circuit a
testRXCircuit = let parameters@[Left a] = [Left (MetaQubitRegister "a")]
                    instructions = [(RX (ComplexConstant (5.0 :+ 10.0)) a)]
                in Circuit "TESTRX" parameters instructions

hadamardRotateXYZ :: (Floating a, Show a, Ord a) => Complex a -> Circuit a
hadamardRotateXYZ rotation
    = let parameters@[Left a, Left b, Right r, Left z] = [Left (MetaQubitRegister "a"), Left (MetaQubitRegister "b"), Right (MetaRegister "r"), Left (MetaQubitRegister "z")]
          instructions = [(Hadamard a),
                          (MeasureOut a r)] ++
                          (ifC "HROTXYZTHEN0" "HROTXYZEND0"
                               r
                               [(RX (ComplexConstant rotation) z)]
                               ([(Hadamard b),
                                 (MeasureOut b r)] ++
                                (ifC "HROTXYZTHEN1" "HROTXYZEND1"
                                     r
                                     [(RY (ComplexConstant rotation) z)]
                                     [(RZ (ComplexConstant rotation) z)])))
      in Circuit "HROTXYZ" parameters instructions

compile :: IO ()
compile = putStrLn "Compiling quantum executable" >>
          --Test registers
          putStrLn (show $ QubitRegister 10) >>
          putStrLn (show $ Register 5) >>
          putStrLn (show $ ComplexConstant (6.0 :+ 7.0)) >>
          putStrLn (show $ ComplexConstant ((-8.0) :+ 1.0)) >>
          putStrLn (show $ ComplexConstant (2.0 :+ (-3.0))) >>
          putStrLn (show $ ComplexConstant ((-4.0) :+ (-6.0))) >>
          --Test instructions
          putStrLn (show $ CNot (QubitRegister 0) (QubitRegister 1)) >>
          putStrLn (show $ PSwap (ComplexConstant (5.0 :+ (-3.2))) (QubitRegister 0) (QubitRegister 1)) >>
          putStrLn (show $ Measure (QubitRegister 4)) >>
          putStrLn (show $ MeasureOut (QubitRegister 4) (Register 5)) >>
          --Test circuits
          putStrLn (show $ DefCircuit bell) >>
          putStrLn (show $ CallCircuit bell [Left (QubitRegister 5), Left (QubitRegister 3)]) >>
          putStrLn (show $ DefCircuit testRXCircuit) >>
          putStrLn (show $ CallCircuit testRXCircuit [Left (QubitRegister 1)]) >>
          putStrLn (show $ DefCircuit xor) >>
          putStrLn (show $ CallCircuit xor [Right (Register 0), Right (Register 1), Right (Register 2)]) >>
          putStrLn (show $ DefCircuit halfAdder) >>
          putStrLn (show $ CallCircuit halfAdder [Right (Register 0), Right (Register 1), Right (Register 2), Right (Register 3)]) >>
          putStrLn (show $ DefCircuit adder) >>
          putStrLn (show $ CallCircuit adder [Right (Register 0), Right (Register 1), Right (Register 2), Right (Register 3), Right (Register 4)]) >>
          --Load the integer 53 into registers [0-31]
          putStrLn (intercalate "\n" (fmap show (loadIntToRegisters 53 0))) >>
          --Load the integer 18 into registers [32-63]
          putStrLn (intercalate "\n" (fmap show (loadIntToRegisters 18 32))) >>
          --Add [0-31] + [32-63] into [64-95]
          let carry = 128
              temp = 129
          in putStrLn (intercalate "\n" (fmap show (addRegisters carry temp 0 32 64))) >>
          putStrLn (show $ DefCircuit (hadamardRotateXYZ (0.70710678118 :+ 0.70710678118))) >>
          putStrLn (show $ CallCircuit (hadamardRotateXYZ (0.70710678118 :+ 0.70710678118)) [Left (QubitRegister 0), Left (QubitRegister 1), Right (Register 0), Left (QubitRegister 2)])


{-# LANGUAGE Arrows              #-}
{-# LANGUAGE GADTs               #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies        #-}
{-# OPTIONS_GHC -Wall #-}
{-# OPTIONS -fplugin=Overloaded -fplugin-opt=Overloaded:Categories #-}
module Main where

import Control.Monad      (when)
import Data.Word          (Word64)
import Numeric            (showFFloat)
import System.Environment (getArgs)
import Data.List (intercalate)

import qualified Control.Category
import qualified Numeric.LinearAlgebra  as LA
import qualified System.Random.SplitMix as SM

import Overloaded.Categories
import VectorSpace

-- | A Function which computes value and derivative at the point.
newtype AD a b = AD (a -> (b, L a b))

instance Category AD where
    id = AD (\x -> (x, L id))

    AD g . AD f = AD $ \a ->
        let (b, L f') = f a
            (c, L g') = g b
        in (c, L (g' . f'))

instance CategoryWith1 AD where
    type Terminal AD = ()

    terminal = AD (const ((), terminal))

instance CartesianCategory AD where
    type Product AD = (,)

    proj1 = AD (\x -> (fst x, proj1))
    proj2 = AD (\x -> (snd x, proj2))

    fanout (AD f) (AD g) = AD $ \a ->
        let (b, f') = f a
            (c, g') = g a
        in ((b, c), fanout f' g')

instance GeneralizedElement AD where
    type Object AD a = a

    konst x = AD (\_ -> (x, L $ \_ -> LZ))

ladd :: LinMap r (a, a) -> LinMap r a
ladd (LH f g) = LA f g
ladd (LV f g) = LV (ladd f) (ladd g)
ladd (LA a b) = LA (ladd a) (ladd b)
ladd LZ       = LZ
ladd (LD k)   = LV (LD k) (LD k)

lmult :: Double -> Double -> LinMap r (a, a) -> LinMap r a
lmult x y (LH f g) = LA (lmul y f) (lmul x g)
lmult x y (LV f g) = LV (lmult x y f) (lmult x y g)
lmult x y (LA f g) = LA (lmult x y f) (lmult x y g)
lmult _ _ LZ       = LZ
lmult x y (LD k)   = LV (LD (k * y)) (LD (k * x))

plus :: AD (Double, Double) Double
plus = AD $ \(x,y) -> (x + y, L ladd)

minus :: AD (Double, Double) Double
minus = AD $ \(x,y) -> (x - y, L $ lmult (-1) 1)

mult :: AD (Double, Double) Double
mult = AD $ \(x,y) -> (x * y, L $ lmult x y)

scale :: Double -> AD Double Double
scale k = AD $ \x -> (k * x, linear k)

evaluateAD :: (HasDim a, HasDim b) => AD a b -> a -> (b, LA.Matrix Double)
evaluateAD (AD f) x = let (y, f') = f x in (y, evalL f')

-------------------------------------------------------------------------------
-- Simple examples
-------------------------------------------------------------------------------

ex1 :: AD Double Double
ex1 = plus %% fanout identity identity

ex2 :: AD Double Double
ex2 = mult %% fanout identity identity

-------------------------------------------------------------------------------
-- Quadratic function
-------------------------------------------------------------------------------

quad :: AD (Double, Double) Double
quad = proc (x, y) -> do
    x2  <- mult    -< (x, x)
    y2  <- mult    -< (y, y)
    tmp <- plus    -< (x2, y2)
    z   <- konst 5 -< ()
    plus -< (tmp, z)

-------------------------------------------------------------------------------
-- Newton
-------------------------------------------------------------------------------

findZero :: AD Double Double -> Double -> [Double]
findZero f x0 = take 10 results
  where
    results = iterate go x0

    go :: Double -> Double
    go x =
        let (y, m) = evaluateAD f x
            [[y']] = LA.toLists m
        in x - gamma * (y / y')

    gamma = 0.1

-------------------------------------------------------------------------------
-- Gradient descent
-------------------------------------------------------------------------------

gradDesc :: forall a. VectorSpace a => AD a Double -> a -> [a]
gradDesc f = iterate go where
    go :: a -> a
    go x =
        let (_, m) = evaluateAD f x
            [grad] = LA.toLists $ LA.tr $ LA.scale gamma m

        in fromVector $ zipWith (-) (toVector x) grad

    gamma = 0.1

-------------------------------------------------------------------------------
-- Random
-------------------------------------------------------------------------------

randomDoubles :: Word64 -> [Double]
randomDoubles seed = go (SM.mkSMGen seed) where
    go g = let (d, g') = SM.nextDouble g in d : go g'

-------------------------------------------------------------------------------
-- Dot
-------------------------------------------------------------------------------

class VectorSpace' a where
    sumN :: AD a Double
    multN :: AD (a, a) a

instance (VectorSpace' a, VectorSpace' b) => VectorSpace' (a, b) where
    sumN = proc (x, y) -> do
        x' <- sumN -< x
        y' <- sumN -< y
        plus -< (x', y')

    multN = proc ((x1, x2), (y1, y2)) -> do
        z1 <- multN -< (x1, y1)
        z2 <- multN -< (x2, y2)
        identity -< (z1, z2)

instance VectorSpace' Double where
    sumN  = identity
    multN = mult

dot :: VectorSpace' a => AD (a, a) Double
dot = sumN %% multN

-------------------------------------------------------------------------------
-- ML stuff
-------------------------------------------------------------------------------

tanhAD :: AD Double Double
tanhAD = AD $ \x ->
    let y = tanh x
    in (y, linear (1 - y * y))

sigmoidAD :: AD Double Double
sigmoidAD = AD $ \x ->
    let y = 1 / (1 + exp (- x))
    in (x, linear (y * (1 - y)))

-- | weights for 2x1 connection. Two weights and bias.
type Weights' = ((Double, Double), Double)

-- | Two internal neurons, and final output
type Weights = ((Weights', Weights'), Weights')

startWeights :: Weights
startWeights = fromVector $ randomDoubles 1337

neuron :: AD (Weights', (Double, Double)) Double
neuron = proc ((ws, bias), i) -> do
    o <- dot -< (ws, i)
    tanhAD %% plus -< (o, bias)

network :: AD (Weights, (Double, Double)) Double
network = proc (((w1, w2), w3), xy) -> do
    u <- neuron  -< (w1, xy)
    v <- neuron  -< (w2, xy)
    neuron -< (w3, (u, v))

networkError :: AD Weights Double
networkError = proc ws -> do
    -- xor!
    s1 <- ex 1 1 0 -< ws
    s2 <- ex 0 0 0 -< ws
    s3 <- ex 1 0 1 -< ws
    s4 <- ex 0 1 1 -< ws

    sumN -< ((s1,s2), (s3, s4))
  where
    ex :: Double -> Double -> Double -> AD Weights Double
    ex x y z = proc ws -> do
         x1 <- konst x -< ()
         y1 <- konst y -< ()
         e1 <- konst z -< ()
         a1 <- network -< (ws, (x1, y1))
         r1 <- minus   -< (e1, a1)
         mult -< (r1, r1)

train :: Weights
train = gradDesc networkError startWeights !! 500

-------------------------------------------------------------------------------
-- Main
-------------------------------------------------------------------------------

main :: IO ()
main = do
    putStrLn $ "quad (2,3) = " ++ show (evaluateAD quad (2,3))
    putStrLn $ "gradDesc quad (2,3) = " ++ show (gradDesc quad (2,3) !! 30)

    print $ evaluateAD tanhAD 1
    print $ evaluateAD sigmoidAD 1

    putStrLn "Training the net (for xor)"
    let ws = train
    putStrLn $ "Parameters = " ++ show (toVector ws)
    putStrLn $ "Error = " ++ show (fst $ evaluateAD networkError ws)
    let example xy =
          putStrLn $ "eval " ++ show xy ++ " = " ++ showFFloat (Just 2) (fst $ evaluateAD network (ws, xy)) ""

    example (0, 0)
    example (0, 1)
    example (1, 0)
    example (1, 1)

    args <- getArgs
    when ("plot" `elem` args) $ do
        putStrLn "Outputting plot data: datafile.dat"

        let n = 20 :: Int
        let points = [ fromIntegral x / fromIntegral n | x <- [0..n] ] :: [Double]

        let output :: String
            output = unlines
                [ intercalate "\t"
                    [ show x
                    , show y
                    , show (fst (evaluateAD network (ws, (x, y))))
                    ]
                | x <- points
                , y <- points
                ]

        writeFile "datafile.dat" output


{-# LANGUAGE CPP                    #-}
{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE DataKinds              #-}
{-# LANGUAGE DefaultSignatures      #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE ParallelListComp       #-}
{-# LANGUAGE PartialTypeSignatures  #-}
{-# LANGUAGE PatternGuards          #-}
{-# LANGUAGE Rank2Types             #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE TypeApplications       #-}
{-# LANGUAGE TypeFamilies           #-}
{-# LANGUAGE TypeSynonymInstances   #-}
{-# LANGUAGE UndecidableInstances   #-} -- See below

-- experiment
-- {-# LANGUAGE PartialTypeSignatures #-}

{-# OPTIONS_GHC -Wall #-}
-- {-# OPTIONS_GHC -fno-warn-unused-imports #-} -- TEMP

-- #define TESTING

#ifdef TESTING
{-# OPTIONS_GHC -fno-warn-unused-binds   #-} -- TEMP
#endif

-- #define GenericPowFFT

----------------------------------------------------------------------
-- | Generic FFT
----------------------------------------------------------------------

module ConCat.FFT
  ( dft, FFT(..), DFTTy, genericFft, GFFT
  -- -- Temporary while debugging
  , twiddle, twiddles, omega, cis
  , o8sq
  ) where

import Prelude hiding (zipWith)

import Data.Complex
-- import Control.Applicative (liftA2)
import GHC.Generics hiding (C,S)

import Data.Key (Zip(..))
import Data.Pointed

#ifdef TESTING
import Data.Foldable (toList)
-- import Test.QuickCheck (quickCheck)
import Test.QuickCheck.All (quickCheckAll)
import ShapedTypes.ApproxEq
#endif

import ConCat.Misc (Unop,transpose,C,inGeneric1,inComp)
import ConCat.Sized
import ConCat.Scan (LScan,lproducts)
import ConCat.Pair
import ConCat.Free.LinearRow (($*))

{--------------------------------------------------------------------
    DFT
--------------------------------------------------------------------}

type AS  h = (Pointed h, Zip h, LScan h)
type ASZ h = (AS h, Sized h)

-- To resolve: Zip vs Applicative. traverse/transpose needs Applicative.

dft :: forall f a. (ASZ f, Foldable f, RealFloat a) => Unop (f (Complex a))
dft xs = omegas (size @f) $* xs
{-# INLINE dft #-}

omegas :: (AS f, AS g, RealFloat a) => Int -> g (f (Complex a))
omegas = fmap powers . powers . omega
-- omegas n = powers <$> powers (omega n)
-- omegas n = powers <$> powers (exp (- i * 2 * pi / fromIntegral n))
{-# INLINE omegas #-}

omegas' :: forall f g a. (ASZ f, ASZ g, RealFloat a) => g (f (Complex a))
omegas' = fmap powers (powers (cis (- 2 * pi / fromIntegral (size @(g :.: f)))))

-- i :: Num a => Complex a
-- i = 0 :+ 1

omega :: RealFloat a => Int -> Complex a
omega n = cis (- 2 * pi / fromIntegral n)
-- omega n = exp (0 :+ (- 2 * pi / fromIntegral n))
-- omega n = exp (- 2 * (0:+1) * pi / fromIntegral n)
{-# INLINE omega #-}

{--------------------------------------------------------------------
    FFT
--------------------------------------------------------------------}

type DFTTy f = forall a. RealFloat a => f (Complex a) -> FFO f (Complex a)

class FFT f where
  type FFO f :: * -> *
  fft :: DFTTy f
  -- default fft :: ( Generic1 f, Generic1 (FFO f)
  --                , FFT (Rep1 f), FFO (Rep1 f) ~ Rep1 (FFO f) ) => DFTTy f
  -- fft = genericFft
  -- Temporary hack to avoid newtype-like representation.
  fftDummy :: f a
  fftDummy = undefined

-- TODO: Eliminate FFO, in favor of fft :: Unop (f (Complex a)).
-- Use dft as spec.

twiddle :: forall g f a. (ASZ g, ASZ f, RealFloat a) => Unop (g (f (Complex a)))
-- twiddle = (zipWith.zipWith) (*) twiddles
twiddle = (zipWith.zipWith) (*) omegas'
-- twiddle = (zipWith.zipWith) (*) (omegas (size @(g :.: f)))
{-# INLINE twiddle #-}

twiddles :: forall g f a. (ASZ g, ASZ f, RealFloat a) => g (f (Complex a))
twiddles = omegas (size @(g :.: f))
{-# INLINE twiddles #-}

o8sq :: C
o8sq = omega (8 :: Int) ^ (2 :: Int)

-- Powers of x, starting x^0. Uses 'LScan' for log parallel time
powers :: (LScan f, Pointed f, Num a) => a -> f a
powers = fst . lproducts . point
{-# INLINE powers #-}

-- TODO: Consolidate with powers in TreeTest and rename sensibly. Maybe use
-- "In" and "Ex" suffixes to distinguish inclusive and exclusive cases.

{--------------------------------------------------------------------
    Generic support
--------------------------------------------------------------------}

instance FFT Par1 where
  type FFO Par1 = Par1
  fft = id

#if 0
inTranspose :: (Traversable f', Traversable g, Applicative g', Applicative f)
            => (f (g a) -> f' (g' a)) -> g (f a) -> g' (f' a)
inTranspose = transpose <-- transpose

ffts' :: ( FFT g, Traversable f, Traversable g
         , Applicative (FFO g), Applicative f, RealFloat a) =>
     g (f (Complex a)) -> FFO g (f (Complex a))
ffts' = transpose . fmap fft . transpose
#endif

#if 0

transpose :: g (f C)     -> f (g C)
fmap fft  :: f (g C)     -> f (FFO g C)
transpose :: f (FFO g C) -> FFO g (f C)

#endif

instance ( Zip f,  Traversable f , Traversable g
         , Applicative f, Applicative (FFO f), Applicative (FFO g), Zip (FFO g)
         , Pointed f, Traversable (FFO g), Pointed (FFO g)
         , FFT f, FFT g, LScan f, LScan (FFO g), Sized f, Sized (FFO g) )
      => FFT (g :.: f) where
  type FFO (g :.: f) = FFO f :.: FFO g
  fft = inComp (traverse fft . twiddle . traverse fft . transpose)
  -- fft = inComp (ffts' . transpose . twiddle . ffts')
  {-# INLINE fft #-}

#if 0
  fft = Comp1 . transpose . fmap fft . twiddle . transpose . fmap fft . transpose . unComp1
  fft = Comp1 . traverse fft . twiddle . traverse fft . transpose . unComp1

-- Types in fft for (g :. f):
  unComp1   :: (g :. f) a -> g  (f  a)
  transpose :: g  (f  a)  -> f  (g  a)
  fmap fft  :: f  (g  a)  -> f  (g' a)
  transpose :: f  (g' a)  -> g' (f  a)
  twiddle   :: g' (f  a)  -> g' (f  a)
  fmap fft  :: g' (f  a)  -> g' (f' a)
  transpose :: g' (f' a)  -> f' (g' a)
  Comp1     :: f' (g' a)  -> (f' :. g') a
#endif

#if 0

--   fft = inComp (ffts' . transpose . twiddle . ffts')

ffts'     :: g (f C)     -> FFO g (f C)
twiddle   :: FFO g (f C) -> FFO g (f C)
transpose :: FFO g (f C) -> f (FFO g C)
ffts'     :: f (FFO g C) -> FFO f (FFO g C)

#endif

-- -- Generalization of 'dft' to traversables. Note that liftA2 should
-- -- work zippily (unlike with lists).
-- dftT :: forall f a. (ASZ f, Traversable f, RealFloat a)
--      => Unop (f (Complex a))
-- dftT xs = out <$> indices
--  where
--    out k   = sum (liftA2 (\ n x -> x * ok^n) indices xs) where ok = om ^ k
--    indices = fst iota :: f Int
--    om      = omega (size @f)
-- {-# INLINE dftT #-}

-- | Generic FFT
genericFft :: ( Generic1 f, Generic1 (FFO f)
              , FFT (Rep1 f), FFO (Rep1 f) ~ Rep1 (FFO f) ) => DFTTy f
genericFft = inGeneric1 fft

type GFFT f = (Generic1 f, Generic1 (FFO f), FFT (Rep1 f), FFO (Rep1 f) ~ Rep1 (FFO f))

#define GenericFFT(f,g) instance GFFT (f) => FFT (f) where { type FFO (f) = (g); fft = genericFft }

-- #define GenericFFT(f,g) instance GFFT (f) => FFT (f) where { type FFO (f) = g; INLINE }

-- TODO: Replace Applicative with Zippable.
-- Can't, because Traversable needs Applicative.

-- Perhaps dftT isn't very useful. Its result and argument types match, unlike fft.

{--------------------------------------------------------------------
    Specialized FFT instances.
--------------------------------------------------------------------}

-- I put the specific instances here in order to avoid an import loop between
-- the LPow and RPow modules. I'd still like to find an elegant FFT that maps f
-- to f, and then move the instances to RPow and LPow.

-- Radix 2 butterfly
instance FFT Pair where
  type FFO Pair = Pair
  -- bogus "non-exhaustive" warning in ghc 8.0.2
  -- fft (a :# b) = (a + b) :# (a - b)
  fft = \ (a :# b) -> (a + b) :# (a - b)
  -- fft = dft
  {-# INLINE fft #-}

#ifdef TESTING

#if 0
twiddles :: (ASZ g, ASZ f, RealFloat a) => g (f (Complex a))

as :: f C
(<.> as) :: f C -> C
twiddles :: f (f C)
(<.> as) <$> twiddles :: f C
#endif

-- -- Binary dot product
-- infixl 7 <.>
-- (<.>) :: (Foldable f, Applicative f, Num a) => f a -> f a -> a
-- u <.> v = sum (liftA2 (*) u v)

{--------------------------------------------------------------------
    Simple, quadratic DFT (for specification & testing)
--------------------------------------------------------------------}

-- Adapted from Dave's definition
dftL :: RealFloat a => Unop [Complex a]
dftL xs = [ sum [ x * ok^n | x <- xs | n <- [0 :: Int ..] ]
          | k <- [0 .. length xs - 1], let ok = om ^ k ]
 where
   om = omega (length xs)

{--------------------------------------------------------------------
    Tests
--------------------------------------------------------------------}

-- > powers 2 :: LTree N2 Int
-- B (B (L ((1 :# 2) :# (4 :# 8))))
-- > powers 2 :: LTree N3 Int
-- B (B (B (L (((1 :# 2) :# (4 :# 8)) :# ((16 :# 32) :# (64 :# 128))))))

fftl :: (FFT f, Foldable (FFO f), RealFloat a) => f (Complex a) -> [Complex a]
fftl = toList . fft

type LTree = L.Pow Pair
type RTree = R.Pow Pair

type LC n = LTree n C
type RC n = RTree n C

p1 :: Pair C
p1 = 1 :# 0

tw1 :: LTree N1 (Pair C)
tw1 = twiddles

tw2 :: RTree N2 (Pair C)
tw2 = twiddles

tw3 :: RTree N2 (RTree N2 C)
tw3 = twiddles

tw3' :: [[C]]
tw3' = toList (toList <$> tw3)

-- Adapted from Dave's testing

-- test :: (FFT f, Foldable f, Foldable (FFO f)) => f C -> IO ()
-- test fx =
--   do ps "\nTesting input" xs
--      ps "Expected output" (dftL xs)
--      ps "Actual output  " (toList (fft fx))
--  where
--    ps label z = putStrLn (label ++ ": " ++ show z)
--    xs = toList fx

#if 0
t0 :: LC N0
t0 = L.fromList [1]

t1 :: LC N1
t1 = L.fromList [1, 0]

t2s :: [LC N2]
t2s = L.fromList <$>
        [ [1,  0,  0,  0]  -- Delta
        , [1,  1,  1,  1]  -- Constant
        , [1, -1,  1, -1]  -- Nyquist
        , [1,  0, -1,  0]  -- Fundamental
        , [0,  1,  0, -1]  -- Fundamental w/ 90-deg. phase lag
       ]

tests :: IO ()
tests = do test p1
           test t0
           test t1
           mapM_ test t2s
#endif

infix 4 ===
(===) :: Eq b => (a -> b) -> (a -> b) -> a -> Bool
(f === g) x = f x == g x

infix 4 =~=
(=~=) :: ApproxEq b => (a -> b) -> (a -> b) -> a -> Bool
(f =~= g) x = f x =~ g x

fftIsDftL :: (FFT f, Foldable f, Foldable (FFO f), RealFloat a, ApproxEq a) =>
             f (Complex a) -> Bool
fftIsDftL = toList . fft =~= dftL . toList

dftTIsDftL :: (ASZ f, Traversable f, RealFloat a, ApproxEq a) =>
              f (Complex a) -> Bool
dftTIsDftL = toList . dftT =~= dftL . toList

dftIsDftL :: (ASZ f, Foldable f, RealFloat a, ApproxEq a) =>
             f (Complex a) -> Bool
dftIsDftL = toList . dft =~= dftL . toList

-- -- TEMP:
-- dftDft :: (ASZ f, Traversable f, RealFloat a, ApproxEq a) =>
--           f (Complex a) -> ([Complex a], [Complex a])
-- dftDft xs = (toList . dft $ xs, dftL . toList $ xs)

{--------------------------------------------------------------------
    Properties to test
--------------------------------------------------------------------}

transposeTwiddleCommutes :: (ASZ g, Traversable g, ASZ f, (ApproxEq (f (g C))))
                         => g (f C) -> Bool
transposeTwiddleCommutes =
 twiddle . transpose =~= transpose . twiddle

prop_transposeTwiddle_L3P :: LTree N3 (Pair C) -> Bool
prop_transposeTwiddle_L3P = transposeTwiddleCommutes

prop_transposeTwiddle_R3P :: RTree N3 (Pair C) -> Bool
prop_transposeTwiddle_R3P = transposeTwiddleCommutes

-- dft tests fail. Hm!

-- prop_dft_R3 :: RTree N3 C -> Bool
-- prop_dft_R3 = dftIsDftL

-- prop_dft_L3 :: LTree N3 C -> Bool
-- prop_dft_L3 = dftIsDftL

prop_dftT_p :: Pair C -> Bool
prop_dftT_p = dftTIsDftL

prop_dftT_L3 :: LTree N3 C -> Bool
prop_dftT_L3 = dftTIsDftL

prop_dftT_R3 :: RTree N3 C -> Bool
prop_dftT_R3 = dftTIsDftL

prop_fft_p :: Pair C -> Bool
prop_fft_p = fftIsDftL

prop_fft_L1 :: LTree N1 C -> Bool
prop_fft_L1 = fftIsDftL

prop_fft_L2 :: LTree N2 C -> Bool
prop_fft_L2 = fftIsDftL

prop_fft_L3 :: LTree N3 C -> Bool
prop_fft_L3 = fftIsDftL

prop_fft_L4 :: LTree N4 C -> Bool
prop_fft_L4 = fftIsDftL

prop_fft_R1 :: RTree N1 C -> Bool
prop_fft_R1 = fftIsDftL

prop_fft_R2 :: RTree N2 C -> Bool
prop_fft_R2 = fftIsDftL

prop_fft_R3 :: RTree N3 C -> Bool
prop_fft_R3 = fftIsDftL

prop_fft_R4 :: RTree N4 C -> Bool
prop_fft_R4 = fftIsDftL

-- TH oddity
return []

runTests :: IO Bool
runTests = $quickCheckAll

-- end of tests
#endif


{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash                  #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeInType                 #-}
{-# LANGUAGE UnboxedTuples              #-}
{-# LANGUAGE UndecidableInstances       #-}
{-# LANGUAGE ViewPatterns               #-}
{-# OPTIONS_GHC -fno-warn-orphans       #-}
module Numeric.Quaternion.Internal.QDouble
    ( QDouble, Quater (..)
    ) where

import qualified Control.Monad.ST                     as ST
import           Data.Coerce                          (coerce)
import           Numeric.Basics
import           Numeric.DataFrame.Internal.PrimArray
import qualified Numeric.DataFrame.ST                 as ST
import           Numeric.DataFrame.Type
import           Numeric.Quaternion.Internal
import           Numeric.Vector.Internal

type QDouble = Quater Double

deriving instance PrimBytes (Quater Double)
deriving instance PrimArray Double (Quater Double)

instance Quaternion Double where
    newtype Quater Double = QDouble (Vector Double 4)
    {-# INLINE packQ #-}
    packQ = coerce (vec4 :: Double -> Double -> Double -> Double -> Vector Double 4)
    {-# INLINE unpackQ# #-}
    unpackQ# = coerce (unpackV4# :: Vector Double 4 -> (# Double, Double, Double, Double #))
    {-# INLINE fromVecNum #-}
    fromVecNum (unpackV3# -> (# x, y, z #))  = packQ x y z
    {-# INLINE fromVec4 #-}
    fromVec4 = coerce
    {-# INLINE toVec4 #-}
    toVec4 = coerce
    {-# INLINE square #-}
    square (unpackQ# -> (# x, y, z, w #)) = (x * x) + (y * y) + (z * z) + (w * w)
    {-# INLINE im #-}
    im (unpackQ# -> (# x, y, z, _ #)) = packQ x y z 0.0
    {-# INLINE re #-}
    re (unpackQ# -> (# _, _, _, w #)) = packQ 0 0 0 w
    {-# INLINE imVec #-}
    imVec (unpackQ# -> (# x, y, z, _ #)) = vec3 x y z
    {-# INLINE taker #-}
    taker (unpackQ# -> (# _, _, _, w #)) = w
    {-# INLINE takei #-}
    takei (unpackQ# -> (# x, _, _, _ #)) = x
    {-# INLINE takej #-}
    takej (unpackQ# -> (# _, y, _, _ #)) = y
    {-# INLINE takek #-}
    takek (unpackQ# -> (# _, _, z, _ #)) = z
    {-# INLINE conjugate #-}
    conjugate (unpackQ# -> (# x, y, z, w #))
      = packQ (negate x) (negate y) (negate z) w
    {-# INLINE rotScale #-}
    rotScale (unpackQ# -> (# i, j, k, t #))
             (unpackV3# -> (# x, y, z #))
      = let l = t*t - i*i - j*j - k*k
            d = 2.0 * ( i*x + j*y + k*z)
            t2 = t * 2.0
        in vec3 (l*x + d*i + t2 * (z*j - y*k))
                (l*y + d*j + t2 * (x*k - z*i))
                (l*z + d*k + t2 * (y*i - x*j))
    {-# INLINE getRotScale #-}
    getRotScale a b = case (# unpackV3# a, unpackV3# b #) of
      (# _, (# 0, 0, 0 #) #) -> packQ 0 0 0 0
      (# (# 0, 0, 0 #), _ #) -> let x = (1 / 0 :: Double) in packQ x x x x
      (# (# a1, a2, a3 #), (# b1, b2, b3 #) #) ->
        let ma = sqrt (a1*a1 + a2*a2 + a3*a3)
            mb = sqrt (b1*b1 + b2*b2 + b3*b3)
            d  = a1*b1 + a2*b2 + a3*b3
            c  = sqrt (ma*mb + d)
            ma2 = ma * 1.4142135623730951 -- sqrt 2.0
            r  = recip (ma2 * c)
            c' = sqrt (mb / ma) -- ratio of a and b for corner cases
            r' = recip (sqrt ( negate (a1*b1 + a2*b2) ))
        in case unpackV3# (cross a b) of
          (# 0, 0, 0 #)
              -- if a and b face the same direction, q is fully real
            | d >= 0       -> packQ 0 0 0 c'
              -- if a and b face opposite directions, find an orthogonal vector
              -- prerequisites: w == 0  and  a·(x,y,z) == 0
              -- corner cases: only one vector component is non-zero
            | b1 == 0      -> packQ c' 0 0 0
              -- otherwise set the last component to zero,
              -- and get an orthogonal vector in 2D.
            | otherwise    -> packQ (-b2*r') (b1*r') 0 0
              -- NB: here we have some precision troubles
              --     when a and b are close to parallel and opposite.
          (# t1, t2, t3 #) -> packQ (t1 * r) (t2 * r) (t3 * r) (c / ma2)
    {-# INLINE axisRotation #-}
    axisRotation v a = case unpackV3# v of
      (# 0, 0, 0 #) -> packQ 0 0 0 (negateUnless (abs a < M_PI) 1)
      (# x, y, z #) ->
        let c = cos (a * 0.5)
            s = sin (a * 0.5)
                / sqrt (x*x + y*y + z*z)
        in packQ (x * s) (y * s) (z * s) c
    {-# INLINE qArg #-}
    qArg (unpackQ# -> (# x, y, z, w #)) = 2 * atan2 (sqrt (x*x + y*y + z*z)) w
    {-# INLINE fromMatrix33 #-}
    fromMatrix33 m = fromM 1
      (ix# 0# m) (ix# 1# m) (ix# 2# m)
      (ix# 3# m) (ix# 4# m) (ix# 5# m)
      (ix# 6# m) (ix# 7# m) (ix# 8# m)

    {-# INLINE fromMatrix44 #-}
    fromMatrix44 m = fromM (ix# 15# m)
      (ix# 0# m) (ix# 1# m) (ix# 2# m)
      (ix# 4# m) (ix# 5# m) (ix# 6# m)
      (ix# 8# m) (ix# 9# m) (ix# 10# m)

    {-# INLINE toMatrix33 #-}
    toMatrix33 (unpackQ# -> (# 0.0, 0.0, 0.0, w #))
      = let x = w * w
            f 0 = (# 3 :: Int , x #)
            f k = (# k-1, 0 #)
        in case gen# (CumulDims [9,3,1]) f 0 of
            (# _, m #) -> m -- diag (scalar (w * w))
    toMatrix33 (unpackQ# -> (# x', y', z', w' #)) =
      let x = scalar x'
          y = scalar y'
          z = scalar z'
          w = scalar w'
          x2 = x * x
          y2 = y * y
          z2 = z * z
          w2 = w * w
          l2 = x2 + y2 + z2 + w2
      in ST.runST $ do
        df <- ST.newDataFrame
        ST.writeDataFrameOff df 0 $ l2 - 2*(z2 + y2)
        ST.writeDataFrameOff df 1 $ 2*(x*y + z*w)
        ST.writeDataFrameOff df 2 $ 2*(x*z - y*w)
        ST.writeDataFrameOff df 3 $ 2*(x*y - z*w)
        ST.writeDataFrameOff df 4 $ l2 - 2*(z2 + x2)
        ST.writeDataFrameOff df 5 $ 2*(y*z + x*w)
        ST.writeDataFrameOff df 6 $ 2*(x*z + y*w)
        ST.writeDataFrameOff df 7 $ 2*(y*z - x*w)
        ST.writeDataFrameOff df 8 $ l2 - 2*(y2 + x2)
        ST.unsafeFreezeDataFrame df
    {-# INLINE toMatrix44 #-}
    toMatrix44 (unpackQ# -> (# 0.0, 0.0, 0.0, w #)) = ST.runST $ do
      df <- ST.newDataFrame
      mapM_ (flip (ST.writeDataFrameOff df) 0) [0..15]
      let w2 = scalar (w * w)
      ST.writeDataFrameOff df 0 w2
      ST.writeDataFrameOff df 5 w2
      ST.writeDataFrameOff df 10 w2
      ST.writeDataFrameOff df 15 1
      ST.unsafeFreezeDataFrame df
    toMatrix44 (unpackQ# -> (# x', y', z', w' #)) =
      let x = scalar x'
          y = scalar y'
          z = scalar z'
          w = scalar w'
          x2 = x * x
          y2 = y * y
          z2 = z * z
          w2 = w * w
          l2 = x2 + y2 + z2 + w2
      in ST.runST $ do
        df <- ST.newDataFrame
        ST.writeDataFrameOff df 0 $ l2 - 2*(z2 + y2)
        ST.writeDataFrameOff df 1 $ 2*(x*y + z*w)
        ST.writeDataFrameOff df 2 $ 2*(x*z - y*w)
        ST.writeDataFrameOff df 3 0
        ST.writeDataFrameOff df 4 $ 2*(x*y - z*w)
        ST.writeDataFrameOff df 5 $ l2 - 2*(z2 + x2)
        ST.writeDataFrameOff df 6 $ 2*(y*z + x*w)
        ST.writeDataFrameOff df 7 0
        ST.writeDataFrameOff df 8 $ 2*(x*z + y*w)
        ST.writeDataFrameOff df 9 $ 2*(y*z - x*w)
        ST.writeDataFrameOff df 10 $ l2 - 2*(y2 + x2)
        ST.writeDataFrameOff df 11 0
        ST.writeDataFrameOff df 12 0
        ST.writeDataFrameOff df 13 0
        ST.writeDataFrameOff df 14 0
        ST.writeDataFrameOff df 15 1
        ST.unsafeFreezeDataFrame df

{- Calculate quaternion from a 3x3 matrix.

   First argument is a constant; it is either 1 for a 3x3 matrix,
   or m44 for a 4x4 matrix. I just need to multiply all components by
   this number.

   Further NB for the formulae:

   d == square q == det m ** (1/3)
   t == trace m  == 4 w w - d
   m01 - m10 == 4 z w
   m20 - m02 == 4 y w
   m12 - m21 == 4 x w
   m01 + m10 == 4 x y
   m20 + m02 == 4 x z
   m12 + m21 == 4 y z
   m00 == + x x - y y - z z + w w
   m11 == - x x + y y - z z + w w
   m22 == - x x - y y + z z + w w
   4 x x == d + m00 - m11 - m22
   4 y y == d - m00 + m11 - m22
   4 z z == d - m00 - m11 + m22
   4 w w == d + m00 + m11 + m22
 -}
fromM :: Double
      -> Double -> Double -> Double
      -> Double -> Double -> Double
      -> Double -> Double -> Double
      -> QDouble
fromM c'
  m00 m01 m02
  m10 m11 m12
  m20 m21 m22
    | t > 0
      = let dd = sqrt ( d + t )
            is = c / dd
        in packQ ((m12 - m21)*is) ((m20 - m02)*is) ((m01 - m10)*is) (c*dd)
    | m00 > m11 && m00 > m22
      = let dd = sqrt ( d + m00 - m11 - m22 )
            is = c / dd
        in packQ (c*dd) ((m01 + m10)*is) ((m02 + m20)*is) ((m12 - m21)*is)
    | m11 > m22
      = let dd = sqrt ( d - m00 + m11 - m22 )
            is = c / dd
        in packQ ((m01 + m10)*is) (c*dd) ((m12 + m21)*is) ((m20 - m02)*is)
    | otherwise
      = let dd = sqrt ( d - m00 - m11 + m22 )
            is = c / dd
        in packQ ((m02 + m20)*is) ((m12 + m21)*is) (c*dd) ((m01 - m10)*is)

  where
    -- normalizing constant
    c = recip $ 2 * sqrt c'
    -- trace
    t = m00 + m11 + m22
    -- cubic root of determinant
    d = ( m00 * ( m11 * m22 - m12 * m21 )
        - m01 * ( m10 * m22 - m12 * m20 )
        + m02 * ( m10 * m21 - m11 * m20 )
        ) ** 0.33333333333333333333333333333333

instance Num QDouble where
    QDouble a + QDouble b
      = QDouble (a + b)
    {-# INLINE (+) #-}
    QDouble a - QDouble b
      = QDouble (a - b)
    {-# INLINE (-) #-}
    (unpackQ# -> (# a1, a2, a3, a4 #)) * (unpackQ# -> (# b1, b2, b3, b4 #))
      = packQ ((a4 * b1) + (a1 * b4) + (a2 * b3) - (a3 * b2))
              ((a4 * b2) - (a1 * b3) + (a2 * b4) + (a3 * b1))
              ((a4 * b3) + (a1 * b2) - (a2 * b1) + (a3 * b4))
              ((a4 * b4) - (a1 * b1) - (a2 * b2) - (a3 * b3))
    {-# INLINE (*) #-}
    negate (QDouble a) = QDouble (negate a)
    {-# INLINE negate #-}
    abs = packQ 0 0 0 . sqrt . square
    {-# INLINE abs #-}
    signum q@(unpackQ# -> (# x, y, z, w #))
      | qd == 0   = q
      | otherwise = case ix + iy + iz + iw + nn of
        0 -> packQ (x * l) (y * l) (z * l) (w * l)
        1 -> packQ (copysign 1 x) 0 0 0
        2 -> packQ 0 (copysign 1 y) 0 0
        4 -> packQ 0 0 (copysign 1 z) 0
        8 -> packQ 0 0 0 (copysign 1 w)
        _ -> packQ n n n n
      where
        n  = 0 / 0 :: Double
        qd = x*x + y*y + z*z + w*w
        ix = if isInfinite x then 1 else 0 :: Int
        iy = if isInfinite y then 2 else 0 :: Int
        iz = if isInfinite z then 4 else 0 :: Int
        iw = if isInfinite w then 8 else 0 :: Int
        nn = if isNaN x || isNaN y || isNaN z || isNaN w then 16 else 0 :: Int
        l  = recip (sqrt qd)
    {-# INLINE signum #-}
    fromInteger = packQ 0 0 0 . fromInteger
    {-# INLINE fromInteger #-}

instance Fractional QDouble where
    {-# INLINE recip #-}
    recip q@(unpackQ# -> (# x, y, z, w #)) = case negate (recip (square q)) of
      c -> packQ (x * c) (y * c) (z * c) (negate (w * c))
    {-# INLINE (/) #-}
    a / b = a * recip b
    {-# INLINE fromRational #-}
    fromRational = packQ 0 0 0 . fromRational

instance Floating QDouble where
    {-# INLINE pi #-}
    pi = packQ 0 0 0 M_PI
    {-# INLINE exp #-}
    exp (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z ew
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        ew  = exp w
        l   = ew * sin mv / mv
        arg = ew * cos mv
    {-# INLINE log #-}
    log = log' (Vec3 1 0 0)
    {-# INLINE sqrt #-}
    sqrt = sqrt' (Vec3 1 0 0)
    {-# INLINE sin #-}
    sin (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z (sin w)
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        l   = cos w * sinh mv / mv
        arg = sin w * cosh mv
    {-# INLINE cos #-}
    cos (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z (cos w)
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        l   = sin w * sinh mv / negate mv
        arg = cos w * cosh mv
    {-# INLINE tan #-}
    tan (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0       = packQ x y z (tan w)
      | isInfinite mv2 = signum (packQ x y z 0)
      | otherwise      = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv = sqrt mv2
        b = 2*mv
        a = 2*w
        sina = sin a
        eb = exp (-b)
        eb2 = eb*eb
        d = 1 + eb2 + 2 * eb * cos a
        rd = recip d
        pa = M_PI - abs a
        rd' = 2 / (b*b + pa*pa)
        (l, arg) =
          if d >= M_EPS
          then ((1 - eb2) * rd / mv, 2 * sina * eb * rd)
          else (2 * rd' , negate sina * rd')
    {-# INLINE sinh #-}
    sinh (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z (sinh w)
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        l   = cosh w * sin mv / mv
        arg = sinh w * cos mv
    {-# INLINE cosh #-}
    cosh (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0  = packQ x y z (cosh w)
      | otherwise = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv  = sqrt mv2
        l   = sinh w * sin mv / mv
        arg = cosh w * cos mv
    {-# INLINE tanh #-}
    tanh (unpackQ# -> (# x, y, z, w #))
      | mv2 == 0       = packQ x y z (tanh w)
      | isInfinite mv2 = packQ 0 0 0 (signum w)
      | otherwise      = packQ (x * l) (y * l) (z * l) arg
      where
        mv2 = (x * x) + (y * y) + (z * z)
        mv = sqrt mv2
        b = 2*w
        a = 2*mv
        eb = exp (- abs b)
        eb2 = eb*eb
        d = 1 + eb2 + 2 * eb * cos a
        rd = recip d
        pa = M_PI - a
        rd' = 2 / (b*b + pa*pa)
        (l, arg) =
          if d >= M_EPS
          then (2 * sin a * eb * rd / mv, copysign (1 - eb2) b * rd)
          else (2 * rd' , b * rd')
    {-# INLINE asin #-}
    -- The original formula:
    -- asin q = -i * log (i*q + sqrt (1 - q*q))
    -- below is a more numerically stable version.
    asin (unpackQ# -> (# x, y, z, w #))
      | v2 == 0   = if w2 <= 1
                    then packQ x y z (asin w)
                    else packQ l 0 0 arg
      | otherwise  = packQ (x*c) (y*c) (z*c) arg
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        w1qq = 0.5 *(1 - w2 + v2)       -- real part of (1 - q*q)/2
        l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (1 - q*q)/2
        sp2 = l1qq + w1qq
        sn2 = l1qq - w1qq
        sp = sqrt sp2
        sn = copysign (sqrt sn2) w
        -- choose a more stable (symbolically equiv) version
        dp = if 2 * v2 <= sp2 then sp - v else v2 / ((sp + v)*(sn2 + v2))
        dn = if 2 * w2 <= sn2 then w - sn else w2 / ((sn + w)*(sp2 + w2))
        (wD, vD) = case compare w1qq 0 of
            GT -> (dp, w * dp / sp)
            LT -> (v * dn / sn, dn)
            EQ -> (-v, w)
        l = -0.5 * log (wD*wD + vD*vD)
        c = l / v
        arg = atan2 vD wD
    {-# INLINE acos #-}
    acos q = M_PI_2 - asin q
    {-# INLINE atan #-}
    -- atan q = i / 2 * log ( (i + q) / (i - q) )
    atan (unpackQ# -> (# x, y, z, w #))
      | v2 == 0   = packQ x y z (atan w)
      | otherwise = packQ (x*c) (y*c) (z*c) arg
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        q2 = w2 + v2
        v' = v - 1
        mzero = w2 + v'*v'
        (c, arg) =
          if mzero == 0
          then ( sqrt maxFinite / v, 0)
          else ( 0.25 * (log (1 + q2 + 2*v) - log mzero) / v
               , 0.5 * atan2 (2*w) (1 - q2) )
    {-# INLINE asinh #-}
    -- The original formula:
    -- asinh q = log (q + sqrt (q*q + 1))
    -- below is a more numerically stable version.
    asinh (unpackQ# -> (# x, y, z, w #))
      | v2 == 0   = packQ x y z (asinh w)
      | otherwise = packQ (x*c) (y*c) (z*c) arg
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        w1qq = 0.5 *(1 + w2 - v2)       -- real part of (1 + q*q)/2
        l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (1 + q*q)/2
        sp2 = l1qq + w1qq
        sn2 = l1qq - w1qq
        sp = sqrt sp2
        sn = copysign (sqrt sn2) w
        -- choose a more stable (symbolically equiv) version
        dp = if 2 * w >= - sp then w + sp else w2 / ((sp - w)*(w2 + sn2))
        dn = if 2 * v <= - sn || sn >= 0
                              then v + sn else v2 / ((v - sn)*(v2 + sp2))
        (wD, vD) = case compare w1qq 0 of
            GT -> (dp, v * dp / sp)
            LT -> (w * dn / sn, dn)
            EQ -> (w, v)
        c = atan2 vD wD / v
        arg = 0.5 * log (wD*wD + vD*vD)
    {-# INLINE acosh #-}
    -- The original formula:
    -- asinh q = log (q + sqrt (q + 1) * sqrt (q - 1))
    -- below is a more numerically stable version.
    -- note, log (q + sqrt (q*q - 1)) would not work, because that would not
    -- be the principal value.
    acosh (unpackQ# -> (# x, y, z, w #))
      | v2 == 0   = packQ x y z (acosh w)
      | otherwise = packQ (x*c) (y*c) (z*c) arg
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        w1qq = 0.5 *(w2 - v2 - 1)       -- real part of (q*q - 1)/2
        l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (q*q - 1)/2
        sp2 = l1qq + w1qq
        sn2 = l1qq - w1qq
        sp = sqrt sp2
        sn = copysign (sqrt sn2) w
        -- choose a more stable (symbolically equiv) version
        dp = if 2 * w >= - sp then w + sp else w2 / ((w - sp)*(w2 + sn2))
        dn = if 2 * v <= - sn || sn >= 0
                              then v + sn else v2 / ((sn - v)*(v2 + sp2))
        (wD, vD) = case compare w1qq 0 of
            GT -> (dp, v * dp / sp)
            LT -> (w * dn / sn, dn)
            EQ -> (w, v)
        c = atan2 vD wD / v
        arg = 0.5 * log (wD*wD + vD*vD)
    {-# INLINE atanh #-}
    -- atanh q = 0.5 * log ( (1 + q) / (1 - q) )
    atanh (unpackQ# -> (# x, y, z, w #))
      | v2 ==  0  = packQ x y z (atanh w)
      | otherwise = packQ (x*c) (y*c) (z*c) (copysign arg w)
      where
        v2 = (x * x) + (y * y) + (z * z)
        v = sqrt v2
        w2 = w*w
        q2 = w2 + v2
        w' = abs w - 1
        c  = 0.5 * atan2 (2*v) (1 - q2) / v
        arg = if w' == 0
              then (1/0)
              else 0.25 * (log (1 + q2 + 2 * abs w) - log (v2 + w'*w'))

-- If q is negative real, provide a fallback axis to align log.
log' :: Vector Double 3 -> QDouble -> QDouble
log' r (unpackQ# -> (# x, y, z, w #))
  = case (x * x) + (y * y) + (z * z) of
    0.0 | w >= 0
           -> packQ 0 0 0 (log w)
        | Vec3 rx ry rz <- r
           ->  packQ (M_PI*rx) (M_PI*ry) (M_PI*rz) (log (negate w))
    mv2 -> case (# mv2 + w * w, sqrt mv2 #) of
      (# q2, mv #) -> case atan2 mv w / mv of
        l -> packQ (x * l) (y * l) (z * l) (0.5 * log q2)

-- If q is negative real, provide a fallback axis to align sqrt.
sqrt' :: Vector Double 3 -> QDouble -> QDouble
sqrt' r (unpackQ# -> (# x, y, z, w #))
  | v2 == 0 && w >= 0
    = packQ x y z (sqrt w)
  | v2 == 0
  , Vec3 rx ry rz <- r
  , sw <- sqrt (negate w)
    = packQ (sw*rx) (sw*ry) (sw*rz) 0
  | otherwise
    = packQ (x * c) (y * c) (z * c) arg
  where
    v2 = (x * x) + (y * y) + (z * z)
    mq = sqrt (v2 + w * w)
    arg = sqrt $ 0.5 * if w >= 0 then mq + w else v2 / (mq - w)
    c = 0.5 / arg

instance Eq QDouble where
    {-# INLINE (==) #-}
    QDouble a == QDouble b = a == b
    {-# INLINE (/=) #-}
    QDouble a /= QDouble b = a /= b


{-# 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 BangPatterns         #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE MagicHash            #-}
{-# LANGUAGE TypeSynonymInstances #-}
{-# LANGUAGE UnboxedTuples        #-}

{-# OPTIONS_GHC -fno-warn-orphans #-}

-- |
-- Module      :  Internal.Vector
-- Copyright   :  (c) Alberto Ruiz 2007-15
-- License     :  BSD3
-- Maintainer  :  Alberto Ruiz
-- Stability   :  provisional
--

module Internal.Vector(
    I,Z,R,C,
    fi,ti,
    Vector, fromList, unsafeToForeignPtr, unsafeFromForeignPtr, unsafeWith,
    createVector, avec, inlinePerformIO,
    toList, dim, (@>), at', (|>),
    vjoin, subVector, takesV, idxs,
    buildVector,
    asReal, asComplex,
    toByteString,fromByteString,
    zipVector, unzipVector, zipVectorWith, unzipVectorWith,
    foldVector, foldVectorG, foldVectorWithIndex, foldLoop,
    mapVector, mapVectorM, mapVectorM_,
    mapVectorWithIndex, mapVectorWithIndexM, mapVectorWithIndexM_
) where

import           Data.Complex
import           Data.Int                      (Int64)
import           Data.Vector.Storable          (Vector, fromList, unsafeFromForeignPtr,
                                                unsafeToForeignPtr, unsafeWith)
import qualified Data.Vector.Storable          as Vector
import           Foreign.C.Types               (CInt)
import           Foreign.ForeignPtr
import           Foreign.Marshal.Array
import           Foreign.Ptr
import           Foreign.Storable
import           GHC.Base                      (IO (IO), realWorld#, when)
import           GHC.ForeignPtr                (mallocPlainForeignPtrBytes)
import           System.IO.Unsafe              (unsafePerformIO)

import           Control.Monad                 (replicateM)
import           Data.Binary
import           Data.Binary.Put
import qualified Data.ByteString.Internal      as BS
import           Data.Vector.Storable.Internal (updPtr)

type I = CInt
type Z = Int64
type R = Float
type C = Complex Float

-- | specialized fromIntegral
fi :: Int -> CInt
fi = fromIntegral

-- | specialized fromIntegral
ti :: CInt -> Int
ti = fromIntegral

-- | Number of elements
dim :: (Storable t) => Vector t -> Int
dim = Vector.length
{-# INLINE dim #-}

-- C-Haskell vector adapter
{-# INLINE avec #-}
avec :: Storable a => Vector a -> (f -> IO r) -> ((CInt -> Ptr a -> f) -> IO r)
avec v f g = unsafeWith v $ \ptr -> f (g (fromIntegral (Vector.length v)) ptr)

-- allocates memory for a new vector
createVector :: Storable a => Int -> IO (Vector a)
createVector n = do
    when (n < 0) $ error ("trying to createVector of negative dim: "++show n)
    fp <- doMalloc undefined
    return $ unsafeFromForeignPtr fp 0 n
  where
    --
    -- Use the much cheaper Haskell heap allocated storage
    -- for foreign pointer space we control
    --
    doMalloc :: Storable b => b -> IO (ForeignPtr b)
    doMalloc dummy = do
        mallocPlainForeignPtrBytes (n * sizeOf dummy)

{- | creates a Vector from a list:

@> fromList [2,3,5,7]
4 |> [2.0,3.0,5.0,7.0]@

-}

safeRead :: Storable a => Vector a -> (Ptr a -> IO c) -> c
safeRead v = inlinePerformIO . unsafeWith v
{-# INLINE safeRead #-}

inlinePerformIO :: IO a -> a
inlinePerformIO (IO m) = case m realWorld# of (# _, r #) -> r
{-# INLINE inlinePerformIO #-}

{- extracts the Vector elements to a list

>>> toList (linspace 5 (1,10))
[1.0,3.25,5.5,7.75,10.0]

-}
toList :: Storable a => Vector a -> [a]
toList v = safeRead v $ peekArray (dim v)

{- | Create a vector from a list of elements and explicit dimension. The input
     list is truncated if it is too long, so it may safely
     be used, for instance, with infinite lists.

>>> 5 |> [1..]
[1.0,2.0,3.0,4.0,5.0]
it :: (Enum a, Num a, Foreign.Storable.Storable a) => Vector a

-}
(|>) :: (Storable a) => Int -> [a] -> Vector a
infixl 9 |>
n |> l
    | length l' == n = fromList l'
    | otherwise      = error "list too short for |>"
  where
    l' = take n l

-- | Create a vector of indexes, useful for matrix extraction using '(??)'
idxs :: [Int] -> Vector I
idxs js = fromList (map fromIntegral js) :: Vector I

{- | takes a number of consecutive elements from a Vector

>>> subVector 2 3 (fromList [1..10])
[3.0,4.0,5.0]
it :: (Enum t, Num t, Foreign.Storable.Storable t) => Vector t

-}
subVector :: Storable t => Int       -- ^ index of the starting element
                        -> Int       -- ^ number of elements to extract
                        -> Vector t  -- ^ source
                        -> Vector t  -- ^ result
subVector = Vector.slice
{-# INLINE subVector #-}

{- | Reads a vector position:

>>> fromList [0..9] @> 7
7.0

-}
(@>) :: Storable t => Vector t -> Int -> t
infixl 9 @>
v @> n
    | n >= 0 && n < dim v = at' v n
    | otherwise = error "vector index out of range"
{-# INLINE (@>) #-}

-- | access to Vector elements without range checking
at' :: Storable a => Vector a -> Int -> a
at' v n = safeRead v $ flip peekElemOff n
{-# INLINE at' #-}

{- | concatenate a list of vectors

>>> vjoin [fromList [1..5::Float], konst 1 3]
[1.0,2.0,3.0,4.0,5.0,1.0,1.0,1.0]
it :: Vector Float

-}
vjoin :: Storable t => [Vector t] -> Vector t
vjoin [] = fromList []
vjoin [v] = v
vjoin as = unsafePerformIO $ do
    let tot = sum (map dim as)
    r <- createVector tot
    unsafeWith r $ \ptr ->
        joiner as tot ptr
    return r
  where joiner [] _ _ = return ()
        joiner (v:cs) _ p = do
            let n = dim v
            unsafeWith v $ \pb -> copyArray p pb n
            joiner cs 0 (advancePtr p n)

{- | Extract consecutive subvectors of the given sizes.

>>> takesV [3,4] (linspace 10 (1,10::Float))
[[1.0,2.0,3.0],[4.0,5.0,6.0,7.0]]
it :: [Vector Float]

-}
takesV :: Storable t => [Int] -> Vector t -> [Vector t]
takesV ms w | sum ms > dim w = error $ "takesV " ++ show ms ++ " on dim = " ++ (show $ dim w)
            | otherwise = go ms w
    where go [] _ = []
          go (n:ns) v = subVector 0 n v
                      : go ns (subVector n (dim v - n) v)

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

-- | transforms a complex vector into a real vector with alternating real and imaginary parts
asReal :: (RealFloat a, Storable a) => Vector (Complex a) -> Vector a
asReal v = unsafeFromForeignPtr (castForeignPtr fp) (2*i) (2*n)
    where (fp,i,n) = unsafeToForeignPtr v

-- | transforms a real vector into a complex vector with alternating real and imaginary parts
asComplex :: (RealFloat a, Storable a) => Vector a -> Vector (Complex a)
asComplex v = unsafeFromForeignPtr (castForeignPtr fp) (i `div` 2) (n `div` 2)
    where (fp,i,n) = unsafeToForeignPtr v

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

-- | map on Vectors
mapVector :: (Storable a, Storable b) => (a-> b) -> Vector a -> Vector b
mapVector f v = unsafePerformIO $ do
    w <- createVector (dim v)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) = return ()
                go !k = do x <- peekElemOff p k
                           pokeElemOff      q k (f x)
                           go (k-1)
            go (dim v -1)
    return w
{-# INLINE mapVector #-}

-- | zipWith for Vectors
zipVectorWith :: (Storable a, Storable b, Storable c) => (a-> b -> c) -> Vector a -> Vector b -> Vector c
zipVectorWith f u v = unsafePerformIO $ do
    let n = min (dim u) (dim v)
    w <- createVector n
    unsafeWith u $ \pu ->
        unsafeWith v $ \pv ->
            unsafeWith w $ \pw -> do
                let go (-1) = return ()
                    go !k = do x <- peekElemOff pu k
                               y <- peekElemOff pv k
                               pokeElemOff      pw k (f x y)
                               go (k-1)
                go (n -1)
    return w
{-# INLINE zipVectorWith #-}

-- | unzipWith for Vectors
unzipVectorWith :: (Storable (a,b), Storable c, Storable d)
                   => ((a,b) -> (c,d)) -> Vector (a,b) -> (Vector c,Vector d)
unzipVectorWith f u = unsafePerformIO $ do
      let n = dim u
      v <- createVector n
      w <- createVector n
      unsafeWith u $ \pu ->
          unsafeWith v $ \pv ->
              unsafeWith w $ \pw -> do
                  let go (-1) = return ()
                      go !k   = do z <- peekElemOff pu k
                                   let (x,y) = f z
                                   pokeElemOff      pv k x
                                   pokeElemOff      pw k y
                                   go (k-1)
                  go (n-1)
      return (v,w)
{-# INLINE unzipVectorWith #-}

foldVector :: Storable a => (a -> b -> b) -> b -> Vector a -> b
foldVector f x v = unsafePerformIO $
    unsafeWith v $ \p -> do
        let go (-1) s = return s
            go !k !s = do y <- peekElemOff p k
                          go (k-1::Int) (f y s)
        go (dim v -1) x
{-# INLINE foldVector #-}

-- the zero-indexed index is passed to the folding function
foldVectorWithIndex :: Storable a => (Int -> a -> b -> b) -> b -> Vector a -> b
foldVectorWithIndex f x v = unsafePerformIO $
    unsafeWith v $ \p -> do
        let go (-1) s = return s
            go !k !s = do y <- peekElemOff p k
                          go (k-1::Int) (f k y s)
        go (dim v -1) x
{-# INLINE foldVectorWithIndex #-}

foldLoop :: (Int -> t -> t) -> t -> Int -> t
foldLoop f s0 d = go (d - 1) s0
     where
       go 0 s   = f (0::Int) s
       go !j !s = go (j - 1) (f j s)

foldVectorG :: Storable t1 => (Int -> (Int -> t1) -> t -> t) -> t -> Vector t1 -> t
foldVectorG f s0 v = foldLoop g s0 (dim v)
    where g !k !s = f k (safeRead v . flip peekElemOff) s
          {-# INLINE g #-} -- Thanks to Ryan Ingram (http://permalink.gmane.org/gmane.comp.lang.haskell.cafe/46479)
{-# INLINE foldVectorG #-}

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

-- | monadic map over Vectors
--    the monad @m@ must be strict
mapVectorM :: (Storable a, Storable b, Monad m) => (a -> m b) -> Vector a -> m (Vector b)
mapVectorM f v = do
    w <- return $! unsafePerformIO $! createVector (dim v)
    mapVectorM' w 0 (dim v -1)
    return w
    where mapVectorM' w' !k !t
              | k == t               = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f x
                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
              | otherwise            = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f x
                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
                                       mapVectorM' w' (k+1) t
{-# INLINE mapVectorM #-}

-- | monadic map over Vectors
mapVectorM_ :: (Storable a, Monad m) => (a -> m ()) -> Vector a -> m ()
mapVectorM_ f v = do
    mapVectorM' 0 (dim v -1)
    where mapVectorM' !k !t
              | k == t            = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    f x
              | otherwise         = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    _ <- f x
                                    mapVectorM' (k+1) t
{-# INLINE mapVectorM_ #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
--    the monad @m@ must be strict
mapVectorWithIndexM :: (Storable a, Storable b, Monad m) => (Int -> a -> m b) -> Vector a -> m (Vector b)
mapVectorWithIndexM f v = do
    w <- return $! unsafePerformIO $! createVector (dim v)
    mapVectorM' w 0 (dim v -1)
    return w
    where mapVectorM' w' !k !t
              | k == t               = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f k x
                                       return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
              | otherwise            = do
                                       x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                       y <- f k x
                                       _ <- return $! inlinePerformIO $! unsafeWith w' $! \q -> pokeElemOff q k y
                                       mapVectorM' w' (k+1) t
{-# INLINE mapVectorWithIndexM #-}

-- | monadic map over Vectors with the zero-indexed index passed to the mapping function
mapVectorWithIndexM_ :: (Storable a, Monad m) => (Int -> a -> m ()) -> Vector a -> m ()
mapVectorWithIndexM_ f v = do
    mapVectorM' 0 (dim v -1)
    where mapVectorM' !k !t
              | k == t            = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    f k x
              | otherwise         = do
                                    x <- return $! inlinePerformIO $! unsafeWith v $! \p -> peekElemOff p k
                                    _ <- f k x
                                    mapVectorM' (k+1) t
{-# INLINE mapVectorWithIndexM_ #-}

mapVectorWithIndex :: (Storable a, Storable b) => (Int -> a -> b) -> Vector a -> Vector b
--mapVectorWithIndex g = head . mapVectorWithIndexM (\a b -> [g a b])
mapVectorWithIndex f v = unsafePerformIO $ do
    w <- createVector (dim v)
    unsafeWith v $ \p ->
        unsafeWith w $ \q -> do
            let go (-1) = return ()
                go !k = do x <- peekElemOff p k
                           pokeElemOff      q k (f k x)
                           go (k-1)
            go (dim v -1)
    return w
{-# INLINE mapVectorWithIndex #-}

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

-- a 64K cache, with a Float taking 13 bytes in Bytestring,
-- implies a chunk size of 5041
chunk :: Int
chunk = 5000

chunks :: Int -> [Int]
chunks d = let c = d `div` chunk
               m = d `mod` chunk
           in if m /= 0 then reverse (m:(replicate c chunk)) else (replicate c chunk)

putVector :: (Storable t, Binary t) => Vector t -> Data.Binary.Put.PutM ()
putVector v = mapM_ put $! toList v

getVector :: (Storable a, Binary a) => Int -> Get (Vector a)
getVector d = do
              xs <- replicateM d get
              return $! fromList xs

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

toByteString :: Storable t => Vector t -> BS.ByteString
toByteString v = BS.PS (castForeignPtr fp) (sz*o) (sz * dim v)
  where
    (fp,o,_n) = unsafeToForeignPtr v
    sz = sizeOf (v@>0)

fromByteString :: Storable t => BS.ByteString -> Vector t
fromByteString (BS.PS fp o n) = r
  where
    r = unsafeFromForeignPtr (castForeignPtr (updPtr (`plusPtr` o) fp)) 0 n'
    n' = n `div` sz
    sz = sizeOf (r@>0)

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

instance (Binary a, Storable a) => Binary (Vector a) where

    put v = do
            let d = dim v
            put d
            mapM_ putVector $! takesV (chunks d) v

    -- put = put . v2bs

    get = do
          d <- get
          vs <- mapM getVector $ chunks d
          return $! vjoin vs

    -- get = fmap bs2v get

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

{- | creates a Vector of the specified length using the supplied function to
     to map the index to the value at that index.

@> buildVector 4 fromIntegral
4 |> [0.0,1.0,2.0,3.0]@

-}
buildVector :: Storable a => Int -> (Int -> a) -> Vector a
buildVector len f =
    fromList $ map f [0 .. (len - 1)]

-- | zip for Vectors
zipVector :: (Storable a, Storable b, Storable (a,b)) => Vector a -> Vector b -> Vector (a,b)
zipVector = zipVectorWith (,)

-- | unzip for Vectors
unzipVector :: (Storable a, Storable b, Storable (a,b)) => Vector (a,b) -> (Vector a,Vector b)
unzipVector = unzipVectorWith id

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


--
-- ConvLayer : convolution layer
--

module CNN.ConvLayer (
  initFilterC
, convolve
, deconvolve
, reverseConvFilter
, updateConvFilter
) where

import Control.Monad hiding (msum)
import Data.List hiding (transpose)
import Data.Maybe
import Debug.Trace
import Numeric.LinearAlgebra
import System.Random.Mersenne as MT

import CNN.Algebra
import CNN.Image
import CNN.LayerType

{- |
initFilterC

  IN : #kernel
       #channel
       input size (x * y)
       kernel size (n * n)
       pool size (m * m)

  OUT: filter of convolution layer

-}

initFilterC :: Int -> Int -> Int -> Int -> Int -> Int -> IO [FilterC]
initFilterC k c x y n m = do
  f <- forM [1..k] $ \i -> do
    w <- initKernel c n
    return (w, 0.0)
  return f
  where
    ri = fromIntegral (n * n * c)   -- resolution (input)
    ro = fromIntegral (n * n * k) / (fromIntegral (m * m))
    a = sqrt (6.0 / (ri + ro))
    initKernel :: Int -> Int -> IO Kernel
    initKernel c n = do
      let
        sz = n * n
      w <- forM [1..c] $ \i -> do
        w' <- forM [1..sz] $ \j -> do
          r <- MT.randomIO :: IO Double
          return ((r * 2.0 - 1.0) * a)
        return w'
      return $ fromLists w

--

{- |
convolve

  IN: kernel size
      filter
      image

test: filter k=3, c=2, s=2
>>> let filter = [(fromLists [[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0]],0.25),(fromLists [[3.0,4.0,5.0,6.0],[7.0,8.0,1.0,2.0]],0.5),(fromLists [[5.0,6.0,7.0,8.0],[1.0,2.0,3.0,4.0]],0.75)] :: [FilterC]
>>> let img1 = [fromLists [[1.0,2.0,3.0],[4.0,5.0,6.0],[7.0,8.0,9.0]], fromLists [[4.0,5.0,6.0],[7.0,8.0,9.0],[1.0,2.0,3.0]]] :: Image
>>> let img2 = [fromLists [[1.0,2.0,3.0,4.0],[4.0,5.0,6.0,7.0],[7.0,8.0,9.0,10.0],[10.0,11.0,12.0,13.0]],fromLists [[4.0,5.0,6.0,7.0],[7.0,8.0,9.0,10.0],[1.0,2.0,3.0,4.0],[4.0,5.0,6.0,7.0]]] :: Image
>>> convolve 2 filter img1
[(2><2)
 [ 200.25, 236.25
 , 173.25, 209.25 ],(2><2)
 [ 152.5, 188.5
 , 233.5, 269.5 ],(2><2)
 [ 152.75, 188.75
 , 197.75, 233.75 ]]
>>> convolve 2 filter img2
[(3><3)
 [ 200.25, 236.25, 272.25
 , 173.25, 209.25, 245.25
 , 182.25, 218.25, 254.25 ],(3><3)
 [ 152.5, 188.5, 224.5
 , 233.5, 269.5, 305.5
 , 206.5, 242.5, 278.5 ],(3><3)
 [ 152.75, 188.75, 224.75
 , 197.75, 233.75, 269.75
 , 278.75, 314.75, 350.75 ]]

-}

--
-- USE corr2 ??
--

convolve :: Int -> [FilterC] -> Image -> Image
convolve s fs im = map (convolveImage x iss) fs
  where
    pl = head im
    x  = cols pl - (s - 1)
    y  = rows pl - (s - 1)
    ps = [(i, j) | i <- [0..(x-1)], j <- [0..(y-1)]]
    iss = map subImage im
    subImage :: Plain -> Matrix R
    subImage i = fromColumns $ map (\p -> flatten $ subMatrix p (s, s) i) ps

convolveImage :: Int -> [Matrix R] -> FilterC -> Plain
convolveImage x iss (k, b) = reshape x $ cmap (+ b) $ vsum vs
  where
    vs = zipWith (<#) (toRows k) iss

-- back prop

deconvolve :: Int -> [FilterC] -> Image -> Delta -> (Delta, Maybe Layer)
deconvolve s fs im d = (delta, Just (ConvLayer s dw))
  where
    delta = convolve s fs $ addBorder s d
    sim = slideImage s im
    dw = map (hadamard sim) d

{- |
addBorder

  IN : kernel size
       delta (Image)
  OUT: delta (Image)

>>> let s = [fromLists [[1.0, 2.0],[3.0,4.0]], fromLists [[5.0,6.0],[7.0,8.0]]]
>>> addBorder 2 s
[(4><4)
 [ 0.0, 0.0, 0.0, 0.0
 , 0.0, 1.0, 2.0, 0.0
 , 0.0, 3.0, 4.0, 0.0
 , 0.0, 0.0, 0.0, 0.0 ],(4><4)
 [ 0.0, 0.0, 0.0, 0.0
 , 0.0, 5.0, 6.0, 0.0
 , 0.0, 7.0, 8.0, 0.0
 , 0.0, 0.0, 0.0, 0.0 ]]

-}

addBorder :: Int -> Delta -> Delta
addBorder s d = map ab d
  where
    s' = s - 1
    ab :: Plain -> Plain
    ab p = yb ||| ((xb === (p === xb)) ||| yb)
      where
        x = cols p
        y = rows p
        xb = (s'><x) $ repeat 0.0
        yb = ((y + 2 * s')><s') $ repeat 0.0

{- |
slideImage

>>> let p = [fromLists [[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0],[9.0,10.0,11.0,12.0],[13.0,14.0,15.0,16.0]]] :: Image
>>> slideImage 2 p
[[(3><3)
 [ 1.0,  2.0,  3.0
 , 5.0,  6.0,  7.0
 , 9.0, 10.0, 11.0 ],(3><3)
 [  2.0,  3.0,  4.0
 ,  6.0,  7.0,  8.0
 , 10.0, 11.0, 12.0 ],(3><3)
 [  5.0,  6.0,  7.0
 ,  9.0, 10.0, 11.0
 , 13.0, 14.0, 15.0 ],(3><3)
 [  6.0,  7.0,  8.0
 , 10.0, 11.0, 12.0
 , 14.0, 15.0, 16.0 ]]]

-}

slideImage :: Int -> Image -> [[Plain]]
slideImage s im = map (slidePlain s) im

{- |
slidePlain

>>> let p = fromLists [[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0],[9.0,10.0,11.0,12.0],[13.0,14.0,15.0,16.0]] :: Plain
>>> slidePlain 2 p
[(3><3)
 [ 1.0,  2.0,  3.0
 , 5.0,  6.0,  7.0
 , 9.0, 10.0, 11.0 ],(3><3)
 [  2.0,  3.0,  4.0
 ,  6.0,  7.0,  8.0
 , 10.0, 11.0, 12.0 ],(3><3)
 [  5.0,  6.0,  7.0
 ,  9.0, 10.0, 11.0
 , 13.0, 14.0, 15.0 ],(3><3)
 [  6.0,  7.0,  8.0
 , 10.0, 11.0, 12.0
 , 14.0, 15.0, 16.0 ]]

>>> slidePlain 3 p
[(2><2)
 [ 1.0, 2.0
 , 5.0, 6.0 ],(2><2)
 [ 2.0, 3.0
 , 6.0, 7.0 ],(2><2)
 [ 3.0, 4.0
 , 7.0, 8.0 ],(2><2)
 [ 5.0,  6.0
 , 9.0, 10.0 ],(2><2)
 [  6.0,  7.0
 , 10.0, 11.0 ],(2><2)
 [  7.0,  8.0
 , 11.0, 12.0 ],(2><2)
 [  9.0, 10.0
 , 13.0, 14.0 ],(2><2)
 [ 10.0, 11.0
 , 14.0, 15.0 ],(2><2)
 [ 11.0, 12.0
 , 15.0, 16.0 ]]

-}

slidePlain :: Int -> Plain -> [Plain]
slidePlain s m
  | s < 1     = []
  | s == 1    = [m]
  | otherwise = map (\p -> subMatrix p (x, y) m) ps
  where
    x = cols m - (s - 1)
    y = rows m - (s - 1)
    ps = [(i, j) | i <- [0..s-1], j <- [0..s-1]]

{- |
hadamard

  IN : slided images
       delta
  OUT: (Kernel, Bias)

>>> let p = [fromLists [[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0],[9.0,10.0,11.0,12.0],[13.0,14.0,15.0,16.0]], fromLists [[16.0,15.0,14.0,13.0],[12.0,11.0,10.0,9.0],[8.0,7.0,6.0,5.0],[4.0,3.0,2.0,1.0]]] :: Image
>>> let ds = fromLists [[0.1,0.2,0.3],[0.4,0.5,0.6],[0.7,0.8,0.9]]
>>> hadamard (slideImage 2 p) ds
((2><4)
 [ 34.800000000000004, 39.3, 52.800000000000004, 57.3
 ,               41.7, 37.2, 23.700000000000006, 19.2 ],4.5)

-}

hadamard :: [[Plain]] -> Plain -> FilterC
hadamard sim ds = (w, b)
  where
    b = sumElements ds
    d' = flatten ds
    w = fromLists $ map (map ((<.>) d' . flatten)) sim

{- |
reverseConvFilter

  IN : kernel size
       original filter
  OUT: reversed filter

>>> let fs = [(fromLists [[0.1,0.2,0.3,0.4],[0.5,0.6,0.7,0.8]],1.0),(fromLists [[0.1,0.3,0.5,0.7],[0.2,0.4,0.6,0.8]],2.0),(fromLists [[0.1,0.2,0.5,0.6],[0.3,0.4,0.7,0.8]],3.0)] :: [FilterC]
>>> let ConvLayer s fs' = reverseConvFilter 2 fs
>>> s
2
>>> fs'
[((3><4)
 [ 0.4, 0.3, 0.2, 0.1
 , 0.7, 0.5, 0.3, 0.1
 , 0.6, 0.5, 0.2, 0.1 ],0.0),((3><4)
 [ 0.8, 0.7, 0.6, 0.5
 , 0.8, 0.6, 0.4, 0.2
 , 0.8, 0.7, 0.4, 0.3 ],0.0)]

-}

reverseConvFilter :: Int -> [FilterC] -> Layer
reverseConvFilter s fs = ConvLayer s (zip rv (repeat 0.0))
  where
    (k, _) = unzip fs
    k'     = map (toLists . fliprl) k
    --r      = transpose $ map (map reverse) k'
    r      = transpose k'
    rv     = map fromLists r

-- update

updateConvFilter :: Int -> [FilterC] -> Double -> [Maybe Layer] -> Layer
updateConvFilter s fs lr dl
  | dl' == [] = ConvLayer s fs
  | otherwise = ConvLayer s $ zip ks' (toList bs')
  where
    dl' = catMaybes dl
    sc = lr / fromIntegral (length dl')
    (ks , bs ) = unzip fs
    (kss, bss) = unzip $ map unzip $ strip dl'
    dbs = vsum $ map fromList bss
    dks = map (mscale sc . msum) $ transpose kss
    bs' = (fromList bs) - (vscale sc dbs)
    ks' = zipWith (-) ks dks

strip :: [Layer] -> [[FilterC]]
strip [] = []
strip (ConvLayer s fs:ds) = fs:strip ds
strip (_:ds) = strip ds  -- if not ConvLayer


{-|
Module : Mnist
|-}
module Mnist
  ( Image(..)
  , readTrainingData
  , readTestData
  , getData
  ) where

import qualified Codec.Compression.GZip as GZip
import qualified Data.ByteString.Lazy as BL
import Data.Binary.Get
import Data.Word
import qualified Data.List.Split as S
import Numeric.LinearAlgebra

-- | data representing an image
data Image = Image {
      iRows :: Int -- ^ Number of rows
    , iColumns :: Int -- ^ Number of columns
    , iPixels :: [Word8] -- ^ Pixels
    } deriving (Eq, Show)

-- | Create column vector of Image pixels
toColVector :: Image
         -> Matrix Double
toColVector image = (rc><1) p :: Matrix Double
  where rc = iRows image * iColumns image
        p = normalisedData image

-- | Normalise Image pixels
normalisedData :: Image -> [Double]
normalisedData image = normalisePixel <$> iPixels image

-- | Normalise pixel
normalisePixel :: Word8 -> Double
normalisePixel p = fromIntegral p / 255.0

{-
[offset] [type]          [value]          [description]
0000     32 bit integer  0x00000801(2049) magic number (MSB first)
0004     32 bit integer  60000            number of items
0008     unsigned byte   ??               label
0009     unsigned byte   ??               label
 ........
xxxx     unsigned byte   ??               label
The labels values are 0 to 9.
-}

-- | Extract information from labels file
deserialiseLabels :: Get (Word32, Word32, [Word8])
deserialiseLabels = do
  magicNumber <- getWord32be
  count <- getWord32be
  labelData <- getRemainingLazyByteString
  let labels = BL.unpack labelData
  return (magicNumber, count, labels)

-- | Read labels from file
readLabels :: FilePath -> IO [Int]
readLabels filename = do
  content <- GZip.decompress <$> BL.readFile filename
  let (_, _, labels) = runGet deserialiseLabels content
  return (fromIntegral <$> labels)

{-
[offset] [type]          [value]          [description]
0000     32 bit integer  0x00000803(2051) magic number
0004     32 bit integer  60000            number of images
0008     32 bit integer  28               number of rows
0012     32 bit integer  28               number of columns
0016     unsigned byte   ??               pixel
0017     unsigned byte   ??               pixel
 ........
xxxx     unsigned byte   ??               pixel
Pixels are organized row-wise. Pixel values are 0 to 255. 0 means background (white), 255 means foreground (black).
-}

-- | Extract information from image file
deserialiseHeader :: Get (Word32, Word32, Word32, Word32, [[Word8]])
deserialiseHeader = do
  magicNumber <- getWord32be
  imageCount <- getWord32be
  r <- getWord32be
  c <- getWord32be
  packedData <- getRemainingLazyByteString
  let len = fromIntegral (r * c)
  let unpackedData = S.chunksOf len (BL.unpack packedData)
  return (magicNumber, imageCount, r, c, unpackedData)

-- | Read Images form file
readImages :: FilePath -> IO [Image]
readImages filename = do
  content <- GZip.decompress <$> BL.readFile filename
  let (_, _, r, c, unpackedData) = runGet deserialiseHeader content
  return (Image (fromIntegral r) (fromIntegral c) <$> unpackedData)

-- | Create column vector of label
vectorizedLabels :: Int -> Matrix Double
vectorizedLabels x = (10><1) p :: Matrix Double
  where p = [i | j <- [0 .. 9], let i = if j == x then 1.0 else 0.0]

-- | Get training data
readTrainingData :: IO [(Matrix Double, Matrix Double)]
readTrainingData = do
  trainingLabels <- readLabels "train-labels-idx1-ubyte.gz"
  trainingImages <- readImages "train-images-idx3-ubyte.gz"
  return (zip (map toColVector trainingImages) (vectorizedLabels <$> trainingLabels))

-- | Get test data
readTestData :: IO [(Matrix Double, Int)]
readTestData = do
  testLabels <- readLabels "t10k-labels-idx1-ubyte.gz"
  testImages <- readImages "t10k-images-idx3-ubyte.gz"
  return (zip (map toColVector testImages) testLabels)

-- | Get training and test data
getData :: IO ([(Matrix Double, Matrix Double)], [(Matrix Double, Int)])
getData = do
  trainingData <- readTrainingData
  testData <- readTestData
  return (trainingData, testData)