module Lib (module Lib) where

import Data.Ratio ((%), numerator, denominator)
import Data.List (inits)

-- | Two implementations of naive rational enumeration

inits1 :: [a] -> [[a]]
inits1 xs = drop 1 $ inits xs

diags :: [a] -> [b] -> [[(a,b)]]
diags xs ys =
  let xs' = inits1 xs
      ys' = map reverse $ inits1 ys
  in zipWith zip xs' ys'

rats1 :: [Rational]
rats1 =
  let pairs = concat $ diags [1..] [1..]
  in map (\(x, y) -> x % y) pairs

rats2 :: [Rational]
rats2 = [m % (d-m) | d <- [2..], m <- [1..d-1]]

-- | Instrumenting and reversing the GCD.

-- | Instrumented GCD with left/right list (as corresponding booleans False/True).
igcd :: Integer -> Integer -> (Integer, [Bool])
igcd m n
  | m < n = step False $ igcd m (n - m)
  | m > n = step True $ igcd (m - n) n
  | otherwise = (m, [])
  where step b (d, bs) = (d, b:bs)

-- | Only the boolean list of the instrumented GCD.
pgcd :: Integer -> Integer -> [Bool]
pgcd m n = snd $ igcd m n

-- | Are two fractional expressions m/n and m'/n' the same rational?
equiv :: (Integer, Integer) -> (Integer, Integer) -> Bool
equiv (m, n) (m', n') = pgcd m n == pgcd m' n'

ungcd :: (Integer, [Bool]) -> (Integer, Integer)
ungcd (d, bs) = foldr undo (d, d) bs
  where undo False (m, n) = (m, n + m)
        undo True (m, n) = (m + n, n)

-- | Lexicographic list of all boolean lists.
boolseqs :: [[Bool]]
boolseqs = [] : [b:bs | bs <- boolseqs, b <- [False, True]]

-- | Using @ungcdg 1@ means that we only obtain co-prime pairs.
rats3 :: [Rational]
rats3 = map (uncurry (%) . curry ungcd 1) boolseqs

-- | The Stern-Brocot tree
--
-- = Structure
-- A binary tree with two additional vertices: the root @1/1@ has two
-- parents, \(0/1\) and \(1/0\). Every other node is of the form \((m
-- + m')/(n + n')\) where \(m/n\) is the /parent/ and \(m'/n'\) is the
-- /other-direction parent/ of higher order. Those two are known as
-- the /ancestors/ of the node \((m + m')/(n + n')\).
-- == Example
-- The number \(4/3\) is represented by @R, L, L@. Its parent \(3/2\)
-- is @R, L@ and its other-direction parent is \(1/1\), represented by
-- the empty list @[]@. Those are its ancestors.

data BTree a = Node a (BTree a) (BTree a)

-- | Reduce a `BTree`.
foldt :: (a -> b -> b -> b) -> BTree a -> b
foldt f (Node a x y) = f a (foldt f x) (foldt f y)

-- | Produce a `BTree` from a transformation.
--
-- The /node/, /left/ and /right/ notions are provided by @f@.
unfoldt:: (t -> (a, t, t)) -> t -> BTree a
unfoldt f t =
  let (a, x, y) = f t
  in Node a (unfoldt f x) (unfoldt f y)

-- | A breadth-first list of all elements of a tree.
--
-- Without the `concat`, the @foldt glue@ creates a list of lists of
-- all the levels.
bf :: BTree a -> [a]
bf = concat . foldt glue
  where glue a xs ys = [a] : zipWith (++) xs ys

adj :: (Integer, Integer) -> (Integer, Integer) -> (Integer, Integer)
adj (m, n) (m', n') = (m + m', n + n')

-- | Breadth-first list of the Stern-Brocot tree.
rats4 :: [Rational]
rats4 = bf $ unfoldt step ((0, 1), (1, 0))
  -- @step@ answers /what is the node value here/, and then /how to
  -- proceed on the left/ and /how to proceed on the right/. The
  -- "numbers" @l@ and @r@ are the ancestors of @m@.
  where step (l, r) =
          let m = adj l r
          in (uncurry (%) m, (l, m), (m, r))

-- | Deforesting the Stern-Brocot tree.


-- | Interleave two lists together.
interleave :: [a] -> [a] -> [a]
interleave [] ys = ys
interleave (x:xs) ys = x : interleave ys xs

-- | Produce a list from a transformation.
unfolds :: (t -> (a, t)) -> t -> [a]
unfolds f a =
  let (b, a') = f a
  in b : unfolds f a'

-- | `infill` produces the next level of the Stern-Brocot tree.
--
-- Given the current level, it produces its rationals and the next
-- level.
infill :: [(Integer, Integer)] -> ([Rational], [(Integer, Integer)])
infill xs = (map (uncurry (%)) ys, interleave xs ys)
  -- @ys@ is the list obtained by applying `adj` for every adjacent
  -- pair (overlapping), e.g. @adj 1 2, adj 2 3, adj 3 4@, and so on.
  where ys = zipWith adj xs (tail xs)

rats5 :: [Rational]
rats5 = concat $ unfolds infill [(0,1), (1,0)]

-- | Stern-Brocot growing memory issue.
--
-- So far, all computations using the Stern-Brocot tree involve
-- keeping a state of increasing memory size. In an observation by
-- Moshe (2003), the Calkin-Wilf tree avoids this issue and also
-- produces the rationals.

-- | Breadth-first iteration of the Calkin-Wilf tree.
rats6 :: [Rational]
rats6 = bf $ unfoldt step (1,1)
  where step (m, n) = (m % n, (m, m + n), (n + m, n))

-- | The Calkin-Wilf tree structure
--
-- The parents of \(x\) are \((x - 1, 1/(1/x - 1))\). The children are
-- \(1/(1/x + 1), x+1\).

inverse :: Rational -> Rational
inverse z = (denominator z) % (numerator z)

-- | Moshe's observation (AMS monthly 110, 642-643, 2003)
next :: Rational -> Rational
next x =
  let (m, n) = (numerator x, denominator x)
      d = div m n
      r = x - (d % 1)
  in inverse $ (d + 1) % 1 - r

rats7 :: [Rational]
rats7 = iterate next (1 % 1)

-- | All rationals between 0 and 1.
between01 :: [Rational]
between01 = iterate (next . next) (1 % 2)