import Data.Searchable
import Control.Applicative
import Data.Ratio

data Nat = Z | S Nat
    deriving (Ord, Show, Eq)

instance Num Nat where
    -- Functions defined in Kohlenbach p 45.  These aren't the exact Kohlenbach definitions
    -- because the operations are *much* faster using this constructor order than using
    -- Kohlenbach's.
    Z + y = y
    S x + y = S (x+y)
    Z * y = Z
    S x * y = y + (x*y)
    x - Z = x
    x - S y = pd (x-y)

    -- Additional items required by Num instance
    abs x = x
    signum Z = Z
    signum (S x) = S Z
    fromInteger 0 = Z
    fromInteger n = S . fromInteger $ n - 1

-- Additional functions defined in Kohlenbach p 45
sg :: Nat -> Nat
sg Z = 1
sg (S x) = 0

pd Z = Z
pd (S x) = x

-- least function, taken from Impossible Functional Programs.  The (1 + n) is
-- much faster than (n + 1) based on the way the S value constructor in + is ordered.
least p = if p 0 then 0 else 1 + least(\n -> p (1 + n))

instance Fractional Nat where
    -- Use truncation for division on Nats.  Use S value constructor for Div 0 protection.
    (/) x (S y) = least(\z -> (S z)*(S y) > x)
    fromRational r = (fromInteger . floor) r

-- Infinity declaration for testing totalness of predicates
inf = S inf

-- Declaration of natural numbers for infinite search monad
natural :: Set Nat
natural = union (singleton Z) (S <$> natural)

-- A type to distinguish between true Nats and those used to encode a rational
newtype QCode = QCode {unQCode :: Nat}
    deriving (Show)

-- Cantor pairing function (p 58)
j :: Nat -> Nat -> QCode
j x y = if k <= t then QCode k else QCode 0
        where   k = least (\u -> 2*u == t)
                t = (x+y)*(x+y) + 3*x + y

-- Projection functions j1 and j2 for converting QCode back into first or second Nat argument (p 58)
j1 :: QCode -> Nat
j1 zq = least (\x -> x <= z && forsome natural (\y -> y <= z && unQCode(j x y) == z))
        where z = unQCode zq

j2 :: QCode -> Nat
j2 zq = least (\y -> y <= z && forsome natural (\x -> x <= z && unQCode(j x y) == z))
        where z = unQCode zq

-- Even/Odd definitions for Nat
evenNat :: Nat -> Bool
evenNat n = (n/2)*2 == n
oddNat = not . evenNat

-- Function for comparing two QCodes for equivalence (1/2 and 2/4 are equivalent rationals but have
-- different QCodes)
instance Eq QCode where
    (==) n1 n2 | evenNat j1n1 && evenNat j1n2 = (j1n1/2)*(j2n2 + 1) == (j1n2/2)*(j2n1 + 1)
               | oddNat  j1n1 && oddNat  j1n2 = ((j1n1+1)/2)*(j2n2 + 1) == ((j1n2+1)/2)*(j2n1 + 1)
               | otherwise                    = False
               where j1n1 = j1 n1
                     j1n2 = j1 n2
                     j2n1 = j2 n1
                     j2n2 = j2 n2

-- Utility function for constructing a QCode from a numerator and denominator
numden2QCode :: Integer -> Integer -> QCode
numden2QCode num den  | (num >= 0 && den > 0) || (num <= 0 && den < 0)  = j (fromInteger $ 2*abs(num)) (fromInteger $ abs(den)-1)
                      | otherwise                                       = j (fromInteger $ 2*abs(num)-1) (fromInteger $ abs(den)-1)

-- Test out some QCode equivalence predicates
main = do
        print $ numden2QCode (-1) 1 == numden2QCode (-2) 2 -- True
        print $ numden2QCode (-1) 1 == numden2QCode (-3) 3 -- True
        print $ numden2QCode (-1) 2 == numden2QCode 2 (-4) -- True
        print $ numden2QCode 1 2    == numden2QCode 2 4    -- True
        print $ numden2QCode 1 2    == numden2QCode 1 3    -- False