{-# LANGUAGE DataKinds, FlexibleContexts,
             NoImplicitPrelude, PolyKinds, RebindableSyntax,
             ScopedTypeVariables, TypeFamilies, TypeOperators #-}

module SHETests (sheTests) where

import Gen
import Harness.SHE
import Tests hiding (hideArgs)
import Utils

import Control.Applicative
import Control.Monad
import Control.Monad.Random
import Control.Monad.State

import Crypto.Lol hiding (CT)
import Crypto.Lol.Applications.SymmSHE
import Crypto.Lol.Cyclotomic.Linear
import qualified Crypto.Lol.Cyclotomic.Tensor.CTensor as CT

import qualified Test.Framework as TF
import Test.Framework.Providers.QuickCheck2
import Test.QuickCheck

v = 1 :: Double

hideArgs :: forall a rnd bnch. 
  (GenArgs (StateT (Maybe (SKOf a)) rnd) bnch, MonadRandom rnd, 
   ShowType a, ResultOf bnch ~ Test a)
  => bnch -> Proxy a -> rnd TF.Test
hideArgs f p = do
  res <- evalStateT (genArgs f) (Nothing :: Maybe (SKOf a))
  case res of
    Test b -> return $ testProperty (showType p) b
    TestM b -> testProperty (showType p) <$> b

sheTests = 
  [testGroupM "Dec . Enc"  $ applyDec (Proxy::Proxy DecParams) $ hideArgs prop_encDec,
   testGroupM "DecU . Enc" $ applyCTFunc (Proxy::Proxy CTParams) $ hideArgs prop_encDecU,
   testGroupM "AddPub"     $ applyCTFunc (Proxy::Proxy CTParams) $ hideArgs prop_addPub,
   testGroupM "MulPub"     $ applyCTFunc (Proxy::Proxy CTParams) $ hideArgs prop_mulPub,
   testGroupM "ScalarPub"  $ applyCTFunc (Proxy::Proxy CTParams) $ hideArgs prop_addScalar,
   testGroupM "CTAdd"      $ applyCTFunc (Proxy::Proxy CTParams) $ hideArgs prop_ctadd,
   testGroupM "CTMul"      $ applyCTFunc (Proxy::Proxy CTParams) $ hideArgs prop_ctmul,
   testGroupM "CT zero"    $ applyCTFunc (Proxy::Proxy CTParams) $ hideArgs prop_ctzero,
   testGroupM "CT one"     $ applyCTFunc (Proxy::Proxy CTParams) $ hideArgs prop_ctone,
   testGroupM "ModSwitch PT" modSwPTTests,
   testGroupM "Tunnel"       tunnelTests,
   testGroupM "Twace"      $ applyCTTwEm (Proxy::Proxy TwoIdxParams) $ hideArgs prop_cttwace,
   testGroupM "Embed"      $ applyCTTwEm (Proxy::Proxy TwoIdxParams) $ hideArgs prop_ctembed,
   testGroupM "KSLin"      $ applyKSQ (Proxy::Proxy KSQParams) $ hideArgs prop_ksLin,
   testGroupM "keySwitch"  $ applyKSQ (Proxy::Proxy KSQParams) $ hideArgs prop_ksQuad
  ]

type CTCombos = '[
  '(F7, F7, Zq 2,Zq (19393921 ** 18869761)),
  '(F7, F21,Zq 2,Zq (19393921 ** 18869761)),
  '(F2, F8, Zq 2,Zq 536871001),
  '(F1, F8, Zq 2,Zq 536871001),
  '(F4, F12,Zq 2,Zq 2148249601),
  '(F4, F8, Zq 3,Zq 2148249601),
  '(F7, F7, Zq 4,Zq (19393921 ** 18869761)),
  '(F7, F21,Zq 4,Zq (19393921 ** 18869761)),
  '(F1, F4, Zq 4,Zq 18869761),
  '(F4, F4, Zq 4,Zq 18869761),
  '(F14,F14,Zq 4,Zq 18869761),
  '(F28,F28,Zq 4,Zq 18869761),
  '(F28,F28,Zq 4,Zq 80221),
  '(F1, F8, Zq 4,Zq 536871001),
  '(F2, F8, Zq 4,Zq 536871001),
  '(F4, F12,Zq 8,Zq 2148249601)
  ]

