{-# LANGUAGE FlexibleInstances #-}
import Data.Set (Set,elems,cartesianProduct,fromList,empty,insert,singleton,
                union,unions,difference,filter)

-- Fix a natural number n, and identify it with the ordinal {0 < ... < n}. Let A be the set of binary relations R on n that satisfy (A1) if R(a, b) then a <= b, (A2) R is reflexive, (A3) R is transitive. Let B be the set of binary relations S on n that satisfy (B1) if S(a, b) then a + b <= n, (B2) S is downward-closed. Count the elements of A and B depending on n.

class Below a where
  below :: a -> Set a

instance (Below a, Below b) => Below (a, b) where
  below (a,b) = cartesianProduct (below a) (below b)

type N = Integer

instance Below N where
  below n = fromList [0..n]

data CombinatorialProblem a = CP
  {
    bound :: Set a,
    close :: Set a -> Set a
  }

enumerate :: Ord a => CombinatorialProblem a -> Set (Set a)
enumerate cp = iterate' f null (empty, singleton (close cp empty))
  where
    f (old, new) = (pretty_old, pretty_new)
      where
        pretty_old = union old new
        fresh = [close cp (insert a s) | a <- elems (bound cp), s <- elems new]
        pretty_new = fromList fresh `difference` pretty_old
        
iterate' :: ((a,a) -> (a,a)) -> (a -> Bool) -> (a,a) -> a
iterate' f p x = let (x0,x1) = iterate'' x in x0
  where
    iterate'' x = if p (let (x0,x1) = x in x1) then x else iterate'' (f x)

count :: Ord a => CombinatorialProblem a -> Int
count cp = length (enumerate cp)


-- PROBLEM B --

downwardClosed :: (Below a, Ord a) => a -> CombinatorialProblem a
downwardClosed a = CP (below a) down
  where
    down s = unions [below n | n <- elems s]

downwardClosedTriangle :: N -> CombinatorialProblem (N,N)
downwardClosedTriangle n = CP bound' (close $ downwardClosed (n,n))
  where
    bound' = Data.Set.filter (\(i,j) -> i+j<=n) (bound $ downwardClosed (n,n))

{-
λ> count . downwardClosed . (\a -> (a,a)) <$> [0..6]
[2,6,20,70,252,924,3432]
OEIS A000984, Central binomial coefficients
(n+1 times as many as the one below, but shifted once)

λ> count . downwardClosedTriangle <$> [0..4]
[2,5,14,42,132]
OEIS A000108, Catalan numbers
-}


-- PROBLEM A --

reflexiveTransitive :: (Below a, Ord a) => a -> CombinatorialProblem (a,a)
reflexiveTransitive a = CP (below (a,a)) (fix free)
  where
    free s = unions [refl,s,trans]
      where
        refl = fromList [(i,i) | i <- elems (below a)]
        trans = fromList [(i,k) | (i,j) <- elems s, (j',k) <- elems s, j == j']

fix :: Eq a => (a -> a) -> a -> a
fix f x = let x' = f x in if x == x' then x else fix f x'

reflexiveTransitiveTriangle :: N -> CombinatorialProblem (N,N)
reflexiveTransitiveTriangle n = CP bound' (close $ reflexiveTransitive n)
  where
    bound' = Data.Set.filter (\(i,j) -> i<=j) (bound $ reflexiveTransitive n)

main = return ()

{-
λ> count . reflexiveTransitive <$> [0..4]
[1,4,29,355,6942]
A000798

λ> count . reflexiveTransitiveTriangle <$> [0..6]
[1,2,7,40,357,4824,96428]
OEIS A006455
-}