content: {-|
  Module      : Crypto.Lol.Cyclotomic.Tensor.Repa.CRT
  Description : Functions to support the chinese remainder transform on Repa arrays.
  Copyright   : (c) Eric Crockett, 2011-2017
                    Chris Peikert, 2011-2017
  License     : GPL-3
  Maintainer  : ecrockett0@email.com
  Stability   : experimental
  Portability : POSIX
  
  Functions to support the chinese remainder transform on Repa arrays.
  -}
  
  {-# LANGUAGE ConstraintKinds       #-}
  {-# LANGUAGE FlexibleContexts      #-}
  {-# LANGUAGE GADTs                 #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE NoImplicitPrelude     #-}
  {-# LANGUAGE ScopedTypeVariables   #-}
  
  module Crypto.Lol.Cyclotomic.Tensor.Repa.CRT
  ( scalarCRT'
  , fCRT, fCRTInv
  , mulGCRT', divGCRT'
  , gCRT, gInvCRT
  ) where
  
  import Crypto.Lol.CRTrans
  import Crypto.Lol.Cyclotomic.Tensor
  import Crypto.Lol.Cyclotomic.Tensor.Repa.RTCommon as RT
  import Crypto.Lol.Prelude                               as LP
  
  import Control.Applicative
  import Data.Coerce
  import Data.Singletons.Prelude
  
  -- | Embeds a scalar into the CRT basis (when it exists).
  scalarCRT' :: forall mon m r . (Fact m, CRTrans mon r, Unbox r)
                => mon (r -> Arr m r)
  {-# INLINABLE scalarCRT' #-}
  scalarCRT'
    = let n = proxy totientFact (Proxy::Proxy m)
          sz = Z :. n
      in pure $ Arr . force . fromFunction sz . const
  
  -- | Multiply by @g_m@ in the CRT basis (when it exists).
  mulGCRT' :: (Fact m, CRTrans mon r, Unbox r)
              => mon (Arr m r -> Arr  m r)
  {-# INLINABLE mulGCRT' #-}
  mulGCRT' = (coerce (\x -> force . RT.zipWith (*) x) `asTypeOf` asTypeOf) <$> gCRT
  
  -- | Divide by @g@ in the CRT basis (when it exists).
  divGCRT' :: (Fact m, CRTrans mon r, Unbox r) => mon (Arr m r -> Arr m r)
  {-# INLINABLE divGCRT' #-}
  divGCRT' = (coerce (\x -> force . RT.zipWith (*) x) `asTypeOf` asTypeOf) <$> gInvCRT
  
  wrapVector :: forall mon m r . (Monad mon, Fact m, Ring r, Unbox r)
                => TaggedT m mon (Kron r) -> mon (Arr m r)
  wrapVector v = do
    vmat <- proxyT v (Proxy::Proxy m)
    let n = proxy totientFact (Proxy::Proxy m)
    return $ coerce $ force $ RT.fromFunction (Z:.n)
      (\(Z:.i) -> indexK vmat i 0)
  
  gCRT, gInvCRT :: (Fact m, CRTrans mon r, Unbox r) => mon (Arr m r)
  {-# INLINABLE gCRT #-}
  {-# INLINABLE gInvCRT #-}
  
  -- | The coefficient vector of @g@ in the CRT basis (when it exists).
  gCRT = wrapVector gCRTK
  -- | The coefficient vector of @g^{ -1 }@ in the CRT basis (when it exists).
  gInvCRT = wrapVector gInvCRTK
  
  fCRT, fCRTInv ::
    forall mon m r . (Fact m, CRTrans mon r, Unbox r, Elt r)
    => mon (Arr m r -> Arr m r)
  
  {-# INLINABLE fCRT #-}
  {-# INLINABLE fCRTInv #-}
  
  -- | The Chinese Remainder Transform.
  -- Exists if and only if CRT exists for all prime powers.
  fCRT = evalM $ fTensor ppCRT
  
  -- divide by mhat after doing crtInv'
  -- | The inverse Chinese Remainder Transform.
  -- Exists if and only if CRT exists for all prime powers.
  fCRTInv = do
    (_, mhatInv) :: (CRTInfo r) <- proxyT crtInfo (Proxy :: Proxy m)
    let totm = proxy totientFact (Proxy :: Proxy m)
        divMhat = trans totm $ RT.map (*mhatInv)
    evalM $ (divMhat .*) <$> fTensor ppCRTInv'
  
  ppDFT, ppDFTInv', ppCRT, ppCRTInv' ::
    forall mon pp r . (PPow pp, CRTrans mon r, Unbox r, Elt r)
    => TaggedT pp mon (Trans r)
  
  {-# INLINABLE ppDFT #-}
  {-# INLINABLE ppDFTInv' #-}
  {-# INLINABLE ppCRT #-}
  {-# INLINABLE ppCRTInv' #-}
  
  ppDFT = case (sing :: SPrimePower pp) of
    (SPP (STuple2 sp SO)) -> tagT $ withWitnessT pDFT sp
    spp@(SPP (STuple2 sp (SS se'))) ->
      tagT $ do
        let spp' = SPP (STuple2 sp se')
        pp'dft <- withWitnessT ppDFT spp'
        pptwid <- withWitnessT (ppTwid False) spp
        pdft <- withWitnessT pDFT sp
        return $ (pp'dft @* Id (dim pdft)) .* pptwid .* (Id (dim pp'dft) @* pdft)
  
  ppDFTInv' = case (sing :: SPrimePower pp) of
    (SPP (STuple2 sp SO)) -> tagT $ withWitnessT pDFTInv' sp
    spp@(SPP (STuple2 sp (SS se'))) ->
      tagT $ do
        let spp' = SPP (STuple2 sp se')
        pp'dftInv' <- withWitnessT ppDFTInv' spp'
        pptwidInv <- withWitnessT (ppTwid True) spp
        pdftInv' <- withWitnessT pDFTInv' sp
        return $ (Id (dim pp'dftInv') @* pdftInv') .* pptwidInv .*
                   (pp'dftInv' @* Id (dim pdftInv'))
  
  ppCRT = case (sing :: SPrimePower pp) of
    (SPP (STuple2 sp SO)) -> tagT $ withWitnessT pCRT sp
    spp@(SPP (STuple2 sp (SS se'))) ->
      tagT $ do
        let spp' = SPP (STuple2 sp se')
        pp'dft <- withWitnessT ppDFT spp'
        pptwid <- withWitnessT (ppTwidHat False) spp
        pcrt <- withWitnessT pCRT sp
        return $
          (pp'dft @* Id (dim pcrt)) .* pptwid .*
          -- save some work when p=2
          (if dim pcrt > 1 then Id (dim pp'dft) @* pcrt else Id (dim pp'dft))
  
  ppCRTInv' = case (sing :: SPrimePower pp) of
    (SPP (STuple2 sp SO)) -> tagT $ withWitnessT pCRTInv' sp
    spp@(SPP (STuple2 sp (SS se'))) ->
      tagT $ do
        let spp' = SPP (STuple2 sp se')
        pp'dftInv' <- withWitnessT ppDFTInv' spp'
        pptwidInv <- withWitnessT (ppTwidHat True) spp
        pcrtInv' <- withWitnessT pCRTInv' sp
        return $
          (Id (dim pp'dftInv') @* pcrtInv') .* pptwidInv .*
          (pp'dftInv' @* Id (dim pcrtInv'))
  
  butterfly :: (Additive r) => Trans r
  butterfly = trans 2 $ \arr ->
              fromFunction (extent arr) $
                               \(sh:.j) -> case j of
                                             0 -> arr ! (sh:.0) +
                                                  arr ! (sh:.1)
                                             1 -> arr ! (sh:.0) -
                                                  arr ! (sh:.1)
  
  -- DFT_p, CRT_p, scaled DFT_p^{ -1 } and CRT_p^{ -1 }
  pDFT, pDFTInv', pCRT, pCRTInv' ::
    forall mon p r . (Prime p, CRTrans mon r, Unbox r, Elt r)
    => TaggedT p mon (Trans r)
  
  {-# INLINABLE pDFT #-}
  {-# INLINABLE pDFTInv' #-}
  {-# INLINABLE pCRT #-}
  {-# INLINABLE pCRTInv' #-}
  
  pDFT = let pval = proxy valuePrime (Proxy::Proxy p)
         in if pval == 2
            then return butterfly
            else do (omegaPPow, _) <- crtInfo
                    return $ trans pval $ mulMat $ force $
                           fromFunction (Z :. pval :. pval)
                                            (\(Z:.i:.j) -> omegaPPow (i*j))
  
  pDFTInv' = let pval = proxy valuePrime (Proxy::Proxy p)
             in if pval == 2
                then return butterfly
                else do (omegaPPow, _) <- crtInfo
                        return $ trans pval $ mulMat $ force $
                               fromFunction (Z :. pval :. pval)
                                                (\(Z:.i:.j) -> omegaPPow (-i*j))
  
  pCRT = let pval = proxy valuePrime (Proxy::Proxy p)
         in if pval == 2
            then return $ Id 1
            else do (omegaPPow, _) <- crtInfo
                    return $ trans (pval-1) $ mulMat $ force $
                           fromFunction (Z :. pval-1 :. pval-1)
                                            (\(Z:.i:.j) -> omegaPPow ((i+1)*j))
  
  -- crt_p * this = \hat{p}*I, for all prime p.
  pCRTInv' =
    let pval = proxy valuePrime (Proxy::Proxy p)
    in if pval == 2 then return $ Id 1
       else do
         (omegaPPow, _) <- crtInfo
         return $ trans (pval-1) $  mulMat $ force $
                fromFunction (Z :. pval-1 :. pval-1)
                                 (\(Z:.i:.j) -> omegaPPow (negate i*(j+1)) -
                                                omegaPPow (j+1))
  
  -- twiddle factors for DFT_pp and CRT_pp decompositions
  ppTwid, ppTwidHat ::
    forall mon pp r . (PPow pp, CRTrans mon r, Unbox r)
    => Bool -> TaggedT pp mon (Trans r)
  
  {-# INLINABLE ppTwid #-}
  {-# INLINABLE ppTwidHat #-}
  
  ppTwid inv =
    let pp@(p,e) = proxy ppPPow (Proxy :: Proxy pp)
        ppval = valuePP pp
    in do
      (omegaPPPow, _) <- crtInfo
      return $ trans ppval $ mulDiag $ force $
                             fromFunction (Z :. ppval)
                             (\(Z:.i) -> let (iq,ir) = i `divMod` p
                                             pow = (if inv then negate else id)
                                                   ir * digitRev (p,e-1) iq
                                         in omegaPPPow pow)
  
  ppTwidHat inv =
    let pp@(p,e) = proxy ppPPow (Proxy :: Proxy pp)
        pptot = totientPP pp
    in do
      (omegaPPPow, _) <- crtInfo
      return $ trans pptot $ mulDiag $ force $
                             fromFunction (Z :. pptot)
                             (\(Z:.i) -> let (iq,ir) = i `divMod` (p-1)
                                             pow = (if inv then negate else id)
                                                   (ir+1) * digitRev (p,e-1) iq
                                         in omegaPPPow pow)
  