{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE FlexibleContexts #-}

-- This file is the new version for C-code FFT algorithm with Nyquist allready on position 1.

-- | This module imports the C-functions for a Fast Fourier Transform and encapsulates them in
--   list-mapping functions. (At first it was a DFT algoritm which I changed, but I didn'd change
--   all the names of the functions, so it is called dft...)
--   The functionality of this module is used for filtering audio data, i. e. for building
--   audio filters.

module Sound.Hommage.DFTFilter
 (

-- | Filtering is done with an FFT algorithm written in C (in the file dft.c).
--   This module imports the c-functions for fft ('analyseDFT') and inverse fft ('syntheseDFT'),
--   which operate on Arrays with 1024 Double values.
--   The array-operations are embedded in a mechanism that does buffering, overlapping and windowing
--   and works like a black box that maps one input value to one output value in the IO monad
--   ('mkAnalyse', 'mkSynthese', 'mkFilterBuffered').
--   Around these IO's some list mapping functions are wrapped ('dftanalyse', 'dftsynthese', 'dftfilter',
--   'dftfilterBy') for convenient use in a Haskell programm.
--
--   Mapping the Fourier coefficients during the filter process means filtering the data.
--   The result depends on the way the coefficients are mapped resp. modified while mapping.
--   The function to do this has the type of 'CoeffMap'.
--
--   /Description of the filtering process:/
--
--   Filtering is done by mapping the coefficients with 'CoeffMap'.
--   For this the wave data is split into parts of 1024 values, where every part overlaps with the
--   preceeding and succeeding part by 512 values. Via the fourier-transform (analyse) every part
--   is mapped  to its frequency spectrum, i. e. its fourier coefficients.
--
--   The 'CoeffMap' action is applied to every frequency spectrum (see 'CoeffArr') in sequence.
--   It reads the coefficients from the first 'CoeffArr' where the input is stored and writes it
--   (or some data derived from it) to the second one. This is the point where the \'real\' filtering
--   happens. Via the inverse fourier-transform (synthese) the resulting coefficients are mapped
--   to their real signal values. These signal values are multiplied with a half cosinus curve for
--   fading in and out (windowing). Then they are mixed in the way their source data was split:
--   Every part overlaps for 512 values with the last and the next one.

 -- * User-Level Filter Functions
   dftanalyse
 , dftsynthese
 , dftfilter
 , dftfilterBy
 , dftfilterBy'

 -- * CoeffArr
 , CoeffArr
 , readCoeffArr
-- , readCoeffArr'
 , writeCoeffArr
-- , writeCoeffArr'
 , storeCoeff
 , unstoreCoeff
 -- * CoeffMap
 , CoeffMap
 , mkCoeffMap
 , coeffmap
 -- * IO Wrapper
 , mkAnalyse
 , mkSynthese
 , mkFilterBuffered
 -- * Interface to C-code
 , analyseDFT
 , syntheseDFT
 , kurveDFT
 )
 where

import GHC.Weak
import GHC.IO
import Foreign.ForeignPtr
import Foreign.Ptr
import Foreign.Storable
import Data.Array.Storable
import Data.Complex
import Data.Array.IO
import Data.IORef

import Sound.Hommage.Misc
---------------------------------------------------------------------------------------------------
-- ([Complex Double] -> [Complex Double]) -> IO CoeffMap

-- | (Fast) Fourier Transformation and Inverse with a buffer-mapping action for filtering.
dftfilterBy :: IO (CoeffMap ()) -> [Double] -> [Double]
dftfilterBy cm = drop 1536 . inList (mkFilterBuffered cm) . (++ replicate 512 0.0)

-- | (Fast) Fourier Transformation and Inverse with a buffer-mapping action for filtering.
dftfilterBy' :: s -> IO (s -> CoeffMap s) -> [Double] -> [Double]
dftfilterBy' s cm = drop 1536 . inList (mkFilterBuffered' s cm) . (++ replicate 512 0.0)

-- | (Fast) Fourier Transformation and Inverse with an additional argument with sublists of
--   factors to weight the coefficients. See 'mkCoeffMap' for a decription of that argument.
dftfilter :: [Double] -> [Double] -> [Double]
dftfilter ds = dftfilterBy (mkCoeffMap (replicate 512 0 ++ ds))

-- | Inverse (Fast) Fourier Transformation
dftsynthese :: [Complex Double] -> [Double]
dftsynthese = inList mkSynthese

-- | (Fast) Fourier Transformation.
dftanalyse :: [Double] -> [(Complex Double)]
dftanalyse = inList mkAnalyse

---------------------------------------------------------------------------------------------------
foreign import ccall "cinit" initfou :: Int

foreign import ccall "canalyse" cAnalyse :: Ptr Double -> Ptr Double -> IO ()

foreign import ccall "csynthese" cSynthese :: Ptr Double -> Ptr Double -> IO ()

foreign import ccall "ckurve" cKurve :: Ptr Double -> IO ()
---------------------------------------------------------------------------------------------------
-- | Maps an array with 1024 real values to its frequency spectrum with 512 complex values.
--   The spectrum is stored in the second array with 1024 Doubles which are the real and imaginary
--   parts of the complex values. See 'CoeffArr' for a detailed frequency spectrum
--   desription.
analyseDFT :: StorableArray Int Double -> StorableArray Int Double -> IO ()
analyseDFT wave coeffs = seq initfou $
 withStorableArray wave $ \ptr_w ->
 withStorableArray coeffs $ \ptr_c ->
 cAnalyse ptr_w ptr_c

-- | Maps a frequency spectrum to its signal value. Both arrays have a length of 1024; the first one
--   contains the 512 complex values of the spectrum. The real signal value will be stored in
--   the second array. See 'CoeffArr' for a detailed frequency spectrum
--   desription.
syntheseDFT :: StorableArray Int Double -> StorableArray Int Double -> IO ()
syntheseDFT coeffs wave = seq initfou $
 withStorableArray wave $ \ptr_w ->
 withStorableArray coeffs $ \ptr_c ->
 cSynthese ptr_c ptr_w

-- | Fades a signal which is stored in an array of length 1024 in and out.
--   The signal is multiplied with a cosinus curve with range 0..pi.
kurveDFT :: StorableArray Int Double -> IO ()
kurveDFT wave = seq initfou $
 withStorableArray wave $ \ptr_w ->
 cKurve ptr_w
---------------------------------------------------------------------------------------------------
--
---------------------------------------------------------------------------------------------------
-- | Represents a frequency spectrum.
--   Range is from 0 to 1023 (N=1024).
--   Index 0 is constant value, index 1 is Nyquest frequency (real part of the N\/2 frequency).
--   Index 2 and 3 are the complex value for the basefrequency, index 4 and 5
--   are the complex value for the double basefrequency, 6 and 7 for 3 * basefrequency and so on:
--
-- > arr [0] = real part of zero frequency (imaginary part is allways 0)
-- > arr [1] = real part of Nyquest frequency (N/2) (imaginary part is allways 0)
-- > arr [2] = real part of base frequency
-- > arr [3] = imag part of base frequency
-- > arr [4] = real part of base frequency * 2
-- > arr [5] = imag part of base frequency * 2
-- > arr [6] = real part of base frequency * 3
-- > arr [7] = imag part of base frequency * 3
-- > ...
-- > arr [N-4] = real part of base frequency * (N/2 - 2)
-- > arr [N-3] = imag part of base frequency * (N/2 - 2)
-- > arr [N-2] = real part of base frequency * (N/2 - 1)
-- > arr [N-1] = imag part of base frequency * (N/2 - 1)
--
-- Pseudocode:
--
-- > c (0)   = (arr [0] :+ 0)
-- > c (512) = (arr [1] :+ 0)
--
--  and for all i=[1..511]:
--
-- > c (i) = (arr [i*2] :+ arr [i*2+1])
--
type CoeffArr = StorableArray Int Double

{-
-- | Reads a Complex value of a CoeffArr as described at 'CoeffArr'.
readCoeffArr :: CoeffArr -> Int -> IO (Complex Double)
readCoeffArr arr n = if n == 0 then readArray arr 0 >>= \x -> readArray arr 1023 >>= \y -> return (x :+ y)
                               else let j  = 2*n
                                        j' = j-1
                                    in readArray arr j' >>= \x -> readArray arr j >>= \y -> return (x :+ y)

-- | Writes a Complex value of a CoeffArr as described at 'CoeffArr'.
writeCoeffArr :: CoeffArr -> Int -> Complex Double -> IO ()
writeCoeffArr arr n c = if n == 0 then writeArray arr 0 (realPart c) >> writeArray arr 1023 (imagPart c)
                                  else let j  = 2*n
                                           j' = j-1
                                  in writeArray arr j' (realPart c) >> writeArray arr j (imagPart c)
-}

-- | Reads a Complex value of a CoeffArr (with Nyquest as imaginary part of coeff 0).
readCoeffArr :: CoeffArr -> Int -> IO (Complex Double)
readCoeffArr arr n =
 let j  = 2*n
     j' = j+1
 in readArray arr j >>= \x -> readArray arr j' >>= \y -> return (x :+ y)

-- | Writes a Complex value of a CoeffArr (with Nyquest as imaginary part of coeff 0).
writeCoeffArr :: CoeffArr -> Int -> Complex Double -> IO ()
writeCoeffArr arr n c =
 let j  = 2*n
     j' = j+1
 in writeArray arr j (realPart c) >> writeArray arr j' (imagPart c)

storeCoeff :: IOArray Int (Complex Double) -> CoeffArr -> IO ()
storeCoeff arr brr = do
 for 0 (<512) (+1) $ \i -> let i1 = i * 2
                               i2 = i1 + 1
                           in do (x :+ y) <- readArray arr i
                                 writeArray brr i1 x
                                 writeArray brr i2 y

unstoreCoeff :: CoeffArr -> IOArray Int (Complex Double) -> IO ()
unstoreCoeff brr arr = do
 for 0 (<512) (+1) $ \i -> let i1 = i * 2
                               i2 = i1 + 1
                           in do x <- readArray brr i1
                                 y <- readArray brr i2
                                 writeArray arr i (x :+ y)
---------------------------------------------------------------------------------------------------
-- | An action that maps the fourier coefficients from the first array to the second.
--   Modifying the data while mapping means filtering. See 'CoeffArr' for a description of the arrays.
type CoeffMap a = CoeffArr -> CoeffArr -> IO a

-- | Constructs a Coeffmap.
--   The sublists contain the 512 real values with which the basefrequency and the 511 complex fourier coefficients
--   are multiplied (the Nyquist frequency will be zero).
--   The first value is the factor for the constant coefficient, the second for the base frequency, the next one
--   for the double base freq and so on.
--   If the sublists have more than 512 elements, these elements are thrown away.
--   If it is shorter, the array will be filled up with zeros.
{-
mkCoeffMap :: [[Double]] -> IO (CoeffMap ())
mkCoeffMap cs = do
 r <- newIORef cs
 b <- newArray (0,511) 0.0
 return $ \a1 a2 -> let loop i (x:xs) = if i >= 512 then return () else
                                        writeArray b i x >> loop (i+1) xs
                        loop i []     = for i (<512) (+1) (\i -> writeArray b i 0.0)
                    in do xs <- readIORef r
                          let (xi,xr) | null xs   = ([], [])
                                      | otherwise = (head xs, tail xs)
                          writeIORef r xr
                          loop 0 xi
                          coeffmap b a1 a2
-}

mkCoeffMap :: [Double] -> IO (CoeffMap ())
mkCoeffMap cs = do
 r <- newIORef cs
 b <- newArray (0,511) 0.0
 return $ \a1 a2 -> let loop i (x:xs) = if i >= 512
                                         then writeIORef r (x:xs)
                                         else writeArray b i x >> loop (i+1) xs
                        loop i []     = for i (<512) (+1) (\i -> writeArray b i 0.0)
                                        >> writeIORef r []
                    in do xs <- readIORef r
                          loop 0 xs
                          coeffmap b a1 a2

-- | 'coeffmap' takes an Array with 512 values and uses them for a weigted map of the fourier
--   coefficients. The Nyquist frequency is muliplied with 0.0.
coeffmap :: StorableArray Int Double -> CoeffMap ()
coeffmap filt coeffs coeffs' = do
 for 0 (<512) (+1) (\i -> let j  = 2*i
                              j' = j + 1
                          in
                          readArray filt i >>= \c ->
                          readArray coeffs j  >>= \c1 ->
                          readArray coeffs j' >>= \c2 ->
                          writeArray coeffs' j  (c * c1) >>
                          writeArray coeffs' j' (c * c2) )
 writeArray coeffs' 1 0.0

