-- Copyright (c) 2011, David Amos. All rights reserved.

{-# LANGUAGE NoMonomorphismRestriction #-}


module Math.Combinatorics.Poset where

import Math.Common.ListSet as LS -- set operations on strictly ascending lists
import Math.Core.Utils -- for set/multiset operations on ordered lists
import Math.Core.Field
import Math.Combinatorics.FiniteGeometry
import Math.Algebra.LinearAlgebra

import Math.Combinatorics.Digraph

import Data.List as L
import qualified Data.Map as M
-- import qualified Data.Set as S


-- |A poset is represented as a pair (set,po), where set is the underlying set of the poset, and po is the partial order relation
newtype Poset t = Poset ([t], t -> t -> Bool)

instance Eq t => Eq (Poset t) where
    Poset (set,po) == Poset (set',po') =
        set == set' && and [po x y == po' x y | x <- set, y <- set]
-- There may be ways to avoid comparing every pair
-- If we could calculate the coverGraph without comparing every pair,
-- then it would be sufficient to test whether their cover graphs are equal

instance Show t => Show (Poset t) where
    show (Poset (set,po)) = "Poset " ++ show set

implies p q = q || not p

isReflexive (set,po) = and [x `po` x | x <- set]
isAntisymmetric (set,po) = and [((x `po` y) && (y `po` x)) `implies` (x == y) | x <- set, y <- set]
isTransitive (set,po) = and [((x `po` y) && (y `po` z)) `implies` (x `po` z) | x <- set, y <- set, z <- set]

isPoset poset = isReflexive poset && isAntisymmetric poset && isTransitive poset

poset (set,po)
    | isPoset (set,po) = Poset (set,po)
    | otherwise = error "poset: Not a partial order"


-- Most of the posets we will deal with are in fact lattices, meaning that any two elements
-- have a meet (greatest lower bound) and join (least upper bound)

intervals (Poset (set,po)) = [(a,b) | a <- set, b <- set, a `po` b]

interval (Poset (set,po)) (x,z) = [y | y <- set, x `po` y, y `po` z]


-- LINEAR ORDER POSET
-- This is of course a lattice, with meet = min, join = max

-- |A chain is a poset in which every pair of elements is comparable (ie either x <= y or y <= x).
-- It is therefore a linear or total order.
-- chainN n is the poset consisting of the numbers [1..n] ordered by (<=)
chainN :: Int -> Poset Int
chainN n = Poset ( [1..n], (<=) )

-- hasseN n = DG [1..n] [(i,i+1) | i <- [1..n-1]]


-- |An antichain is a poset in which distinct elements are incomparable.
-- antichainN n is the poset consisting of [1..n], with x <= y only when x == y.
antichainN :: Int -> Poset Int
antichainN n = Poset ( [1..n], (==) )


-- LATTICE OF (POSITIVE) DIVISORS OF N

divides a b = b `rem` a == 0

divisors n = toSet [ d' | d <- takeWhile (\d -> d*d <= n) [1..],
                          let (q,r) = n `quotRem` d, r == 0,
                          d' <- [d,q] ]
-- The toSet call sorts, and deduplicates if n is a square
-- divisors n | n >= 1 = [a | a <- [1..n], a `divides` n]

-- |posetD n is the lattice of (positive) divisors of n
posetD :: Int -> Poset Int
posetD n | n >= 1 = Poset ( divisors n, divides )


-- LATTICE OF SUBSETS OF [1..N] ORDERED BY INCLUSION
-- (Boolean lattice)

powerset [] = [[]]
powerset (x:xs) = let p = powerset xs in p ++ map (x:) p

{-
-- subset test for sorted lists
isSubset (x:xs) (y:ys) =
    case compare x y of
    LT -> False
    EQ -> isSubset xs ys
    GT -> isSubset (x:xs) ys
isSubset [] _ = True
isSubset _ [] = False
-}

-- |posetB n is the lattice of subsets of [1..n] ordered by inclusion
posetB :: Int -> Poset [Int]
posetB n = Poset ( powerset [1..n], LS.isSubset )


-- LATTICE OF SET PARTITIONS OF [1..N] ORDERED BY REFINEMENT

partitions [] = [[]]
partitions [x] = [[[x]]]
partitions (x:xs) = let ps = partitions xs in
    map ([x]:) ps ++ [ (x:cell):p' | p <- ps, (cell,p') <- picks p]
-- if the input is sorted, then so is the output

isRefinement a b = and [or [acell `isSubset` bcell | bcell <- b] | acell <- a]
-- if we know that a and b are appropriately sorted, then this can probably be done more efficiently

-- |posetP n is the lattice of set partitions of [1..n], ordered by refinement
posetP :: Int -> Poset [[Int]]
posetP n = Poset ( partitions [1..n], isRefinement )

