-- |
-- Module:      Data.Euclidean
-- Copyright:   (c) 2019 Andrew Lelechenko
-- Licence:     BSD3
-- Maintainer:  Andrew Lelechenko <andrew.lelechenko@gmail.com>
--

{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DefaultSignatures          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MagicHash                  #-}

module Data.Euclidean
  ( Euclidean(..)
  , Field
  , GcdDomain(..)
  , WrappedIntegral(..)
  , WrappedFractional(..)
  , gcdExt
  ) where

import Prelude hiding (quotRem, quot, rem, divMod, div, mod, gcd, lcm, negate, (*))
import qualified Prelude as P
import Control.Exception
import Data.Bits
import Data.Complex
import Data.Int (Int, Int8, Int16, Int32, Int64)
import Data.Maybe
import Data.Ratio
import Data.Semiring
import Data.Word (Word, Word8, Word16, Word32, Word64)
import Foreign.C.Types
import GHC.Exts
import GHC.Integer.GMP.Internals

import Numeric.Natural

---------------------------------------------------------------------
-- Classes
---------------------------------------------------------------------

-- | 'GcdDomain' represents a
-- <https://en.wikipedia.org/wiki/GCD_domain GCD domain>.
-- This is a domain, where GCD can be defined,
-- but which does not necessarily allow a well-behaved
-- division with remainder (as in 'Euclidean' domains).
--
-- For example, there is no way to define 'rem' over
-- polynomials with integer coefficients such that
-- remainder is always "smaller" than divisor. However,
-- 'gcd' is still definable, just not by means of
-- Euclidean algorithm.
--
-- All methods of 'GcdDomain' have default implementations
-- in terms of 'Euclidean'. So most of the time
-- it is enough to write:
--
-- > instance GcdDomain Foo
-- > instance Euclidean Foo where
-- >   quotRem = ...
-- >   degree  = ...
class Semiring a => GcdDomain a where

  -- | Division without remainder.
  --
  -- prop> \x y -> (x * y) `divide` y == Just x
  -- prop> \x y -> maybe True (\z -> x == z * y) (x `divide` y)
  divide :: a -> a -> Maybe a

  default divide :: (Eq a, Euclidean a) => a -> a -> Maybe a
  divide x y = let (q, r) = quotRem x y in
    if isZero r then Just q else Nothing

  -- | Greatest common divisor. Must satisfy
  --
  -- prop> \x y -> isJust (x `divide` gcd x y) && isJust (y `divide` gcd x y)
  -- prop> \x y z -> isJust (gcd (x * z) (y * z) `divide` z)
  gcd :: a -> a -> a

  default gcd :: (Eq a, Euclidean a) => a -> a -> a
  gcd a b
    | isZero b  = a
    | otherwise = gcd b (a `rem` b)

  -- | Lowest common multiple. Must satisfy
  --
  -- prop> \x y -> isJust (lcm x y `divide` x) && isJust (lcm x y `divide` y)
  -- prop> \x y z -> isNothing (z `divide` x) || isNothing (z `divide` y) || isJust (z `divide` lcm x y)
  lcm :: a -> a -> a

  default lcm :: Eq a => a -> a -> a
  lcm a b
    | isZero a || isZero b = zero
    | otherwise = case a `divide` gcd a b of
      Nothing -> error "lcm: violated gcd invariant"
      Just c  -> c * b

  -- | Test whether two arguments are
  -- <https://en.wikipedia.org/wiki/Coprime_integers coprime>.
  -- Must match its default definition:
  --
  -- prop> \x y -> coprime x y == isJust (1 `divide` gcd x y)
  coprime :: a -> a -> Bool

  default coprime :: Eq a => a -> a -> Bool
  coprime x y = isJust (one `divide` gcd x y)

infixl 7 `divide`

-- | Informally speaking, 'Euclidean' is a superclass of 'Integral',
-- lacking 'toInteger', which allows to define division with remainder
-- for a wider range of types, e. g., complex integers
-- and polynomials with rational coefficients.
--
-- 'Euclidean' represents a
-- <https://en.wikipedia.org/wiki/Euclidean_domain Euclidean domain>
-- endowed by a given Euclidean function 'degree'.
class GcdDomain a => Euclidean a where
  -- | Division with remainder.
  --
  -- prop> \x y -> y == 0 || let (q, r) = x `quotRem` y in x == q * y + r
  quotRem :: a -> a -> (a, a)

  -- | Division. Must match its default definition:
  --
  -- prop> \x y -> quot x y == fst (quotRem x y)
  quot :: a -> a -> a
  quot x y = fst (quotRem x y)

  -- | Remainder. Must match its default definition:
  --
  -- prop> \x y -> rem x y == snd (quotRem x y)
  rem :: a -> a -> a
  rem x y = snd (quotRem x y)

  -- | Euclidean (aka degree, valuation, gauge, norm) function on 'a'. Usually 'fromIntegral' . 'abs'.
  --
  -- 'degree' is rarely used by itself. Its purpose
  -- is to provide an evidence of soundness of 'quotRem'
  -- by testing the following property:
  --
  -- prop> \x y -> y == 0 || let (q, r) = x `quotRem` y in (r == 0 || degree r < degree y)
  degree :: a -> Natural

infixl 7 `quot`
infixl 7 `rem`

coprimeIntegral :: Integral a => a -> a -> Bool
coprimeIntegral x y = (odd x || odd y) && P.gcd x y == 1

-- | Execute the extended Euclidean algorithm.
-- For elements @a@ and @b@, compute their greatest common divisor @g@
-- and the coefficient @s@ satisfying @as + bt = g@ for some @t@.
gcdExt :: (Eq a, Euclidean a, Ring a) => a -> a -> (a, a)
gcdExt = go one zero
  where
    go s s' r r'
      | r' == zero = (r, s)
      | otherwise  = case quotRem r r' of
        (q, r'') -> go s' (minus s (times q s')) r' r''
{-# INLINABLE gcdExt #-}

-- | 'Field' represents a
-- <https://en.wikipedia.org/wiki/Field_(mathematics) field>,
-- a ring with a multiplicative inverse for any non-zero element.
class (Euclidean a, Ring a) => Field a

---------------------------------------------------------------------
-- Instances
---------------------------------------------------------------------

instance GcdDomain () where
  divide  = const $ const (Just ())
  gcd     = const $ const ()
  lcm     = const $ const ()
  coprime = const $ const True

instance Euclidean () where
  degree  = const 0
  quotRem = const $ const ((), ())
  quot    = const $ const ()
  rem     = const $ const ()

instance Field ()

instance GcdDomain Mod2 where

instance Euclidean Mod2 where
  degree = const 0
  quotRem x y
    | isZero y  = throw DivideByZero
    | otherwise = (x, zero)

instance Field Mod2

-- | Wrapper around 'Integral' with 'GcdDomain'
-- and 'Euclidean' instances.
newtype WrappedIntegral a = WrapIntegral { unwrapIntegral :: a }
  deriving (Eq, Ord, Show, Num, Integral, Real, Enum, Bits)

instance Num a => Semiring (WrappedIntegral a) where
  plus  = (P.+)
  zero  = 0
  times = (P.*)
  one   = 1
  fromNatural = P.fromIntegral

instance Num a => Ring (WrappedIntegral a) where
  negate = P.negate

instance Integral a => GcdDomain (WrappedIntegral a) where
  gcd     = P.gcd
  lcm     = P.lcm
  coprime = coprimeIntegral

instance Integral a => Euclidean (WrappedIntegral a) where
  degree  = P.fromIntegral . abs . unwrapIntegral
  quotRem = P.quotRem
  quot    = P.quot
  rem     = P.rem

instance GcdDomain Int where
#if MIN_VERSION_integer_gmp(0,5,1)
  gcd (I# x) (I# y) = I# (gcdInt x y)
#else
  gcd     = P.gcd
#endif
  lcm     = P.lcm
  coprime = coprimeIntegral

instance GcdDomain Word where
#if MIN_VERSION_integer_gmp(1,0,0)
  gcd (W# x) (W# y) = W# (gcdWord x y)
#else
  gcd     = P.gcd
#endif
  lcm     = P.lcm
  coprime = coprimeIntegral

instance GcdDomain Integer where
  gcd     = gcdInteger
  lcm     = lcmInteger
  coprime = coprimeIntegral

#define deriveGcdDomain(ty)     \
instance GcdDomain (ty) where { \
;  gcd     = P.gcd              \
;  lcm     = P.lcm              \
;  coprime = coprimeIntegral    \
}

deriveGcdDomain(Int8)
deriveGcdDomain(Int16)
deriveGcdDomain(Int32)
deriveGcdDomain(Int64)
deriveGcdDomain(Word8)
deriveGcdDomain(Word16)
deriveGcdDomain(Word32)
deriveGcdDomain(Word64)
deriveGcdDomain(Natural)

#define deriveEuclidean(ty)       \
instance Euclidean (ty) where {   \
;  degree  = P.fromIntegral . abs \
;  quotRem = P.quotRem            \
;  quot    = P.quot               \
;  rem     = P.rem                \
}

deriveEuclidean(Int)
deriveEuclidean(Int8)
deriveEuclidean(Int16)
deriveEuclidean(Int32)
deriveEuclidean(Int64)
deriveEuclidean(Integer)
deriveEuclidean(Word)
deriveEuclidean(Word8)
deriveEuclidean(Word16)
deriveEuclidean(Word32)
deriveEuclidean(Word64)
deriveEuclidean(Natural)

-- | Wrapper around 'Fractional'
-- with trivial 'GcdDomain'
-- and 'Euclidean' instances.
newtype WrappedFractional a = WrapFractional { unwrapFractional :: a }
  deriving (Eq, Ord, Show, Num, Fractional)

instance Num a => Semiring (WrappedFractional a) where
  plus  = (P.+)
  zero  = 0
  times = (P.*)
  one   = 1
  fromNatural = P.fromIntegral

instance Num a => Ring (WrappedFractional a) where
  negate = P.negate

instance Fractional a => GcdDomain (WrappedFractional a) where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Fractional a => Euclidean (WrappedFractional a) where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Fractional a => Field (WrappedFractional a)

instance Integral a => GcdDomain (Ratio a) where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Integral a => Euclidean (Ratio a) where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Integral a => Field (Ratio a)

instance GcdDomain Float where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean Float where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field Float

instance GcdDomain Double where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean Double where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field Double

instance GcdDomain CFloat where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean CFloat where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field CFloat

instance GcdDomain CDouble where
  divide x y = Just (x / y)
  gcd        = const $ const 1
  lcm        = const $ const 1
  coprime    = const $ const True

instance Euclidean CDouble where
  degree      = const 0
  quotRem x y = (x / y, 0)
  quot        = (/)
  rem         = const $ const 0

instance Field CDouble

conjQuotAbs :: Field a => Complex a -> Complex a
conjQuotAbs (x :+ y) = x `quot` norm :+ (negate y) `quot` norm
  where
    norm = (x `times` x) `plus` (y `times` y)

instance Field a => GcdDomain (Complex a) where
  divide x y = Just (x `times` conjQuotAbs y)
  gcd        = const $ const one
  lcm        = const $ const one
  coprime    = const $ const True

instance Field a => Euclidean (Complex a) where
  degree      = const 0
  quotRem x y = (quot x y, zero)
  quot x y    = x `times` conjQuotAbs y
  rem         = const $ const zero

instance Field a => Field (Complex a)