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

    module Math.Combinatorics.LatinSquares where

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

    -- import Math.Combinatorics.FiniteGeometry
    import Math.Combinatorics.Design
    import Math.Algebra.Field.Base
    import Math.Algebra.Field.Extension
    import Math.Algebra.LinearAlgebra (fMatrix')
    import Math.Combinatorics.Graph
    import Math.Combinatorics.GraphAuts
    import Math.Combinatorics.StronglyRegularGraph
    import Math.Core.Utils (combinationsOf)

    -- LATIN SQUARES

    findLatinSqs :: (Eq a) => [a] -> [[[a]]]
    findLatinSqs :: [a] -> [[[a]]]
    findLatinSqs xs :: [a]
    xs = Int -> [[a]] -> [[[a]]]
    findLatinSqs' 1 [[a]
    xs] where
        n :: Int
    n = [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [a]
    xs
        findLatinSqs' :: Int -> [[a]] -> [[[a]]]
    findLatinSqs' i :: Int
    i rows :: [[a]]
    rows
            | Int
    i Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    n    = [[[a]] -> [[a]]
    forall a. [a] -> [a]
    reverse [[a]]
    rows]
            | Bool
    otherwise = [[[[a]]]] -> [[[a]]]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [Int -> [[a]] -> [[[a]]]
    findLatinSqs' (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    +1) ([a]
    row[a] -> [[a]] -> [[a]]
    forall a. a -> [a] -> [a]
    :[[a]]
    rows)
                                 | [a]
    row <- [[a]] -> [a] -> [a] -> [[a]]
    forall (t :: * -> *) a.
    (Foldable t, Eq a) =>
    [t a] -> [a] -> [a] -> [[a]]
    findRows ([[a]] -> [[a]]
    forall a. [[a]] -> [[a]]
    L.transpose [[a]]
    rows) [] [a]
    xs]
        findRows :: [t a] -> [a] -> [a] -> [[a]]
    findRows (col :: t a
    col:cols :: [t a]
    cols) ls :: [a]
    ls rs :: [a]
    rs = [[[a]]] -> [[a]]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[t a] -> [a] -> [a] -> [[a]]
    findRows [t a]
    cols (a
    ra -> [a] -> [a]
    forall a. a -> [a] -> [a]
    :[a]
    ls) (a -> [a] -> [a]
    forall a. Eq a => a -> [a] -> [a]
    L.delete a
    r [a]
    rs)
                                        | a
    r <- [a]
    rs, a
    r a -> t a -> Bool
    forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
    `notElem` t a
    col]
        findRows [] ls :: [a]
    ls _ = [[a] -> [a]
    forall a. [a] -> [a]
    reverse [a]
    ls]

    isLatinSq :: (Ord a) => [[a]] -> Bool
    isLatinSq :: [[a]] -> Bool
    isLatinSq rows :: [[a]]
    rows = ([a] -> Bool) -> [[a]] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all [a] -> Bool
    forall a. Ord a => [a] -> Bool
    isOneOfEach [[a]]
    rows Bool -> Bool -> Bool
    && ([a] -> Bool) -> [[a]] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all [a] -> Bool
    forall a. Ord a => [a] -> Bool
    isOneOfEach [[a]]
    cols where
        cols :: [[a]]
    cols = [[a]] -> [[a]]
    forall a. [[a]] -> [[a]]
    L.transpose [[a]]
    rows

    isOneOfEach :: [a] -> Bool
    isOneOfEach xs :: [a]
    xs = [a] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [a]
    xs Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Set a -> Int
    forall a. Set a -> Int
    S.size ([a] -> Set a
    forall a. Ord a => [a] -> Set a
    S.fromList [a]
    xs)

    -- The incidence graph of a latin square
    -- It is distance-regular
    -- Godsil & Royle p69
    incidenceGraphLS :: [[c]] -> Graph (Int, Int, c)
    incidenceGraphLS l :: [[c]]
    l = ([(Int, Int, c)], [[(Int, Int, c)]]) -> Graph (Int, Int, c)
    forall t. Ord t => ([t], [[t]]) -> Graph t
    graph ([(Int, Int, c)]
    vs,[[(Int, Int, c)]]
    es) where
        n :: Int
    n = [[c]] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [[c]]
    l -- the order of the latin square
        vs :: [(Int, Int, c)]
    vs = [ (Int
    i, Int
    j, [[c]]
    l [[c]] -> (Int, Int) -> c
    forall a. [[a]] -> (Int, Int) -> a
    ! (Int
    i,Int
    j)) | Int
    i <- [1..Int
    n], Int
    j <- [1..Int
    n] ]
        es :: [[(Int, Int, c)]]
    es = [ [(Int, Int, c)
    v1,(Int, Int, c)
    v2] | [v1 :: (Int, Int, c)
    v1@(i :: Int
    i,j :: Int
    j,lij :: c
    lij), v2 :: (Int, Int, c)
    v2@(i' :: Int
    i',j' :: Int
    j',lij' :: c
    lij')] <- Int -> [(Int, Int, c)] -> [[(Int, Int, c)]]
    forall a. Int -> [a] -> [[a]]
    combinationsOf 2 [(Int, Int, c)]
    vs, Int
    i Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    i' Bool -> Bool -> Bool
    || Int
    j Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    j' Bool -> Bool -> Bool
    || c
    lij c -> c -> Bool
    forall a. Eq a => a -> a -> Bool
    == c
    lij' ]
        m :: [[a]]
    m ! :: [[a]] -> (Int, Int) -> a
    ! (i :: Int
    i,j :: Int
    j) = [[a]]
    m [[a]] -> Int -> [a]
    forall a. [a] -> Int -> a
    !! (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -1) [a] -> Int -> a
    forall a. [a] -> Int -> a
    !! (Int
    jInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -1)

    incidenceGraphLS' :: [[a]] -> Graph (Int, Int)
    incidenceGraphLS' l :: [[a]]
    l = ([(Int, Int)], [[(Int, Int)]]) -> Graph (Int, Int)
    forall t. Ord t => ([t], [[t]]) -> Graph t
    graph ([(Int, Int)]
    vs,[[(Int, Int)]]
    es) where
        n :: Int
    n = [[a]] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [[a]]
    l -- the order of the latin square
        vs :: [(Int, Int)]
    vs = [ (Int
    i, Int
    j) | Int
    i <- [1..Int
    n], Int
    j <- [1..Int
    n] ]
        es :: [[(Int, Int)]]
    es = [ [(Int, Int)
    v1,(Int, Int)
    v2] | [v1 :: (Int, Int)
    v1@(i :: Int
    i,j :: Int
    j), v2 :: (Int, Int)
    v2@(i' :: Int
    i',j' :: Int
    j')] <- Int -> [(Int, Int)] -> [[(Int, Int)]]
    forall a. Int -> [a] -> [[a]]
    combinationsOf 2 [(Int, Int)]
    vs, Int
    i Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    i' Bool -> Bool -> Bool
    || Int
    j Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    j' Bool -> Bool -> Bool
    || Map (Int, Int) a
    l' Map (Int, Int) a -> (Int, Int) -> a
    forall k a. Ord k => Map k a -> k -> a
    M.! (Int
    i,Int
    j) a -> a -> Bool
    forall a. Eq a => a -> a -> Bool
    == Map (Int, Int) a
    l' Map (Int, Int) a -> (Int, Int) -> a
    forall k a. Ord k => Map k a -> k -> a
    M.! (Int
    i',Int
    j') ]
        l' :: Map (Int, Int) a
    l' = [((Int, Int), a)] -> Map (Int, Int) a
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList [ ( (Int
    i,Int
    j), [[a]]
    l [[a]] -> Int -> [a]
    forall a. [a] -> Int -> a
    !! (Int
    iInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -1) [a] -> Int -> a
    forall a. [a] -> Int -> a
    !! (Int
    jInt -> Int -> Int
    forall a. Num a => a -> a -> a
    -1) ) | Int
    i <- [1..Int
    n], Int
    j <- [1..Int
    n] ]
    -- vertices are grid positions
    -- adjacent if they're in the same row, same column, or have the same entry

    -- ORTHOGONAL AND MUTUALLY ORTHOGONAL LATINS SQUARES

    -- |Are the two latin squares orthogonal?
    isOrthogonal :: (Ord a, Ord b) => [[a]] -> [[b]] -> Bool
    isOrthogonal :: [[a]] -> [[b]] -> Bool
    isOrthogonal greeks :: [[a]]
    greeks latins :: [[b]]
    latins = [(a, b)] -> Bool
    forall a. Ord a => [a] -> Bool
    isOneOfEach [(a, b)]
    pairs
        where pairs :: [(a, b)]
    pairs = [a] -> [b] -> [(a, b)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip ([[a]] -> [a]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[a]]
    greeks) ([[b]] -> [b]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [[b]]
    latins)

    findMOLS :: t -> [[[a]]] -> [[[[a]]]]
    findMOLS k :: t
    k lsqs :: [[[a]]]
    lsqs = t -> [[[a]]] -> [[[a]]] -> [[[[a]]]]
    forall t a.
    (Num t, Ord a, Eq t) =>
    t -> [[[a]]] -> [[[a]]] -> [[[[a]]]]
    findMOLS' t
    k [] [[[a]]]
    lsqs where
        findMOLS' :: t -> [[[a]]] -> [[[a]]] -> [[[[a]]]]
    findMOLS' 0 ls :: [[[a]]]
    ls _ = [[[[a]]] -> [[[a]]]
    forall a. [a] -> [a]
    reverse [[[a]]]
    ls]
        findMOLS' i :: t
    i ls :: [[[a]]]
    ls (r :: [[a]]
    r:rs :: [[[a]]]
    rs) =
            if ([[a]] -> Bool) -> [[[a]]] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ([[a]] -> [[a]] -> Bool
    forall a b. (Ord a, Ord b) => [[a]] -> [[b]] -> Bool
    isOrthogonal [[a]]
    r) [[[a]]]
    ls
            then t -> [[[a]]] -> [[[a]]] -> [[[[a]]]]
    findMOLS' (t
    it -> t -> t
    forall a. Num a => a -> a -> a
    -1) ([[a]]
    r[[a]] -> [[[a]]] -> [[[a]]]
    forall a. a -> [a] -> [a]
    :[[[a]]]
    ls) [[[a]]]
    rs [[[[a]]]] -> [[[[a]]]] -> [[[[a]]]]
    forall a. [a] -> [a] -> [a]
    ++ t -> [[[a]]] -> [[[a]]] -> [[[[a]]]]
    findMOLS' t
    i [[[a]]]
    ls [[[a]]]
    rs
            else t -> [[[a]]] -> [[[a]]] -> [[[[a]]]]
    findMOLS' t
    i [[[a]]]
    ls [[[a]]]
    rs
        findMOLS' _ _ [] = []

    -- |Are the latin squares mutually orthogonal (ie each pair is orthogonal)?
    isMOLS :: (Ord a) => [[[a]]] -> Bool
    isMOLS :: [[[a]]] -> Bool
    isMOLS (greek :: [[a]]
    greek:latins :: [[[a]]]
    latins) = ([[a]] -> Bool) -> [[[a]]] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ([[a]] -> [[a]] -> Bool
    forall a b. (Ord a, Ord b) => [[a]] -> [[b]] -> Bool
    isOrthogonal [[a]]
    greek) [[[a]]]
    latins Bool -> Bool -> Bool
    && [[[a]]] -> Bool
    forall a. Ord a => [[[a]]] -> Bool
    isMOLS [[[a]]]
    latins
    isMOLS [] = Bool
    True

    -- |MOLS from a projective plane
    fromProjectivePlane :: (Ord k, Num k) => Design [k] -> [[[Int]]]
    fromProjectivePlane :: Design [k] -> [[[Int]]]
    fromProjectivePlane (D xs :: [[k]]
    xs bs :: [[[k]]]
    bs) = ([[[k]]] -> [[Int]]) -> [[[[k]]]] -> [[[Int]]]
    forall a b. (a -> b) -> [a] -> [b]
    map [[[k]]] -> [[Int]]
    forall a. (Enum a, Num a) => [[[k]]] -> [[a]]
    toLS [[[[k]]]]
    parallelClasses where
        k :: [k]
    k = [k
    x | [0,1,x :: k
    x] <- [[k]]
    xs] -- the field we're working over
        n :: Int
    n = [k] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [k]
    k            -- the order of the projective plane
        parallelClasses :: [[[[k]]]]
    parallelClasses = Int -> [[[[k]]]] -> [[[[k]]]]
    forall a. Int -> [a] -> [a]
    drop 2 ([[[[k]]]] -> [[[[k]]]]) -> [[[[k]]]] -> [[[[k]]]]
    forall a b. (a -> b) -> a -> b
    $ ([[k]] -> [[k]] -> Bool) -> [[[k]]] -> [[[[k]]]]
    forall a. (a -> a -> Bool) -> [a] -> [[a]]
    L.groupBy (\l1 :: [[k]]
    l1 l2 :: [[k]]
    l2 -> [[k]] -> [k]
    forall a. [a] -> a
    head [[k]]
    l1 [k] -> [k] -> Bool
    forall a. Eq a => a -> a -> Bool
    == [[k]] -> [k]
    forall a. [a] -> a
    head [[k]]
    l2) [[[k]]]
    bs
        -- The classes of parallel lines
        -- Each line has its ideal point at its head
        -- The first two classes have [0,0,1] and [0,1,0] as ideal points, and hence consist of horizontal and vertical lines
        toLS :: [[[k]]] -> [[a]]
    toLS ls :: [[[k]]]
    ls = let grid :: Map (k, k) a
    grid = [((k, k), a)] -> Map (k, k) a
    forall k a. Ord k => [(k, a)] -> Map k a
    M.fromList [ ((k
    x,k
    y),a
    i) | (i :: a
    i, [0,1,mu :: k
    mu]:ps :: [[k]]
    ps) <- [a] -> [[[k]]] -> [(a, [[k]])]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [1..] [[[k]]]
    ls, [1,x :: k
    x,y :: k
    y] <- [[k]]
    ps]
                  in Int -> (Int -> Int -> a) -> [[a]]
    forall t a. (Num t, Enum t) => t -> (t -> t -> a) -> [[a]]
    fMatrix' Int
    n (\i :: Int
    i j :: Int
    j -> Map (k, k) a
    grid Map (k, k) a -> (k, k) -> a
    forall k a. Ord k => Map k a -> k -> a
    M.! ([k]
    k [k] -> Int -> k
    forall a. [a] -> Int -> a
    !! Int
    i, [k]
    k [k] -> Int -> k
    forall a. [a] -> Int -> a
    !! Int
    j))

    -- ORTHOGONAL ARRAYS
    -- Godsil & Royle p224

    isOA :: (Int, Int) -> [[b]] -> Bool
    isOA (k :: Int
    k,n :: Int
    n) rows :: [[b]]
    rows =
        [[b]] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length [[b]]
    rows Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    k Bool -> Bool -> Bool
    &&
        ([b] -> Bool) -> [[b]] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all ( (Int -> Int -> Bool
    forall a. Eq a => a -> a -> Bool
    == Int
    nInt -> Integer -> Int
    forall a b. (Num a, Integral b) => a -> b -> a
    ^2) (Int -> Bool) -> ([b] -> Int) -> [b] -> Bool
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . [b] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length ) [[b]]
    rows Bool -> Bool -> Bool
    &&
        ([(b, b)] -> Bool) -> [[(b, b)]] -> Bool
    forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
    all [(b, b)] -> Bool
    forall a. Ord a => [a] -> Bool
    isOneOfEach [[b] -> [b] -> [(b, b)]
    forall a b. [a] -> [b] -> [(a, b)]
    zip [b]
    ri [b]
    rj | [ri :: [b]
    ri,rj :: [b]
    rj] <- Int -> [[b]] -> [[[b]]]
    forall a. Int -> [a] -> [[a]]
    combinationsOf 2 [[b]]
    rows ]

    -- An OA(3,n) from a latin square
    fromLS :: t [Int] -> [[Int]]
    fromLS l :: t [Int]
    l =
        [ [[Int]] -> [Int]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [Int -> Int -> [Int]
    forall a. Int -> a -> [a]
    replicate Int
    n Int
    i | Int
    i <- [1..Int
    n] ] -- row numbers
        , [[Int]] -> [Int]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat (Int -> [Int] -> [[Int]]
    forall a. Int -> a -> [a]
    replicate Int
    n [1..Int
    n])           -- column numbers
        , t [Int] -> [Int]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat t [Int]
    l                              -- entries
        ]
        where n :: Int
    n = t [Int] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length t [Int]
    l -- the order of the latin square

    fromMOLS :: [t [Int]] -> [[Int]]
    fromMOLS mols :: [t [Int]]
    mols =
        ([[Int]] -> [Int]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [Int -> Int -> [Int]
    forall a. Int -> a -> [a]
    replicate Int
    n Int
    i | Int
    i <- [1..Int
    n] ]) [Int] -> [[Int]] -> [[Int]]
    forall a. a -> [a] -> [a]
    : -- row numbers
        ([[Int]] -> [Int]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat (Int -> [Int] -> [[Int]]
    forall a. Int -> a -> [a]
    replicate Int
    n [1..Int
    n]) ) [Int] -> [[Int]] -> [[Int]]
    forall a. a -> [a] -> [a]
    :          -- column numbers
        (t [Int] -> [Int]) -> [t [Int]] -> [[Int]]
    forall a b. (a -> b) -> [a] -> [b]
    map t [Int] -> [Int]
    forall (t :: * -> *) a. Foldable t => t [a] -> [a]
    concat [t [Int]]
    mols                           -- entries for each lsq
        where n :: Int
    n = t [Int] -> Int
    forall (t :: * -> *) a. Foldable t => t a -> Int
    length (t [Int] -> Int) -> t [Int] -> Int
    forall a b. (a -> b) -> a -> b
    $ [t [Int]] -> t [Int]
    forall a. [a] -> a
    head [t [Int]]
    mols -- the order of the latin squares

    -- The graph defined by an OA(k,n)
    -- It is strongly regular with parameters ( n^2, (n-1)k, n-2+(k-1)(k-2), k(k-1) )
    -- Godsil & Royle p225
    graphOA :: [[a]] -> Graph [a]
    graphOA rows :: [[a]]
    rows = ([[a]], [[[a]]]) -> Graph [a]
    forall t. Ord t => ([t], [[t]]) -> Graph t
    graph ([[a]]
    vs,[[[a]]]
    es) where
        vs :: [[a]]
    vs = [[a]] -> [[a]]
    forall a. [[a]] -> [[a]]
    L.transpose [[a]]
    rows -- the vertices are the columns of the OA
        es :: [[[a]]]
    es = [ [[a]
    v1,[a]
    v2] | [v1 :: [a]
    v1,v2 :: [a]
    v2] <- Int -> [[a]] -> [[[a]]]
    forall a. Int -> [a] -> [[a]]
    combinationsOf 2 [[a]]
    vs, [Bool] -> Bool
    forall (t :: * -> *). Foldable t => t Bool -> Bool
    or ((a -> a -> Bool) -> [a] -> [a] -> [Bool]
    forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
    zipWith a -> a -> Bool
    forall a. Eq a => a -> a -> Bool
    (==) [a]
    v1 [a]
    v2) ]
        -- two vertices are adjacent if they agree in any position

    -- Expected SRG parameters
    srgParamsOA :: (d, d) -> Maybe (d, d, d, d)
    srgParamsOA (k :: d
    k,n :: d
    n) =  (d, d, d, d) -> Maybe (d, d, d, d)
    forall a. a -> Maybe a
    Just ( d
    nd -> Integer -> d
    forall a b. (Num a, Integral b) => a -> b -> a
    ^2, (d
    nd -> d -> d
    forall a. Num a => a -> a -> a
    -1)d -> d -> d
    forall a. Num a => a -> a -> a
    *d
    k, d
    nd -> d -> d
    forall a. Num a => a -> a -> a
    -2d -> d -> d
    forall a. Num a => a -> a -> a
    +(d
    kd -> d -> d
    forall a. Num a => a -> a -> a
    -1)d -> d -> d
    forall a. Num a => a -> a -> a
    *(d
    kd -> d -> d
    forall a. Num a => a -> a -> a
    -2), d
    kd -> d -> d
    forall a. Num a => a -> a -> a
    *(d
    kd -> d -> d
    forall a. Num a => a -> a -> a
    -1) )

    -- eg srgParams (4,4) == srgParams $ graphOA $ init $ fromMOLS $ fromProjectivePlane $ pg2 f4

    -- Todo:
    -- Would like a way to find out to what extent two sets of MOLS are really the same,
    -- eg can one be obtained from the other by row or column reordering (with renumbering)
    -- This might provide a proof of the distinctness of phi, omega, omegaD, psi