content: -- |
  -- Module:      Math.NumberTheory.Moduli.Singleton
  -- Copyright:   (c) 2019 Andrew Lelechenko
  -- Licence:     MIT
  -- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
  --
  -- Singleton data types.
  --
  
  {-# LANGUAGE DataKinds           #-}
  {-# LANGUAGE DeriveGeneric       #-}
  {-# LANGUAGE FlexibleInstances   #-}
  {-# LANGUAGE GADTs               #-}
  {-# LANGUAGE LambdaCase          #-}
  {-# LANGUAGE PatternSynonyms     #-}
  {-# LANGUAGE PolyKinds           #-}
  {-# LANGUAGE RankNTypes          #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TupleSections       #-}
  {-# LANGUAGE TypeOperators       #-}
  {-# LANGUAGE ViewPatterns        #-}
  {-# OPTIONS_GHC -Wno-unrecognised-warning-flags -Wno-pattern-namespace-specifier #-}
  
  module Math.NumberTheory.Moduli.Singleton
    ( -- * SFactors singleton
      SFactors
    , sfactors
    , someSFactors
    , unSFactors
    , proofFromSFactors
      -- * CyclicGroup singleton
    , CyclicGroup
    , cyclicGroup
    , cyclicGroupFromFactors
    , cyclicGroupFromModulo
    , proofFromCyclicGroup
    , pattern CG2
    , pattern CG4
    , pattern CGOddPrimePower
    , pattern CGDoubleOddPrimePower
      -- * SFactors \<=\> CyclicGroup
    , cyclicGroupToSFactors
    , sfactorsToCyclicGroup
      -- * Some wrapper
    , Some(..)
    ) where
  
  import Control.DeepSeq
  import Data.Constraint
  import Data.Kind
  import Data.List (sort)
  import qualified Data.Map as M
  import Data.Proxy
  import GHC.Generics
  import GHC.TypeNats (KnownNat, Nat, natVal)
  import Numeric.Natural
  import Unsafe.Coerce
  
  import Math.NumberTheory.Roots (highestPower)
  import Math.NumberTheory.Primes
  import Math.NumberTheory.Primes.Types
  import Math.NumberTheory.Utils.FromIntegral
  
  -- | Wrapper to hide an unknown type-level natural.
  data Some (a :: Nat -> Type) where
    Some :: a m -> Some a
  
  -- | From "Data.Constraint.Nat".
  newtype Magic n = Magic (KnownNat n => Dict (KnownNat n))
  
  -- | This singleton data type establishes a correspondence
  -- between a modulo @m@ on type level
  -- and its factorisation on term level.
  newtype SFactors a (m :: Nat) = SFactors
    { unSFactors :: [(Prime a, Word)]
    -- ^ Factors of @m@.
    } deriving (Show, Generic)
  
  instance Eq (SFactors a m) where
    _ == _ = True
  
  instance Ord (SFactors a m) where
    _ `compare` _ = EQ
  
  instance NFData a => NFData (SFactors a m)
  
  instance Ord a => Eq (Some (SFactors a)) where
    Some (SFactors xs) == Some (SFactors ys) =
      xs == ys
  
  instance Ord a => Ord (Some (SFactors a)) where
    Some (SFactors xs) `compare` Some (SFactors ys) =
      xs `compare` ys
  
  instance Show a => Show (Some (SFactors a)) where
    showsPrec p (Some x) = showsPrec p x
  
  instance NFData a => NFData (Some (SFactors a)) where
    rnf (Some x) = rnf x
  
  -- | Create a singleton from a type-level positive modulo @m@,
  -- passed in a constraint.
  --
  -- >>> :set -XDataKinds
  -- >>> sfactors :: SFactors Integer 13
  -- SFactors {unSFactors = [(Prime 13,1)]}
  sfactors :: forall a m. (Ord a, UniqueFactorisation a, KnownNat m) => SFactors a m
  sfactors = if m == 0
    then error "sfactors: modulo must be positive"
    else SFactors (sort (factorise m))
    where
      m = fromIntegral (natVal (Proxy :: Proxy m))
  
  -- | Create a singleton from factors of @m@.
  -- Factors must be distinct, as in output of 'factorise'.
  --
  -- >>> import Math.NumberTheory.Primes
  -- >>> someSFactors (factorise 98)
  -- SFactors {unSFactors = [(Prime 2,1),(Prime 7,2)]}
  someSFactors :: (Ord a, Num a) => [(Prime a, Word)] -> Some (SFactors a)
  someSFactors
    = Some
    . SFactors
    -- Just a precaution against ill-formed lists of factors
    . M.assocs
    . M.fromListWith (+)
  
  -- | Convert a singleton to a proof that @m@ is known. Usage example:
  --
  -- > toModulo :: SFactors Integer m -> Natural
  -- > toModulo t = case proofFromSFactors t of Sub Dict -> natVal t
  proofFromSFactors :: Integral a => SFactors a m -> (() :- KnownNat m)
  proofFromSFactors (SFactors fs) = Sub $ unsafeCoerce (Magic Dict) (fromIntegral' (factorBack fs) :: Natural)
  
  -- | This singleton data type establishes a correspondence
  -- between a modulo @m@ on type level
  -- and a cyclic group of the same order on term level.
  data CyclicGroup a (m :: Nat)
    = CG2' -- ^ Residues modulo 2.
    | CG4' -- ^ Residues modulo 4.
    | CGOddPrimePower'       (Prime a) Word
    -- ^ Residues modulo @p@^@k@ for __odd__ prime @p@.
    | CGDoubleOddPrimePower' (Prime a) Word
    -- ^ Residues modulo 2@p@^@k@ for __odd__ prime @p@.
    deriving (Show, Generic)
  
  instance Eq (CyclicGroup a m) where
    _ == _ = True
  
  instance Ord (CyclicGroup a m) where
    _ `compare` _ = EQ
  
  instance NFData a => NFData (CyclicGroup a m)
  
  instance Eq a => Eq (Some (CyclicGroup a)) where
    Some CG2' == Some CG2' = True
    Some CG4' == Some CG4' = True
    Some (CGOddPrimePower' p1 k1) == Some (CGOddPrimePower' p2 k2) =
      p1 == p2 && k1 == k2
    Some (CGDoubleOddPrimePower' p1 k1) == Some (CGDoubleOddPrimePower' p2 k2) =
      p1 == p2 && k1 == k2
    _ == _ = False
  
  instance Ord a => Ord (Some (CyclicGroup a)) where
    compare (Some x) (Some y) = case x of
      CG2' -> case y of
        CG2' -> EQ
        _    -> LT
      CG4' -> case y of
        CG2' -> GT
        CG4' -> EQ
        _    -> LT
      CGOddPrimePower' p1 k1 -> case y of
        CGDoubleOddPrimePower'{} -> LT
        CGOddPrimePower' p2 k2 ->
          p1 `compare` p2 <> k1 `compare` k2
        _ -> GT
      CGDoubleOddPrimePower' p1 k1 -> case y of
        CGDoubleOddPrimePower' p2 k2 ->
          p1 `compare` p2 <> k1 `compare` k2
        _ -> GT
  
  instance Show a => Show (Some (CyclicGroup a)) where
    showsPrec p (Some x) = showsPrec p x
  
  instance NFData a => NFData (Some (CyclicGroup a)) where
    rnf (Some x) = rnf x
  
  -- | Create a singleton from a type-level positive modulo @m@,
  -- passed in a constraint.
  --
  -- >>> :set -XDataKinds
  -- >>> import Data.Maybe
  -- >>> cyclicGroup :: Maybe (CyclicGroup Integer 169)
  -- Just (CGOddPrimePower' (Prime 13) 2)
  --
  -- >>> :set -XTypeOperators -XNoStarIsType
  -- >>> import GHC.TypeNats
  -- >>> sfactorsToCyclicGroup (sfactors :: SFactors Integer 4)
  -- Just CG4'
  -- >>> sfactorsToCyclicGroup (sfactors :: SFactors Integer (2 * 13 ^ 3))
  -- Just (CGDoubleOddPrimePower' (Prime 13) 3)
  -- >>> sfactorsToCyclicGroup (sfactors :: SFactors Integer  (4 * 13))
  -- Nothing
  cyclicGroup
    :: forall a m.
       (Integral a, UniqueFactorisation a, KnownNat m)
    => Maybe (CyclicGroup a m)
  cyclicGroup = fromModuloInternal m
    where
      m = fromIntegral (natVal (Proxy :: Proxy m))
  
  -- | Create a singleton from factors.
  -- Factors must be distinct, as in output of 'factorise'.
  cyclicGroupFromFactors
    :: (Eq a, Num a)
    => [(Prime a, Word)]
    -> Maybe (Some (CyclicGroup a))
  cyclicGroupFromFactors = \case
    [(unPrime -> 2, 1)] -> Just $ Some CG2'
    [(unPrime -> 2, 2)] -> Just $ Some CG4'
    [(unPrime -> 2, _)] -> Nothing
    [(p, k)] -> Just $ Some $ CGOddPrimePower' p k
    [(unPrime -> 2, 1), (p, k)] -> Just $ Some $ CGDoubleOddPrimePower' p k
    [(p, k), (unPrime -> 2, 1)] -> Just $ Some $ CGDoubleOddPrimePower' p k
    _ -> Nothing
  
  -- | Similar to 'cyclicGroupFromFactors' . 'factorise',
  -- but much faster, because it
  -- but performes only one primality test instead of full
  -- factorisation.
  cyclicGroupFromModulo
    :: (Integral a, UniqueFactorisation a)
    => a
    -> Maybe (Some (CyclicGroup a))
  cyclicGroupFromModulo = fmap Some . fromModuloInternal
  
  fromModuloInternal
    :: (Integral a, UniqueFactorisation a)
    => a
    -> Maybe (CyclicGroup a m)
  fromModuloInternal = \case
    2 -> Just CG2'
    4 -> Just CG4'
    n
      | even n -> uncurry CGDoubleOddPrimePower' <$> isOddPrimePower (n `div` 2)
      | otherwise -> uncurry CGOddPrimePower' <$> isOddPrimePower n
  
  isOddPrimePower
    :: (Integral a, UniqueFactorisation a)
    => a
    -> Maybe (Prime a, Word)
  isOddPrimePower n
    | even n    = Nothing
    | otherwise = (, k) <$> isPrime p
    where
      (p, k) = highestPower n
  
  -- | Convert a cyclic group to a proof that @m@ is known. Usage example:
  --
  -- > toModulo :: CyclicGroup Integer m -> Natural
  -- > toModulo t = case proofFromCyclicGroup t of Sub Dict -> natVal t
  proofFromCyclicGroup :: Integral a => CyclicGroup a m -> (() :- KnownNat m)
  proofFromCyclicGroup = proofFromSFactors . cyclicGroupToSFactors
  
  -- | Check whether a multiplicative group of residues,
  -- characterized by its modulo, is cyclic and, if yes, return its form.
  --
  -- >>> :set -XTypeOperators -XNoStarIsType
  -- >>> import GHC.TypeNats
  -- >>> sfactorsToCyclicGroup (sfactors :: SFactors Integer 4)
  -- Just CG4'
  -- >>> sfactorsToCyclicGroup (sfactors :: SFactors Integer (2 * 13 ^ 3))
  -- Just (CGDoubleOddPrimePower' (Prime 13) 3)
  -- >>> sfactorsToCyclicGroup (sfactors :: SFactors Integer  (4 * 13))
  -- Nothing
  sfactorsToCyclicGroup :: (Eq a, Num a) => SFactors a m -> Maybe (CyclicGroup a m)
  sfactorsToCyclicGroup (SFactors fs) = case fs of
    [(unPrime -> 2, 1)]         -> Just CG2'
    [(unPrime -> 2, 2)]         -> Just CG4'
    [(unPrime -> 2, _)]         -> Nothing
    [(p, k)]                    -> Just $ CGOddPrimePower' p k
    [(p, k), (unPrime -> 2, 1)] -> Just $ CGDoubleOddPrimePower' p k
    [(unPrime -> 2, 1), (p, k)] -> Just $ CGDoubleOddPrimePower' p k
    _ -> Nothing
  
  -- | Invert 'sfactorsToCyclicGroup'.
  --
  -- >>> import Data.Maybe
  -- >>> cyclicGroupToSFactors (fromJust (sfactorsToCyclicGroup (sfactors :: SFactors Integer 4)))
  -- SFactors {unSFactors = [(Prime 2,2)]}
  cyclicGroupToSFactors :: Num a => CyclicGroup a m -> SFactors a m
  cyclicGroupToSFactors = SFactors . \case
    CG2' -> [(Prime 2, 1)]
    CG4' -> [(Prime 2, 2)]
    CGOddPrimePower' p k -> [(p, k)]
    CGDoubleOddPrimePower' p k -> [(Prime 2, 1), (p, k)]
  
  -- | Unidirectional pattern for residues modulo 2.
  pattern CG2 :: CyclicGroup a m
  pattern CG2 <- CG2'
  
  -- | Unidirectional pattern for residues modulo 4.
  pattern CG4 :: CyclicGroup a m
  pattern CG4 <- CG4'
  
  -- | Unidirectional pattern for residues modulo @p@^@k@ for __odd__ prime @p@.
  pattern CGOddPrimePower :: Prime a -> Word -> CyclicGroup a m
  pattern CGOddPrimePower p k <- CGOddPrimePower' p k
  
  -- | Unidirectional pattern for residues modulo 2@p@^@k@ for __odd__ prime @p@.
  pattern CGDoubleOddPrimePower :: Prime a -> Word -> CyclicGroup a m
  pattern CGDoubleOddPrimePower p k <- CGDoubleOddPrimePower' p k
  
  {-# COMPLETE CG2, CG4, CGOddPrimePower, CGDoubleOddPrimePower #-}
  