content: {-|
  Module      : Crypto.Lol.Cyclotomic.Tensor.Repa.Extension
  Description : RT-specific functions for embedding/twacing in various bases.
  Copyright   : (c) Eric Crockett, 2011-2017
                    Chris Peikert, 2011-2017
  License     : GPL-3
  Maintainer  : ecrockett0@email.com
  Stability   : experimental
  Portability : POSIX
  
  RT-specific functions for embedding/twacing in various bases.
  -}
  
  {-# LANGUAGE BangPatterns          #-}
  {-# LANGUAGE ConstraintKinds       #-}
  {-# LANGUAGE DataKinds             #-}
  {-# LANGUAGE FlexibleContexts      #-}
  {-# LANGUAGE FlexibleInstances     #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE NoImplicitPrelude     #-}
  {-# LANGUAGE PolyKinds             #-}
  {-# LANGUAGE ScopedTypeVariables   #-}
  {-# LANGUAGE TemplateHaskell       #-}
  {-# LANGUAGE TypeFamilies          #-}
  {-# LANGUAGE TypeOperators         #-}
  
  module Crypto.Lol.Cyclotomic.Tensor.Repa.Extension
  ( twacePowDec', twaceCRT', embedPow', embedDec', embedCRT'
  , coeffs', powBasisPow', crtSetDec'
  ) where
  
  import           Crypto.Lol.CRTrans
  import qualified Crypto.Lol.Cyclotomic.Tensor                     as T
  import           Crypto.Lol.Cyclotomic.Tensor.Repa.CRT
  import           Crypto.Lol.Cyclotomic.Tensor.Repa.RTCommon as RT
  import           Crypto.Lol.Prelude                               as LP
  
  import Crypto.Lol.Types.FiniteField
  import Crypto.Lol.Types.ZmStar
  
  import Control.Applicative
  import Control.Arrow       (first, second)
  
  import           Data.Coerce
  import           Data.Default
  import           Data.Maybe
  import           Data.Reflection              (reify)
  import qualified Data.Vector                  as V
  import qualified Data.Vector.Unboxed          as U
  import           Data.Vector.Unboxed.Deriving
  
  -- Default instances
  instance Default Z where def = Z
  instance (Default a, Default b) => Default (a:.b) where def = def:.def
  
  -- derived Unbox instances
  derivingUnbox "DIM1"
    [t| (Z:.Int) -> Int |]
    [| \(Z:.i) -> i |]
    [| (Z :.) |]
  
  -- | The "tweaked trace" function in either the powerful or decoding
  -- basis of the m'th cyclotomic ring to the mth cyclotomic ring when
  -- @m | m'@.
  twacePowDec' :: forall m m' r . (m `Divides` m', Unbox r)
                   => Arr m' r -> Arr m r
  twacePowDec'
    = let indices = proxy extIndicesPowDec (Proxy::Proxy '(m, m'))
      in coerce $ \ !arr -> force $ backpermute (extent indices) (indices !) arr
  
  -- | The "tweaked trace" function in the CRT
  -- basis of the m'th cyclotomic ring to the mth cyclotomic ring when
  -- @m | m'@.
  twaceCRT' :: forall mon m m' r .
               (m `Divides` m', CRTrans mon r, Unbox r, Elt r)
               => mon (Arr m' r -> Arr m r)
  twaceCRT' = do
    g' :: Arr m' r <- gCRT
    gInv <- gInvCRT
    embed :: Arr m r -> Arr m' r <- embedCRT'
    (_, m'hatinv) <- proxyT crtInfo (Proxy::Proxy m')
    let hatRatioInv = m'hatinv * fromIntegral (proxy valueHatFact (Proxy::Proxy m))
        -- tweak = mhat * g' / (m'hat * g)
        tweak = (coerce $ \x -> force . RT.map (* hatRatioInv) . RT.zipWith (*) x) (embed gInv) g' :: Arr m' r
        indices = proxy extIndicesCRT (Proxy::Proxy '(m, m'))
    return $
      -- take true trace after mul-by-tweak
      coerce (\ !arr -> sumS . backpermute (extent indices) (indices !) . RT.zipWith (*) arr) tweak
  
  embedPow', embedDec' :: forall m m' r .
               (m `Divides` m', Unbox r, Additive r)
               => Arr m r -> Arr m' r
  -- | Embeds an array in the powerful basis of the the mth cyclotomic ring
  -- to an array in the powerful basis of the m'th cyclotomic ring when @m | m'@
  embedPow'
    = let indices = proxy baseIndicesPow (Proxy::Proxy '(m, m'))
      in coerce $ \ !arr -> force $ fromFunction (extent indices)
                         (\idx -> let (j0,j1) = (indices ! idx)
                                  in if j0 == 0 then arr ! j1 else zero)
  -- | Embeds an array in the decoding basis of the the mth cyclotomic ring
  -- to an array in the decoding basis of the m'th cyclotomic ring when @m | m'@
  embedDec'
    = let indices = proxy baseIndicesDec (Proxy::Proxy '(m, m'))
      in coerce $ \ !arr -> force $
                         fromFunction (extent indices)
                           (\idx -> maybe zero
                                    (\(sh,b) -> if b then negate (arr ! sh)
                                                else arr ! sh)
                                    (indices ! idx))
  
  -- | Embeds an array in the CRT basis of the the mth cyclotomic ring
  -- to an array in the CRT basis of the m'th cyclotomic ring when @m | m'@
  embedCRT' :: forall mon m m' r . (m `Divides` m', CRTrans mon r, Unbox r)
               => mon (Arr m r -> Arr m' r)
  embedCRT' = do
    -- first check existence of CRT transform of index m'
    _ <- proxyT crtInfo (Proxy::Proxy m') :: mon (CRTInfo r)
    let idxs = proxy baseIndicesCRT (Proxy::Proxy '(m,m'))
    return $ coerce $ \ !arr -> (force $ backpermute (extent idxs) (idxs !) arr)
  
  -- | maps an array in the powerful/decoding basis, representing an
  -- O_m' element, to an array of arrays representing O_m elements in
  -- the same type of basis
  coeffs' :: forall m m' r . (m `Divides` m', Unbox r)
               => Arr m' r -> [Arr m r]
  coeffs' =
    let indices = proxy extIndicesCoeffs (Proxy::Proxy '(m, m'))
    in coerce $ \ !arr -> V.toList $
    V.map (\idxs -> force $ backpermute (extent idxs) (idxs !) arr) indices
  
  -- | The powerful extension basis, wrt the powerful basis.
  -- Outputs a list of arrays in O_m' that are an O_m basis for O_m'
  powBasisPow' :: forall m m' r . (m `Divides` m', Ring r, Unbox r)
                  => Tagged m [Arr m' r]
  powBasisPow' = return $
    let (_, phi, phi', _) = proxy T.indexInfo (Proxy::Proxy '(m,m'))
        idxs = proxy T.baseIndicesPow (Proxy::Proxy '(m,m'))
    in LP.map (\k -> Arr $ force $ fromFunction (Z :. phi')
                           (\(Z:.j) -> let (j0,j1) = idxs U.! j
                                       in if j0==k && j1==0 then one else zero))
        [0..phi' `div` phi - 1]
  
  -- | A list of arrays representing the mod-p CRT set of the
  -- extension O_m'/O_m
  crtSetDec' :: forall m m' fp .
                (m `Divides` m', PrimeField fp, Coprime (PToF (CharOf fp)) m',
                 Unbox fp)
                => Tagged m [Arr m' fp]
  crtSetDec' = return $
    let m'p = Proxy :: Proxy m'
        p = proxy valuePrime (Proxy::Proxy (CharOf fp))
        phi = proxy totientFact m'p
  
        d = proxy (order p) m'p
        h :: Int = proxy valueHatFact m'p
        hinv = recip $ fromIntegral h
    in reify d $ \(_::Proxy d) ->
         let twCRTs' :: T.Kron (GF fp d)
               = fromMaybe (error "internal error: crtSetDec': twCRTs") $ proxyT T.twCRTs m'p
             zmsToIdx = proxy T.zmsToIndexFact m'p
             elt j i = T.indexK twCRTs' j (zmsToIdx i)
             trace' = trace :: GF fp d -> fp
             cosets = proxy (partitionCosets p) (Proxy::Proxy '(m,m'))
         in LP.map (\is -> Arr $ force $ fromFunction (Z :. phi)
                            (\(Z:.j) -> hinv * trace'
                                        (sum $ LP.map (elt j) is))) cosets
  
  
  -- convert memoized reindexing Vectors to Arrays, for convenience and speed
  
  extIndicesPowDec :: forall m m' . (m `Divides` m')
                      => Tagged '(m, m') (Array U DIM1 DIM1)
  extIndicesPowDec = do
    idxs <- T.extIndicesPowDec
    return $ fromUnboxed (Z :. U.length idxs) $ U.map (Z:.) idxs
  
  extIndicesCRT :: forall m m' . (m `Divides` m')
                   => Tagged '(m, m') (Array U DIM2 DIM1)
  extIndicesCRT =
    let phi = proxy totientFact (Proxy::Proxy m)
        phi' = proxy totientFact (Proxy::Proxy m')
    in do
      idxs <- T.extIndicesCRT
      return $ fromUnboxed (Z :. phi :. phi' `div` phi) $ U.map (Z:.) idxs
  
  baseIndicesPow :: forall m m' . (m `Divides` m')
                    => Tagged '(m, m') (Array U DIM1 (Int,DIM1))
  
  baseIndicesDec :: forall m m' . (m `Divides` m')
                    => Tagged '(m, m') (Array U DIM1 (Maybe (DIM1, Bool)))
  
  baseIndicesCRT :: forall m m' . (m `Divides` m')
                    => Tagged '(m, m') (Array U DIM1 DIM1)
  
  baseIndicesPow = do
    idxs <- T.baseIndicesPow
    return $ fromUnboxed (Z :. U.length idxs) $ U.map (second (Z:.)) idxs
  
  baseIndicesDec = do
    idxs <- T.baseIndicesDec
    return $ fromUnboxed (Z :. U.length idxs) $ U.map (liftA (first (Z:.))) idxs
  
  baseIndicesCRT = do
    idxs <- T.baseIndicesCRT
    return $ fromUnboxed (Z :. U.length idxs) $ U.map (Z:.) idxs
  
  extIndicesCoeffs :: forall m m' . (m `Divides` m')
                      => Tagged '(m, m') (V.Vector (Array U DIM1 DIM1))
  extIndicesCoeffs =
    V.map (\arr -> fromUnboxed (Z :. U.length arr) $
                   U.map (Z:.) arr) <$> T.extIndicesCoeffs
  