{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE EmptyDataDecls #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

{-|
This module defines a type for binary fixed-point arithmetic. The main
advantage this provides over decimal fixed-point arithmetic is that
the point is maintained using fast bit shifts instead of slow 'div'
operations. This type is also polymorphic on the underlying
representation, so you can use whatever size and signedness you
want. You just have to be mindful of overflows if you use a fixed-size
representation, especially with operations like multiplication.
-}
module Data.Fixed.Binary
       ( div'
       , mod'
       , divMod'
       , Fixed ()
       , HasResolution (..)
       , E0, E1, E2, E4, E8, E16, E32, E64
       , S, P
       , fixedRadix, fixedSize
       , (:+), (*.)
       , (:-), (/.)
       ) where

import Data.Bits
import Data.Fixed (div', mod', divMod')
import Data.Int
import Data.Ratio
import Data.Typeable
import Data.Word
import Text.Read

-- | The first type parameter represents the number of bits to devote
-- to the fractional part of the number. The second type parameter is
-- the underlying representation. For example, @Fixed E8 Int16@ uses
-- eight bits for the integer component (of which one bit is used for
-- the sign) and eight bits for the fractional component.
newtype Fixed r a = Fixed { unFixed :: a }
                  deriving (Enum, Eq, Ord, Typeable)

inFixed :: (a -> b) -> (Fixed r a -> Fixed s b)
{-# INLINE inFixed #-}
inFixed = (Fixed .) . (. unFixed)

inFixed2 :: (a -> b -> c) -> (Fixed r a -> Fixed s b -> Fixed t c)
{-# INLINE inFixed2 #-}
inFixed2 = (inFixed .) . (. unFixed)

-- | Instances of this class are useful as the first parameter of
-- 'Fixed'.
class HasResolution r where
  -- | Given a fixed-point number, give the number of bits used to
  -- represent its fractional part.
  resolution :: Num a => Fixed r a -> Int

withResolution :: (HasResolution r, Num a) => (Int -> Fixed r a) -> Fixed r a
{-# INLINE withResolution #-}
withResolution f = withType (f . resolution)
  where withType :: (Fixed r a -> Fixed r a) -> Fixed r a
        withType g = g undefined

instance (HasResolution r, Bits a, Integral a) => Read (Fixed r a) where
  readPrec = fmap fromRational readPrec

instance (HasResolution r, Bits a, Real a) => Show (Fixed r a) where
  show = show . toRational

instance (HasResolution r, Bits a, Real a) => Num (Fixed r a) where
  {-# INLINABLE (+) #-}
  (+) = inFixed2 (+)
  {-# INLINABLE (-) #-}
  (-) = inFixed2 (-)
  {-# INLINABLE (*) #-}
  Fixed x * Fixed y = withResolution $ Fixed . shiftR (x*y)
  {-# INLINABLE negate #-}
  negate = inFixed negate
  {-# INLINABLE abs #-}
  abs = inFixed abs
  {-# INLINABLE signum #-}
  signum (Fixed x) = withResolution $ \s -> Fixed $ signum x `shiftL` s
  {-# INLINABLE fromInteger #-}
  fromInteger i = withResolution $ Fixed . shiftL (fromInteger i)

instance (HasResolution r, Bits a, Real a) => Real (Fixed r a) where
  {-# INLINABLE toRational #-}
  toRational x = toRational (unFixed x) / toRational (2 ^ resolution x :: Int)

instance (HasResolution r, Bits a, Integral a) => Fractional (Fixed r a) where
  {-# INLINABLE (/) #-}
  a / b = Fixed $ (unFixed a `shiftL` resolution a) `div` unFixed b
  {-# INLINABLE recip #-}
  recip x = Fixed $ (1 `shiftL` (2 * resolution x)) `div` unFixed x
  {-# INLINABLE fromRational #-}
  fromRational r = withResolution $ \s ->
    Fixed . floor $ (numerator r `shiftL` s) % denominator r

instance (HasResolution r, Bits a, Integral a) => RealFrac (Fixed r a) where
  {-# INLINABLE properFraction #-}
  properFraction a = let i = truncate a in (i, a - fromIntegral i)
  {-# INLINABLE truncate #-}
  truncate = truncate . toRational
  {-# INLINABLE round #-}
  round = round . toRational
  {-# INLINABLE ceiling #-}
  ceiling = ceiling . toRational
  {-# INLINABLE floor #-}
  floor = floor . toRational

-- | Fast conversion between fixed-point numbers with the same
-- fractional size.
fixedRadix :: (Integral a, Num b) => Fixed r a -> Fixed r b
{-# INLINABLE fixedRadix #-}
fixedRadix = inFixed fromIntegral

-- TODO Can't I write this as one awesome, polymorphic rule?
{-# RULES
"realToFrac/fixedRadixInt"     forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Int
"realToFrac/fixedRadixInt8"    forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Int8
"realToFrac/fixedRadixInt16"   forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Int16
"realToFrac/fixedRadixInt32"   forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Int32
"realToFrac/fixedRadixInt64"   forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Int64
"realToFrac/fixedRadixWord"    forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Word
"realToFrac/fixedRadixWord8"   forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Word8
"realToFrac/fixedRadixWord16"  forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Word16
"realToFrac/fixedRadixWord32"  forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Word32
"realToFrac/fixedRadixWord64"  forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Word64
"realToFrac/fixedRadixInteger" forall (x :: Integral a => Fixed r a). realToFrac x = fixedRadix x :: Fixed r Integer
  #-}

-- | Fast conversion between fixed-point numbers with the same
-- representation size.
fixedSize :: (HasResolution r, HasResolution s, Bits a) => Fixed r a -> Fixed s a
{-# INLINABLE fixedSize #-}
fixedSize x = withResolution $ \s -> Fixed $ unFixed x `shift` (s - resolution x)
-- TODO Rewrite rules?

-- | Multiplication without throwing away fractional information. Note
-- that this doesn't help against losing significant information from
-- the integer component. If you are concerned about preventing
-- overflow then convert to a larger representation first.
(*.) :: Num a => Fixed r a -> Fixed s a -> Fixed (r :+ s) a
{-# INLINABLE (*.) #-}
(*.) = inFixed2 (*)

-- | Division without throwing away fractional information. Same
-- caveats apply as with '(*.)'.
(/.) :: Integral a => Fixed r a -> Fixed s a -> Fixed (r :- s) a
{-# INLINABLE (/.) #-}
(/.) = inFixed2 div

data E0

-- | Increment a resolution
data S n

-- | Add resolutions
type family a :+ b
type instance E0 :+ b = b
type instance S a :+ b = S (a :+ b)

-- | Subtract resolutions
type family a :- b
type instance a :- E0 = a
type instance S a :- S b = a :- b

-- | Decrement a resolution
type family P a
type instance P (S a) = a

type E1  = S E0
type E2  = E1 :+ E1
type E4  = E2 :+ E2
type E8  = E4 :+ E4
type E16 = E8 :+ E8
type E32 = E16 :+ E16
type E64 = E32 :+ E32

instance HasResolution n => HasResolution (S n) where
  {-# INLINE resolution #-}
  resolution = succ . resolution' undefined
    where resolution' :: (HasResolution n, Num a) =>
                         Fixed n a -> Fixed (S n) a -> Int
          {-# INLINE resolution' #-}
          resolution' dummy = const $ resolution dummy

instance HasResolution E0 where
  {-# INLINE resolution #-}
  resolution = const 0