-- |
-- Module:      Math.NumberTheory.Prefactored
-- Copyright:   (c) 2017 Andrew Lelechenko
-- Licence:     MIT
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--
-- Type for numbers, accompanied by their factorisation.
--

{-# LANGUAGE TypeFamilies  #-}

{-# OPTIONS_GHC -fno-warn-unused-imports #-}

module Math.NumberTheory.Prefactored
  ( Prefactored(prefValue, prefFactors)
  , fromValue
  , fromFactors
  ) where

import Prelude hiding ((^), gcd)
import Control.Arrow
import Data.Euclidean
import Data.Semigroup
import Data.Semiring (Semiring(..), Mul(..), (^))
import qualified Data.Semiring as Semiring
import Unsafe.Coerce

import Math.NumberTheory.Euclidean.Coprimes
import Math.NumberTheory.Primes
import Math.NumberTheory.Primes.Types

-- | A container for a number and its pairwise coprime (but not necessarily prime)
-- factorisation.
-- It is designed to preserve information about factors under multiplication.
-- One can use this representation to speed up prime factorisation
-- and computation of arithmetic functions.
--
-- For instance, let @p@ and @q@ be big primes:
--
-- >>> let p = 1000000000000000000000000000057 :: Integer
-- >>> let q = 2000000000000000000000000000071 :: Integer
--
-- It would be difficult to compute the totient function
-- of their product as is, because once we multiplied them
-- the information of factors is lost and
-- 'Math.NumberTheory.ArithmeticFunctions.totient' (@p@ * @q@)
-- would take ages. Things become different if we simply
-- change types of @p@ and @q@ to prefactored ones:
--
-- >>> let p = 1000000000000000000000000000057 :: Prefactored Integer
-- >>> let q = 2000000000000000000000000000071 :: Prefactored Integer
--
-- Now the 'Math.NumberTheory.ArithmeticFunctions.totient' function
-- can be computed instantly:
--
-- >>> import Math.NumberTheory.ArithmeticFunctions
-- >>> prefValue $ totient (p^2 * q^3)
-- 8000000000000000000000000001752000000000000000000000000151322000000000000000000000006445392000000000000000000000135513014000000000000000000001126361040
-- >>> prefValue $ totient $ totient (p^2 * q^3)
-- 2133305798262843681544648472180210822742702690942899511234131900112583590230336435053688694839034890779375223070157301188739881477320529552945446912000
--
-- Let us look under the hood:
--
-- >>> import Math.NumberTheory.ArithmeticFunctions
-- >>> prefFactors $ totient (p^2 * q^3)
-- Coprimes {unCoprimes = [(1000000000000000000000000000057,1),(41666666666666666666666666669,1),(2000000000000000000000000000071,2),(111111111111111111111111111115,1),(2,4),(3,3)]}
-- >>> prefFactors $ totient $ totient (p^2 * q^3)
-- Coprimes {unCoprimes = [(39521,1),(227098769,1),(22222222222222222222222222223,1),(2000000000000000000000000000071,1),(361696272343,1),(85331809838489,1),(6046667,1),(199937,1),(5,3),(41666666666666666666666666669,1),(2,22),(3,8)]}
--
-- Pairwise coprimality of factors is crucial, because it allows
-- us to process them independently, possibly even
-- in parallel or concurrent fashion.
--
-- Following invariant is guaranteed to hold:
--
-- > abs (prefValue x) = abs (product (map (uncurry (^)) (prefFactors x)))
data Prefactored a = Prefactored
  { Prefactored a -> a
prefValue   :: a
    -- ^ Number itself.
  , Prefactored a -> Coprimes a Word
prefFactors :: Coprimes a Word
    -- ^ List of pairwise coprime (but not necessarily prime) factors,
    -- accompanied by their multiplicities.
  } deriving (Prefactored a -> Prefactored a -> Bool
(Prefactored a -> Prefactored a -> Bool)
-> (Prefactored a -> Prefactored a -> Bool) -> Eq (Prefactored a)
forall a. Eq a => Prefactored a -> Prefactored a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Prefactored a -> Prefactored a -> Bool
$c/= :: forall a. Eq a => Prefactored a -> Prefactored a -> Bool
== :: Prefactored a -> Prefactored a -> Bool
$c== :: forall a. Eq a => Prefactored a -> Prefactored a -> Bool
Eq, Int -> Prefactored a -> ShowS
[Prefactored a] -> ShowS
Prefactored a -> String
(Int -> Prefactored a -> ShowS)
-> (Prefactored a -> String)
-> ([Prefactored a] -> ShowS)
-> Show (Prefactored a)
forall a. Show a => Int -> Prefactored a -> ShowS
forall a. Show a => [Prefactored a] -> ShowS
forall a. Show a => Prefactored a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Prefactored a] -> ShowS
$cshowList :: forall a. Show a => [Prefactored a] -> ShowS
show :: Prefactored a -> String
$cshow :: forall a. Show a => Prefactored a -> String
showsPrec :: Int -> Prefactored a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Prefactored a -> ShowS
Show)

-- | Create 'Prefactored' from a given number.
--
-- >>> fromValue 123
-- Prefactored {prefValue = 123, prefFactors = Coprimes {unCoprimes = [(123,1)]}}
fromValue :: (Eq a, GcdDomain a) => a -> Prefactored a
fromValue :: a -> Prefactored a
fromValue a
a = a -> Coprimes a Word -> Prefactored a
forall a. a -> Coprimes a Word -> Prefactored a
Prefactored a
a (a -> Word -> Coprimes a Word
forall a b.
(Eq a, GcdDomain a, Eq b, Num b) =>
a -> b -> Coprimes a b
singleton a
a Word
1)

-- | Create 'Prefactored' from a given list of pairwise coprime
-- (but not necessarily prime) factors with multiplicities.
--
-- >>> fromFactors (splitIntoCoprimes [(140, 1), (165, 1)])
-- Prefactored {prefValue = 23100, prefFactors = Coprimes {unCoprimes = [(28,1),(33,1),(5,2)]}}
-- >>> fromFactors (splitIntoCoprimes [(140, 2), (165, 3)])
-- Prefactored {prefValue = 88045650000, prefFactors = Coprimes {unCoprimes = [(28,2),(33,3),(5,5)]}}
fromFactors :: Semiring a => Coprimes a Word -> Prefactored a
fromFactors :: Coprimes a Word -> Prefactored a
fromFactors Coprimes a Word
as = a -> Coprimes a Word -> Prefactored a
forall a. a -> Coprimes a Word -> Prefactored a
Prefactored (Mul a -> a
forall a. Mul a -> a
getMul (Mul a -> a) -> Mul a -> a
forall a b. (a -> b) -> a -> b
$ ((a, Word) -> Mul a) -> [(a, Word)] -> Mul a
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\(a
a, Word
k) -> a -> Mul a
forall a. a -> Mul a
Mul (a -> Mul a) -> a -> Mul a
forall a b. (a -> b) -> a -> b
$ a
a a -> Word -> a
forall a b. (Semiring a, Integral b) => a -> b -> a
^ Word
k) (Coprimes a Word -> [(a, Word)]
forall a b. Coprimes a b -> [(a, b)]
unCoprimes Coprimes a Word
as)) Coprimes a Word
as

