{-# 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, E10, E16, E20, E30, E32, E64
       , S, P
       , fixedRadix, fixedSize
       , (:+), (*.)
       , (:-), (/.)
       ) where

import Control.Applicative
import Control.Arrow
import Control.Monad
import Data.Bits
import Data.Fixed (div', mod', divMod')
import Data.List
import Data.Int
import Data.Maybe
import Data.Ratio
import Data.Typeable
import Data.Word
import Text.Read
import qualified Text.Read.Lex as L

-- | 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

-- Read a signed number. Stolen from GHC.Read.
readNumber :: Num a => (L.Lexeme -> ReadPrec a) -> ReadPrec a
readNumber convert =
  parens
  ( do x <- lexP
       case x of
         L.Symbol "-" -> do y <- lexP
                            n <- convert y
                            return (negate n)

         _   -> convert x
  )

-- Stolden from GHC.Read.
convertFrac :: Fractional a => L.Lexeme -> ReadPrec a
convertFrac (L.Int i) = return (fromInteger i)
convertFrac (L.Rat r) = return (fromRational r)
convertFrac _         = pfail

instance (HasResolution r, Bits a, Integral a) => Read (Fixed r a) where
  readPrec = readNumber convertFrac

instance (HasResolution r, Bits a, Integral a) => Show (Fixed r a) where
  show = reverse . uncurry (++) . second ('.':) .
         join (***) (\str -> if null str then "0" else str) . uncurry splitAt .
         second (reverse . show . numerator) . fromJust .
         find ((==1) . denominator . snd) . iterate (succ *** (*10)) . (,) 0 .
         toRational

instance (HasResolution r, Bits a, Integral 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, Integral a) => Real (Fixed r a) where
  {-# INLINABLE toRational #-}
  toRational x = toRational (unFixed x) / toRational (2 ^ resolution x :: Integer)

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

-- TODO Don't assume it's a binary float so that we can actually
-- expose this function.
toRealFloat :: (HasResolution r, Integral a, RealFloat b) => Fixed r a -> b
{-# INLINABLE toRealFloat #-}
toRealFloat = liftA2 encodeFloat (fromIntegral . unFixed) (negate . resolution)
{-# SPECIALIZE toRealFloat :: (HasResolution r, Integral a) => Fixed r a -> Float #-}
{-# SPECIALIZE toRealFloat :: (HasResolution r, Integral a) => Fixed r a -> Double #-}

{-# RULES
"realToFrac/Float" forall (x :: (HasResolution r, Integral a) => Fixed r a). realToFrac x = toRealFloat x :: Float
"realToFrac/Double" forall (x :: (HasResolution r, Integral a) => Fixed r a). realToFrac x = toRealFloat x :: Double
  #-}

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 E10 = S (S E8)
type E16 = E8 :+ E8
type E20 = E10 :+ E10
type E30 = E20 :+ E10
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