    {-# LANGUAGE FlexibleInstances, MultiParamTypeClasses, DeriveDataTypeable, PatternGuards #-}
    -----------------------------------------------------------------------------
    -- |
    -- Module      :  XMonad.Layout.Mosaic
    -- Copyright   :  (c) 2009 Adam Vogt, 2007 James Webb
    -- License     :  BSD-style (see xmonad/LICENSE)
    --
    -- Maintainer  :  vogt.adam<at>gmail.com
    -- Stability   :  unstable
    -- Portability :  unportable
    --
    -- Based on MosaicAlt, but aspect ratio messages always change the aspect
    -- ratios, and rearranging the window stack changes the window sizes.
    --
    -----------------------------------------------------------------------------

    module XMonad.Layout.Mosaic (
        -- * Usage
        -- $usage
         Aspect(..)
        ,mosaic
        ,changeMaster
        ,changeFocused
        )
        where

    import Prelude hiding (sum)

    import XMonad(Typeable,
                  LayoutClass(doLayout, handleMessage, pureMessage, description),
                  Message, X, fromMessage, withWindowSet, Resize(..),
                  splitHorizontallyBy, splitVerticallyBy, sendMessage, Rectangle)
    import qualified XMonad.StackSet as W
    import Control.Arrow(second, first)
    import Control.Monad(mplus)
    import Data.Foldable(Foldable,foldMap, sum)
    import Data.Function(on)
    import Data.List(sortBy)
    import Data.Monoid(Monoid,mempty, mappend)

    -- $usage
    -- You can use this module with the following in your @~\/.xmonad\/xmonad.hs@:
    --
    -- > import XMonad.Layout.Mosaic
    --
    -- Then edit your @layoutHook@ by adding the Mosaic layout:
    --
    -- > myLayout = mosaic 2 [3,2] ||| Full ||| etc..
    -- > main = xmonad $ defaultConfig { layoutHook = myLayout }
    --
    -- Unfortunately, infinite lists break serialization, so don't use them. And if
    -- the list is too short, it is extended with @++ repeat 1@, which covers the
    -- main use case.
    --
    -- To change the choice in aspect ratio and the relative sizes of windows, add
    -- to your keybindings:
    --
    --  > , ((modm, xK_a), sendMessage Taller)
    --  > , ((modm, xK_z), sendMessage Wider)
    --
    --  > , ((modm, xK_r), sendMessage Reset)
    --
    -- For more detailed instructions on editing the layoutHook see:
    --
    -- "XMonad.Doc.Extending#Editing_the_layout_hook"

    data Aspect
        = Taller
        | Wider
        | Reset
        | SlopeMod ([Rational] -> [Rational])
        deriving (Typeable)

    instance Message Aspect

    -- | The relative magnitudes (the sign is ignored) of the rational numbers in
    -- the second argument determine the relative areas that the windows receive.
    -- The first number represents the size of the master window, the second is for
    -- the next window in the stack, and so on.
    --
    -- The list is extended with @++ repeat 1@, so @mosaic 1.5 []@ is like a
    -- resizable grid.
    --
    -- The first parameter is the multiplicative factor to use when responding to
    -- the 'Expand' message.
    mosaic :: Rational -> [Rational] -> Mosaic a
    mosaic = Mosaic Nothing

    data Mosaic a = -- | True to override the aspect, current index, maximum index
                    Mosaic (Maybe(Bool,Rational,Int)) Rational [Rational] deriving (Read,Show)

    instance LayoutClass Mosaic a where
        description = const "Mosaic"

        pureMessage (Mosaic Nothing _ _) _ = Nothing
        pureMessage (Mosaic (Just(_,ix,mix)) delta ss) ms = fromMessage ms >>= ixMod
            where ixMod Taller | round ix >= mix = Nothing
                               | otherwise = Just $ Mosaic (Just(False,succ ix,mix)) delta ss
                  ixMod Wider  | round ix <= (0::Integer) = Nothing
                               | otherwise = Just $ Mosaic (Just(False,pred ix,mix)) delta ss
                  ixMod Reset              = Just $ Mosaic Nothing delta ss
                  ixMod (SlopeMod f)       = Just $ Mosaic (Just(False,ix,mix)) delta (f ss)

        handleMessage l@(Mosaic _ delta _) ms
            | Just Expand <- fromMessage ms = changeFocused (*delta) >> return Nothing
            | Just Shrink <- fromMessage ms = changeFocused (/delta) >> return Nothing
            | otherwise = return $ pureMessage l ms

        doLayout (Mosaic state delta ss) r st = let
            ssExt = zipWith const (ss ++ repeat 1) $ W.integrate st
            rects = splits r ssExt
            nls = length rects
            fi = fromIntegral
            nextIx (ov,ix,mix)
                    | mix <= 0 || ov = fromIntegral $ nls `div` 2
                    | otherwise = max 0 $ (*fi (pred nls)) $ min 1 $ ix / fi mix
            rect = rects !! maybe (nls `div` 2) round (nextIx `fmap` state)
            state' = fmap (\x@(ov,_,_) -> (ov,nextIx x,pred nls)) state
                        `mplus` Just (True,fromIntegral nls / 2,pred nls)
            ss' = maybe ss (const ss `either` const ssExt) $ zipRemain ss ssExt
            in return (zip (W.integrate st) rect, Just $ Mosaic state' delta ss')

    zipRemain :: [a] -> [b] -> Maybe (Either [a] [b])
    zipRemain (_:xs) (_:ys) = zipRemain xs ys
    zipRemain [] [] = Nothing
    zipRemain [] y = Just (Right y)
    zipRemain x [] = Just (Left x)

    -- | These sample functions are meant to be applied to the list of window sizes
    -- through the 'SlopeMod' message.
    changeMaster :: (Rational -> Rational) -> X ()
    changeMaster = sendMessage . SlopeMod . onHead

    -- | Apply a function to the Rational that represents the currently focused
    -- window.
    --
    -- 'Expand' and 'Shrink' messages are responded to with @changeFocused
    -- (*delta)@ or @changeFocused (delta/)@ where @delta@ is the first argument to
    -- 'mosaic'.
    --
    -- This is exported because other functions (ex. @const 1@, @(+1)@) may be
    -- useful to apply to the current area.
    changeFocused :: (Rational -> Rational) -> X ()
    changeFocused f = withWindowSet $ sendMessage . SlopeMod
                        . maybe id (mulIx . length . W.up)
                        . W.stack . W.workspace . W.current
        where mulIx i = uncurry (++) . second (onHead f) . splitAt i

    onHead :: (a -> a) -> [a] -> [a]
    onHead f = uncurry (++) . first (fmap f) . splitAt 1

    splits :: Rectangle -> [Rational] -> [[Rectangle]]
    splits rect = map (reverse . map snd . sortBy (compare `on` fst))
                    . splitsL rect . makeTree snd . zip [1..]
                    . normalize . reverse . map abs

    splitsL :: Rectangle -> Tree (Int,Rational) -> [[(Int,Rectangle)]]
    splitsL _rect Empty = []
    splitsL rect (Leaf (x,_)) = [[(x,rect)]]
    splitsL rect (Branch l r) = do
        let mkSplit f = f ((sumSnd l /) $ sumSnd l + sumSnd r) rect
            sumSnd = sum . fmap snd
        (rl,rr) <- map mkSplit [splitVerticallyBy,splitHorizontallyBy]
        splitsL rl l `interleave` splitsL rr r

    -- like zipWith (++), but when one list is shorter, its elements are duplicated
    -- so that they match
    interleave :: [[a]] -> [[a]] -> [[a]]
    interleave xs ys | lx > ly = zc xs (extend lx ys)
                     | otherwise = zc (extend ly xs) ys
      where lx = length xs
            ly = length ys
            zc = zipWith (++)

            extend :: Int -> [a] -> [a]
            extend n pat = do
                (p,e) <- zip pat $ replicate m True ++ repeat False
                [p | e] ++ replicate d p
                where (d,m) = n `divMod` length pat

    normalize :: Fractional a => [a] -> [a]
    normalize x = let s = sum x in map (/s) x

    data Tree a = Branch (Tree a) (Tree a) | Leaf a | Empty

    instance Foldable Tree where
       foldMap _f Empty = mempty
       foldMap f (Leaf x) = f x
       foldMap f (Branch l r) = foldMap f l `mappend` foldMap f r

    instance Functor Tree where
       fmap f (Leaf x) = Leaf $ f x
       fmap f (Branch l r) = Branch (fmap f l) (fmap f r)
       fmap _ Empty = Empty

    instance Monoid (Tree a) where
        mempty = Empty
        mappend Empty x = x
        mappend x Empty = x
        mappend x y = Branch x y

    makeTree ::  (Num a1, Ord a1) => (a -> a1) -> [a] -> Tree a
    makeTree _ [] = Empty
    makeTree _ [x] = Leaf x
    makeTree f xs = Branch (makeTree f a) (makeTree f b)
        where ((a,b),_) = foldr go (([],[]),(0,0)) xs
              go n ((ls,rs),(l,r))
                | l > r     = ((ls,n:rs),(l,f n+r))
                | otherwise = ((n:ls,rs),(f n+l,r))