-- muP n = ...
-- see van Lint and Wilson p336


-- LATTICE OF INTERVAL PARTITIONS OF [1..N] ORDERED BY REFINEMENT
-- Interval partitions of [1..n] correspond to compositions of n

intervalPartitions xs = filter (all isInterval) (partitions xs)

isInterval (x1:x2:xs) = x1+1 == x2 && isInterval (x2:xs)
isInterval _ = True

intervalPartitions2 [] = [[]]
intervalPartitions2 [x] = [[[x]]]
intervalPartitions2 (x:xs) = let ips = intervalPartitions2 xs in
    map ([x]:) ips ++ [ (x:head):tail | (head:tail) <- ips]
-- we're guaranteed that x+1 is at the head of the head


-- LATTICE OF INTEGER PARTITIONS OF N ORDERED BY REFINEMENT

-- For example, integerPartitions 5 -> [ [5], [4,1], [3,2], [3,1,1], [2,2,1], [2,1,1,1], [1,1,1,1,1] ]
integerPartitions n | n >= 0 = ips n n where
    ips 0 _ = [[]]
    ips _ 0 = []
    ips n m | m <= n = map (m:) (ips (n-m) m) ++ ips n (m-1)
            | otherwise = ips n n

isIPRefinement ys xs = dfs xs ys
    where dfs (x:xs) (y:ys) | x < y = False
                            | x == y = dfs xs ys
                            | otherwise = or [dfs xs' ys' | (y', ys') <- picks (y:ys),
                                                            let xs' = insertDesc (x-y') xs]
          dfs [] [] = True
          insertDesc = L.insertBy (flip compare)

{-
-- In theory it feels like this ought to be faster for large n, but in practice it's unclear
isIPRefinement2 ys xs = isIPRefinement (ys \\ xs) (xs \\ ys)
    where (\\) = diffDesc
-}

-- |posetIP n is the poset of integer partitions of n, ordered by refinement
posetIP :: Int -> Poset [Int]
posetIP n = Poset (integerPartitions n, isIPRefinement)


-- Could also implement the Young lattice (or the part up to n)
-- Of integer partitions <= n, ordered by inclusion
-- Kassel, Turaev, Braid Groups, p202


-- INTEGER COMPOSITIONS


-- LATTICE OF SUBSPACES OF Fq^n

subspaces fq n = [] : concatMap (flatsPG (n-1) fq) [0..n-1]
-- note that flatsPG returns the subspaces as a matrix of row vectors in reduced row echelon form

-- inSpanRE m v returns whether the vector v is in the span of the rows of the matrix m, where m is required to be in row echelon form
isSubspace s1 s2 = all (inSpanRE s2) s1

-- This is the projective geometry PG(n,q)
-- |posetL n fq is the lattice of subspaces of the vector space Fq^n, ordered by inclusion.
-- Subspaces are represented by their reduced row echelon form.
-- Example usage: posetL 2 f3
posetL :: (Eq fq, Num fq) => Int -> [fq] -> Poset [[fq]]
posetL n fq = Poset ( subspaces fq n, isSubspace )


-- choose n k = product [n-k+1..n] `div` product [1..k]


-- NEW FROM OLD CONSTRUCTIONS


-- |The subposet of a poset satisfying a predicate
subposet :: Poset a -> (a -> Bool) -> Poset a
subposet (Poset (set,po)) p = Poset (filter p set, po)

-- |The direct sum of two posets
dsum :: Poset a -> Poset b -> Poset (Either a b)
dsum (Poset (setA,poA)) (Poset (setB,poB)) = Poset (set,po)
    where set = map Left setA ++ map Right setB
          po (Left a1) (Left a2) = poA a1 a2
          po (Right b1) (Right b2) = poB b1 b2
          po _ _ = False

-- |The direct product of two posets
dprod :: Poset a -> Poset b -> Poset (a,b)
dprod (Poset (setA,poA)) (Poset (setB,poB)) =
    Poset ( [(a,b) | a <- setA, b <- setB], \(a1,b1) (a2,b2) -> (a1 `poA` a2) && (b1 `poB` b2) )