{-
coeffmap' :: ([Complex Double] -> [Complex Double]) -> IO (CoeffMap ())
coeffmap' f = do
 let toDest arr 1024 _   = return ()
     toDest arr n (x:xs) = writeArray arr n x >> toDest arr (n+1) xs
     toDest arr n []     = writeArray arr n 0 >> toDest arr (n+1) []
 return $ \arr brr -> getElems arr >>= toDest brr 0 . f
-}
---------------------------------------------------------------------------------------------------
-- | Constructs an action that maps wave-data to coefficient-data.
--   Has a delay of 512, i. e. the first 512 elements are zero and the result has (these)
--   512 elemets more than the input.
mkAnalyse :: IO (Maybe Double -> IO (Maybe (Complex Double)))
mkAnalyse = do
 warr <- newArray (0, 1023) 0.0
 carr <- newArray (0, 1023) 0.0
 oarr <- newArray (0, 511) (0.0 :+ 0.0)
-- SIGNALSTREAM reada closea <- openSignalStream sa
 rpos <- newIORef 0
 rcnt <- newIORef Nothing
 let read ma = checkpos >> readIORef rcnt >>= maybe (more ma) rest
     more ma = maybe irest snext ma
     rest k | k <= 0    = return Nothing
            | otherwise =        do writeIORef rcnt $ Just (k-1)
                                    pos <- readIORef rpos
                                    x <- readArray oarr pos
                                    writeIORef rpos (pos + 1)
                                    return $ Just x
     irest =        do writeIORef rcnt (Just 511)
                       pos <- readIORef rpos
                       fillrest pos
                       x <- readArray oarr pos
                       writeIORef rpos (pos + 1)
                       return $ Just x
     snext a =        do pos <- readIORef rpos
                         x <- readArray oarr pos
                         writeArray warr (pos+512) a
                         writeIORef rpos (pos + 1)
                         return $ Just x
     checkpos = readIORef rpos >>= \p -> if p < 512 then return () else do
                writeIORef rpos 0
                analyseDFT warr carr
                unstoreCoeff carr oarr
                for 0 (<512) (+1) (\i -> readArray warr (i+512) >>= writeArray warr i)
     fillrest k = for (512+k) (<1024) (+1) (\i -> writeArray warr i 0.0)
 return read
---------------------------------------------------------------------------------------------------
-- | Constructs an action that maps coefficient-data to wave-data.
--   Has a delay of 512, i. e. the first 512 elements are zero and the result has (these)
--   512 elemets more than the input.
mkSynthese :: IO (Maybe (Complex Double) -> IO (Maybe Double))
mkSynthese = do
 iarr <- newArray (0, 511) (0.0 :+ 0.0)
 carr <- newArray (0, 1023) 0.0
 warr <- newArray (0, 1023) 0.0
 (oarr :: IOArray Int Double) <- newArray (0, 511) 0.0
-- SIGNALSTREAM reada closea <- openSignalStream sa
 rpos <- newIORef 0
 rcnt <- newIORef Nothing
 let read ma = checkpos >> readIORef rcnt >>= maybe (more ma) rest
     more ma = maybe irest snext ma
     rest k | k <= 0    = return Nothing
            | otherwise =        do writeIORef rcnt $ Just (k-1)
                                    pos <- readIORef rpos
                                    x <- readArray oarr pos
                                    writeIORef rpos (pos + 1)
                                    return $ Just x
     irest =        do writeIORef rcnt (Just 511)
                       pos <- readIORef rpos
                       fillrest pos
                       x <- readArray oarr pos
                       writeIORef rpos (pos + 1)
                       return $ Just x
     snext a =        do pos <- readIORef rpos
                         x <- readArray oarr pos
                         writeArray iarr pos a
                         writeIORef rpos (pos + 1)
                         return $ Just x
     checkpos = readIORef rpos >>= \p -> if p < 512 then return () else do
                writeIORef rpos 0
                storeCoeff iarr carr
                for 0 (<512) (+1) (\i -> readArray warr (i+512) >>= writeArray oarr i)
                syntheseDFT carr warr
                kurveDFT warr
                for 0 (<512) (+1) (\i -> readArray warr i >>= \x ->
                                         readArray oarr i >>= \y ->  writeArray oarr i (x+y))
     fillrest k = for k (<512) (+1) (\i -> writeArray iarr i (0.0 :+ 0.0))
 return read
---------------------------------------------------------------------------------------------------
-- | Constructs an action that maps wave-data to wave-data via a Fast Foutrier Transform and inverse,
--   filtered by the given 'CoeffMap'.
--   Has a delay of 1024.
mkFilterBuffered :: IO (CoeffMap ()) -> IO (Maybe Double -> IO (Maybe Double))
mkFilterBuffered mkf = do
 f <- mkf
 warr <- newArray (0, 1023) 0.0
 carr <- newArray (0, 1023) 0.0
 carr' <- newArray (0, 1023) 0.0
 warr' <- newArray (0, 1023) 0.0
 (oarr :: IOArray Int Double) <- newArray (0, 511) 0.0
