module Monocle.Markup where

import Monocle.Core
import Control.Monad.State


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 ""

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

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

modifLab s lf op = transLab lf () $ \x ->
    let xlab = getLabel x in
        if xlab == s then
            return $ op x
        else return x

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

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

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"

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