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

module Math.Combinatorics.Design where

import Data.Maybe (fromJust, isJust)
import qualified Data.List as L
import qualified Data.Map as M
import qualified Data.Set as S

import Math.Common.ListSet (symDiff)
import Math.Algebra.Field.Base
import Math.Algebra.Field.Extension
import Math.Algebra.Group.PermutationGroup hiding (elts, order, isMember)
import Math.Algebra.Group.SchreierSims as SS
import Math.Combinatorics.Graph as G hiding (to1n, combinationsOf)
import Math.Combinatorics.GraphAuts (refine, isSingleton, graphAuts, removeGens)
import Math.Combinatorics.FiniteGeometry

-- Cameron & van Lint, Designs, Graphs, Codes and their Links


{-
set xs = map head $ group $ sort xs

-- subsets of size k (returned in ascending order)
combinationsOf 0 _ = [[]]
combinationsOf _ [] = []
combinationsOf k (x:xs) = map (x:) (combinationsOf (k-1) xs) ++ combinationsOf k xs
-}
isSubset xs ys = all (`elem` ys) xs


-- DESIGNS

data Design a = D [a] [[a]] deriving (Eq,Ord,Show)

design (xs,bs) | isValid d = d where d = D xs bs

toDesign (xs,bs) = D xs' bs' where
    xs' = L.sort xs
    bs' = L.sort $ map L.sort bs -- in fact don't require that the blocks are in order

isValid (D xs bs) = (xs == L.sort xs || error "design: points are not in order")
                 && (all (\b -> b == L.sort b) bs || error "design: blocks do not have points in order")
-- could also check that each block is a subset of xs, etc

points (D xs bs) = xs

blocks (D xs bs) = bs


-- FINDING DESIGN PARAMETERS

noRepeatedBlocks (D xs bs) = all ( (==1) . length ) $ L.group $ L.sort bs


-- Note that the design parameters functions don't check no repeated blocks, so they're also valid for t-structures

-- given t and a t-(v,k,lambda) design, return (v,k,lambda)
tDesignParams t d =
    case findvk d of
    Nothing -> Nothing
    Just (v,k) ->
        case findlambda t d of
        Nothing -> Nothing
        Just lambda -> Just (v,k,lambda)

findvk (D xs bs) =
    let k:ls = map length bs
    in if all (==k) ls then Just (v,k) else Nothing
    where v = length xs

findlambda t (D xs bs) =
    let lambda:ls = [length [b | b <- bs, ts `isSubset` b] | ts <- combinationsOf t xs]
    in if all (==lambda) ls then Just lambda else Nothing

-- given (xs,bs), return design parameters t-(v,k,lambda) with t maximal
designParams d =
    case findvk d of
    Nothing -> Nothing
    Just (v,k) ->
        case reverse (takeWhile (isJust . snd) [(t, findlambda t d) | t <- [0..k] ]) of
        [] -> Nothing
        (t,Just lambda):_ -> Just (t,(v,k,lambda))


isStructure t d = isJust $ tDesignParams t d

isDesign t d = noRepeatedBlocks d && isStructure t d

is2Design d = isDesign 2 d

-- square 2-design (more often called "symmetric" in the literature)
isSquare d@(D xs bs) = is2Design d && length xs == length bs


-- incidence matrix of a design
-- (rows and columns indexed by blocks and points respectively)
-- (this follows Cameron & van Lint, though elsewhere in the literature it is sometimes the other way round)
incidenceMatrix (D xs bs) = [ [if x `elem` b then 1 else 0 | x <- xs] | b <- bs]


-- SOME FAMILIES OF DESIGNS

-- the following is trivially a k-(v,k,lambda) design
subsetDesign v k = design (xs,bs) where
    xs = [1..v]
    bs = combinationsOf k xs

-- Cameron & van Lint, p30
-- the pair design on n points is the complete graph on n points considered as a 2-(n,2,1) design
pairDesign n = D vs es where
    graph = G.k n
    vs = vertices graph
    es = edges graph

-- the affine plane AG(2,Fq) - a 2-(q^2,q,1) design
ag2 fq = design (points, lines) where
    points = ptsAG 2 fq
    lines = map line $ tail $ ptsPG 2 fq
    line [a,b,c] = [ [x,y] | [x,y] <- points, a*x+b*y+c==0 ]

