-- 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 xs = findLatinSqs' 1 [xs] where
    n = length xs
    findLatinSqs' i rows
        | i == n    = [reverse rows]
        | otherwise = concat [findLatinSqs' (i+1) (row:rows)
                             | row <- findRows (L.transpose rows) [] xs]
    findRows (col:cols) ls rs = concat [findRows cols (r:ls) (L.delete r rs)
                                    | r <- rs, r `notElem` col]
    findRows [] ls _ = [reverse ls]

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

isOneOfEach xs = length xs == S.size (S.fromList xs)


-- The incidence graph of a latin square
-- It is distance-regular
-- Godsil & Royle p69
incidenceGraphLS l = graph (vs,es) where
    n = length l -- the order of the latin square
    vs = [ (i, j, l ! (i,j)) | i <- [1..n], j <- [1..n] ]
    es = [ [v1,v2] | [v1@(i,j,lij), v2@(i',j',lij')] <- combinationsOf 2 vs, i == i' || j == j' || lij == lij' ]
    m ! (i,j) = m !! (i-1) !! (j-1)

incidenceGraphLS' l = graph (vs,es) where
    n = length l -- the order of the latin square
    vs = [ (i, j) | i <- [1..n], j <- [1..n] ]
    es = [ [v1,v2] | [v1@(i,j), v2@(i',j')] <- combinationsOf 2 vs, i == i' || j == j' || l' M.! (i,j) == l' M.! (i',j') ]
    l' = M.fromList [ ( (i,j), l !! (i-1) !! (j-1) ) | i <- [1..n], j <- [1..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 greeks latins = isOneOfEach pairs
    where pairs = zip (concat greeks) (concat latins)

findMOLS k lsqs = findMOLS' k [] lsqs where
    findMOLS' 0 ls _ = [reverse ls]
    findMOLS' i ls (r:rs) =
        if all (isOrthogonal r) ls
        then findMOLS' (i-1) (r:ls) rs ++ findMOLS' i ls rs
        else findMOLS' i ls rs
    findMOLS' _ _ [] = []

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

-- |MOLS from a projective plane
fromProjectivePlane :: (Ord k, Num k) => Design [k] -> [[[Int]]]
fromProjectivePlane (D xs bs) = map toLS parallelClasses where
    k = [x | [0,1,x] <- xs] -- the field we're working over
    n = length k            -- the order of the projective plane
    parallelClasses = drop 2 $ L.groupBy (\l1 l2 -> head l1 == head l2) 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 ls = let grid = M.fromList [ ((x,y),i) | (i, [0,1,mu]:ps) <- zip [1..] ls, [1,x,y] <- ps]
              in fMatrix' n (\i j -> grid M.! (k !! i, k !! j))


-- ORTHOGONAL ARRAYS
-- Godsil & Royle p224

isOA (k,n) rows =
    length rows == k &&
    all ( (== n^2) . length ) rows &&
    all isOneOfEach [zip ri rj | [ri,rj] <- combinationsOf 2 rows ]

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

fromMOLS mols =
    (concat [replicate n i | i <- [1..n] ]) : -- row numbers
    (concat (replicate n [1..n]) ) :          -- column numbers
    map concat mols                           -- entries for each lsq
    where n = length $ head 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 rows = graph (vs,es) where
    vs = L.transpose rows -- the vertices are the columns of the OA
    es = [ [v1,v2] | [v1,v2] <- combinationsOf 2 vs, or (zipWith (==) v1 v2) ]
    -- two vertices are adjacent if they agree in any position

-- Expected SRG parameters
srgParamsOA (k,n) =  Just ( n^2, (n-1)*k, n-2+(k-1)*(k-2), k*(k-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