{-|
Module      : Crypto.Lol.Types.Unsafe.ZqBasic
Description : An implementation of modular arithmetic over the integers.
Copyright   : (c) Eric Crockett, 2011-2017
                  Chris Peikert, 2011-2017
License     : GPL-3
Maintainer  : ecrockett0@email.com
Stability   : experimental
Portability : POSIX

  \( \def\Z{\mathbb{Z}} \)
  \( \def\C{\mathbb{C}} \)

An implementation of the quotient ring \(\Z_q = \Z/(q\Z)\).  This
module is "unsafe" because it exports the 'ZqBasic' constructor.  This
module should only be used to make tensor-specific instances for
'ZqBasic'.  The safe way to use this type is to import
"Crypto.Lol.Types".
-}

{-# LANGUAGE AllowAmbiguousTypes        #-}
{-# LANGUAGE ConstraintKinds            #-}
{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE PolyKinds                  #-}
{-# LANGUAGE RebindableSyntax           #-}
{-# LANGUAGE RoleAnnotations            #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TypeApplications           #-}
{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances       #-}

module Crypto.Lol.Types.Unsafe.ZqBasic
( ZqBasic(..) -- export constructor from this (unsafe) module
, goodQs
) where

import Crypto.Lol.CRTrans
import Crypto.Lol.Gadget
import Crypto.Lol.Prelude  as LP
import Crypto.Lol.Reflects

import Math.NumberTheory.Primes

import Control.Applicative
import Control.Arrow
import Control.DeepSeq        (NFData, force)
import Data.Maybe
import NumericPrelude.Numeric as NP (round)
import System.Random

import qualified Data.Vector as V

import qualified Algebra.Additive       as Additive (C)
import qualified Algebra.Field          as Field (C)
import qualified Algebra.IntegralDomain as IntegralDomain (C)
import qualified Algebra.Ring           as Ring (C)
import qualified Algebra.ZeroTestable   as ZeroTestable (C)

-- | An infinite list of primes greater than @lower@ and congruent to
-- 1 mod @m@.
goodQs :: (ToInteger a) => a -> a -> [a]
goodQs m lower = filter (isJust . isPrime . toInteger) $
  iterate (+m) $ lower + ((m-lower) `mod` m) + 1
{-# SPECIALIZE goodQs :: Int64 -> Int64 -> [Int64] #-}

-- | The ring \(\Z_q\) of integers modulo 'q', using underlying integer
-- type 'z'.
newtype ZqBasic q z = ZqB { unZqB :: z }
    deriving (Eq, Ord, ZeroTestable.C, Show, NFData)

-- the q argument, though phantom, matters for safety
type role ZqBasic nominal representational

--deriving instance (U.Unbox i) => G.Vector U.Vector (ZqBasic q i)
--deriving instance (U.Unbox i) => M.MVector U.MVector (ZqBasic q i)
--deriving instance (U.Unbox i) => U.Unbox (ZqBasic q i)

reduce' :: forall q z . (Reflects q z, IntegralDomain z) => z -> ZqBasic q z
reduce' = ZqB . (`mod` value @q)
{-# INLINABLE reduce' #-}

-- puts value in range [-q/2, q/2)
decode' :: forall q z . (Reflects q z, Ring z, Ord z) => ZqBasic q z -> z
decode' = let qval = value @q
          in \(ZqB x) -> if 2 * x < qval then x else x - qval
{-# INLINABLE decode' #-}

instance (Reflects q z, Ring z, Enum z) => Enumerable (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Enumerable (ZqBasic q Int64) #-}
  values = ZqB <$> [0..(value @q - 1)]

instance (Reflects q z, ToInteger z) => Mod (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Mod (ZqBasic q Int64) #-}
  type ModRep (ZqBasic q z) = z

  modulus = value @q

type instance CharOf (ZqBasic p z) = p

instance (Reflects q z, IntegralDomain z) => Reduce z (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Reduce Int64 (ZqBasic q Int64) #-}
  reduce = reduce'

instance {-# OVERLAPPING #-} (Reflects q z, ToInteger z) => Reduce Integer (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Reduce Integer (ZqBasic q Int64) #-}
  reduce = fromInteger

type instance LiftOf (ZqBasic q z) = z

instance (Reflects q z, Ring z, Ord z, IntegralDomain z)
  => Lift' (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Lift' (ZqBasic q Int64) #-}
  lift = decode'

instance (Reflects q z, ToInteger z, Reflects q' z, Ring z)
         => Rescale (ZqBasic q z) (ZqBasic q' z) where
  {-# SPECIALIZE instance (Reflects q Int64, Reflects q' Int64) => Rescale (ZqBasic q Int64) (ZqBasic q' Int64) #-}

  rescale = rescaleMod

instance (Reflects p z, Reflects q z, IntegralDomain z,
          Field (ZqBasic q z), Field (ZqBasic p z))
         => Encode (ZqBasic p z) (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects p Int64, Reflects q Int64) => Encode (ZqBasic p Int64) (ZqBasic q Int64) #-}

  lsdToMSD = let pval = value @p
                 negqval = negate $ value @q
             in (reduce' negqval, recip $ reduce' pval)

-- | Yield a /principal/ \(m\)th root of unity \(\omega_m \in \Z_q^*\).
-- The implementation requires \(q\) to be prime.  It works by finding a
-- generator of \(\Z_q^*\) and raising it to the \( (q-1)/m\) power.
-- Therefore, outputs for different values of \(m\) are consistent,
-- i.e., \(\omega_{m'}^(m'/m) = \omega_m\).
principalRootUnity :: forall m q z .
  (Reflects m Int, Reflects q z, ToInteger z, Enum z, NFData z)
  => TaggedT m Maybe (Int -> ZqBasic q z)
principalRootUnity =        -- use Integers for all intermediate calcs
  let qval = fromIntegral (value @q :: z)
      mval = value @m
      -- order of Zq^* (assuming q prime)
      order = qval-1
      -- the primes dividing the order of Zq^*
      pfactors = unPrime . fst <$> (factorise @Integer) order
      -- the powers we need to check
      exps = div order <$> pfactors
      -- whether an element is a generator of Zq^*
      isGen x = (x^order == one) && all (\e -> x^e /= one) exps
   in tagT $ if isJust (isPrime qval) -- for simplicity, require q to be prime
            then let (mq,mr) = order `divMod` fromIntegral mval
                 in if mr == 0
                    then let omega = head (filter isGen values) ^ mq
                             omegaPows = force $ V.iterateN mval (*omega) one
                         in Just $ (omegaPows V.!) . (`mod` mval)
                    else Nothing
            else Nothing       -- fail if q composite

mhatInv :: forall m q z . (Reflects m Int, Reflects q z, ToInteger z, PID z)
           => TaggedT m Maybe (ZqBasic q z)
mhatInv = tagT $ reduce' <$>
          (`modinv` value @q) (fromIntegral $ valueHat (value @m :: Int))

instance (Reflects q z, ToInteger z, PID z, Enum z, NFData z)
         => CRTrans Maybe (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => CRTrans Maybe (ZqBasic q Int64) #-}

  {-# INLINABLE crtInfo #-}
  crtInfo = (,) <$> principalRootUnity <*> mhatInv

-- | Embeds into the complex numbers \( \C \).
instance (Reflects q z, ToInteger z, Ring (ZqBasic q z))
  => CRTEmbed (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => CRTEmbed (ZqBasic q Int64) #-}

  type CRTExt (ZqBasic q z) = Complex Double

  toExt (ZqB x) = fromReal $ fromIntegral x
  fromExt = reduce' . NP.round . real

-- instance of Additive
instance (Reflects q z, IntegralDomain z) => Additive.C (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Additive.C (ZqBasic q Int64) #-}

  {-# INLINABLE zero #-}
  zero = ZqB zero

  {-# INLINABLE (+) #-}
  (ZqB x) + (ZqB y) = reduce' $ x + y

  {-# INLINABLE negate #-}
  negate (ZqB x) = reduce' $ negate x

-- instance of Ring
instance (Reflects q z, ToInteger z) => Ring.C (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Ring.C (ZqBasic q Int64) #-}

  {-# INLINABLE (*) #-}
  (ZqB x) * (ZqB y) = reduce' $ x * y

  {-# INLINABLE fromInteger #-}
  fromInteger =
    let qval = toInteger (value @q :: z)
    -- this is safe as long as type z can hold the value of q
    in \x -> ZqB $ fromInteger $ x `mod` qval

-- instance of Field
instance (Reflects q z, ToInteger z, PID z, Show z) => Field.C (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Field.C (ZqBasic q Int64) #-}

  {-# INLINABLE recip #-}
  recip = let qval = value @q
              -- safe because modinv returns in range 0..qval-1
          in \(ZqB x) -> ZqB $
               fromMaybe (error $ "ZqB.recip fail: " ++
                         show x ++ "\t" ++ show qval) $ modinv x qval

-- (canonical) instance of IntegralDomain, needed for Cyclotomics
instance (Reflects q z, ToInteger z, PID z, Show z)
  => IntegralDomain.C (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => IntegralDomain.C (ZqBasic q Int64) #-}
  divMod a b = (a/b, zero)

-- Gadget-related instances
instance (Reflects q z, ToInteger z) => Gadget TrivGad (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Gadget TrivGad (ZqBasic q Int64) #-}
  gadget = [one]

instance (Reflects q z, ToInteger z) => Decompose TrivGad (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Decompose TrivGad (ZqBasic q Int64) #-}
  type DecompOf (ZqBasic q z) = z
  decompose x = [lift x]

instance (Reflects q z, ToInteger z, Ring z) => Correct TrivGad (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Correct TrivGad (ZqBasic q Int64) #-}
  correct [b] = (b, [zero])
  correct _   = error "Correct TrivGad: wrong length"

-- BaseBGad instances

gadlen :: (RealIntegral z) => z -> z -> Int
gadlen _ q | isZero q = 0
gadlen b q = 1 + gadlen b (q `div` b)

-- | The base-\(b\) gadget for modulus \(q\), over integers (not mod
-- anything).
gadgetZ :: (RealIntegral z) => z -> z -> [z]
gadgetZ b q = take (gadlen b q) $ iterate (*b) one

instance (Reflects q z, ToInteger z, RealIntegral z, Reflects b z)
         => Gadget (BaseBGad b) (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64, Reflects b Int64) => Gadget (BaseBGad b) (ZqBasic q Int64) #-}

  gadget = let qval = value @q
               bval = value @b
           in reduce' <$> gadgetZ bval qval

instance (Reflects q z, ToInteger z, Reflects b z)
    => Decompose (BaseBGad b) (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64, Reflects b Int64) => Decompose (BaseBGad b) (ZqBasic q Int64) #-}

  type DecompOf (ZqBasic q z) = z
  decompose = let qval = value @q
                  bval = value @b
                  k = gadlen bval qval
                  radices = replicate (k-1) bval
              in decomp radices . lift

-- | Yield the error vector for a noisy multiple of the gadget (all
-- over the integers).
correctZ :: forall z . (RealIntegral z)
            => z                   -- ^ modulus @q@
            -> z                   -- ^ base @b@
            -> [z]                 -- ^ input vector @v = s \cdot g^t + e@
            -> [z]                 -- ^ error @e@
correctZ q b =
  let gadZ = gadgetZ b q
      k = length gadZ
      gadlast = last gadZ
  in \v ->
    if length v /= k
    then error $ "correctZ: wrong length: was " ++ show (length v) ++", expected " ++ show k
    else let (w, x) = barBtRnd (q `div` b) v
             (v', v'l) = subLast v $ qbarD w x
             s = fst $ v'l `divModCent` gadlast
         in zipWith (-) v' $ (s*) <$> gadZ

    where
      -- | Yield @w = round(\bar{B}^t \cdot v / q)@, along with the inner
      -- product of @w@ with the top row of @q \bar{D}@.
      barBtRnd :: z -> [z] -> ([z], z)
      barBtRnd _ [_] = ([], zero)
      barBtRnd q' (v1:vs@(v2:_)) = let quo = fst $ divModCent (b*v1-v2) q
                                   in (quo:) *** (quo*q' +) $
                                      barBtRnd (q' `div` b) vs

      -- | Yield @(q \bar{D}) \cdot w@, given precomputed first entry
      qbarD :: [z] -> z -> [z]
      qbarD [] x      = [x]
      qbarD (w0:ws) x = x : qbarD ws (b*x - q*w0)

      -- | Yield the difference between the input vectors, along with
      -- their final entry.
      subLast :: [z] -> [z] -> ([z], z)
      subLast [v0] [v'0]        = let y = v0-v'0 in ([y], y)
      subLast (v0:vs) (v'0:v's) = first ((v0-v'0):) $ subLast vs v's

instance (Reflects q z, ToInteger z, Reflects b z)
    => Correct (BaseBGad b) (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64, Reflects b Int64) => Correct (BaseBGad b) (ZqBasic q Int64) #-}

  correct =
    let qval = value @q
        bval = value @b
        correct' = correctZ qval bval
    in \tv -> let es = correct' $ lift <$> tv
              in (head tv - reduce (head es), es)

-- instance of Random
instance (Reflects q z, Ring z, Random z) => Random (ZqBasic q z) where
  {-# SPECIALIZE instance (Reflects q Int64) => Random (ZqBasic q Int64) #-}
  random = let high = value @q - 1
           in \g -> let (x,g') = randomR (0,high) g
                    in (ZqB x, g')

  randomR _ = error "randomR non-sensical for Zq types"
  {-# INLINABLE random #-}