module Monocle.Markup (
    Lab (..),
    markup,
    unmark,
    modif,
    modif',
    modifLab,
    choose,
    getLabel
    )where

import Monocle.Core
import Control.Monad.State

-- | Labelled arrow data type.
data Lab a =
    MArrow (Mor a) String |
    MId (Mor a) String |
    MTensor [Lab a] String |
    MComposition [Lab a] String |
    MFunc String [Lab a] FuncT String |
    MTransform String (Lab a) [Mor a] String
    deriving (Eq, Ord)

instance (Eq a) => Morphism (Lab a) where
    dom f = markup $ dom $ unmark f
    cod f = markup $ cod $ unmark f
    isId f = isId $ unmark f
    f \. g = markup $ (unmark f) \. (unmark g)
    f \* g = markup $ (unmark f) \* (unmark g)

makeLab f = case f of
    Arrow _ _ -> MArrow f ""
    Id _ -> MId f ""
    Tensor xs -> MTensor (map makeLab xs) ""
    Composition xs -> MComposition (map makeLab xs) ""
    Func nm xs t -> MFunc nm (map makeLab xs) t ""
    Transform nm x xs -> MTransform nm (makeLab x) xs ""

-- | Removes labels and returns corresponding 'Mor'.
unmark :: (Eq a) => Lab a -> Mor a
unmark lf = nrm $ case lf of
    MArrow f _ -> f
    MId f _ -> f
    MTensor xs _ -> Tensor (map unmark xs)
    MComposition xs _ -> Composition (map unmark xs)
    MFunc nm xs t _ -> Func nm (map unmark xs) t
    MTransform nm x xs _ -> Transform nm (unmark x) xs

mapLabM prep func lf = case lf of
    MTensor xs@(_:_) lab -> do
        prep; xs' <- mapM (mapLabM prep func) xs
        func $ MTensor xs' lab
    MComposition xs lab -> do
        prep;  xs' <- mapM (mapLabM prep func) xs
        func $ MComposition xs' lab
    MFunc nm xs t lab -> do
        prep;  xs' <- mapM (mapLabM prep func) xs
        func $ MFunc nm xs' t lab
    MTransform nm x xs lab -> do
        prep;  x' <- mapLabM prep func x
        func $ MTransform nm x' xs lab
    _ -> func lf

-- | Returns the label of the given marked morphism.
getLabel :: Lab a -> String
getLabel lf = case lf of
    MArrow _ lab -> lab
    MId _ lab -> lab
    MTensor _ lab -> lab
    MComposition _ lab -> lab
    MFunc _ _ _ lab -> lab
    MTransform _ _ _ lab -> lab

setLabel nlab lf = case lf of
    MArrow f _ -> MArrow f nlab
    MId f _ -> MId f nlab
    MTensor xs _ -> MTensor xs nlab
    MComposition xs _ -> MComposition xs nlab
    MFunc nm xs t _ -> MFunc nm xs t nlab
    MTransform nm x xs _ -> MTransform nm x xs nlab


transLab mor inits wlk       = evalState ((mapLabM (return ()) wlk) mor) inits
calcLab mor inits wlk        = execState ((mapLabM (return ()) wlk) mor) inits
transLabP mor inits prep wlk = evalState ((mapLabM prep wlk) mor) inits
calcLabP mor inits prep wlk  = execState ((mapLabM prep wlk) mor) inits

-- | Applies operation to the marked subterm of the given morphism.
modifLab :: String -> Lab a -> (Lab a -> Lab a) -> Lab a
modifLab s lf op = transLab lf () $ \x ->
    let xlab = getLabel x in
        if xlab == s then
            return $ op x
        else return x

-- | Applies operation to the marked subterm of the given morphism.
modif' :: (Eq a) => String -> Lab a -> (Mor a -> Mor a) -> Lab a
modif' s lf op = transLab lf () $ \x ->
    let xlab = getLabel x in
        if xlab == s then
            return $ setLabel xlab $ makeLab $ nrm $ op $ unmark x
        else return x

-- | @modif s lf op == 'unmark' $ 'modif'' s lf op@
modif :: (Eq a) => String -> Lab a -> (Mor a -> Mor a) -> Mor a
modif s lf op = unmark $ modif' s lf op

getByLab s f = calcLab f Nothing $ \x ->
    let xlab = getLabel x in
        if xlab == s then do
            ls <- get
            put $ Just (unmark x)
            return x
        else return x

labels f = calcLab f [] $ \x -> do
    ls <- get
    put $ (getLabel x, unmark x):ls
    return x

-- | Chooses subterm of an associative operation (composition or tensor product).
choose :: String -> Int -> Int -> Lab a -> Lab a
choose nlab start end f = case f of
    MTensor xs lab -> let ((xs1, xs2), xs3) = (let (x1, x2) = splitAt end xs in (splitAt (start-1) x1, x2)) in
            MTensor (xs1 ++ [MTensor xs2 nlab] ++ xs3) lab
    MComposition xs lab -> let s' = length xs - end; e' = length xs - start + 1 in
        let ((xs1, xs2), xs3) = (let (x1, x2) = splitAt e' xs in (splitAt s' x1, x2)) in
            MComposition (xs1 ++ [MComposition xs2 nlab] ++ xs3) lab
    _ -> error "choose: arrow is not composition or tensor"

-- | Returns the given morphism marked up.
markup :: Mor a -> Lab a
markup f = transLab (makeLab f) 1 $ \x ->
    case x of
        MComposition _ _ -> do
            n <- get; put (n+1)
            return $ setLabel ("lab:" ++ show n) x
        MTensor _ _ -> do
            n <- get; put (n+1)
            return $ setLabel ("lab:" ++ show n) x
        MFunc _ _ _ _ -> do
            n <- get; put (n+1)
            return $ setLabel ("lab:" ++ show n) x
        MTransform _ _ _ _ -> do
            n <- get; put (n+1)
            return $ setLabel ("lab:" ++ show n) x
        _ -> return x