-- SIGNALSTREAM reada closea <- openSignalStream sa
 rpos <- newIORef 0
 rcnt <- newIORef Nothing
 let read ma = checkpos >> readIORef rcnt >>= maybe (more ma) rest
     more ma = maybe irest snext ma
     rest k | k <= 0    = return Nothing
            | otherwise =        do writeIORef rcnt $ Just (k-1)
                                    pos <- readIORef rpos
                                    x <- readArray oarr pos
                                    writeIORef rpos (pos + 1)
                                    return $ Just x
     irest =        do writeIORef rcnt (Just 511)
                       pos <- readIORef rpos
                       fillrest pos
                       x <- readArray oarr pos
                       writeIORef rpos (pos + 1)
                       return $ Just x
     snext a =        do pos <- readIORef rpos
                         x <- readArray oarr pos
                         writeArray warr (pos+512) a
                         writeIORef rpos (pos + 1)
                         return $ Just x
     checkpos = readIORef rpos >>= \p -> if p < 512 then return () else do
                writeIORef rpos 0
                analyseDFT warr carr
                for 0 (<512) (+1) (\i -> readArray warr (i+512) >>= writeArray warr i >>
                                         readArray warr' (i+512) >>= writeArray oarr i)
                f carr carr'
                syntheseDFT carr' warr'
                kurveDFT warr'
                for 0 (<512) (+1) (\i -> readArray warr' i >>= \x ->
                                         readArray oarr i >>= \y ->  writeArray oarr i (x+y))

     fillrest k = for (512+k) (<1024) (+1) (\i -> writeArray warr i 0.0)
 return read
---------------------------------------------------------------------------------------------------

-- | Constructs an action that maps wave-data to wave-data via a Fast Foutrier Transform and inverse,
--   filtered by the given 'CoeffMap'.
--   Has a delay of 1024.
mkFilterBuffered' :: s -> IO (s -> CoeffMap s) -> IO (Maybe Double -> IO (Maybe Double))
mkFilterBuffered' st mkf = do
 stref <- newIORef st
 f <- mkf
 warr <- newArray (0, 1023) 0.0
 carr <- newArray (0, 1023) 0.0
 carr' <- newArray (0, 1023) 0.0
 warr' <- newArray (0, 1023) 0.0
 (oarr :: IOArray Int Double) <- newArray (0, 511) 0.0
-- SIGNALSTREAM reada closea <- openSignalStream sa
 rpos <- newIORef 0
 rcnt <- newIORef Nothing
 let read ma = checkpos >> readIORef rcnt >>= maybe (more ma) rest
     more ma = maybe irest snext ma
     rest k | k <= 0    = return Nothing
            | otherwise =        do writeIORef rcnt $ Just (k-1)
                                    pos <- readIORef rpos
                                    x <- readArray oarr pos
                                    writeIORef rpos (pos + 1)
                                    return $ Just x
     irest =        do writeIORef rcnt (Just 511)
                       pos <- readIORef rpos
                       fillrest pos
                       x <- readArray oarr pos
                       writeIORef rpos (pos + 1)
                       return $ Just x
     snext a =        do pos <- readIORef rpos
                         x <- readArray oarr pos
                         writeArray warr (pos+512) a
                         writeIORef rpos (pos + 1)
                         return $ Just x
     checkpos = readIORef rpos >>= \p -> if p < 512 then return () else do
                writeIORef rpos 0
                analyseDFT warr carr
                for 0 (<512) (+1) (\i -> readArray warr (i+512) >>= writeArray warr i >>
                                         readArray warr' (i+512) >>= writeArray oarr i)
                s <- readIORef stref
                s' <- f s carr carr'
                writeIORef stref s'
                syntheseDFT carr' warr'
                kurveDFT warr'
                for 0 (<512) (+1) (\i -> readArray warr' i >>= \x ->
                                         readArray oarr i >>= \y ->  writeArray oarr i (x+y))

     fillrest k = for (512+k) (<1024) (+1) (\i -> writeArray warr i 0.0)
 return read
---------------------------------------------------------------------------------------------------

{-
-- | Reads a Complex value of a CoeffArr
readCoeffArr :: CoeffArr -> Int -> IO (Complex Double)
readCoeffArr arr n = if n == 0 then readArray arr 0 >>= \x -> readArray arr 1023 >>= \y -> return (x :+ y)
                               else let j  = 2*n
                                        j' = j-1
                                    in readArray arr j' >>= \x -> readArray arr j >>= \y -> return (x :+ y)

-- | Writes a Complex value of a CoeffArr
writeCoeffArr :: CoeffArr -> Int -> Complex Double -> IO ()
writeCoeffArr arr n c = if n == 0 then writeArray arr 0 (realPart c) >> writeArray arr 1023 (imagPart c)
                                  else let j  = 2*n
                                           j' = j-1
                                  in writeArray arr j' (realPart c) >> writeArray arr j (imagPart c)

-- | Copies coefficient data from an IOArray with Complex values (512 elements) to a
--   StorableArray with Doubles (1024 elements)
storeCoeff :: IOArray Int (Complex Double) -> CoeffArr -> IO ()
storeCoeff arr brr = do
 for 1 (<512) (+1) $ \i -> let i2 = i * 2
                               i1 = i2 - 1
                           in do (x :+ y) <- readArray arr i
                                 writeArray brr i1 x
                                 writeArray brr i2 y
 (x :+ y) <- readArray arr 0
 writeArray brr 0 x
 writeArray brr 1023 y

-- | Copies coefficient data from a StorableArray with Doubles (1024 elements) to
--   an IOArray with Complex values (512 elements).
unstoreCoeff :: CoeffArr -> IOArray Int (Complex Double) -> IO ()
unstoreCoeff brr arr = do
 for 1 (<512) (+1) $ \i -> let i2 = i * 2
                               i1 = i2 - 1
                           in do x <- readArray brr i1
                                 y <- readArray brr i2
                                 writeArray arr i (x :+ y)
 x <- readArray brr 0
 y <- readArray brr 1023
 writeArray arr 0 (x :+ y)

coeffmap :: StorableArray Int Double -> CoeffMap ()
coeffmap filt coeffs coeffs' = do
 for 1 (<512) (+1) (\i -> let j' = 2*i
                              j  = j' - 1
                          in
                          readArray filt i >>= \c ->
                          readArray coeffs j  >>= \c1 ->
                          readArray coeffs j' >>= \c2 ->
                          writeArray coeffs' j  (c * c1) >>
                          writeArray coeffs' j' (c * c2) )
 readArray filt 0 >>= \c -> readArray coeffs 0 >>= \c' -> writeArray coeffs' 0 (c * c')
 writeArray coeffs' 1023 0.0
-}


import System.Environment
import Data.Complex
import Paraiso

main = do
  args <- getArgs
  let arch = if "--cuda" `elem` args then
               CUDA 128 128
             else
               X86
  putStrLn $ compile arch code
    where
      code = do
         parallel 16384 $ do
           r <- allocate
           x <- allocate
           r =$ Rand 0.0 (4.0::Double)
           x =$ Rand 0.0 (1.0::Double)
           cuda $ do
             sequential 65536 $ do
               x =$ r * x * (1-x)
           output [r,x]


{-| Provides 'hash' function for many data types
-}
module Hash(Hash, combineHash, emptyHash, hashToInt, hashToMax, Hashable(..)) where
--
-- Hash a value.  Hashing produces an Int of
-- unspecified range.
--
import Data.Array
import Data.Complex
import Data.Ratio

newtype Hash = H Int deriving (Eq)

----instance Show Hash where
--    showsType _ = showString "Hash"

combineHash :: Hash -> Hash -> Hash
combineHash (H x) (H y) = H (x+y)

emptyHash :: Hash
emptyHash = H 0

class Hashable a where
    hash :: a -> Hash

instance Hashable Char where
    hash x = H $ fromEnum x

instance Hashable Int where
    hash x = H $ x

instance Hashable Integer where
    hash x = H $ fromInteger x

instance Hashable Float where
    hash x = H $ truncate x

instance Hashable Double where
    hash x = H $ truncate x

instance Hashable (IO a) where
    hash x = H 0

instance Hashable () where
    hash x = H 0

instance Hashable (a -> b) where
    hash x = H 0

instance (Hashable a) => Hashable (Maybe a) where
    hash Nothing = H 0
    hash (Just x) = hash x

instance (Hashable a, Hashable b) => Hashable (Either a b) where
    hash (Left x) = hash x
    hash (Right y) = hash y

-- Denna bör inte vara bortkommenterad men jag kunde inte
-- göra en instans med String nedan
--instance Hashable a => Hashable [a] where
--    hash l = H $ f l 0
--      where f :: (Hashable a) => [a] -> Int -> Int
--            f [] r = r
--            f (c:cs) r = f cs (3*r + (case hash ( c ) of H h -> h) )

instance (Hashable a,Enum a) => Hashable [a] where
    hash l = H $ f l 0
        where f :: Enum b => [b] -> Int -> Int
              f [] r = r
              f (c:cs) r = f cs (3*r + fromEnum c)

instance (Hashable a, Hashable b) => Hashable (a,b) where
    hash (a,b) = H $ (case hash ( a ) of H h -> h)  + 3 * (case hash ( b ) of H h -> h)

instance (Hashable a, Hashable b, Hashable c) => Hashable (a,b,c) where
    hash (a,b,c) = H $ (case hash ( a ) of H h -> h)  + 3 * (case hash ( b ) of H h -> h)  + 5 * (case hash ( c ) of H h -> h)

instance (Hashable a, Hashable b, Hashable c, Hashable d) => Hashable (a,b,c,d) where
    hash (a,b,c,d) = H $ (case hash ( a ) of H h -> h)  + 3 * (case hash ( b ) of H h -> h)  + 5 * (case hash ( c ) of H h -> h)  + 7 * (case hash ( d ) of H h -> h)

instance (Hashable a, Hashable b, Hashable c, Hashable d, Hashable e) => Hashable (a,b,c,d,e) where
    hash (a,b,c,d,e) = H $ (case hash ( a ) of H h -> h)  + 3 * (case hash ( b ) of H h -> h)  + 5 * (case hash ( c ) of H h -> h)  + 7 * (case hash ( d ) of H h -> h)  + 9 * (case hash ( e ) of H h -> h)

instance Hashable Bool where
    hash False = H 0
    hash True = H 1

instance (Integral a, Hashable a) => Hashable (Ratio a) where
    hash x = H $ (case hash ( denominator x ) of H h -> h)  + (case hash ( numerator x ) of H h -> h)

instance (RealFloat a, Hashable a) => Hashable (Complex a) where
    hash (x :+ y) = H $ (case hash ( x ) of H h -> h)  + (case hash ( y ) of H h -> h)

instance (Ix a) => Hashable (Array a b) where
    hash x = H $ 0 -- !!!

hashToInt :: Int -> Hash -> Int
hashToInt maxhash x =
    case x of
    H h ->
        if h < 0 then
            if -h < 0 then 0
            else (-h) `rem` maxhash
        else h `rem` maxhash

hashToMax maxhash x =
    case hash x of
    H h ->
        if h < 0 then
            if -h < 0 then 0
            else (-h) `rem` maxhash
        else h `rem` maxhash


{-# LANGUAGE ForeignFunctionInterface #-}
module Data.Matrix.Static.Internal
    ( c_dd_mul
    , c_ds_mul
    , c_sd_mul
    , c_ss_mul
    , c_ss_cmul
    , c_sd_plus
    , c_ss_plus
    , c_inverse
    , c_cholesky
    , c_eig
    , c_eigs
    , c_eigsh
    , c_seigs
    , c_seigsh
    , c_geigsh
    , c_bdcsvd

    , computationInfo
    ) where

import Data.Complex (Complex)
import Foreign
import Foreign.C.Types
import Foreign.C.String

-------------------------------------------------------------------------------
-- Arithmetic
-------------------------------------------------------------------------------
foreign import ccall "eigen_dd_mul"
    c_dd_mul :: CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_ds_mul"
    c_ds_mul :: CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_sd_mul"
    c_sd_mul :: CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_sd_plus"
    c_sd_plus :: CInt
          -> Ptr a -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_ss_mul"
    c_ss_mul :: CInt
          -> Ptr (Ptr a) -> Ptr CInt -> Ptr (Ptr CInt) -> CInt -> CInt -> Ptr CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_ss_cmul"
    c_ss_cmul :: CInt
          -> Ptr (Ptr a) -> Ptr CInt -> Ptr (Ptr CInt) -> CInt -> CInt -> Ptr CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_ss_plus"
    c_ss_plus :: CInt
          -> Ptr (Ptr a) -> Ptr CInt -> Ptr (Ptr CInt) -> CInt -> CInt -> Ptr CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> Ptr a -> Ptr CInt -> Ptr CInt -> CInt -> CInt -> CInt
          -> IO CString

foreign import ccall "eigen_inverse"
    c_inverse :: CInt
              -> Ptr a -> CInt -> CInt
              -> Ptr a -> CInt -> CInt
              -> IO CString

foreign import ccall "eigen_cholesky"
    c_cholesky :: CInt
               -> Ptr a -> Ptr a ->  CInt -> IO CString

foreign import ccall "eigen_eig"
    c_eig :: Ptr (Complex Double) -> Ptr (Complex Double)
          -> Ptr Double -> CInt -> IO CString

foreign import ccall "spectral_eigs"
    c_eigs :: CInt      -- ^ The number of eigenvectors to search
           -> Ptr (Complex Double)   -- ^
           -> Ptr (Complex Double)   -- ^
           -> Ptr Double   -- ^ Input matrix
           -> CInt         -- ^ Matrix size
           -> CInt   -- ^ ncv
           -> CInt   -- ^ max iterations
           -> Double  -- ^ tolerance
           -> CInt
           -> Double
           -> CInt
           -> IO CInt

foreign import ccall "spectral_eigsh"
    c_eigsh :: CInt
            -> Ptr Double -> Ptr Double
            -> Ptr Double -> CInt
            -> CInt -> CInt -> Double -> CInt -> Double -> CInt
            -> IO CInt

foreign import ccall "spectral_seigs"
    c_seigs :: CInt
            -> Ptr (Complex Double) -> Ptr (Complex Double)
            -> Ptr Double -> Ptr CInt -> Ptr CInt -> CInt -> CInt
            -> CInt -> CInt -> Double -> CInt -> Double -> CInt
            -> IO CInt

foreign import ccall "spectral_seigsh"
    c_seigsh :: CInt
             -> Ptr Double -> Ptr Double
             -> Ptr Double -> Ptr CInt -> Ptr CInt -> CInt -> CInt
             -> CInt -> CInt -> Double -> CInt -> Double -> CInt
             -> IO CInt

foreign import ccall "spectral_geigsh"
    c_geigsh :: CInt
             -> Ptr Double -> Ptr Double
             -> Ptr Double -> CInt
             -> Ptr Double -> Ptr CInt -> Ptr CInt -> CInt
             -> CInt -> CInt -> Double -> CInt
             -> IO CInt

foreign import ccall "eigen_bdcsvd"
    c_bdcsvd :: CInt -> Ptr a -> Ptr b -> Ptr a
             -> Ptr a -> CInt -> CInt -> IO CString

computationInfo :: CInt -> Maybe String
computationInfo 0 = Nothing
computationInfo 1 = Just "NOT_COMPUTED"
computationInfo 2 = Just "NOT_CONVERGING"
computationInfo 3 = Just "NUMERICAL_ISSUE"
computationInfo _ = Just "UNKNOWN ERROR"
{-# INLINE computationInfo #-}


-- This file is part of Quipper. Copyright (C) 2011-2016. Please see the
-- file COPYRIGHT for a list of authors, copyright holders, licensing,
-- and other details. All rights reserved.
--
-- ======================================================================

{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE TypeSynonymInstances #-}

{-# OPTIONS -fcontext-stack=50 #-}

-- | This module contains the Quipper implementation of the Quantum
-- Linear Systems Algorithm.
--
-- The algorithm estimates the radar cross section for a FEM
-- scattering problem by using amplitude estimation to calculate
-- probability amplitudes.
--
-- The notations are based on the paper
--
-- * B. D. Clader, B. C. Jacobs, C. R. Sprouse. Quantum algorithm to
-- calculate electromagnetic scattering cross
-- sections. <http://arxiv.org/abs/1301.2340>.
module Algorithms.QLS.QLS where

import Quipper

import QuipperLib.QFT
import QuipperLib.Arith
import QuipperLib.Decompose

import Data.Complex
import qualified Data.Map as Map

import qualified Algorithms.QLS.TemplateOracle as Template
import Algorithms.QLS.QDouble
import Algorithms.QLS.QSignedInt
import Algorithms.QLS.CircLiftingImport
import Algorithms.QLS.Utils

import Libraries.Auxiliary(boollist_of_int_bh)

-- | The type of 'oracle_A' input arguments during runtime.
type OracleARunTime = Double  -- ^Value resolution.
       -> Int                 -- ^Band.
       -> Bool                -- ^Argflag.
       -> ([Qubit],[Qubit],[Qubit])  -- ^(x=index,y+node,z+value).
       -> Circ ([Qubit],[Qubit],[Qubit])

-- | The type of 'oracle_b' and 'oracle_r' input arguments during runtime.
type OracleBRRunTime = Double  -- ^Magnitude resolution.
        -> Double              -- ^Phase resolution.
        -> ([Qubit],[Qubit],[Qubit]) -- ^(x=index,m+magnitude,p+phase).
        -> Circ ([Qubit],[Qubit],[Qubit])

-- | A type to encapsulate all three oracles.
data Oracle = Oracle {
  oracle_A :: RunTimeParam -> OracleARunTime,
  oracle_b :: RunTimeParam -> OracleBRRunTime,
  oracle_r :: RunTimeParam -> OracleBRRunTime
}

-- | A set of oracles using only blackboxes.
dummy_oracle :: Oracle
dummy_oracle = Oracle {
  oracle_A = \d r i b x -> named_gate "Oracle A" x,
  oracle_b = \d1 d2 r x -> named_gate "Oracle b" x,
  oracle_r = \d1 d2 r x -> named_gate "Oracle r" x
}

-- | A type to hold the runtime parameters.
data RunTimeParam = RT_param {
  k :: Double,       -- ^ Wave number.
  theta :: Double,   -- ^ Incident wave angle.
  phi :: Double,     -- ^ Direction of desired far field radiation pattern.
  e0 :: Double,      -- ^ Incident wave amplitude.
  lambda :: Double,  -- ^ Wavelength.
  xlength :: Double, -- ^ /x/-length of square scattering region.
  ylength :: Double, -- ^ /y/-length of square scattering region.

  scatteringnodes :: [(Int,Int)], -- ^ Metallic region.

  nx :: Int,
  ny :: Int,
  lx :: Double,
  ly :: Double,

  kappa :: Double,
  epsilon :: Double,
  t0 :: Double,
  r :: Double,
  b_max :: Double,
  r_max :: Double,
  d  :: Int,
  nb :: Int,
  p2 :: Double,
  n0 :: Int,
  n1 :: Int,
  n2 :: Int,
  n4 :: Int,

  -- argflags for oracle_A
  magnitudeArgflag :: Bool,
  phaseArgflag :: Bool
} deriving (Show)

-- | A convenient set of runtime parameters for testing.
dummy_RT_param :: RunTimeParam
dummy_RT_param = RT_param {

  k = 2.0*pi*1.0,
  theta = 0.0*pi/4.0,
  phi = 0.0,
  e0 = 1.0,
  lambda = (k dummy_RT_param)/(2.0*pi),
  xlength = 2.0*(lambda dummy_RT_param),
  ylength = 2.0*(lambda dummy_RT_param),

  scatteringnodes =
    let rt = dummy_RT_param in
    let xul = round((fromIntegral $ nx rt)/2)
              - round((xlength rt)/(2.0*(lx rt))) in -- Upper left x index
    let yul = round((fromIntegral $ ny rt)/2)
              - round((ylength rt)/(2.0*(ly rt))) in -- Upper left y index
    let xlr = round((fromIntegral $ nx rt)/2)
              + round((xlength rt)/(2.0*(lx rt))) in  -- Lower right x index
    let ylr = round((fromIntegral $ ny rt)/2)
              + round((ylength rt)/(2.0*(ly rt))) in -- Lower right y in
      [(xul, yul), (xlr, ylr)],

  nx = 12885,
  ny = 12885,
  lx = 0.1,
  ly = 0.1,

  kappa = 1.0,
  epsilon = 1.0,
  t0 = 1.0,
  r = 1.0,
  b_max = 1.0,
  r_max = 1.0,
  d  = 3,
  nb = 2,
  p2 = 3,
  n0 = 3,
  n1 = 3,
  n2 = 3,
  n4 = 3,

  -- argflags for oracle_A
  magnitudeArgflag = False,
  phaseArgflag = True
}

-- | A set of larger values, for testing scalability.
large_RT_param :: RunTimeParam
large_RT_param = RT_param {

  k = 2.0*pi*1.0,
  theta = 0.0*pi/4.0,
  phi = 0.0,
  e0 = 1.0,
  lambda = (k large_RT_param)/(2.0*pi),
  xlength = 2.0*(lambda large_RT_param),
  ylength = 2.0*(lambda large_RT_param),

  scatteringnodes =
    let rt = large_RT_param in
    let xul = round((fromIntegral $ nx rt)/2)
              - round((xlength rt)/(2.0*(lx rt))) in -- Upper left x index
    let yul = round((fromIntegral $ ny rt)/2)
              - round((ylength rt)/(2.0*(ly rt))) in -- Upper left y index
    let xlr = round((fromIntegral $ nx rt)/2)
              + round((xlength rt)/(2.0*(lx rt))) in  -- Lower right x index
    let ylr = round((fromIntegral $ ny rt)/2)
              + round((ylength rt)/(2.0*(ly rt))) in -- Lower right y in
      [(xul, yul), (xlr, ylr)],

  nx = 12885,
  ny = 12885,
  lx = 0.1,
  ly = 0.1,

  kappa = 1e4,
  epsilon = 0.01,
  t0 = 7.0e6,
  r = 2.5e12,
  b_max = 5.0,
  r_max = 1.01,
  d  = 7,
  nb = 9,
  p2 = (  1.0 / (4-(4**(1/3)))  ),
  n0 = 14,
  n1 = 24,
  n2 = 30,
  n4 = 65,

  -- argflags for oracle_A
  magnitudeArgflag = False,
  phaseArgflag = True
}

-- | A set of smaller values, for manageable yet meaningful output.
small_RT_param :: RunTimeParam
small_RT_param = RT_param {

  k = 2.0*pi*1.0,
  theta = 0.0*pi/4.0,
  phi = 0.0,
  e0 = 1.0,
  lambda = (k small_RT_param)/(2.0*pi),
  xlength = 2.0*(lambda small_RT_param),
  ylength = 2.0*(lambda small_RT_param),

  scatteringnodes =
    let xul = 2 in
    let yul = 2 in
    let xlr = 3 in
    let ylr = 3 in
      [(xul, yul), (xlr, ylr)],

  nx = 4,
  ny = 4,
  lx = 0.1,
  ly = 0.1,

  kappa = 1e4,
  epsilon = 0.01,
  t0 = 7.0e6,
  r = 2.5e12,
  b_max = 5.0,
  r_max = 1.01,
  d  = 7,
  nb = 9,
  p2 = (  1.0 / (4-(4**(1/3)))  ),
  n0 = 14,
  n1 = 24,
  n2 = 6,
  n4 = 65,

  -- argflags for oracle_A
  magnitudeArgflag = False,
  phaseArgflag = True
}

-- | Apply an [exp −/iYt/] gate. The timestep /t/ is a parameter.
expYt :: Timestep -> Qubit -> Circ Qubit
expYt = named_rotation "exp(-i%Y)"

-- | Apply an [exp −/iYt/] gate. The timestep /t/ is a parameter.
expYt_at :: Timestep -> Qubit -> Circ ()
expYt_at = named_rotation_at "exp(-i%Y)"

-- | Read a list of bits and make it into a 'Double', by multiplying
-- its integer value by the provided factor.
dynamic_lift_double :: Double -> [Bit] -> Circ Double
dynamic_lift_double factor cl = do
      cdiscard cl

-- Implementation note: removed dynamic_lift as for now it breaks all
-- output formats except ASCII.

--      bl <- dynamic_lift cl
      let sign = 1 -- if (head $ reverse bl) then 1 else -1
      let unsigned_value = 1 -- integer_of_intm_unsigned $
                              -- intm_of_boollist_bh (tail $ reverse bl)
      return (sign * factor * (fromIntegral unsigned_value))

-- | A black box gate to stand in as a replacement for QFT.
qft_for_show :: [Qubit] -> Circ [Qubit]
qft_for_show qs = named_gate "QFT" qs

-- | Main function: for estimating the radar cross section for a FEM
-- scattering problem. The problem can be reduced to the calculation
-- of four angles: φ[sub /b/], φ[sub /bx/], φ[sub /r/1] and φ[sub /r/0].
qlsa_FEM_main :: RunTimeParam -> Oracle -> Circ Double
qlsa_FEM_main param oracle = do
     comment "FEM_main"
     phi_b  <- qlsa_AmpEst_phi_b param oracle
     phi_bx <- qlsa_AmpEst_phi_bx param oracle
     phi_r1 <- qlsa_AmpEst_phi_bxr param oracle True
     phi_r0 <- qlsa_AmpEst_phi_bxr param oracle False
     let sigma = ((((fromIntegral $ nb param) ^ 2) * ((b_max param) ^ 2) * ((r_max param) ^ 2) * ((sin phi_b) ^ 2)) / ( 4 * pi))
     comment "FEM_main"
     return sigma

-- * Amplitude Estimation Functions

-- | Estimates φ[sub /b/], related to the probability of success for the
-- preparation of the known state /b/, using amplitude amplification.
qlsa_AmpEst_phi_b :: RunTimeParam -> Oracle -> Circ Double
qlsa_AmpEst_phi_b param oracle = do
    g <- qinit $ replicate (n0 param) False
    with_ancilla_init (replicate (n2 param) False) $ \x -> do
       label (g,x) ("g","x")
       with_ancilla $ \a -> do
           label (a) ("anc. a")
           with_ancilla $ \b -> do
               label (b) ("anc. b")
               g <- map_hadamard g
               u_b (x,b)
               loop g u_g (x,b,a)
               return ()
           return ()
    g' <- qft_big_endian g -- QFT : Is it really big-endian?
    value_bits  <- measure g'
    value_double <- dynamic_lift_double 1.0 value_bits
    return (pi * value_double / (2 ** (fromIntegral $ n0 param)))
    where
        loop :: [Qubit] -> (a -> Circ ()) -> a -> Circ ()
        loop [] f x = return ()
        loop (h:t) f x = do
            f x `controlled` h
            loop t f' x;
            where
               f' x = do f x; f x

        u_b :: ([Qubit],Qubit) -> Circ ()
        u_b xb = qlsa_StatePrep param xb (oracle_b oracle param) (1.0/(b_max param))

        u_g :: ([Qubit],Qubit,Qubit) -> Circ ()
        u_g (x,b,a) = do
            comment "U_g starts"
            gate_Z_at b
            -- For unitrary linear transformation, adjoint == inverse, hence reverse_...
            (reverse_generic_imp u_b) (x,b)
            qnot_at a `controlled` x .==. (map (\x -> False) x)
            gate_X_at a
            gate_Z_at a
            gate_X_at a
            qnot_at a `controlled` x .==. (map (\x -> False) x)
            u_b (x,b)
            comment "U_g ends"
            return ()

-- | Testing function for 'qlsa_AmpEst_phi_b'.
test_qlsa_AmpEst_phi_b :: Bool -> IO ()
test_qlsa_AmpEst_phi_b dummyRTParamFlag = do
    let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
    print_simple GateCount (qlsa_AmpEst_phi_b param dummy_oracle)
    print_simple Preview (qlsa_AmpEst_phi_b param dummy_oracle)

-- | Estimates φ[sub /bx/], related to the probability of success in
-- computing solution value /x/.
qlsa_AmpEst_phi_bx :: RunTimeParam -> Oracle -> Circ Double
qlsa_AmpEst_phi_bx param oracle  = do
    g <- qinit (take (n0 param) (repeat False))
    with_ancilla_init (take (n2 param) (repeat False)) $ \x -> do
      with_ancilla $ \a -> do
          with_ancilla $ \b -> do
              with_ancilla $ \s -> do
                  g <- map_hadamard g
                  u_bx (x,b,s)
                  loop g u_g (x,b,s,a)
                  return ()
              return ()
          return ()
    g' <- qft_big_endian g -- QFT : Is it really big-endian?
    value_bits  <- measure g'
    value_double <- dynamic_lift_double 1.0 value_bits
    return (pi * value_double / (2 ** (fromIntegral $ n0 param)))
    where
        loop :: [Qubit] -> (a -> Circ ()) -> a -> Circ ()
        loop [] f x = return ()
        loop (h:t) f x = do
            f x `controlled` h
            loop t f' x
            where
            f' x = do f x; f x

        u_bx :: ([Qubit],Qubit,Qubit) -> Circ ()
        u_bx (x,b,s) = do
            qlsa_StatePrep param (x,b) (oracle_b oracle param) (1.0/(b_max param))
            qlsa_Solve_x param (x,s) oracle
            return ()

        u_g :: ([Qubit],Qubit,Qubit,Qubit) -> Circ ()
        u_g (x,b,s,a) = do --named_gate_at "Ug_phi_bx" (x,b,s,a)
            qnot_at a `controlled` b .&&. s
            gate_Z_at a
            qnot_at a `controlled` b .&&. s
            (reverse_generic_imp u_bx) (x,b,s)
            qnot_at a `controlled` [ q .==. 0 | q <- x ] .&&. b .==. 0 .&&. s .==. 0
            gate_X_at a
            gate_Z_at a
            gate_X_at a
            qnot_at a `controlled` [ q .==. 0 | q <- x ] .&&. b .==. 0 .&&. s .==. 0
            u_bx (x,b,s)
            return ()

-- | Estimates φ[sub /r/0] and φ[sub /r/1] (depending on the boolean
-- parameter), related to the overlap of the solution with the
-- arbitrary state /r/.
qlsa_AmpEst_phi_bxr :: RunTimeParam -> Oracle -> Bool -> Circ Double
qlsa_AmpEst_phi_bxr param oracle target = do
    g <- qinit (take (n0 param) (repeat False))
    with_ancilla_init (take (n2 param) (repeat False)) $ \x -> do
      with_ancilla_init (take (n2 param) (repeat False)) $ \y -> do
        with_ancilla $ \a -> do
          with_ancilla $ \b -> do
             with_ancilla $ \s -> do
                with_ancilla $ \r -> do
                    with_ancilla $ \c -> do
                        g <- map_hadamard g
                        u_r (x,y,b,s,r,c)
                        loop g u_g (x,y,b,s,r,c,a)
                        return ()
                    return ()
                return ()
             return ()
          return ()
    g' <- qft_big_endian g -- QFT : Is it really big-endian?
    value_bits  <- measure g'
    value_double <- dynamic_lift_double 1.0 value_bits
    return (pi * value_double / (2 ** (fromIntegral $ n0 param)))
    where
    loop :: [Qubit] -> (a -> Circ ()) -> a -> Circ ()
    loop [] f x = return ()
    loop (h:t) f x = do
       f x `controlled` h
       loop t f' x
       where
         f' x = do f x; f x

    u_r :: ([Qubit],[Qubit],Qubit,Qubit,Qubit,Qubit) -> Circ ()
    u_r (x,y,b,s,r,c) = do
        qlsa_Solve_xr param (x,y,b,s,r,c) oracle
        return ()

    u_g :: ([Qubit],[Qubit],Qubit,Qubit,Qubit,Qubit,Qubit) -> Circ ()
    u_g (x,y,b,s,r,c,a) = do --named_gate_at "Ug_phi_bxr" (x,y,b,s,r,c,a)
         qnot_at a `controlled` (b .==. 1 .&&. s .==. 1 .&&. r .==. 1 .&&. c .==. target)
         gate_Z_at a
         qnot_at a `controlled` (b .==. 1 .&&. s .==. 1 .&&. r .==. 1 .&&. c .==. target)
         (reverse_generic_imp u_r) (x,y,b,s,r,c)
         qnot_at a `controlled` [ q .==. 0 | q <- x ] .&&. [ q .==. 0 | q <- y ] .&&. b .==. 0 .&&. s .==. 0 .&&. r .==. 0 .&&. c .==. 0
         gate_X_at a
         gate_Z_at a
         gate_X_at a
         qnot_at a `controlled` [ q .==. 0 | q <- x ] .&&. [ q .==. 0 | q <- y ] .&&. b .==. 0 .&&. s .==. 0 .&&. r .==. 0 .&&. c .==. 0
         u_r (x,y,b,s,r,c)
         return ()

-- * State Preparation.

-- | Prepares a quantum state /x/, as specified by an oracle function,
-- entangled with a single qubit flag /q/ marking the desired state.
qlsa_StatePrep ::
    RunTimeParam
    -> ([Qubit], Qubit) -- x & qare handles to wires to be changed by StatePrep
    -> OracleBRRunTime -- Common type of (oracle_b oracle) and (oracle_r oracle), if oracle is typed Oracle
    -> Double
    -> Circ ()
qlsa_StatePrep param (x, q) oracle phi0 = do
  _ <- (flip $ box ("qlsa_StatePrep_" ++ (show phi0))) (x,q) $ \(x,q) -> do
    -- comment "StatePrep starts"
    label (x,q) ("x", "q")
    with_ancilla_list (n4 param) $ \m -> do
        label (m) ("anc. m")
        with_ancilla_list (n4 param) $ \p -> do
            label (p) ("anc. p")
            x <- map_hadamard x
            (x, m, p) <- oracle phi0 (epsilon param) (x, m, p)
            qlsa_ControlledPhase p (epsilon param) False
            qlsa_ControlledRotation (m, q) phi0 False
            (x, m, p) <- oracle phi0 (epsilon param) (x, m, p)
            return (x,q)
    -- comment "StatePrep ends"
  return ()

-- | Testing function for 'qlsa_StatePrep'.
test_qlsa_StatePrep :: Bool -> IO ()
test_qlsa_StatePrep dummyRTParamFlag = do
          let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
          let oraclebRunTime = (oracle_b dummy_oracle param)
          let testCirc = do
              x <- qinit $ replicate (n2 param) False
              q <- qinit False
              qlsa_StatePrep param (x,q) oraclebRunTime 1.0
          print_simple GateCount testCirc
          print_simple Preview testCirc

-- * Linear System Solver Functions

-- | Implements the QLSA procedure to generate the solution state |/x/〉.
qlsa_Solve_x :: RunTimeParam ->([Qubit],Qubit) -> Oracle -> Circ ()
qlsa_Solve_x param (x,s) oracle = do
  _ <- (flip $ box "qlsa_Solve_x") (x,s) $ \(x,s) -> do
    with_ancilla_list (n1 param) $ \t -> do
        with_ancilla_list (length t) $ \f -> do
            let phi0 = 2 * pi  / (2 ** (fromIntegral $ n1 param) - (epsilon param))
            t <- map_hadamard t
            u_hs (t,x)
            t <- qft_big_endian t
            integer_inverse (t,f)
            qlsa_ControlledRotation (f,s) phi0 False
            integer_inverse (t,f)
            (reverse_generic_endo qft_big_endian) t
            (reverse_generic_imp u_hs) (t,x)
            t <- map_hadamard t
            return()
        return ()
    return (x,s)
  return ()
    where
        u_hs :: ([Qubit],[Qubit]) -> Circ ()
        u_hs (t,x) = do
            qlsa_HamiltonianSimulation param (t,x) (oracle_A oracle param)
            return ()

-- | Implementation of the integer division. The two registers are
-- supposed to be of the same size and represent little-headian
-- unsigned integers, i.e., the head of the list holds the least
-- significant bit.
integer_inverse :: ([Qubit],[Qubit]) -> Circ ()
integer_inverse (t,f) = do
  _ <- (flip $ box "integer_inverse") (t,f) $ \(t,f) -> do
    -- sanity check
    if (length t /= length f)
      then error "integer_inverse: registers of distinct sizes"
      else return ()
    -- initialize an unsigned integer to 2^(length t) - 1
    with_ancilla_init (map (\_ -> True) t) $ \num -> do
      -- perform the division (encapsulated in a subroutine)
      let d = classical_to_reversible $ \(t,num) -> do
                let x = ((qdint_of_qulist_lh num),(qdint_of_qulist_lh t))
                (_,_,f') <- uncurry q_div_unsigned x
                return $ qulist_of_qdint_lh f'
      d ((t,num),f)
      return (t,f)
  return ()

-- | Implements the complete QLSA procedure to find the
-- solution state |/x/〉 and then implements the swap protocol
-- required for estimation of 〈/x/|/r/〉.
qlsa_Solve_xr :: RunTimeParam ->([Qubit],[Qubit],Qubit,Qubit,Qubit,Qubit) -> Oracle -> Circ ()
-- qlsa_Solve_xr param (x,y,b,s,r,c) oracle  = named_gate_at "Solve_xr" (x,y,b,s,r,c)
qlsa_Solve_xr param (x,y,b,s,r,c) oracle = do
  _ <- (flip $ box "qlsa_Solve_xr") (x,y,b,s,r,c) $ \(x,y,b,s,r,c) -> do
    qlsa_StatePrep param (x,b) (oracle_b oracle param) (1.0/(b_max param))
    qlsa_Solve_x param (x,s) oracle
    qlsa_StatePrep param (y,r) (oracle_r oracle param) (1.0/(r_max param))
    hadamard_at c
    swap_at y x  `controlled` c
    hadamard_at c
    return (x,y,b,s,r,c)
  return ()

-- * Hamiltonian Simulation Functions.

-- | Uses a quantum register |/t/〉 to control the
-- implementation of the Suzuki method for simulating a Hamiltonian
-- specified by an oracle function.
qlsa_HamiltonianSimulation ::
    RunTimeParam
    -> ([Qubit], [Qubit])
    -> OracleARunTime
    -> Circ ()
qlsa_HamiltonianSimulation param (t, x) oracleA = do
  _ <- (flip $ box "qlsa_HamiltonianSimulation") (t,x) $ \(t,x) -> do
    -- Code the first line in a way that depends on the length of t rather
    -- explicitly on (n1 param) which is supposed to be the length of t
    -- Hence replaced the line below by the line after it.
    -- let denom = 2 * (r param) * (fromIntegral ((n1 param) - 1))
    let denom = 2 * (r param) * ( 2^((length t) - 1) )
    let t1 = (p2 param) * (t0 param) / denom
    let t2 = (1.0 - 4.0 * (p2 param)) * (t0 param) / denom
    (t,x) <- box_loopM  "TrotterLoop" (round $ r param) (t,x) (hs_loop t1 t2)
    return (t,x)
  return ()
    where
          hs_loop :: Double -> Double -> ([Qubit], [Qubit]) -> Circ ([Qubit], [Qubit])
          hs_loop t1 t2 (t,x) = do
            u_z_at (t,x) t1
            u_z_at (t,x) t1
            u_z_at (t,x) t2
            u_z_at (t,x) t1
            u_z_at (t,x) t1
            return (t,x)

          u_z_at :: ([Qubit], [Qubit]) -> Double -> Circ ()
          u_z_at (t, x) timeStep = do
              for (nb param) 1 (-1) $ \jj -> do
                  qlsa_HsimKernel param (t, x) jj timeStep oracleA
                  endfor
              for 1 (nb param) 1 $ \jj -> do
                  qlsa_HsimKernel param (t, x) jj timeStep oracleA
                  endfor
              return ()

-- | Testing function for 'qlsa_HamiltonianSimulation'.
test_qlsa_HamiltonianSimulation :: Bool -> IO ()
test_qlsa_HamiltonianSimulation dummyRTParamFlag = do
    let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
    let oracleARunTime = (oracle_A dummy_oracle param)
    let testCirc = do
              t <- qinit (take (n0 param) (repeat False))
              x <- qinit (take (n4 param) (repeat False))
              label(t,x) ("t", "x")
              qlsa_HamiltonianSimulation param (t,x) oracleARunTime
    print_simple GateCount testCirc
--    print_simple Preview testCirc

-- | Uses an oracle function and timestep control register (/t/) to
-- apply 1-sparse Hamiltonian to the input state |/t/, /x/〉.
qlsa_HsimKernel ::
    RunTimeParam
    -> ([Qubit], [Qubit])
    -> Int
    -> Double
    -> OracleARunTime
    -> Circ ()
-- qlsa_HsimKernel param tx band timeStep oracleA = named_gate_at "HsimKernel" tx

qlsa_HsimKernel param (t, x) band timeStep oracleA = do
  _ <- (flip $ box ("qlsa_HsimKernel" ++ (show band) ++ (show timeStep))) (t,x) $ \(t,x) -> do
    let phiP = (epsilon param)
    with_ancilla_list (n2 param) $ \y -> do
       with_ancilla_list (n4 param) $ \m -> do
           with_ancilla_list (n4 param) $ \p -> do
               label (y,m,p) ("y","m","p")
               -- phases
               oracleA phiP band (phaseArgflag param) (x,y,p)
               qlsa_ControlledPhase p phiP False
               oracleA phiP band (phaseArgflag param) (x,y,p)
               -- magnitudes
               let phiMag = 2 ** (negate $ fromIntegral $ after_radix_length)
               oracleA phiMag band (magnitudeArgflag param) (x,y,m)
               for 0 ((length t) - 1) 1 $ \ii -> do
                   let phi_mt = timeStep * phiMag * (2^ii)
                   qlsa_ApplyHmag param (x,y,m) phi_mt `controlled` (t !! ii)
                   endfor
               oracleA phiMag band (magnitudeArgflag param) (x,y,m)
               -- phases again
               oracleA phiP band (phaseArgflag param) (x,y,p)
               qlsa_ControlledPhase p phiP True
               oracleA phiP band (phaseArgflag param) (x,y,p)
               return ()
           return ()
       return ()
    return (t,x)
  return ()

-- | Testing function for 'qlsa_HsimKernel'.
test_qlsa_HsimKernel :: Bool -> IO ()
test_qlsa_HsimKernel dummyRTParamFlag = do
    let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
    let oracleARunTime = (oracle_A dummy_oracle param)
    let testCirc = do
              t <- qinit (take (n0 param) (repeat False))
              x <- qinit (take (n4 param) (repeat False))
              label(t,x) ("t", "x")
              qlsa_HsimKernel param (t,x) 1 0.1 oracleARunTime
    print_simple GateCount testCirc
--    print_simple Preview testCirc

-- | Applies the magnitude component of coupling elements in a
-- 1-sparse Hamiltonian.
qlsa_ApplyHmag :: RunTimeParam -> ([Qubit], [Qubit], [Qubit]) -> Double -> Circ ()
qlsa_ApplyHmag param (x,y,m) phi0 = do
  _ <- (flip $ box ("qlsa_ApplyHmag " ++ (show phi0))) (x,y,m) $ \(x,y,m) -> do
    let (onOne, onZero) = (True, False)
    with_ancilla $ \a -> do -- Assume initialized to False (0)
        label (a) ("anc. a")
        if (length x /= length y)
           then error "qlsa_ApplyHmag: Input registers x and y have different lengths."
           else return ()
        let length_xy = length x
        for 0 (length_xy - 1) 1 $ \ii -> do
            let (xi, yi) = (x !! ii, y !! ii)
            w (xi, yi)
            qnot_at a `controlled` (xi .==. onOne .&&. yi .==. onZero)
            endfor
        qlsa_ControlledPhase (m ++ [a]) phi0 False
        for (length_xy - 1) 0 (-1) $ \ii -> do
            let (xi, yi) = (x !! ii, y !! ii)
            qnot_at a `controlled` (xi .==. onOne .&&. yi .==. onZero)
            w (xi, yi)
            endfor
        return ()
    return (x,y,m)
  return ()

-- | Testing function for 'qlsa_ApplyHmag'.
test_qlsa_ApplyHmag :: Bool -> IO ()
test_qlsa_ApplyHmag dummyRTParamFlag = do
    let param = if dummyRTParamFlag then dummy_RT_param else large_RT_param
    let testCirc = do
              x <- qinit (take (n2 param) (repeat False))
              y <- qinit (take (n2 param) (repeat False))
              m <- qinit (take (n4 param) (repeat False))
              label(x,y,m) ("x", "y", "m")
              qlsa_ApplyHmag param (x,y,m) 0.1
    print_simple GateCount testCirc
    print_simple Preview testCirc

-- | Auxiliary function: the /W/-gate.
w :: (Qubit, Qubit) -> Circ ()
-- w xy = named_gate_at "W" xy
w (xi,yi) = do
  _ <- box "w" (\(xi, yi) -> do
    label (xi,yi) ("x[i]","y[i]")
    gate_X_at xi `controlled` yi
    gate_X_at yi `controlled` xi
    hadamard_at yi `controlled` xi
    gate_X_at yi `controlled` xi
    gate_X_at xi `controlled` yi
    return (xi,yi)) (xi,yi)
  return ()

-- | Testing function for 'w'.
test_w :: IO ()
test_w = do
    let testCirc = do
        xi <- qinit False
        yi <- qinit False
        w (xi, yi)
    print_simple GateCount testCirc
    print_simple Preview testCirc

-- * Controlled Logic Operations

-- | Applies a phase shift of φ\/2 to the signed input register |φ〉.

-- For c, bit locations are counted starting from 0 (least significant) to (length c - 2) (most significant)
-- last c (or c !! (n-1)) is the sign bit
qlsa_ControlledPhase :: [Qubit] -> Double -> Bool -> Circ ()
-- qlsa_ControlledPhase c phi0 f = named_gate_at "CPhase" c
qlsa_ControlledPhase c phi0 f = do
  _ <- (flip $ box ("qlsa_ControlledPhase " ++ (show phi0) ++ " " ++ (show f))) c $ \c -> do
    with_ancilla $ \a -> do -- ancilla is initialized to False
        if f then (qnot_at a) else return ()
        let signQubit = last c
        qnot_at a `controlled` signQubit
        for 0 (length c - 2) 1 $ \ii -> do
            let theta = ( (2.0^(ii)) * phi0 / 2.0) -- Note the divide by 2
            -- The following three lines are equivalent
            expZt_at theta a `controlled` (c !! ii)
            -- a <- expZt theta a `controlled` (c !! ii); return ()
            -- qlsa_ControlledU (c !! ii, a) (expZt theta) False
            endfor
        qnot_at a `controlled` signQubit
        return c
  return ()

-- | Applies a rotation of φ\/2 to the signed input register |φ〉.
qlsa_ControlledRotation :: ([Qubit], Qubit) -> Double -> Bool -> Circ ()
-- qlsa_ControlledRotation ct phi0 f = named_gate_at "CRotate" ct
qlsa_ControlledRotation (c, t) phi0 f = do
  _ <- (flip $ box ("qlsa_ControlledRotation " ++ (show phi0) ++ " " ++ (show f))) (c,t) $ \(c,t) -> do
    if f then (qnot_at t) else return ()
    let signQubit = last c
    qnot_at t `controlled` signQubit
    for 0 (length c - 2) 1 $ \ii -> do
        let theta = (2.0^(ii)) * phi0 / 2.0
        expYt_at theta t `controlled` (c !! ii)
        endfor
    qnot_at t `controlled` signQubit
    if f then (qnot_at t) else return ()
    return (c,t)
  return ()

----------------------------------------------------------------------
-- * Oracles

-- | Map a 'QDouble' into an integer, understood as being scaled by
-- the given factor. Take the factor and the size of the output
-- register as parameter.
make_factor_rep :: Double -> Int -> QDouble -> Circ [Qubit]
make_factor_rep factor size p = do
     -- number of high bits of p
     let p_int_size = (xdouble_length p) - (xdouble_exponent p)
     -- get the required number of bits for doing the encoding
     let auxsize = max (size - 1) p_int_size
     -- do the operation:
     --    (1) build a QDouble from the factor: need (size-1) bits of integer part
     qfactor <- qinit $ fdouble_of_double (xdouble_exponent p)
                                          (auxsize + xdouble_exponent p) factor
     --    (2) scale p
     p_large <- qdouble_pad_to_extent (auxsize,- (xdouble_exponent p)) p
     --    (3) divide p by factor
     qreal_multiples <- unpack template_symb_slash_ p_large qfactor
     --    (4) get the floor of the result
     q_floor <- template_floor
     qmultiples <- q_floor qreal_multiples
     --    (5) set them in the right format (it has size elements)
     let (SInt tp bp) = qmultiples
     let new_p = (take (size - 1) $ reverse tp) ++ [bp]
     return new_p

-- | Implements the oracle for the arbitrary state |/r/〉, using the
-- Template Haskell implementation of 'calcRweights'.
inline_oracle_r :: RunTimeParam -> Double -> Double -> ([Qubit],[Qubit],[Qubit]) -> Circ ([Qubit],[Qubit],[Qubit])
inline_oracle_r rt factor_m factor_p =  box ("Or " ++ (show factor_m) ++ " " ++ (show factor_p)) $ decompose_generic Toffoli $ \(x',m',p') ->
    with_ancilla_init (fromIntegral $ nx rt :: FSignedInt) $ \qnx ->
     with_ancilla_init (fromIntegral $ ny rt :: FSignedInt) $ \qny ->
      with_ancilla_init (fdouble $ lx rt) $ \qlx ->
       with_ancilla_init (fdouble $ ly rt) $ \qly ->
        with_ancilla_init (fdouble $ k rt) $ \qk ->
         with_ancilla_init (fdouble $ theta rt) $ \qtheta ->
          with_ancilla_init (fdouble $ phi rt) $ \qphi ->
                -- the x register is smaller than the size of a QSInt,
                -- and x is an unsigned integer: make up a positive
                -- sign and a pad

                with_ancilla_init False $ \bx -> do
                  with_ancilla_init (replicate (fixed_int_register_length - (n2 rt))
                                               False) $ \pad_x -> do

                    let x = SInt (pad_x ++ (reverse x')) bx
                    let f = classical_to_reversible $
                              \(x,nx,ny,lx,ly,k,theta,phi) -> do
                                 (m,p) <- unpack Template.template_calcRweights
                                            x nx ny lx ly k theta phi
                                 new_p <- make_factor_rep factor_p (n4 rt) p
                                 new_m <- make_factor_rep factor_m (n4 rt) m
                                 return (new_m, new_p)
                    f ((x,qnx,qny,qlx,qly,qk,qtheta,qphi),(m',p'))
                    return (x',m',p')

-- | Implements the oracle for the known state |/b/〉, using the
-- Template Haskell implementation of 'getKnownWeights'.
inline_oracle_b :: RunTimeParam -> Double -> Double -> ([Qubit],[Qubit],[Qubit]) -> Circ ([Qubit],[Qubit],[Qubit])
inline_oracle_b rt factor_m factor_p = box ("Ob " ++ (show factor_m) ++ " " ++ (show factor_p)) $ decompose_generic Toffoli $ \(x',m',p') ->
    -- Make ancillas for constant values

    with_ancilla_init (listpair_fmap fromIntegral $ scatteringnodes rt :: [(FDouble,FDouble)]) $ \qscatteringnodes ->
     with_ancilla_init (fromIntegral $ ny rt :: FSignedInt) $ \qny ->
      with_ancilla_init (fdouble $ lx rt) $ \qlx ->
       with_ancilla_init (fdouble $ ly rt) $ \qly ->
        with_ancilla_init (fdouble $ k rt) $ \qk ->
         with_ancilla_init (fdouble $ theta rt) $ \qtheta ->
          with_ancilla_init (fdouble $ e0 rt) $ \qe0 ->
           with_ancilla_init (fromIntegral $ nx rt :: FSignedInt) $ \qnx -> do

                -- the x register is smaller than the size of a QSInt,
                -- and x is an unsigned integer: make up a positive
                -- sign and a pad

                with_ancilla_init False $ \bx -> do
                  with_ancilla_init (replicate (fixed_int_register_length - (n2 rt))
                                               False) $ \pad_x -> do

                    let x = SInt (pad_x ++ (reverse x')) bx

                    let f = classical_to_reversible $
                              \(y,nx,ny,scatteringnodes,lx,ly,k,theta,e0) -> do
                                  -- get some QSInt and QDouble
                                  (m,p) <- unpack Template.template_getKnownWeights
                                              y nx ny scatteringnodes lx ly k theta e0 7
                                  new_p <- make_factor_rep factor_p (n4 rt) p
                                  new_m <- make_factor_rep factor_m (n4 rt) m

                                  return (new_m, new_p)

                    f ((x,qnx,qny,qscatteringnodes,qlx,qly,qk,qtheta,qe0),(m',p'))
                    return (x',m',p')

-- | Implementation of the oracle calculating the matrix /A/
-- corresponding to the discretization of the scattering problem,
-- using the Template Haskell implementation of
-- 'getNodeValuesMoreOutputs'.
inline_oracle_A ::  RunTimeParam -> Double -> Int -> Bool -> ([Qubit],[Qubit],[Qubit]) -> Circ ([Qubit],[Qubit],[Qubit])
inline_oracle_A rt factor band argflag =  box ("OA " ++ (show band) ++ " " ++ (show argflag)) $ decompose_generic Toffoli $ \(x',y',p') -> do

    let argflag' = if argflag then PTrue else PFalse

    -- Make ancillas for constant values

    with_ancilla_init (listpair_fmap fromIntegral $ scatteringnodes rt :: [(FDouble, FDouble)]) $ \qscatteringnodes ->
      with_ancilla_init (fromIntegral $ ny rt :: FSignedInt) $ \qny ->
        with_ancilla_init (fromRational $ toRational $ lx rt) $ \qlx ->
          with_ancilla_init (fromRational $ toRational $ ly rt) $ \qly ->
            with_ancilla_init (fromRational $ toRational $ k rt) $ \qk ->
             with_ancilla_init (fromIntegral $ nx rt :: FSignedInt) $ \qnx -> do

                -- the x register is smaller than the size of a QSInt,
                -- and x is an unsigned integer: make up a positive
                -- sign and a pad

                with_ancilla_init False $ \bx -> do
                  with_ancilla_init (replicate (fixed_int_register_length - (n2 rt))
                                               False) $ \pad_x -> do

                    let x = SInt (pad_x ++ (reverse x')) bx

                    let f = classical_to_reversible $
                              \(v,nx,ny,scatteringnodes,lx,ly,k) -> do
                               -- get some QSInt and QDouble
                               (y,p) <- unpack Template.template_getNodeValuesMoreOutputs
                                           v band nx ny scatteringnodes lx ly k argflag' 7

                               new_p <- make_factor_rep factor (n4 rt) p

                               -- set y in the correct format (it has n2 elements)
                               -- we can assume that the sign is positive.
                               -- we also assume that n2 < size of QSInt register
                               let (SInt ty _) = y
                               let new_y = (take (n2 rt) $ reverse ty)

                               -- return the values in the right format.
                               return (new_y,new_p)

                    f ((x,qnx,qny,qscatteringnodes,qlx,qly,qk),(y',p'))

                    return (x',y',p')

-- | Encapsulate the inline oracles in Template Haskell into an object
-- of type 'Oracle'.
inline_oracle :: Oracle
inline_oracle = Oracle {
  oracle_A = inline_oracle_A,
  oracle_b = inline_oracle_b,
  oracle_r = inline_oracle_r
}


{-# LANGUAGE ConstraintKinds  #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE QuasiQuotes      #-}

module Layers
  ( ForwardLayer(..)
  , OutputLayer(..)
  , ForwardNN(..)
  , BackwardLayer(..)
  , BackputLayer(..)
  , BackwardNN(..)
  , TrainBatch(..)
  , NElement(..)
  , (<~)
  , (~>)
  , forward
  , output
  , backward
  , backput
  , learnForward
  , learnBackward
  , learn
  , learnAll
  , predict
  , evaluate
  ) where

import           Control.Arrow
import           Control.Lens            hiding ((<~))
import           Data.String.Interpolate as S (i)
import           Debug.Trace
import           Numeric
import           Numeric.LinearAlgebra
import           Synapse

type Weight a = Matrix a

type Bias a = Vector a

type SignalX a = Matrix a

type SignalY a = Matrix a

type Diff a = Matrix a

type TeacherBatch a = Matrix a

type InputBatch a = Matrix a

newtype TrainBatch a =
  TrainBatch (TeacherBatch a, InputBatch a)
  deriving (Show)

data ForwardLayer a
  = AffineForward (Weight a) (Bias a)
  | SigmoidForward
  | ReLUForward
  | JoinedForwardLayer (ForwardLayer a) (ForwardLayer a)
  deriving (Show)

instance (Numeric a, Eq a) => Eq (ForwardLayer a) where
  AffineForward w b == AffineForward w' b' = w == w' && b == b'
  SigmoidForward == SigmoidForward = True
  ReLUForward == ReLUForward = True
  JoinedForwardLayer a b == JoinedForwardLayer a' b' = a == a' && b == b'
  _ == _ = False

infixl 4 ~>

(~>) :: ForwardLayer a -> ForwardLayer a -> ForwardLayer a
a ~> (JoinedForwardLayer x y) = (a ~> x) ~> y
a ~> b = JoinedForwardLayer a b

data OutputLayer a =
  SoftmaxWithCrossForward
  deriving (Show)

data ForwardNN a =
  ForwardNN (ForwardLayer a) (OutputLayer a)
  deriving (Show)

data BackwardLayer a
  = AffineBackward (Weight a) (Bias a) (SignalX a)
  | SigmoidBackward (SignalY a)
  | ReLUBackward (SignalX a)
  | JoinedBackwardLayer (BackwardLayer a) (BackwardLayer a)
  deriving (Show)

instance (Numeric a, Eq a) => Eq (BackwardLayer a) where
  AffineBackward w b d == AffineBackward w' b' d' =
    w == w' && b == b' && d == d'
  SigmoidBackward y == SigmoidBackward y' = y == y'
  ReLUBackward x == ReLUBackward x' = x == x'
  JoinedBackwardLayer a b == JoinedBackwardLayer a' b' = a == a' && b == b'
  _ == _ = False

infixr 4 <~

(<~) :: BackwardLayer a -> BackwardLayer a -> BackwardLayer a
(JoinedBackwardLayer x y) <~ b = x <~ (y <~ b)
a <~ b = JoinedBackwardLayer a b

data BackputLayer a =
  SoftmaxWithCrossBackward (TeacherBatch a) (SignalY a)

data BackwardNN a =
  BackwardNN (BackwardLayer a) (BackputLayer a)

type NElement a = (Ord a, Floating a, Numeric a, Num (Vector a), Show a)

forward ::
     NElement a => ForwardLayer a -> SignalX a -> (BackwardLayer a, SignalY a)
forward (AffineForward w b) x = (AffineBackward w b x, affinem w b x)
forward SigmoidForward x = (SigmoidBackward y, y)
  where
    y = sigmoidm x
forward ReLUForward x = (ReLUBackward x, relum x)
forward (JoinedForwardLayer a b) x0 = (a' <~ b', x2)
  where
    (a', x1) = forward a x0
    (b', x2) = forward b x1

backward ::
     NElement a => a -> BackwardLayer a -> Diff a -> (ForwardLayer a, Diff a)
backward rate (AffineBackward w b x) d =
  (AffineForward (w - scale rate w') (b - scale rate b'), x')
  where
    (x', w', b') = affinemBackward w x d
backward _ (SigmoidBackward y) d = (SigmoidForward, sigmoidBackward y d)
backward _ (ReLUBackward x) d = (ReLUForward, relumBackward x d)
backward r (JoinedBackwardLayer a b) d0 = (a' ~> b', d2)
  where
    (b', d1) = backward r b d0
    (a', d2) = backward r a d1

output ::
     NElement a
  => OutputLayer a
  -> TeacherBatch a
  -> SignalY a
  -> (BackputLayer a, a)
output SoftmaxWithCrossForward t y = (SoftmaxWithCrossBackward t y', loss)
  where
    (y', loss) = softmaxWithCross t y

backput :: NElement a => BackputLayer a -> (OutputLayer a, Diff a)
backput (SoftmaxWithCrossBackward t y) =
  (SoftmaxWithCrossForward, softmaxWithCrossBackward t y)

learnForward :: NElement a => ForwardNN a -> TrainBatch a -> (BackwardNN a, a)
learnForward (ForwardNN layers loss) (TrainBatch (t, x)) =
  result `seq` (BackwardNN layers' loss', result)
  where
    (layers', y) = forward layers x
    (loss', result) = output loss t y

learnBackward :: NElement a => a -> BackwardNN a -> ForwardNN a
learnBackward rate (BackwardNN layers loss) = ForwardNN layers' loss'
  where
    (loss', d) = backput loss
    (layers', _) = backward rate layers d

learn :: NElement a => a -> ForwardNN a -> TrainBatch a -> (ForwardNN a, a)
learn rate a = first (learnBackward rate) . learnForward a

learnAll ::
     NElement a => a -> ForwardNN a -> [TrainBatch a] -> (ForwardNN a, [a])
learnAll rate origin batches = ls `seq` (nn, ls)
  where
    (nn, ls) = foldr f (origin, []) batches
    f batch (a, ls) = ls `seq` second (~: ls) $ learn rate a batch
    x ~: xs = trace [i|[#{length xs + 1}/#{n}] #{show x}|] (x : xs)
    n = length batches

predict :: NElement a => ForwardLayer a -> Vector a -> Int
predict layers = maxIndex . flatten . snd . forward layers . asRow

evaluate :: NElement a => ForwardLayer a -> [(Int, Vector a)] -> Double
evaluate layers samples = fromIntegral nOk / fromIntegral (length samples)
  where
    nOk = length $ filter (uncurry (==)) results
    results = map (second $ predict layers) samples


{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-#LANGUAGE OverloadedLists#-}

module Lib
    ( app
    , main'
    ) where

import Network.Wai
import Network.Wai.Handler.Warp
import Servant
import Servant.API
import Data.Aeson
import GHC.Generics
import Data.String.Conversions
import Text.Read (readMaybe)
import Servant.HTML.Lucid
import Lucid
import Data.Time.Calendar
import Data.List (sortBy)
import Data.Complex

import Data.HashMap.Strict
import Data.Text(Text(..))

------------------------ API SPEC ---------------------------
type UserAPI =
           "users" :> QueryParam "sortby" SortBy :> Get '[JSON] [User]
      :<|> "members" :> Get '[JSON] (Headers '[Header "Access-Control-Allow-Origin" String] Members)
      :<|> "albert" :> Get '[JSON] User
      :<|> "isaac" :> Get '[JSON] User

type MathAPI =
          "sum" :> Capture "x" Int :> Capture "y" Int :> Get '[JSON] Int
     :<|> "conjugate" :> Capture "x" (Complex Int)  :> Get '[JSON] String

type MyAPI = UserAPI :<|> MathAPI

data User = User
    { name :: String
    , age :: Int
    , img :: String
    } deriving (Eq, Show, Generic)

newtype Members = Members [User]

--------------------- INSTANCES ---------------------
instance ToJSON User

instance ToJSON Members where
   toJSON (Members us) = Object [("members", toJSON us)]

instance ToHtml User where
    toHtml u = tr_ $ do
                td_ $ toHtml (name u)
                td_ $ toHtml (show $ age u)

    toHtmlRaw = toHtml

instance ToHtml [User] where
    toHtml u = table_ $ do
        tr_ $ do
            th_ $ toHtml ("Name"::String)
            th_ $ toHtml ("Age"::String)
        foldMap toHtml u

    toHtmlRaw = toHtml

data SortBy = Age | Name deriving (Generic, Show,Eq)

instance FromHttpApiData SortBy where
    parseQueryParam "age" = Right Age
    parseQueryParam "name" = Right Name
    parseQueryParam _ = Left "Not valid query param"

instance (Num a,Read a) => FromHttpApiData (Complex a) where
    --parseUrlPiece :: Text -> Either Text a
    parseUrlPiece txt = maybe (Left "Cant parse complex num") (Right) (readMaybe $ cs txt)

---------------------- ENDPOINT HANDLERS -------------------
users :: [User]
users =
  [ User "Isaac Newton"    372 "http://placehold.it/250x250"
  , User "Albert Einstein" 136 "http://placehold.it/250x250"
  , User "Haskell Curry" 200 "http://placehold.it/250x250"
  ]

members :: Monad m => m (Headers '[Header "Access-Control-Allow-Origin" String] Members)
members = return $ addHeader "*" (Members users)

albert = users !! 0
isaac = users !! 1

sortedUsers Age = sortBy (compareWith age)
sortedUsers Name = sortBy (compareWith name)

-- m is actually a EitherT ServantErr IO, in the abscense of SortBy query param users gets returned
sortUsers :: Monad m =>  Maybe SortBy -> m [User]
sortUsers = return . maybe users (flip sortedUsers users)

server1 :: Server UserAPI
server1 =
          sortUsers
     :<|> members
     :<|> return isaac
     :<|> return albert

server2 :: Server MathAPI
server2 = sum
      :<|> showConjugate
      where
       sum = return <.. (+)

       showConjugate :: Monad m => Complex Int -> m String
       showConjugate = return . show . conjugate

server :: Server MyAPI
server = server1 :<|> server2

myAPI :: Proxy MyAPI
myAPI = Proxy

---------------------------------- SERVER --------------------------------
{- 'serve' comes from servant and hands you a WAI Application,
which you can think of as an "abstract" web application,
not yet a webserver. -}

app :: Application
app = serve myAPI server

main' = do
    let hostIP = "0.0.0.0"
    let port = 5000
    putStrLn $ "Server running in " ++ (show hostIP) ++ " port " ++ (show port)
    let settings = setPort port $ setHost hostIP defaultSettings
    runSettings settings app

----------------------------- UTILITIES -----------------------------
(<..) = (.) . (.)
compareWith f x y = compare (f x) (f y)


{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE GADTs                    #-}
{-# LANGUAGE Strict                   #-}
module Grenade.Layers.Internal.Update (
    descendMatrix
  , descendVector
  , MatrixInputValues (..)
  , MatrixResult (..)
  , VectorInputValues (..)
  , VectorResult (..)
  , descendMatrixV
  , descendVectorV
  , MatrixInputValuesV (..)
  , MatrixResultV (..)
  , VectorInputValuesV (..)
  , VectorResultV (..)
  ) where

import           Control.Parallel.Strategies
import           Data.Maybe                   (fromJust, fromMaybe)
import qualified Data.Vector.Storable         as U (unsafeFromForeignPtr0,
                                                    unsafeToForeignPtr0)
import qualified Data.Vector.Storable         as V
import           Foreign                      (mallocForeignPtrArray, withForeignPtr)
import           Foreign.Ptr                  (Ptr)
import           GHC.TypeLits
import           Numeric.LinearAlgebra        (Vector, flatten)
import qualified Numeric.LinearAlgebra.Devel  as U
import           Numeric.LinearAlgebra.Static
import           System.IO.Unsafe             (unsafePerformIO)

import           Grenade.Core.Optimizer
import           Grenade.Layers.Internal.CUDA
import           Grenade.Types

data MatrixInputValues rows columns
  = MatrixValuesSGD
      !(L rows columns) -- ^ current weights
      !(L rows columns) -- ^ gradients
      !(L rows columns) -- ^ last update (old momentum)
  | MatrixValuesAdam
      !Int              -- ^ Step
      !(L rows columns) -- ^ current weights
      !(L rows columns) -- ^ gradients
      !(L rows columns) -- ^ current m
      !(L rows columns) -- ^ current v

data MatrixInputValuesV
  = MatrixValuesSGDV
      !(V.Vector RealNum) -- ^ current weights
      !(V.Vector RealNum) -- ^ gradients
      !(V.Vector RealNum) -- ^ old momentum
  | MatrixValuesAdamV
      !Int                -- ^ Step
      !(V.Vector RealNum) -- ^ current weights
      !(V.Vector RealNum) -- ^ gradients
      !(V.Vector RealNum) -- ^ current m
      !(V.Vector RealNum) -- ^ current v

data MatrixResult rows columns
  = MatrixResultSGD
      { matrixActivations :: !(L rows columns) -- ^ new activations (weights)
      , matrixMomentum    :: !(L rows columns) -- ^ new momentum
      }
  | MatrixResultAdam
      { matrixActivations :: !(L rows columns) -- ^ new activations (weights)
      , matrixM           :: !(L rows columns) -- ^ new m
      , matrixV           :: !(L rows columns) -- ^ new v
      }

data MatrixResultV
  = MatrixResultSGDV
      { matrixActivationsV :: !(V.Vector RealNum) -- ^ new activations (weights)
      , matrixMV           :: !(V.Vector RealNum) -- ^ new m
      }
  | MatrixResultAdamV
      { matrixActivationsV :: !(V.Vector RealNum) -- ^ new activations (weights)
      , matrixMV           :: !(V.Vector RealNum) -- ^ new m
      , matrixVV           :: !(V.Vector RealNum) -- ^ new v
      }

data VectorInputValues r
  = VectorValuesSGD
      !(R r) -- ^ current weights
      !(R r) -- ^ gradients
      !(R r) -- ^ last update (old momentum)
  | VectorValuesAdam
      !Int   -- ^ Step
      !(R r) -- ^ current weights
      !(R r) -- ^ current gradients
      !(R r) -- ^ current m
      !(R r) -- ^ current v

data VectorInputValuesV
  = VectorValuesSGDV
      !(V.Vector RealNum) -- ^ current weights
      !(V.Vector RealNum) -- ^ current gradients
      !(V.Vector RealNum) -- ^ current m
  | VectorValuesAdamV
      !Int                -- ^ Step
      !(V.Vector RealNum) -- ^ current weights
      !(V.Vector RealNum) -- ^ current gradients
      !(V.Vector RealNum) -- ^ current m
      !(V.Vector RealNum) -- ^ current v

data VectorResult r
  = VectorResultSGD
      { vectorBias     :: !(R r) -- ^ new activations (bias)
      , vectorMomentum :: !(R r) -- ^ new momentum
      }
  | VectorResultAdam
      { vectorBias :: !(R r) -- ^ new activations (bias)
      , vectorM    :: !(R r) -- ^ new m
      , vectorV    :: !(R r) -- ^ new v
      }

data VectorResultV
  = VectorResultSGDV
      { vectorBiasV :: !(V.Vector RealNum) -- ^ new activations (bias)
      , vectorMV    :: !(V.Vector RealNum) -- ^ new m
      }
  | VectorResultAdamV
      { vectorBiasV :: !(V.Vector RealNum) -- ^ new activations (bias)
      , vectorMV    :: !(V.Vector RealNum) -- ^ new m
      , vectorVV    :: !(V.Vector RealNum) -- ^ new v
      }

descendMatrix :: (KnownNat rows, KnownNat columns) => Optimizer o -> MatrixInputValues rows columns -> MatrixResult rows columns
descendMatrix (OptSGD rate momentum regulariser) (MatrixValuesSGD weights gradient lastUpdate) =
  let (rows, cols) = size weights
      len          = rows * cols
      -- Most gradients come in in ColumnMajor,
      -- so we'll transpose here before flattening them
      -- into a vector to prevent a copy.
      --
      -- This gives ~15% speed improvement for LSTMs.
      weights'     = flatten . tr . extract $ weights
      gradient'    = flatten . tr . extract $ gradient
      lastUpdate'  = flatten . tr . extract $ lastUpdate
      (vw, vm)     = descendUnsafeSGD len rate momentum regulariser weights' gradient' lastUpdate'

      -- Note that it's ColumnMajor, as we did a transpose before
      -- using the internal vectors.
      mw           = U.matrixFromVector U.ColumnMajor rows cols vw
      mm           = U.matrixFromVector U.ColumnMajor rows cols vm
  in  MatrixResultSGD (fromJust . create $ mw) (fromJust . create $ mm)
descendMatrix (OptAdam alpha beta1 beta2 epsilon lambda) (MatrixValuesAdam step weights gradient m v) =
  let (rows, cols) = size weights
      len          = rows * cols
      -- Most gradients come in in ColumnMajor,
      -- so we'll transpose here before flattening them
      -- into a vector to prevent a copy.
      --
      -- This gives ~15% speed improvement for LSTMs.
      weights'  = flatten . tr . extract $ weights
      gradient' = flatten . tr . extract $ gradient
      m'        = flatten . tr . extract $ m
      v'        = flatten . tr . extract $ v
      (vw, vm, vv)     = descendUnsafeAdamGPU len step alpha beta1 beta2 epsilon lambda weights' gradient' m' v'

      -- Note that it's ColumnMajor, as we did a transpose before
      -- using the internal vectors.
      mw           = U.matrixFromVector U.ColumnMajor rows cols vw
      mm           = U.matrixFromVector U.ColumnMajor rows cols vm
      mv           = U.matrixFromVector U.ColumnMajor rows cols vv
  in  MatrixResultAdam (fromJust . create $ mw) (fromJust . create $ mm) (fromJust . create $ mv)
descendMatrix opt _ = error $ "optimzer does not match to MatrixInputValues in implementation! Optimizer: " ++ show opt

descendMatrixV :: Optimizer o -> MatrixInputValuesV -> MatrixResultV
descendMatrixV (OptSGD rate momentum regulariser) (MatrixValuesSGDV weights gradient lastUpdate) =
  let len          = V.length weights
      (vw, vm)     = descendUnsafeSGD len rate momentum regulariser weights gradient lastUpdate
  in  MatrixResultSGDV vw vm
descendMatrixV (OptAdam alpha beta1 beta2 epsilon lambda) (MatrixValuesAdamV step weights gradient m v) =
  let len          = V.length weights
      (vw, vm, vv) = descendUnsafeAdamGPU len step alpha beta1 beta2 epsilon lambda weights gradient m v
  in  MatrixResultAdamV vw vm vv
descendMatrixV opt _ = error $ "optimzer does not match to MatrixInputValues in implementation! Optimizer: " ++ show opt

descendVector :: (KnownNat r) => Optimizer o -> VectorInputValues r -> VectorResult r
descendVector (OptSGD rate momentum regulariser) (VectorValuesSGD weights gradient lastUpdate) =
  let len          = size weights
      weights'     = extract weights
      gradient'    = extract gradient
      lastUpdate'  = extract lastUpdate
      (vw, vm)     = descendUnsafeSGD len rate momentum regulariser weights' gradient' lastUpdate'
  in  VectorResultSGD (fromJust $ create vw) (fromJust $ create vm)
descendVector (OptAdam alpha beta1 beta2 epsilon lambda) (VectorValuesAdam step weights gradient m v) =
  let len       = size weights
      weights'  = extract weights
      gradient' = extract gradient
      m'        = extract m
      v'        = extract v
      (vw, vm, vv)     = descendUnsafeAdamGPU len step alpha beta1 beta2 epsilon lambda weights' gradient' m' v'
  in  VectorResultAdam (fromJust $ create vw) (fromJust $ create vm) (fromJust $ create vv)
descendVector opt _ = error $ "optimzer does not match to VectorInputValues in implementation! Optimizer: " ++ show opt

descendVectorV :: Optimizer o -> VectorInputValuesV -> VectorResultV
descendVectorV (OptSGD rate momentum regulariser) (VectorValuesSGDV weights gradient lastUpdate) =
  let len          = V.length weights
      (vw, vm)     = descendUnsafeSGD len rate momentum regulariser weights gradient lastUpdate
  in  VectorResultSGDV vw vm
descendVectorV (OptAdam alpha beta1 beta2 epsilon lambda) (VectorValuesAdamV step weights gradient m v) =
  let len = V.length weights
      (vw, vm, vv) = descendUnsafeAdamGPU len step alpha beta1 beta2 epsilon lambda weights gradient m v
   in VectorResultAdamV vw vm vv
descendVectorV opt _ = error $ "optimzer does not match to VectorInputValues in implementation! Optimizer: " ++ show opt

-- -- | Caching of data
-- type CacheKey = (LookupType, ProxyType, StateFeatures)

-- cacheMVar :: MVar (M.Map CacheKey [Values])
-- cacheMVar = unsafePerformIO $ newMVar mempty
-- {-# NOINLINE cacheMVar #-}

-- emptyCache :: MonadIO m => m ()
-- emptyCache = liftIO $ modifyMVar_ cacheMVar (const mempty)

-- addCache :: (MonadIO m) => CacheKey -> [Values] -> m ()
-- addCache k val = liftIO $ modifyMVar_ cacheMVar (return . M.insert k val)

-- lookupCache :: (MonadIO m) => CacheKey -> m (Maybe [Values])
-- lookupCache k = liftIO $ (M.lookup k =<<) <$> tryReadMVar cacheMVar

-- -- | Get output of function f, if possible from cache according to key (st).
-- cached :: (MonadIO m) => (LookupType, ProxyType, StateFeatures) -> m [Values] -> m [Values]
-- cached st ~f = do
--   c <- lookupCache st
--   case c of
--     Nothing -> do
--       res <- f
--       res `seq` addCache st res
--       return res
--     Just res -> return res

descendUnsafeSGD :: Int -> RealNum -> RealNum -> RealNum -> Vector RealNum -> Vector RealNum -> Vector RealNum -> (Vector RealNum, Vector RealNum)
descendUnsafeSGD len rate momentum regulariser weights gradient lastUpdate =
  unsafePerformIO $ do

    outWPtr <- mallocForeignPtrArray len
    outMPtr <- mallocForeignPtrArray len
    let (wPtr, _) = U.unsafeToForeignPtr0 weights
    let (gPtr, _) = U.unsafeToForeignPtr0 gradient
    let (lPtr, _) = U.unsafeToForeignPtr0 lastUpdate

    withForeignPtr wPtr $ \wPtr' ->
      withForeignPtr gPtr $ \gPtr' ->
        withForeignPtr lPtr $ \lPtr' ->
          withForeignPtr outWPtr $ \outWPtr' ->
            withForeignPtr outMPtr $ \outMPtr' ->
              descend_sgd_cpu len rate momentum regulariser wPtr' gPtr' lPtr' outWPtr' outMPtr'

    return (U.unsafeFromForeignPtr0 outWPtr len, U.unsafeFromForeignPtr0 outMPtr len)
{-# NOINLINE descendUnsafeSGD #-}

descendUnsafeAdam ::
     Int -- Len
  -> Int -- Step
  -> RealNum -- Alpha
  -> RealNum -- Beta1
  -> RealNum -- Beta2
  -> RealNum -- Epsilon
  -> RealNum -- Lambda
  -> Vector RealNum -- Weights
  -> Vector RealNum -- Gradient
  -> Vector RealNum -- M
  -> Vector RealNum -- V
  -> (Vector RealNum, Vector RealNum, Vector RealNum)
descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights gradient m v =
  unsafePerformIO $ do
  V.unsafeWith weights $ \wPtr' ->
    V.unsafeWith gradient $ \gPtr' ->
      V.unsafeWith m $ \mPtr' ->
        V.unsafeWith v $ \vPtr' ->
           descend_adam_cpu len step alpha beta1 beta2 epsilon lambda wPtr' gPtr' mPtr' vPtr'
  return (weights, m, v)
{-# NOINLINE descendUnsafeAdam #-}

descendUnsafeAdamGPU ::
     Int -- Len
  -> Int -- Step
  -> RealNum -- Alpha
  -> RealNum -- Beta1
  -> RealNum -- Beta2
  -> RealNum -- Epsilon
  -> RealNum -- Lambda
  -> Vector RealNum -- Weights
  -> Vector RealNum -- Gradient
  -> Vector RealNum -- M
  -> Vector RealNum -- V
  -> (Vector RealNum, Vector RealNum, Vector RealNum)
descendUnsafeAdamGPU len step alpha beta1 beta2 epsilon lambda weights gradient m v
  --  | useCuda len = -- We need some size, otherwise the overhead is too big and processing using BLAS is faster
  --   fromMaybe
  --     (descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights gradient m v)
  --     (cudaDescendUnsafeAdamGPU len step alpha beta1 beta2 epsilon lambda weights gradient m v)
  | otherwise = descendUnsafeAdam len step alpha beta1 beta2 epsilon lambda weights gradient m v

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

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


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

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveAnyClass        #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-|
Module      : Grenade.Layers.SamePadPooling
Description : Pooling layer that supports SAME_UPPER padding
Maintainer  : Theo Charalambous
License     : BSD2
Stability   : experimental
-}
module Grenade.Layers.SamePadPooling (
  -- * Layer Definition
    SamePadPooling (..)
  ) where

import           Control.DeepSeq
import           Data.Function                   ((&))
import           Data.Kind                       (Type)
import           Data.Maybe
import           Data.Proxy
import           Data.Serialize
import           Data.Singletons.TypeLits        hiding (natVal)
import           GHC.Generics
import           GHC.TypeLits
import           Numeric.LinearAlgebra.Static    (create, extract)

import           Grenade.Core
import           Grenade.Layers.Internal.Pooling
import           Grenade.Onnx

-- | A pooling layer for a neural network, for when auto_pad is SAME_UPPER or SAME_LOWER
--
--   Pads on the X and Y dimension of an image.
data SamePadPooling  :: Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Nat -> Type where
  SamePadPooling :: SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom
  deriving (NFData, Generic)

instance Show (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) where
  show SamePadPooling = "SamePadPooling"

instance UpdateLayer (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) where
  type Gradient (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) = ()
  runUpdate _ x _ = x
  reduceGradient _ = ()

instance RandomLayer (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom)  where
  createRandomWith _ _ = return SamePadPooling

instance Serialize (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) where
  put _ = return ()
  get = return SamePadPooling

-- | A two dimentional image can be padded.
instance ( KnownNat padLeft
         , KnownNat padTop
         , KnownNat padRight
         , KnownNat padBottom
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , (padLeft + padRight) ~ ((outputColumns - 1) * strideColumns + kernelColumns - inputColumns)
         , (padTop + padBottom) ~ ((outputRows - 1) * strideRows + kernelRows - inputRows)
         ) => Layer (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns) where
  type Tape (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) ('D2 inputRows inputColumns) ('D2 outputRows outputColumns)  = ()
  runForwards SamePadPooling (S2D input) =
    let padl  = fromIntegral $ natVal (Proxy :: Proxy padLeft)
        padt  = fromIntegral $ natVal (Proxy :: Proxy padTop)
        padr  = fromIntegral $ natVal (Proxy :: Proxy padRight)
        padb  = fromIntegral $ natVal (Proxy :: Proxy padBottom)
        kr    = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        kc    = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sr    = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sc    = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        h     = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        w     = fromIntegral $ natVal (Proxy :: Proxy inputColumns)

        m     = extract input

        r     = validPadPoolForwards 1 h w kr kc sr sc padl padt padr padb m

    in  ((), S2D . fromJust . create $ r)
  runBackwards _ _ _ = error "backward pass for SamePadPooling not implemented"

-- | A two dimentional image can be padded.
instance ( KnownNat padLeft
         , KnownNat padTop
         , KnownNat padRight
         , KnownNat padBottom
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat kernelRows
         , KnownNat kernelColumns
         , KnownNat inputRows
         , KnownNat inputColumns
         , KnownNat outputRows
         , KnownNat outputColumns
         , KnownNat channels
         , KnownNat (inputRows * channels)
         , KnownNat (outputRows * channels)
         , (padLeft + padRight) ~ ((outputColumns - 1) * strideColumns + kernelColumns - inputColumns)
         , (padTop + padBottom) ~ ((outputRows - 1) * strideRows + kernelRows - inputRows)
         ) => Layer (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels) where
  type Tape (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) ('D3 inputRows inputColumns channels) ('D3 outputRows outputColumns channels)  = ()
  runForwards SamePadPooling (S3D input) =
    let padl  = fromIntegral $ natVal (Proxy :: Proxy padLeft)
        padt  = fromIntegral $ natVal (Proxy :: Proxy padTop)
        padr  = fromIntegral $ natVal (Proxy :: Proxy padRight)
        padb  = fromIntegral $ natVal (Proxy :: Proxy padBottom)
        kr    = fromIntegral $ natVal (Proxy :: Proxy kernelRows)
        kc    = fromIntegral $ natVal (Proxy :: Proxy kernelColumns)
        sr    = fromIntegral $ natVal (Proxy :: Proxy strideRows)
        sc    = fromIntegral $ natVal (Proxy :: Proxy strideColumns)
        h     = fromIntegral $ natVal (Proxy :: Proxy inputRows)
        w     = fromIntegral $ natVal (Proxy :: Proxy inputColumns)
        c     = fromIntegral $ natVal (Proxy :: Proxy channels)

        m     = extract input

        r     = validPadPoolForwards c h w kr kc sr sc padl padt padr padb m
    in  ((), S3D . fromJust . create $ r)

  runBackwards _ _ _ = error "backward pass for SamePadPooling not implemented"

instance OnnxOperator (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) where
  onnxOpTypeNames _ = ["MaxPool"]

instance ( KnownNat padLeft
         , KnownNat padTop
         , KnownNat padRight
         , KnownNat padBottom
         , KnownNat strideRows
         , KnownNat strideColumns
         , KnownNat kernelRows
         , KnownNat kernelColumns
         ) => OnnxLoadable (SamePadPooling kernelRows kernelColumns strideRows strideColumns padLeft padTop padRight padBottom) where

  loadOnnxNode _ node = do
    node & hasSupportedDilations

    (node `hasMatchingShape` "kernel_shape") kernelShape
    (node `hasMatchingShape` "strides")      strideShape

    -- todo: check that attribute is one of: SAME_UPPER or SAME_LOWER

    return SamePadPooling
      where
        kernelShape = [natVal (Proxy :: Proxy kernelRows), natVal (Proxy :: Proxy kernelColumns)]
        strideShape = [natVal (Proxy :: Proxy strideRows), natVal (Proxy :: Proxy strideColumns)]