instance (Eq a, GcdDomain a) => Semiring (Prefactored a) where
  Prefactored a
v1 Coprimes a Word
_ plus :: Prefactored a -> Prefactored a -> Prefactored a
`plus` Prefactored a
v2 Coprimes a Word
_
    = a -> Prefactored a
forall a. (Eq a, GcdDomain a) => a -> Prefactored a
fromValue (a
v1 a -> a -> a
forall a. Semiring a => a -> a -> a
`plus` a
v2)
  Prefactored a
v1 Coprimes a Word
f1 times :: Prefactored a -> Prefactored a -> Prefactored a
`times` Prefactored a
v2 Coprimes a Word
f2
    = a -> Coprimes a Word -> Prefactored a
forall a. a -> Coprimes a Word -> Prefactored a
Prefactored (a
v1 a -> a -> a
forall a. Semiring a => a -> a -> a
`times` a
v2) (Coprimes a Word
f1 Coprimes a Word -> Coprimes a Word -> Coprimes a Word
forall a. Semigroup a => a -> a -> a
<> Coprimes a Word
f2)
  fromNatural :: Natural -> Prefactored a
fromNatural Natural
n = a -> Prefactored a
forall a. (Eq a, GcdDomain a) => a -> Prefactored a
fromValue (Natural -> a
forall a. Semiring a => Natural -> a
fromNatural Natural
n)

instance (Eq a, Num a, GcdDomain a) => Num (Prefactored a) where
  Prefactored a
v1 Coprimes a Word
_ + :: Prefactored a -> Prefactored a -> Prefactored a
+ Prefactored a
v2 Coprimes a Word
_
    = a -> Prefactored a
forall a. (Eq a, GcdDomain a) => a -> Prefactored a
fromValue (a
v1 a -> a -> a
forall a. Num a => a -> a -> a
+ a
v2)
  Prefactored a
v1 Coprimes a Word
_ - :: Prefactored a -> Prefactored a -> Prefactored a
- Prefactored a
v2 Coprimes a Word
_
    = a -> Prefactored a
forall a. (Eq a, GcdDomain a) => a -> Prefactored a
fromValue (a
v1 a -> a -> a
forall a. Num a => a -> a -> a
- a
v2)
  Prefactored a