-- the projective plane PG(2,Fq) - a square 2-(q^2+q+1,q+1,1) design
pg2 fq = design (points, lines) where
    points = ptsPG 2 fq
    lines = map line points
    line u = [v | v <- points, u <.> v == 0]
    u <.> v = sum (zipWith (*) u v)
-- Remember that the points and lines of PG(2,Fp) are really the lines and planes of AG(3,Fp).
-- A line in AG(3,Fp) defines a plane orthogonal to it.


-- The points and i-flats of PG(n,fq), 1<=i<=n-1, form a 2-design
-- For i==1, this is a 2-((q^(n+1)-1)/(q-1),q+1,1) design
-- For i==n-1, this is a 2-((q^(n+1)-1)/(q-1),(q^n-1)/(q-1),(q^(n-1)-1)/(q-1)) design
-- Cameron & van Lint, p8
flatsDesignPG n fq k = design (points, blocks) where
    points = ptsPG n fq
    blocks = map closurePG $ flatsPG n fq k -- the closurePG replaces the generators of the flat by the list of points of the flat

-- The projective point-hyperplane design is also denoted PG(n,q)
pg n fq = flatsDesignPG n fq (n-1)

-- (Cameron & van Lint don't actually state that this is a design except when k == n-1)
flatsDesignAG n fq k = design (points, blocks) where
    points = ptsAG n fq
    blocks = map closureAG $ flatsAG n fq k -- the closureAG replaces the generators of the flat by the list of points of the flat

-- The affine point-hyperplane design is also denoted AG(n,q)
-- It a 2-(q^n,q^(n-1),(q^(n-1)-1)/(q-1)) design
-- Cameron & van Lint, p17
ag n fq = flatsDesignAG n fq (n-1)



-- convert a design to be defined over the set [1..n]
to1n (D xs bs) = (D xs' bs') where
    mapping = M.fromList $ zip xs [1..] -- the mapping from vs to [1..n]
    xs' = M.elems mapping
    bs' = [map (mapping M.!) b | b <- bs] -- the blocks will already be sorted correctly by construction


-- Cameron & van Lint p10
paleyDesign fq | length fq `mod` 4 == 3 = design (xs,bs) where
    xs = fq
    qs = set [x^2 | x <- xs] L.\\ [0] -- the non-zero squares in Fq
    bs = [L.sort (map (x+) qs) | x <- xs]

fanoPlane = paleyDesign f7
-- isomorphic to PG(2,F2)


-- NEW DESIGNS FROM OLD

-- Dual of a design. Cameron & van Lint p11
dual (D xs bs) = design (bs, map beta xs) where
    beta x = filter (x `elem`) bs

-- Derived design relative to a point. Cameron & van Lint p11
-- Derived design of a t-(v,k,lambda) is a t-1-(v-1,k-1,lambda) design.
derivedDesign (D xs bs) p = design (xs L.\\ [p], [b L.\\ [p] | b <- bs, p `elem` b])

-- Residual design relative to a point. Cameron & van Lint p13
-- Point-residual of a t-(v,k,lambda) is a t-1-(v-1,k,mu).
pointResidual (D xs bs) p = design (xs L.\\ [p], [b | b <- bs, p `notElem` b])

-- Complementary design relative to a point. Cameron & van Lint p13
-- Complement of a t-(v,k,lambda) is a t-(v,v-k,mu).
complementaryDesign (D xs bs) = design (xs, [xs L.\\ b | b <- bs])

-- Residual design relative to a block. Cameron & van Lint p13
-- This is only a design if (xs,bs) is a square design
-- It may have repeated blocks - but if so, residuals of the complement will not
-- Block-residual of a 2-(v,k,lambda) is a 2-(v-k,k-lambda,lambda).
blockResidual d@(D xs bs) b | isSquare d = design (xs L.\\ b, [b' L.\\ b | b' <- bs, b' /= b])


-- DESIGN AUTOMORPHISMS

isDesignAut (D xs bs) g | supp g `isSubset` xs = all (`S.member` bs') [b -^ g | b <- bs]
    where bs' = S.fromList bs

incidenceGraph (D xs bs) = graph (vs,es) where
    vs = L.sort $ map Left xs ++ map Right bs
    es = L.sort [ [Left x, Right b] | x <- xs, b <- bs, x `elem` b ]

-- We find design auts by finding graph auts of the incidence graph of the design
-- In a square design, we need to watch out for graph auts which are mapping points <-> blocks
designAuts1 d = filter (/=1) $ map points $ graphAuts $ incidenceGraph d where
    points (P m) = P $ M.fromList [(x,y) | (Left x, Left y) <- M.toList m]
     -- This implicitly filters out (Right x, Right y) action on blocks,
     -- and also (Left x, Right y) auts taking points to blocks.
     -- The filter (/=1) is to remove points <-> blocks auts

-- The incidence graph is a bipartite graph, so the distance function naturally partitions points from blocks


-- !! Not strictly correct to stop when a level is empty
-- For example, Fano plane, then if you fix 1,2, you can't move 3, but you can move 4
-- However, the reason it doesn't seem to cause any problem is that in fact it's almost certainly okay to just find the first level
-- (since we know the group is transitive on the points)
-- (In the graph case, non-transitive graphs like kb 2 3 require us to continue after null levels)

-- !! No, not okay to stop at first level
-- The strong generators of one level don't necessarily generate the whole group
-- For example, suppose the group is Sn. We might by chance find a transversal consisting entirely of elts of An
-- We might only finally get Sn when looking for transversal for n-1, and finding (n-1 n)

-- The distance partition of the incidence graph of a design will always look like this:
-- [x], [b | x `elem` b], xs \\ [x], [b | x `notElem` b],
-- When we refine it, we're crossing it with a similar partition but with a different x
-- This won't have much effect on the point cell, but will hopefully split the block cells in half


-- based on graphAuts as applied to the incidence graph, but modified to avoid point-block crossover auts
designAuts d@(D xs bs) = map points (designAuts' [] [vs]) where
    n = length xs
    g@(G vs es) = incidenceGraph d
    points (P m) = P $ M.fromList [(k,v) | (Left k, Left v) <- M.toList m]
    designAuts' us p@((x@(Left _):ys):pt) =
        let p' = L.sort $ filter (not . null) $ refine (ys:pt) (dps M.! x)
        in level us p x ys []
        ++ designAuts' (x:us) p'
    designAuts' us ([]:pt) = designAuts' us pt
    designAuts' _ (((Right _):_):_) = [] -- if we fix all the points, then the blocks must be fixed too 
    -- designAuts' _ [] = []
    level us p@(ph:pt) x (y@(Left _):ys) hs =
        let px = refine (L.delete x ph : pt) (dps M.! x)
            py = refine (L.delete y ph : pt) (dps M.! y)
            uus = zip us us
        in case dfs ((x,y):uus) px py of
           []  -> level us p x ys hs
           h:_ -> let hs' = h:hs in h : level us p x (ys L.\\ (x .^^ hs')) hs'
    -- level _ _ _ [] _ = []
    level _ _ _ _ _ = [] -- includes the case where y matches Right _, which can only occur on first level, before we've distance partitioned
    dfs xys p1 p2
        | map length p1 /= map length p2 = []
        | otherwise =
            let p1' = filter (not . null) p1
                p2' = filter (not . null) p2
            in if all isSingleton p1'
               then let xys' = xys ++ zip (concat p1') (concat p2')
                    in if isCompatible xys' then [fromPairs' xys'] else []
               else let (x:xs):p1'' = p1'
                        ys:p2'' = p2'
                    in concat [dfs ((x,y):xys)
                                   (refine (xs : p1'') (dps M.! x))
                                   (refine ((L.delete y ys):p2'') (dps M.! y))
                                   | y <- ys]
    isCompatible xys = and [([x,x'] `S.member` es') == (L.sort [y,y'] `S.member` es') | (x,y) <- xys, (x',y') <- xys, x < x']
    dps = M.fromList [(v, G.distancePartition g v) | v <- vs]
    es' = S.fromList es


designAutsNew d@(D xs bs) = map points (designAuts' [] [] [vs]) where
    n = length xs
    g@(G vs es) = incidenceGraph d
    -- xs' = take n vs -- the points of the design as vertices of the incidence graph
    points (P m) = P $ M.fromList [(k,v) | (Left k, Left v) <- M.toList m]
    designAuts' us hs p@((x@(Left _):ys):pt) =
        let p' = L.sort $ filter (not . null) $ refine (ys:pt) (dps M.! x)
            hs' = level us p x (ys L.\\ (x .^^ hs)) []
            reps = cosetRepsGx (hs'++hs) x
            schreierGens = removeGens x $ schreierGeneratorsGx (x,reps) (hs'++hs)
        in hs' ++ designAuts' (x:us) schreierGens p'
    designAuts' us hs ([]:pt) = designAuts' us hs pt
    designAuts' us hs p@((x@(Right _):ys):pt) = []
    designAuts' _ _ [] = []
    level us p@(ph:pt) x (y@(Left _):ys) hs =
        let px = refine (L.delete x ph : pt) (dps M.! x)
            py = refine (L.delete y ph : pt) (dps M.! y)
            uus = zip us us
        in case dfs ((x,y):uus) px py of
                []  -> level us p x ys hs
                h:_ -> let hs' = h:hs in h : level us p x (ys L.\\ (x .^^ hs')) hs'
    -- level _ _ _ [] _ = []
    level _ _ _ _ _ = [] -- includes the case where y matches Right _, which can only occur on first level, before we've distance partitioned
    dfs xys p1 p2
        | map length p1 /= map length p2 = []
        | otherwise =
            let p1' = filter (not . null) p1
                p2' = filter (not . null) p2
            in if all isSingleton p1'
               then let xys' = xys ++ zip (concat p1') (concat p2')
                    in if isCompatible xys' then [fromPairs' xys'] else []
               else let (x:xs):p1'' = p1'
                        ys:p2'' = p2'
                    in concat [dfs ((x,y):xys)
                                   (refine (xs : p1'') (dps M.! x))
                                   (refine ((L.delete y ys):p2'') (dps M.! y))
                                   | y <- ys]
    isCompatible xys = and [([x,x'] `S.member` es') == (L.sort [y,y'] `S.member` es') | (x,y) <- xys, (x',y') <- xys, x < x']
    dps = M.fromList [(v, G.distancePartition g v) | v <- vs]
    es' = S.fromList es


-- MATHIEU GROUPS AND WITT DESIGNS

alphaL2_23 = p [[-1],[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22]]                      -- t -> t+1
betaL2_23  = p [[-1],[0],[1,2,4,8,16,9,18,13,3,6,12],[5,10,20,17,11,22,21,19,15,7,14]]                  -- t -> 2*t
gammaL2_23 = p [[-1,0],[1,22],[2,11],[3,15],[4,17],[5,9],[6,19],[7,13],[8,20],[10,16],[12,21],[14,18]]  -- t -> -1/t

l2_23 = [alphaL2_23, betaL2_23, gammaL2_23]

-- Mathieu group M24
-- Conway and Sloane p274ff
-- This is the automorphism group of the extended binary Golay code G24
-- or alternatively of the unique Steiner system S(5,8,24) (which consists of the weight 8 codewords of the above)

deltaM24 = p [[-1],[0],[1,18,4,2,6],[3],[5,21,20,10,7],[8,16,13,9,12],[11,19,22,14,17],[15]]
-- this is t -> t^3 / 9 (for t a quadratic residue), t -> 9 t^3 (t a non-residue)

m24 = [alphaL2_23, betaL2_23, gammaL2_23, deltaM24]

m24sgs = sgs m24
-- order 244823040

m23sgs = filter (\g -> (-1).^g == -1) m24sgs
-- order 10200960

m22sgs = filter (\g -> 0.^g == 0) m23sgs
-- order 443520

-- !! The above assume that the base is [-1,0,..], which isn't guaranteed


-- Steiner system S(5,8,24)

octad = [0,1,2,3,4,7,10,12]
-- Conway&Sloane p276 - this is a weight 8 codeword from Golay code G24


s_5_8_24 = design ([-1..22], octad -^^ l2_23)
-- S(5,8,24) constructed as the image of a single octad under the action of PSL(2,23)
-- 759 blocks ( (24 `choose` 5) `div` (8 `choose` 5) )

s_4_7_23 = derivedDesign s_5_8_24 (-1)
-- 253 blocks ( (23 `choose` 4) `div` (7 `choose` 4) )

s_3_6_22 = derivedDesign s_4_7_23 0
-- 77 blocks

-- Could test that m22sgs are all designAuts of s_3_6_22

{-
s_5_8_24' = octad -^^ m24 -- [alphaL2_23, betaL2_23, gammaL2_23]
-- S(5,8,24) constructed as the image of a single octad under the action of PSL(2,23)
-- 759 blocks ( (24 `choose` 5) `div` (8 `choose` 5) )


s_4_7_23' = [xs | (-1):xs <- s_5_8_24]
-- 253 blocks ( (23 `choose` 4) `div` (7 `choose` 4) )

s_3_6_22' = [xs | 0:xs <- s_4_7_23]
-- 77 blocks
-}


{-
-- WITT DESIGNS
-- S(5,8,24) AND S(5,6,12)

-- Let D be a square 2-design.
-- An n-arc is a set of n points of D, no three of which are contained in a block
arcs n (D xs bs) = map reverse $ dfs n [] xs where
    dfs 0 ys _ = [ys]
    dfs i ys xs = concat [dfs (i-1) (x:ys) (dropWhile (<=x) xs) | x <- xs, isCompatible (x:ys)]
    isCompatible ys = all ((<=2) . length) [ys `L.intersect` b | b <- bs]

tangents (D xs bs) arc = [b | b <- bs, length (arc `L.intersect` b) == 1]

-- !! NOT QUITE AS EXPECTED
-- Cameron van Lint implies that ovals should have n = 1+(k-1)/lambda, whereas I'm finding that they're one bigger than that
-- eg length $ ovals $ ag2 f3 should be 54
-- But ag2 f3 isn't a *square* design
ovals d =
    let Just (_,k,lambda) = tDesignParams 2 d
        (q,r) = (k-1) `quotRem` lambda
        n = 2+q -- == 1+(k-1)/lambda
    in if r == 0
       then [arc | arc <- arcs n d, arc == L.sort (concat $ map (L.intersect arc) $ tangents d arc)] -- each point has a unique tangent
       else []

hyperovals d =
    let Just (_,k,lambda) = tDesignParams 2 d
        (q,r) = k `quotRem` lambda
        n = 1+q -- == 1+k/lambda
    in if r == 0
       then filter (null . tangents d) $ arcs n d
       else []

-- Cameron & van Lint, p22
-- s_5_8_24 = [length (intersect (head h) (head s)) | h <- [h1,h2,h3], s <- [s1,s2,s3]] where
s_5_8_24 = design (points,lines) where
    points = map Left xs ++ map Right [1,2,3]
    lines = [map Left b ++ map Right [1,2,3] | b <- bs] ++ -- line plus three points at infinity
            [map Left h ++ map Right [2,3] | h <- h1] ++ -- hyperoval plus two points at infinity
            [map Left h ++ map Right [1,3] | h <- h2] ++
            [map Left h ++ map Right [1,2] | h <- h3] ++
            [map Left s ++ map Right [1] | s <- s1] ++ -- Baer subplanes plus one point at infinity
            [map Left s ++ map Right [2] | s <- s2] ++
            [map Left s ++ map Right [3] | s <- s3] ++
            [map Left (l1 `symDiff` l2) | l1 <- bs, l2 <- bs, l1 < l2]
    d@(D xs bs) = pg2 f4
    hs = hyperovals d
    [h1,h2,h3] = evenClasses hs
    [s2,s1,s3] = oddClasses baerSubplanes -- we have to number the ss so that if h <- hi, s <- sj, then |h intersect s| is even <=> i == j
    evenClasses (h:hs) = let (ys,ns) = partition (even . length . L.intersect h) hs in (h:ys) : evenClasses ns
    evenClasses [] = []
    oddClasses (h:hs) = let (ys,ns) = partition (odd . length . L.intersect h) hs in (h:ys) : oddClasses ns
    oddClasses [] = []
    baerSubplanes = [s | s <- baerSubplanes', and [length (L.intersect s b) `elem` [1,3] | b <- bs] ] 
    baerSubplanes' = map reverse $ dfs 7 [] xs where
        dfs 0 ys _ = [ys]
        dfs i ys xs = concat [dfs (i-1) (x:ys) (dropWhile (<=x) xs) | x <- xs, isCompatible (x:ys)]
        isCompatible ys = all ((<=3) . length) [ys `L.intersect` b | b <- bs]
-}