type Gadgets = '[TrivGad, BaseBGad 2]
type Tensors = '[CT.CT,RT]
type MM'PQCombos = 
  '[ '(F1, F7, Zq 2, Zq (19393921 ** 18869761)),
     '(F2, F4, Zq 8, Zq (2148854401 ** 2148249601)),
     '(F4, F12, Zq 2, Zq (2148854401 ** 2148249601)),
     '(F8, F64, Zq 2, Zq (2148854401 ** 2148249601)),
     '(F3, F27, Zq 2, Zq (2148854401 ** 2148249601)),
     '(F2, F4, Zq 8, Zq (2148854401 ** 2148249601 ** 2150668801)),
     '(F4, F12, Zq 2, Zq (2148854401 ** 2148249601 ** 2150668801)),
     '(F8, F64, Zq 2, Zq (2148854401 ** 2148249601 ** 2150668801)),
     '(F3, F27, Zq 2, Zq (2148854401 ** 2148249601 ** 2150668801))]


type CTParams  = ( '(,) <$> Tensors) <*> CTCombos
type DecParams = ( '(,) <$> Tensors) <*> (Nub (Filter Liftable CTCombos))
type Zq'Params = ( '(,) <$> Tensors) <*> (Map AddZq (Filter NonLiftable MM'PQCombos))
type KSQParams = ( '(,) <$> Gadgets) <*> Zq'Params
type TwoIdxParams = ( '(,) <$> Tensors) <*> '[ '(F1, F7, F3, F21, Zq 2, Zq 18869761)]

prop_ksLin :: (DecryptUCtx t m m' z zp zq, Eq (Cyc t m zp))
  => SK (Cyc t m' z) 
     -> KSLinear t m m' z zp zq zq' gad 
     -> PTCT m zp (Cyc t m' zq) 
     -> Test '(t,m,m',zp,zq,zq',gad)
prop_ksLin skin (KSL kswlin skout) (PTCT x' x) =
  let y = kswlin x
      y' = decryptUnrestricted skout y
  in test $ x' == y'

prop_ksQuad :: (Ring (CT m zp (Cyc t m' zq)),
                DecryptUCtx t m m' z zp zq, 
                Eq (Cyc t m zp))
  => SK (Cyc t m' z) 
     -> KSHint m zp t m' zq gad zq' 
     -> PTCT m zp (Cyc t m' zq)
     -> PTCT m zp (Cyc t m' zq)
     -> Test '(t,m,m',zp,zq,zq',gad)
prop_ksQuad sk (KeySwitch kswq) (PTCT y1 x1) (PTCT y2 x2) = 
  let x' = kswq $ x1*x2
      y = y1*y2
      x = decryptUnrestricted sk x'
  in test $ y == x

prop_addPub :: forall t m m' z zp zq . 
  (DecryptUCtx t m m' z zp zq,
   AddPublicCtx t m m' zp zq,
   Eq (Cyc t m zp))
  => SK (Cyc t m' z) 
     -> Cyc t m zp 
     -> PTCT m zp (Cyc t m' zq) 
     -> Test '(t,m,m',zp,zq)
prop_addPub sk x (PTCT y' y) = 
  let xy = addPublic x y
      xy' = decryptUnrestricted sk xy
  in test $ xy' == (x+y')

prop_mulPub :: (DecryptUCtx t m m' z zp zq,
                MulPublicCtx t m m' zp zq,
                Eq (Cyc t m zp))
  => SK (Cyc t m' z) 
     -> Cyc t m zp 
     -> PTCT m zp (Cyc t m' zq)
     -> Test '(t,m,m',zp,zq)
prop_mulPub sk x (PTCT y' y) = 
  let xy = mulPublic x y
      xy' = decryptUnrestricted sk xy
  in test $ xy' == (x*y')

prop_addScalar :: (DecryptUCtx t m m' z zp zq,
                   AddScalarCtx t m' zp zq,
                   Eq (Cyc t m zp))
  => SK (Cyc t m' z) -> zp -> PTCT m zp (Cyc t m' zq) -> Test '(t,m,m',zp,zq)
prop_addScalar sk c (PTCT x' x) =
  let cx = addScalar c x
      cx' = decryptUnrestricted sk cx
  in test $ cx' == ((scalarCyc c)+x')

prop_ctadd :: (DecryptUCtx t m m' z zp zq,
               Additive (CT m zp (Cyc t m' zq)),
               Eq (Cyc t m zp))
  => SK (Cyc t m' z) 
     -> PTCT m zp (Cyc t m' zq)
     -> PTCT m zp (Cyc t m' zq)
     -> Test '(t,m,m',zp,zq)
prop_ctadd sk (PTCT x1' x1) (PTCT x2' x2) = 
  let y = x1+x2
      y' = decryptUnrestricted sk y
  in test $ x1'+x2' == y'

prop_ctmul :: (DecryptUCtx t m m' z zp zq,
               Ring (CT m zp (Cyc t m' zq)),
               Eq (Cyc t m zp))
  => SK (Cyc t m' z) 
     -> PTCT m zp (Cyc t m' zq)
     -> PTCT m zp (Cyc t m' zq)
     -> Test '(t,m,m',zp,zq)
prop_ctmul sk (PTCT x1' x1) (PTCT x2' x2) = 
  let y = x1*x2
      y' = decryptUnrestricted sk y
  in test $ x1'*x2' == y'

prop_ctzero :: forall t m m' z zp zq .
 (DecryptUCtx t m m' z zp zq,
  Additive (CT m zp (Cyc t m' zq)),
  Eq (Cyc t m zp)) 
  => SK (Cyc t m' z) -> Test '(t,m,m',zp,zq)
prop_ctzero sk =
  let z = decryptUnrestricted sk (zero :: CT m zp (Cyc t m' zq))
  in test $ zero == z

prop_ctone :: forall t m m' z zp zq .
  (DecryptUCtx t m m' z zp zq,
   Ring (CT m zp (Cyc t m' zq)),
   Eq (Cyc t m zp))
  => SK (Cyc t m' z) -> Test '(t,m,m',zp,zq)
prop_ctone sk = 
  let z = decryptUnrestricted sk (one :: CT m zp (Cyc t m' zq))
  in test $ one == z

prop_ctembed :: forall t r r' s s' z zp zq . 
  (DecryptUCtx t r r' z zp zq,
   DecryptUCtx t s s' z zp zq,
   r `Divides` r', 
   s `Divides` s', 
   r `Divides` s, 
   r' `Divides` s',
   Eq (Cyc t s zp))
  => SK (Cyc t r' z) -> PTCT r zp (Cyc t r' zq) -> Test '(t,r,r',s,s',zp,zq)
prop_ctembed sk (PTCT x' x) = 
  let y = embedCT x :: CT s zp (Cyc t s' zq)
      y' = decryptUnrestricted (embedSK sk) y
  in test $ (embed x' :: Cyc t s zp) == y'

-- CT must be encrypted with key from small ring
prop_cttwace :: forall t r r' s s' z zp zq . 
  (Eq zp,
   EncryptCtx t s s' z zp zq, 
   DecryptUCtx t r r' z zp zq,
   r `Divides` s,
   r' `Divides` s',
   s `Divides` s',
   r ~ (FGCD r' s))
  => SK (Cyc t r' z) -> Cyc t s zp -> Test '(t,r,r',s,s',zp,zq)
prop_cttwace sk x = testIO $ do
  y :: CT s zp (Cyc t s' zq) <- encrypt (embedSK sk) x
  let y' = twaceCT y :: CT r zp (Cyc t r' zq)
      x' = decryptUnrestricted sk y'
  return $ (twace x :: Cyc t r zp) == x'

prop_encDecU :: forall t m m' z zp zq . 
  (GenSKCtx t m' z Double, 
   EncryptCtx t m m' z zp zq, 
   DecryptUCtx t m m' z zp zq,
   Eq (Cyc t m zp))
  => SK (Cyc t m' z) -> Cyc t m zp -> Test '(t,m,m',zp,zq)
prop_encDecU sk x = testIO $ do
  y :: CT m zp (Cyc t m' zq) <- encrypt sk x
  let x' = decryptUnrestricted sk $ y
  return $ x == x'

prop_encDec :: forall t m m' z zp zq . 
  (GenSKCtx t m' z Double, 
   EncryptCtx t m m' z zp zq, 
   DecryptCtx t m m' z zp zq,
   Eq (Cyc t m zp))
  => SK (Cyc t m' z) -> Cyc t m zp -> Test '(t,m,m',zp,zq)
prop_encDec sk x = testIO $ do
  y :: CT m zp (Cyc t m' zq) <- encrypt sk x
  let x' = decrypt sk $ y
  return $ x == x'

helper :: (Proxy '(t,b) -> a) -> Proxy t -> Proxy b -> a
helper f _ _ = f Proxy

-- one-off tests, no hideArgsper
prop_modSwPT :: forall t m m' z zp zp' zq .
  (Eq zp, Eq zp',
   DecryptUCtx t m m' z zp zq,
   DecryptUCtx t m m' z zp' zq,
   ModSwitchPTCtx t m' zp zp' zq,
   RescaleCyc (Cyc t) zp zp',
   Ring (Cyc t m zp),
   Mod zp, Mod zp',
   ModRep zp ~ ModRep zp') 
  => SK (Cyc t m' z) -> CT m zp (Cyc t m' zq) -> Test '(t, '(m,m',zp',zp,zq))
prop_modSwPT sk y =
  let p = proxy modulus (Proxy::Proxy zp)
      p' = proxy modulus (Proxy::Proxy zp')
      z = (fromIntegral $ p `div` p')*y
      x = decryptUnrestricted sk z
      y' = modSwitchPT z :: CT m zp' (Cyc t m' zq)
      x'' = decryptUnrestricted sk y'
  in test $ x'' == rescaleCyc Dec x

modSwPTTests = (modSwPTTests' (Proxy::Proxy CT.CT)) ++ (modSwPTTests' (Proxy::Proxy RT))

modSwPTTests' p = 
  [helper (hideArgs prop_modSwPT) p (Proxy::Proxy '(F7,F21,Zq 4,Zq 8,Zq 18869761)),
   helper (hideArgs prop_modSwPT) p (Proxy::Proxy '(F7,F42,Zq 2,Zq 4,Zq (18869761 ** 19393921)))]


tunnelTests = (tunnelTests' (Proxy::Proxy CT.CT)) ++ (tunnelTests' (Proxy::Proxy RT))

tunnelTests' p = 
  [helper (hideArgs prop_ringTunnel) p 
    (Proxy::Proxy '(F8,F40,F20,F60,Zq 4,Zq (18869761 ** 19393921),TrivGad))]

prop_ringTunnel :: forall t e r s e' r' s' z zp zq gad . 
  (TunnelCtx t e r s e' r' s' z zp zq gad,
   EncryptCtx t r r' z zp zq,
   GenSKCtx t r' z Double,
   GenSKCtx t s' z Double,
   DecryptUCtx t s s' z zp zq,
   Random zp, Eq zp,
   e ~ FGCD r s, Fact e) 
  => Cyc t r zp -> Test '(t,'(r,r',s,s',zp,zq,gad))
prop_ringTunnel x = testIO $ do
  let totr = proxy totientFact (Proxy::Proxy r)
      tote = proxy totientFact (Proxy::Proxy e)
      basisSize = totr `div` tote
  -- choose a random linear function of the appropriate size
  bs :: [Cyc t s zp] <- replicateM basisSize getRandom
  let f = (linearDec bs) \\ (gcdDivides (Proxy::Proxy r) (Proxy::Proxy s)) :: Linear t zp e r s 
      expected = evalLin f x \\ (gcdDivides (Proxy::Proxy r) (Proxy::Proxy s))
  skin :: SK (Cyc t r' (LiftOf zp)) <- genSK v
  skout :: SK (Cyc t s' (LiftOf zp)) <- genSK v
  y :: CT r zp (Cyc t r' zq) <- encrypt skin x
  tunn <- proxyT (tunnelCT f skout skin) (Proxy::Proxy gad)
  let y' = tunn y :: CT s zp (Cyc t s' zq)
      actual = decryptUnrestricted skout y' :: Cyc t s zp
  return $ expected == actual