{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE TypeFamilies #-}
module Test.Synthesizer.LLVM.Utility where

import qualified Synthesizer.LLVM.Parameterized.SignalPacked as SigPS
import qualified Synthesizer.LLVM.Parameterized.Signal as SigP
import qualified Synthesizer.LLVM.Parameter as Param
import qualified Synthesizer.LLVM.Plug.Output as POut
import qualified Synthesizer.LLVM.Frame.SerialVector as Serial

import qualified Synthesizer.State.Signal as SigS
import qualified Synthesizer.Zip as Zip

import Control.Monad (liftM, liftM2)
import Control.Applicative ((<$>))

import qualified Data.StorableVector.Lazy as SVL
import qualified Data.StorableVector as SV
import Data.StorableVector.Lazy (ChunkSize)
import Foreign.Storable (Storable)

import qualified LLVM.Extra.Storable as Storable
import qualified LLVM.Extra.Tuple as Tuple
import qualified LLVM.Core as LLVM

import qualified Type.Data.Num.Decimal as TypeNum

import System.Random (Random, randomRs, StdGen, mkStdGen)

import Data.Tuple.HT (mapPair)

import qualified Test.QuickCheck as QC

import qualified Algebra.RealRing as RealRing
import qualified Algebra.Absolute as Absolute

import NumericPrelude.Numeric
import NumericPrelude.Base


genRandomVectorParam :: QC.Gen (Int, StdGen)
genRandomVectorParam =
   liftM2 (,) (QC.choose (1,100)) (mkStdGen <$> QC.arbitrary)

randomStorableVector ::
   (Storable a, Random a) =>
   (a, a) -> (Int, StdGen) -> SV.Vector a
randomStorableVector range (len, seed) =
   fst $ SV.packN len $ randomRs range seed

randomStorableVectorLoop ::
   (Storable a, Random a) =>
   (a, a) -> (Int, StdGen) -> SVL.Vector a
randomStorableVectorLoop range param =
   SVL.cycle $ SVL.fromChunks [randomStorableVector range param]

randomSignal ::
   (Storable.C a, Tuple.ValueOf a ~ al, Random a) =>
   (a, a) -> Param.T p (Int, StdGen) -> SigP.T p al
randomSignal range p =
   SigP.fromStorableVectorLazy (randomStorableVectorLoop range <$> p)


render ::
   (Storable.C a, Tuple.ValueOf a ~ al) =>
   (SVL.Vector a -> sig) ->
   SigP.T p al -> IO (ChunkSize -> p -> sig)
render limit sig =
   fmap (\func chunkSize -> limit . func chunkSize) $ SigP.runChunky sig

render2 ::
   (Storable.C a, Tuple.ValueOf a ~ al) =>
   (Storable.C b, Tuple.ValueOf b ~ bl) =>
   ((SVL.Vector a, SVL.Vector b) -> sig) ->
   SigP.T p (al, bl) -> IO (ChunkSize -> p -> sig)
render2 limit sig =
   fmap (\func chunkSize ->
            limit . mapPair (SVL.fromChunks, SVL.fromChunks) .
            unzip . map (\(Zip.Cons a b) -> (a,b)) . func chunkSize) $
   SigP.runChunkyPlugged sig POut.deflt


data CheckSimilarityState a =
   CheckSimilarityState a (SVL.Vector a) (SigS.T a)

instance (Storable a, Ord a, Absolute.C a) =>
      QC.Testable (CheckSimilarityState a) where
   property (CheckSimilarityState tol xs ys) =
      QC.property $
         SigS.foldR (&&) True $
         -- dangerous, since shortened signals would be tolerated
         SigS.zipWith (\x y -> abs(x-y) < tol)
            (SigS.fromStorableSignal xs) ys

{-# INLINE checkSimilarityState #-}
checkSimilarityState ::
   (RealRing.C a, Storable.C a, Tuple.ValueOf a ~ av) =>
   a ->
   (SVL.Vector a -> SVL.Vector a) ->
   SigP.T p av ->
   (p -> SigS.T a) ->
   IO (ChunkSize -> p -> CheckSimilarityState a)
checkSimilarityState tol limit gen0 sig1 =
   liftM
      (\sig0 chunkSize p ->
         CheckSimilarityState tol (sig0 chunkSize p) (sig1 p))
      (render limit gen0)


data CheckSimilarity a =
   CheckSimilarity a (SVL.Vector a) (SVL.Vector a)

instance (Storable a, Ord a, Absolute.C a) => QC.Testable (CheckSimilarity a) where
   property (CheckSimilarity tol xs ys) =
      QC.property $
         SigS.foldR (&&) True $
         -- dangerous, since shortened signals would be tolerated
         SigS.zipWith (\x y -> abs(x-y) < tol)
            (SigS.fromStorableSignal xs)
            (SigS.fromStorableSignal ys)

{-# INLINE checkSimilarity #-}
checkSimilarity ::
   (RealRing.C b, Storable.C b,
    Storable.C a, Tuple.ValueOf a ~ av) =>
   b ->
   (SVL.Vector a -> SVL.Vector b) ->
   SigP.T p av -> SigP.T p av ->
   IO (ChunkSize -> p -> CheckSimilarity b)
checkSimilarity tol limit gen0 gen1 =
   liftM2
      (\sig0 sig1 chunkSize p ->
         CheckSimilarity tol (sig0 chunkSize p) (sig1 chunkSize p))
      (render limit gen0)
      (render limit gen1)

checkSimilarityPacked ::
   Float ->
   (SVL.Vector Float -> SVL.Vector Float) ->
   SigP.T p (LLVM.Value Float) ->
   SigP.T p (Serial.Value TypeNum.D4 Float) ->
   IO (ChunkSize -> p -> CheckSimilarity Float)
checkSimilarityPacked tol limit scalar vector =
   checkSimilarity tol limit scalar (SigPS.unpack vector)


{- |
Instead of testing on equality immediately
we use this interim data type.
This allows us to inspect the signals that are compared.
-}
data CheckEqualityGen a = CheckEqualityGen a a

type CheckEquality a = CheckEqualityGen (SVL.Vector a)
type CheckEquality2 a b = CheckEqualityGen (SVL.Vector a, SVL.Vector b)

instance (Eq a) => QC.Testable (CheckEqualityGen a) where
   property (CheckEqualityGen x y) = QC.property (x==y)

checkEquality ::
   (Eq a, Storable.C a, Tuple.ValueOf a ~ av) =>
   (SVL.Vector a -> SVL.Vector a) ->
   SigP.T p av -> SigP.T p av ->
   IO (ChunkSize -> p -> CheckEquality a)
checkEquality limit gen0 gen1 =
   liftM2
      (\sig0 sig1 chunkSize p ->
         CheckEqualityGen (sig0 chunkSize p) (sig1 chunkSize p))
      (render limit gen0)
      (render limit gen1)

checkEquality2 ::
   (Eq a, Storable.C a, Tuple.ValueOf a ~ al) =>
   (Eq b, Storable.C b, Tuple.ValueOf b ~ bl) =>
   (SVL.Vector a -> SVL.Vector a) ->
   (SVL.Vector b -> SVL.Vector b) ->
   SigP.T p (al,bl) -> SigP.T p (al,bl) ->
   IO (ChunkSize -> p -> CheckEquality2 a b)
checkEquality2 limitA limitB gen0 gen1 =
   liftM2
      (\sig0 sig1 chunkSize p ->
         CheckEqualityGen (sig0 chunkSize p) (sig1 chunkSize p))
      (render2 (mapPair (limitA, limitB)) gen0)
      (render2 (mapPair (limitA, limitB)) gen1)