v1 Coprimes a Word
f1 * :: Prefactored a -> Prefactored a -> Prefactored a
* Prefactored a
v2 Coprimes a Word
f2
    = a -> Coprimes a Word -> Prefactored a
forall a. a -> Coprimes a Word -> Prefactored a
Prefactored (a
v1 a -> a -> a
forall a. Num a => a -> a -> a
* a
v2) (Coprimes a Word
f1 Coprimes a Word -> Coprimes a Word -> Coprimes a Word
forall a. Semigroup a => a -> a -> a
<> Coprimes a Word
f2)
  negate :: Prefactored a -> Prefactored a
negate (Prefactored a
v Coprimes a Word
f) = a -> Coprimes a Word -> Prefactored a
forall a. a -> Coprimes a Word -> Prefactored a
Prefactored (a -> a
forall a. Num a => a -> a
negate a
v) Coprimes a Word
f
  abs :: Prefactored a -> Prefactored a
abs (Prefactored a
v Coprimes a Word
f)    = a -> Coprimes a Word -> Prefactored a
forall a. a -> Coprimes a Word -> Prefactored a
Prefactored (a -> a
forall a. Num a => a -> a
abs a
v) Coprimes a Word
f
  signum :: Prefactored a -> Prefactored a
signum (Prefactored a
v Coprimes a Word
_) = a -> Coprimes a Word -> Prefactored a
forall a. a -> Coprimes a Word -> Prefactored a
Prefactored (a -> a
forall a. Num a => a -> a
signum a
v) Coprimes a Word
forall a. Monoid a => a
mempty
  fromInteger :: Integer -> Prefactored a
fromInteger Integer
n = a -> Prefactored a
forall a. (Eq a, GcdDomain a) => a -> Prefactored a
fromValue (Integer -> a
forall a. Num a => Integer -> a
fromInteger Integer
n)

instance (Eq a, GcdDomain a, UniqueFactorisation a) => UniqueFactorisation (Prefactored a) where
  factorise :: Prefactored a -> [(Prime (Prefactored a), Word)]
factorise (Prefactored a
_ Coprimes a Word
f)
    = ((a, Word) -> [(Prime (Prefactored a), Word)])
-> [(a, Word)] -> [(Prime (Prefactored a), Word)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(a
x, Word
xm) -> ((Prime a, Word) -> (Prime (Prefactored a), Word))
-> [(Prime a, Word)] -> [(Prime (Prefactored a), Word)]
forall a b. (a -> b) -> [a] -> [b]
map (\(Prime a
p, Word
k) -> (Prefactored a -> Prime (Prefactored a)
forall a. a -> Prime a
Prime (Prefactored a -> Prime (Prefactored a))
-> Prefactored a -> Prime (Prefactored a)
forall a b. (a -> b) -> a -> b
$ a -> Prefactored a
forall a. (Eq a, GcdDomain a) => a -> Prefactored a
fromValue (a -> Prefactored a) -> a -> Prefactored a
forall a b. (a -> b) -> a -> b
$ Prime a -> a
forall a. Prime a -> a
unPrime Prime a
p, Word
k Word -> Word -> Word
forall a. Num a => a -> a -> a
* Word
xm)) (a -> [(Prime a, Word)]
forall a. UniqueFactorisation a => a -> [(Prime a, Word)]
factorise a
x)) (Coprimes a Word -> [(a, Word)]
forall a b. Coprimes a b -> [(a, b)]
unCoprimes Coprimes a Word
f)
  isPrime :: Prefactored a -> Maybe (Prime (Prefactored a))
isPrime (Prefactored a
_ Coprimes a Word
f) = case Coprimes a Word -> [(a, Word)]
forall a b. Coprimes a b -> [(a, b)]
unCoprimes Coprimes a Word
f of
    [(a
n, Word
1)] -> Prefactored a -> Prime (Prefactored a)
forall a. a -> Prime a
Prime (Prefactored a -> Prime (Prefactored a))
-> (Prime a -> Prefactored a) -> Prime a -> Prime (Prefactored a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Prefactored a
forall a. (Eq a, GcdDomain a) => a -> Prefactored a
fromValue (a -> Prefactored a) -> (Prime a -> a) -> Prime a -> Prefactored a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Prime a -> a
forall a. Prime a -> a
unPrime (Prime a -> Prime (Prefactored a))
-> Maybe (Prime a) -> Maybe (Prime (Prefactored a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> Maybe (Prime a)
forall a. UniqueFactorisation a => a -> Maybe (Prime a)
isPrime a
n
    [(a, Word)]
_        -> Maybe (Prime (Prefactored a))
forall a. Maybe a
Nothing