-- | Plane partitions. See eg. <http://en.wikipedia.org/wiki/Plane_partition>
--
-- Plane partitions are encoded as lists of lists of Z heights. For example the plane 
-- partition in the picture
-- 
-- <<svg/plane_partition.svg>>
--
-- is encoded as
--
-- > PlanePart [ [5,4,3,3,1]
-- >           , [4,4,2,1]
-- >           , [3,2]
-- >           , [2,1]
-- >           , [1]
-- >           , [1]
-- >           ]
-- 
{-# LANGUAGE BangPatterns #-}
module Math.Combinat.Partitions.Plane where

--------------------------------------------------------------------------------

import Data.List
import Data.Array

import Math.Combinat.Classes
import Math.Combinat.Partitions
import Math.Combinat.Tableaux as Tableaux
import Math.Combinat.Helper

--------------------------------------------------------------------------------
-- * the type of plane partitions

-- | A plane partition encoded as a tablaeu (the \"Z\" heights are the numbers)
newtype PlanePart = PlanePart [[Int]] deriving (PlanePart -> PlanePart -> Bool
(PlanePart -> PlanePart -> Bool)
-> (PlanePart -> PlanePart -> Bool) -> Eq PlanePart
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: PlanePart -> PlanePart -> Bool
$c/= :: PlanePart -> PlanePart -> Bool
== :: PlanePart -> PlanePart -> Bool
$c== :: PlanePart -> PlanePart -> Bool
Eq,Eq PlanePart
Eq PlanePart
-> (PlanePart -> PlanePart -> Ordering)
-> (PlanePart -> PlanePart -> Bool)
-> (PlanePart -> PlanePart -> Bool)
-> (PlanePart -> PlanePart -> Bool)
-> (PlanePart -> PlanePart -> Bool)
-> (PlanePart -> PlanePart -> PlanePart)
-> (PlanePart -> PlanePart -> PlanePart)
-> Ord PlanePart
PlanePart -> PlanePart -> Bool
PlanePart -> PlanePart -> Ordering
PlanePart -> PlanePart -> PlanePart
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: PlanePart -> PlanePart -> PlanePart
$cmin :: PlanePart -> PlanePart -> PlanePart
max :: PlanePart -> PlanePart -> PlanePart
$cmax :: PlanePart -> PlanePart -> PlanePart
>= :: PlanePart -> PlanePart -> Bool
$c>= :: PlanePart -> PlanePart -> Bool
> :: PlanePart -> PlanePart -> Bool
$c> :: PlanePart -> PlanePart -> Bool
<= :: PlanePart -> PlanePart -> Bool
$c<= :: PlanePart -> PlanePart -> Bool
< :: PlanePart -> PlanePart -> Bool
$c< :: PlanePart -> PlanePart -> Bool
compare :: PlanePart -> PlanePart -> Ordering
$ccompare :: PlanePart -> PlanePart -> Ordering
$cp1Ord :: Eq PlanePart
Ord,Int -> PlanePart -> ShowS
[PlanePart] -> ShowS
PlanePart -> String
(Int -> PlanePart -> ShowS)
-> (PlanePart -> String)
-> ([PlanePart] -> ShowS)
-> Show PlanePart
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PlanePart] -> ShowS
$cshowList :: [PlanePart] -> ShowS
show :: PlanePart -> String
$cshow :: PlanePart -> String
showsPrec :: Int -> PlanePart -> ShowS
$cshowsPrec :: Int -> PlanePart -> ShowS
Show)

fromPlanePart :: PlanePart -> [[Int]]
fromPlanePart :: PlanePart -> [[Int]]
fromPlanePart (PlanePart [[Int]]
xs) = [[Int]]
xs

isValidPlanePart :: [[Int]] -> Bool
isValidPlanePart :: [[Int]] -> Bool
isValidPlanePart [[Int]]
pps = 
  [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [ Array (Int, Int) Int
tableArray (Int, Int) Int -> (Int, Int) -> Int
forall i e. Ix i => Array i e -> i -> e
!(Int
i,Int
j) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Array (Int, Int) Int
tableArray (Int, Int) Int -> (Int, Int) -> Int
forall i e. Ix i => Array i e -> i -> e
!(Int
i  ,Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Bool -> Bool -> Bool
&&
        Array (Int, Int) Int
tableArray (Int, Int) Int -> (Int, Int) -> Int
forall i e. Ix i => Array i e -> i -> e
!(Int
i,Int
j) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Array (Int, Int) Int
tableArray (Int, Int) Int -> (Int, Int) -> Int
forall i e. Ix i => Array i e -> i -> e
!(Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1,Int
j  )
      | Int
i<-[Int
0..Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] , Int
j<-[Int
0..Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] 
      ]
  where
    table :: Array (Int,Int) Int
    table :: Array (Int, Int) Int
table = (Int -> Int -> Int)
-> Int
-> ((Int, Int), (Int, Int))
-> [((Int, Int), Int)]
-> Array (Int, Int) Int
forall i e a.
Ix i =>
(e -> a -> e) -> e -> (i, i) -> [(i, a)] -> Array i e
accumArray Int -> Int -> Int
forall a b. a -> b -> a
const Int
0 ((Int
0,Int
0),(Int
y,Int
x)) [ ((Int
i,Int
j),Int
k) | (Int
i,[Int]
ps) <- [Int] -> [[Int]] -> [(Int, [Int])]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [[Int]]
pps , (Int
j,Int
k) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0..] [Int]
ps ]
    y :: Int
y = [[Int]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Int]]
pps
    x :: Int
x = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum (([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [[Int]]
pps)

instance CanBeEmpty PlanePart where
  isEmpty :: PlanePart -> Bool
isEmpty = [[Int]] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([[Int]] -> Bool) -> (PlanePart -> [[Int]]) -> PlanePart -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PlanePart -> [[Int]]
fromPlanePart
  empty :: PlanePart
empty   = [[Int]] -> PlanePart
PlanePart []

-- | Throws an exception if the input is not a plane partition
toPlanePart :: [[Int]] -> PlanePart
toPlanePart :: [[Int]] -> PlanePart
toPlanePart [[Int]]
pps = if [[Int]] -> Bool
isValidPlanePart [[Int]]
pps
  then [[Int]] -> PlanePart
PlanePart ([[Int]] -> PlanePart) -> [[Int]] -> PlanePart
forall a b. (a -> b) -> a -> b
$ ([Int] -> Bool) -> [[Int]] -> [[Int]]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> ([Int] -> Bool) -> [Int] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null) ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ ([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map ((Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
0)) ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> a -> b
$ [[Int]]
pps
  else String -> PlanePart
forall a. HasCallStack => String -> a
error String
"toPlanePart: not a plane partition"

-- | The XY projected shape of a plane partition, as an integer partition
planePartShape :: PlanePart -> Partition
planePartShape :: PlanePart -> Partition
planePartShape = [[Int]] -> Partition
forall a. Tableau a -> Partition
Tableaux.tableauShape ([[Int]] -> Partition)
-> (PlanePart -> [[Int]]) -> PlanePart -> Partition
forall b c a. (b -> c) -> (a -> b) -> a -> c
. PlanePart -> [[Int]]
fromPlanePart

-- | The Z height of a plane partition
planePartZHeight :: PlanePart -> Int
planePartZHeight :: PlanePart -> Int
planePartZHeight (PlanePart [[Int]]
xs) = 
  case [[Int]]
xs of
    ((Int
h:[Int]
_):[[Int]]
_) -> Int
h
    [[Int]]
_         -> Int
0

planePartWeight :: PlanePart -> Int
planePartWeight :: PlanePart -> Int
planePartWeight (PlanePart [[Int]]
xs) = [Int] -> Int
forall a. Num a => [a] -> a
sum' (([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall a. Num a => [a] -> a
sum' [[Int]]
xs)

instance HasWeight PlanePart where
  weight :: PlanePart -> Int
weight = PlanePart -> Int
planePartWeight

--------------------------------------------------------------------------------
-- * constructing plane partitions

singleLayer :: Partition -> PlanePart
singleLayer :: Partition -> PlanePart
singleLayer = [[Int]] -> PlanePart
PlanePart ([[Int]] -> PlanePart)
-> (Partition -> [[Int]]) -> Partition -> PlanePart
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> [Int]) -> [Int] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
k -> Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
k Int
1) ([Int] -> [[Int]]) -> (Partition -> [Int]) -> Partition -> [[Int]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Partition -> [Int]
fromPartition 

-- |  Stacks layers of partitions into a plane partition.
-- Throws an exception if they do not form a plane partition.
stackLayers :: [Partition] -> PlanePart
stackLayers :: [Partition] -> PlanePart
stackLayers [Partition]
layers = if [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [ Partition -> Partition -> Bool
isSubPartitionOf Partition
p Partition
q | (Partition
q,Partition
p) <- [Partition] -> [(Partition, Partition)]
forall a. [a] -> [(a, a)]
pairs [Partition]
layers ]
  then [Partition] -> PlanePart
unsafeStackLayers [Partition]
layers
  else String -> PlanePart
forall a. HasCallStack => String -> a
error String
"stackLayers: the layers do not form a plane partition"

-- | Stacks layers of partitions into a plane partition.
-- This is unsafe in the sense that we don't check that the partitions fit on the top of each other.
unsafeStackLayers :: [Partition] -> PlanePart
unsafeStackLayers :: [Partition] -> PlanePart
unsafeStackLayers []            = [[Int]] -> PlanePart
PlanePart []
unsafeStackLayers (Partition
bottom:[Partition]
rest) = [[Int]] -> PlanePart
PlanePart ([[Int]] -> PlanePart) -> [[Int]] -> PlanePart
forall a b. (a -> b) -> a -> b
$ ([[Int]] -> Partition -> [[Int]])
-> [[Int]] -> [Partition] -> [[Int]]
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl [[Int]] -> Partition -> [[Int]]
addLayer (PlanePart -> [[Int]]
fromPlanePart (PlanePart -> [[Int]]) -> PlanePart -> [[Int]]
forall a b. (a -> b) -> a -> b
$ Partition -> PlanePart
singleLayer Partition
bottom) [Partition]
rest where
  addLayer :: [[Int]] -> Partition -> [[Int]]
  addLayer :: [[Int]] -> Partition -> [[Int]]
addLayer [[Int]]
xxs (Partition [Int]
ps) = [ (Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(+) [Int]
xs (Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate Int
p Int
1 [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int]
forall a. a -> [a]
repeat Int
0) | ([Int]
xs,Int
p) <- [[Int]] -> [Int] -> [([Int], Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [[Int]]
xxs ([Int]
ps [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int]
forall a. a -> [a]
repeat Int
0) ] 

-- | The \"layers\" of a plane partition (in direction @Z@). We should have
--
-- > unsafeStackLayers (planePartLayers pp) == pp
-- 
planePartLayers :: PlanePart -> [Partition]
planePartLayers :: PlanePart -> [Partition]
planePartLayers pp :: PlanePart
pp@(PlanePart [[Int]]
xs) = [ Int -> Partition
layer Int
h | Int
h<-[Int
1..PlanePart -> Int
planePartZHeight PlanePart
pp] ] where
  layer :: Int -> Partition
layer Int
h = [Int] -> Partition
Partition ([Int] -> Partition) -> [Int] -> Partition
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>Int
0) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ ([Int] -> Int) -> [[Int]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [Int] -> Int
forall a. Num a => [a] -> a
sum' ([[Int]] -> [Int]) -> [[Int]] -> [Int]
forall a b. (a -> b) -> a -> b
$ (([Int] -> [Int]) -> [[Int]] -> [[Int]]
forall a b. (a -> b) -> [a] -> [b]
map (([Int] -> [Int]) -> [[Int]] -> [[Int]])
-> ((Int -> Int) -> [Int] -> [Int])
-> (Int -> Int)
-> [[Int]]
-> [[Int]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map) (Int -> Int -> Int
forall a p. (Ord a, Num p) => a -> a -> p
f Int
h) [[Int]]
xs
  f :: a -> a -> p
f a
h = \a
k -> if a
ka -> a -> Bool
forall a. Ord a => a -> a -> Bool
>=a
h then p
1 else p
0

--------------------------------------------------------------------------------
-- * generating plane partitions

-- | Plane partitions of a given weight
planePartitions :: Int -> [PlanePart]
planePartitions :: Int -> [PlanePart]
planePartitions Int
d 
  | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
0     = []
  | Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0     = [[[Int]] -> PlanePart
PlanePart []]
  | Bool
otherwise  = [[PlanePart]] -> [PlanePart]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ Int -> [Partition] -> [PlanePart]
go (Int
dInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
n) [Partition
p] | Int
n<-[Int
1..Int
d] , Partition
p<-Int -> [Partition]
partitions Int
n ]
  where
    go :: Int -> [Partition] -> [PlanePart]
    go :: Int -> [Partition] -> [PlanePart]
go  Int
0   [Partition]
acc       = [[Partition] -> PlanePart
unsafeStackLayers ([Partition] -> [Partition]
forall a. [a] -> [a]
reverse [Partition]
acc)]
    go !Int
rem acc :: [Partition]
acc@(Partition
h:[Partition]
_) = [[PlanePart]] -> [PlanePart]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ Int -> [Partition] -> [PlanePart]
go (Int
remInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
k) (Partition
thisPartition -> [Partition] -> [Partition]
forall a. a -> [a] -> [a]
:[Partition]
acc) | Int
k<-[Int
1..Int
rem] , Partition
this <- Int -> Partition -> [Partition]
subPartitions Int
k Partition
h ]

--------------------------------------------------------------------------------