-- |
-- Module:      Math.NumberTheory.Primes.Types
-- Copyright:   (c) 2017 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- This is an internal module, defining types for primes.
-- Should not be exposed to users.
--

{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}

module Math.NumberTheory.Primes.Types
  ( Prime(..)
  , toPrimeIntegral
  ) where

import Data.Bits
import GHC.Generics
import Control.DeepSeq
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed as U

import Math.NumberTheory.Utils.FromIntegral

-- $setup
-- >>> import Math.NumberTheory.Primes (nextPrime, precPrime)

-- | Wrapper for prime elements of @a@. It is supposed to be constructed
-- by 'Math.NumberTheory.Primes.nextPrime' / 'Math.NumberTheory.Primes.precPrime'.
-- and eliminated by 'unPrime'.
--
-- One can leverage 'Enum' instance to generate lists of primes.
-- Here are some examples.
--
-- *  Generate primes from the given interval:
--
--    >>> :set -XFlexibleContexts
--    >>> [nextPrime 101 .. precPrime 130]
--    [Prime 101,Prime 103,Prime 107,Prime 109,Prime 113,Prime 127]
--
-- *  Generate an infinite list of primes:
--
--    > [nextPrime 101 ..]
--    > [Prime 101,Prime 103,Prime 107,Prime 109,Prime 113,Prime 127...
--
-- *  Generate primes from the given interval of form p = 6k+5:
--
--    >>> [nextPrime 101, nextPrime 107 .. precPrime 150]
--    [Prime 101,Prime 107,Prime 113,Prime 131,Prime 137,Prime 149]
--
-- *  Get next prime:
--
--    >>> succ (nextPrime 101)
--    Prime 103
--
-- *  Get previous prime:
--
--    >>> pred (nextPrime 101)
--    Prime 97
--
-- *  Count primes less than a given number (cf. 'Math.NumberTheory.Primes.Counting.approxPrimeCount'):
--
--    >>> fromEnum (precPrime 100)
--    25
--
-- *  Get 25-th prime number (cf. 'Math.NumberTheory.Primes.Counting.nthPrimeApprox'):
--
--    >>> toEnum 25 :: Prime Int
--    Prime 97
--
newtype Prime a = Prime
  { unPrime :: a -- ^ Unwrap prime element.
  }
  deriving (Eq, Ord, Generic)

instance NFData a => NFData (Prime a)

instance Show a => Show (Prime a) where
  showsPrec d (Prime p) r = (if d > 10 then "(" ++ s ++ ")" else s) ++ r
    where
      s = "Prime " ++ show p

newtype instance U.MVector s (Prime a) = MV_Prime (U.MVector s a)
newtype instance U.Vector    (Prime a) = V_Prime  (U.Vector    a)

instance U.Unbox a => U.Unbox (Prime a)

instance M.MVector U.MVector a => M.MVector U.MVector (Prime a) where
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicOverlaps #-}
  {-# INLINE basicUnsafeNew #-}
  {-# INLINE basicInitialize #-}
  {-# INLINE basicUnsafeReplicate #-}
  {-# INLINE basicUnsafeRead #-}
  {-# INLINE basicUnsafeWrite #-}
  {-# INLINE basicClear #-}
  {-# INLINE basicSet #-}
  {-# INLINE basicUnsafeCopy #-}
  {-# INLINE basicUnsafeGrow #-}
  basicLength (MV_Prime v) = M.basicLength v
  basicUnsafeSlice i n (MV_Prime v) = MV_Prime $ M.basicUnsafeSlice i n v
  basicOverlaps (MV_Prime v1) (MV_Prime v2) = M.basicOverlaps v1 v2
  basicUnsafeNew n = MV_Prime <$> M.basicUnsafeNew n
  basicInitialize (MV_Prime v) = M.basicInitialize v
  basicUnsafeReplicate n x = MV_Prime <$> M.basicUnsafeReplicate n (unPrime x)
  basicUnsafeRead (MV_Prime v) i = Prime <$> M.basicUnsafeRead v i
  basicUnsafeWrite (MV_Prime v) i x = M.basicUnsafeWrite v i (unPrime x)
  basicClear (MV_Prime v) = M.basicClear v
  basicSet (MV_Prime v) x = M.basicSet v (unPrime x)
  basicUnsafeCopy (MV_Prime v1) (MV_Prime v2) = M.basicUnsafeCopy v1 v2
  basicUnsafeMove (MV_Prime v1) (MV_Prime v2) = M.basicUnsafeMove v1 v2
  basicUnsafeGrow (MV_Prime v) n = MV_Prime <$> M.basicUnsafeGrow v n

instance G.Vector U.Vector a => G.Vector U.Vector (Prime a) where
  {-# INLINE basicUnsafeFreeze #-}
  {-# INLINE basicUnsafeThaw #-}
  {-# INLINE basicLength #-}
  {-# INLINE basicUnsafeSlice #-}
  {-# INLINE basicUnsafeIndexM #-}
  {-# INLINE elemseq #-}
  basicUnsafeFreeze (MV_Prime v) = V_Prime <$> G.basicUnsafeFreeze v
  basicUnsafeThaw (V_Prime v) = MV_Prime <$> G.basicUnsafeThaw v
  basicLength (V_Prime v) = G.basicLength v
  basicUnsafeSlice i n (V_Prime v) = V_Prime $ G.basicUnsafeSlice i n v
  basicUnsafeIndexM (V_Prime v) i = Prime <$> G.basicUnsafeIndexM v i
  basicUnsafeCopy (MV_Prime mv) (V_Prime v) = G.basicUnsafeCopy mv v
  elemseq _ = seq

-- | Convert between primes of different types, similar in spirit to 'toIntegralSized'.
--
-- A simpler version of this function is:
--
-- > toPrimeIntegral :: (Integral a, Integral b) => a -> Maybe b
-- > toPrimeIntegral (Prime a)
-- >   | toInteger a == b = Just (Prime (fromInteger b))
-- >   | otherwise        = Nothing
-- >   where
-- >     b = toInteger a
--
-- The point of 'toPrimeIntegral' is to avoid redundant conversions and conditions,
-- when it is safe to do so, determining type sizes statically with 'bitSizeMaybe'.
-- For example, 'toPrimeIntegral' from 'Prime' 'Int' to 'Prime' 'Word' boils down to
-- 'Just' . 'fromIntegral'.
--
toPrimeIntegral :: (Integral a, Integral b, Bits a, Bits b) => Prime a -> Maybe (Prime b)
toPrimeIntegral (Prime a) = case unsignedWidth b of
  Nothing -> res
  Just bW -> case unsignedWidth a of
    Just aW
      | aW <= bW -> res
    _
      | a <= bit bW - 1 -> res
      | otherwise       -> Nothing
  where
    b = fromIntegral' a
    res = Just (Prime b)
{-# INLINE toPrimeIntegral #-}

unsignedWidth :: Bits a => a -> Maybe Int
unsignedWidth t
  | isSigned t = subtract 1 <$> bitSizeMaybe t
  | otherwise  =                bitSizeMaybe t
{-# INLINE unsignedWidth #-}