-- |The dual of a poset
dual :: Poset a -> Poset a
dual (Poset (set, po)) = Poset (set, po')
    where po' x y = po y x


-- ANALYSIS OF POSETS

-- |Given a poset (X,<=), we say that y covers x, written x -< y, if x < y and there is no z in X with x < z < y.
-- The Hasse digraph of a poset is the digraph whose vertices are the elements of the poset,
-- with an edge between every pair (x,y) with x -< y.
-- The Hasse digraph can be represented diagrammatically as a Hasse diagram, by drawing x below y whenever x -< y.
hasseDigraph :: (Eq a) => Poset a -> Digraph a
hasseDigraph (Poset (set,po)) = DG set [(x,y) | x <- set, y <- set, x -< y]
    where x -< y = x /= y && x `po` y && null [z | z <- set, x `po` z, x /= z, z `po` y, z /= y]
-- The partial order can be recovered as the transitive closure of the covers relation
-- !! Can we construct the cover graph without having to compare every pair ??

-- If we know in advance the poset we're interested in,
-- then we're probably better off constructing the Hasse digraph directly

-- (In effect, the transitive closure of the edge relation)
-- |Given a DAG (directed acyclic graph), return the poset consisting of the vertices of the DAG, ordered by reachability.
-- This can be used to recover a poset from its Hasse digraph.
reachabilityPoset :: (Ord a) => Digraph a -> Poset a
reachabilityPoset (DG vs es) = Poset (vs,tc') -- \u v -> tc M.! (u,v)
    where tc = M.fromList [ ((u,v), tc' u v) | u <- vs, v <- vs]
          tc' u v | u == v = True
                  | otherwise = or [tc M.! (w,v) | w <- successors u]
          successors u = [v | (u',v) <- es, u' == u]
-- !! looks like we could memoise more than we are doing


{-
-- For example:
> let poset = posetB 3 in poset == reachabilityPoset (hasseDigraph poset)
True
-}


-- This definition is incorrect. This is an order embedding. Order preserving only requires the rightward implication.
isOrderPreserving :: (a -> b) -> Poset a -> Poset b -> Bool
isOrderPreserving f (Poset (seta,poa)) (Poset (setb,pob)) =
    and [ x `poa` y == f x `pob` f y | x <- seta, y <- seta ]

-- Find all order isomorphisms between two posets
-- This is the most naive algorithm, and should not be used on larger posets
-- For example, already on the following, this takes forever compared to almost instant for orderIsos:
-- > head $ orderIsos01 (posetD $ 8*9*25*49) (posetD $ 4*27*25*121)
orderIsos01 (Poset (seta,poa)) (Poset (setb,pob))
    | length seta /= length setb = []
    | otherwise = orderIsos' [] seta setb
    where orderIsos' xys [] [] = [xys]
          orderIsos' xys (x:xs) ys =
              concat [ orderIsos' ((x,y):xys) xs ys'
                     | (y,ys') <- picks ys,
                       and [ (x `poa` x', x' `poa` x) == (y `pob` y', y' `pob` y) | (x',y') <- xys ] ]

-- |Are the two posets order-isomorphic?
isOrderIso :: (Ord a, Ord b) => Poset a -> Poset b -> Bool
isOrderIso poseta posetb = (not . null) (orderIsos poseta posetb)

-- For small posets, it may be that the up-front cost of calculating the hasseDigraph and heightPartitionDAG
-- are not justified.
-- |Find all order isomorphisms between two posets
orderIsos :: (Ord a, Ord b) => Poset a -> Poset b -> [[(a,b)]]
orderIsos posetA@(Poset (_,poa)) posetB@(Poset (_,pob))
    | map length heightPartA /= map length heightPartB = []
    | otherwise = dfs [] heightPartA heightPartB
    where heightPartA = heightPartitionDAG (hasseDigraph posetA)
          heightPartB = heightPartitionDAG (hasseDigraph posetB)
          dfs xys [] [] = [xys]
          dfs xys ([]:las) ([]:lbs) = dfs xys las lbs
          dfs xys ((x:xs):las) (ys:lbs) =
              concat [ dfs ((x,y):xys) (xs:las) (ys' : lbs)
                     | (y,ys') <- picks ys,
                       and [ (x `poa` x', x' `poa` x) == (y `pob` y', y' `pob` y) | (x',y') <- xys ] ]
-- A variant on this algorithm would use the Hasse digraph rather than the partial order in the test on the last line
-- This might be faster, depending how expensive the partial order comparison function is
-- In effect though, it would then be a DAG isomorphism function

-- The order automorphisms of a poset
orderAuts1 poset = orderIsos poset poset
-- This returns all automorphisms
-- What we really want is to return generators of the permutation group


-- |A linear extension of a poset is a linear ordering of the elements which extends the partial order.
-- Equivalently, it is an ordering [x1..xn] of the underlying set, such that if xi <= xj then i <= j.
isLinext (Poset (set,po)) set' = all (\(x,y) -> not (y `po` x)) (pairs set')


-- |Linear extensions of a poset
linexts (Poset (set,po)) = linexts' [[]] set
    where linexts' lss (r:rs) =
              let lss' = [ lts ++ [r] ++ gts
                         | ls <- lss,
                           let ls' = takeWhile (not . (r `po`)) ls,
                           (lts,gts) <- zip (inits ls') (tails ls),
                           all (not . (`po` r)) gts ]
              in linexts' lss' rs
          linexts' lss [] = lss