{-| Data types for functorially lifting sequence positions and
locations onto named sequences.  These are useful for taking functions
that work with sequence positions and locations and associating them
specific, named sequences.

-}

module Bio.Location.OnSeq ( 
  -- * Data types
  SeqName , OnSeq(..)

  -- * Utility functions
  , withSeqData, andSameSeq, onSameSeq

  -- * Sequence collections indexed by name
  , OnSeqs, perSeq, perSeqUpdate, withNameAndSeq
  ) where 

import Control.Monad.Error
import qualified Data.ByteString.Lazy.Char8 as LBS
import Data.List
import qualified Data.Map as M
import Data.Monoid

import Bio.Sequence.SeqData

-- | Sequence name, as in a 'Sequence'
type SeqName = SeqData

-- | Data type for an object associated with a specific, named sequence
data OnSeq a = OnSeq { onSeqName :: !SeqName
                     , onSeqObj :: !a
                     } deriving (Eq, Ord, Show)

instance Functor OnSeq where
    fmap f (OnSeq seqname x) = OnSeq seqname (f x)

-- | Looks up a sequence by name and applies a function to it
withSeqData :: (Monad m) =>
               (SeqData -> a -> m b) -- ^ Function using sequence data
               -> (SeqName -> m SeqData) -- ^ Lookup sequence by name 
               -> OnSeq a -- ^ Object with named sequence 
               -> m b
withSeqData f lookupSeq (OnSeq seqname x) = lookupSeq seqname >>= flip f x

-- | Tests a predicate when two objects are on the same sequence,
-- returning @False@ if they are on different sequences.
andSameSeq :: (a -> b -> Bool) -> OnSeq a -> OnSeq b -> Bool
andSameSeq f (OnSeq xname x) (OnSeq yname y) | xname == yname = f x y
                                             | otherwise = False

-- | Performs an action when two objects are on the same sequence and
-- produces an error otherwise.
onSameSeq :: (Error e, MonadError e m) => (a -> b -> m c) -> OnSeq a -> OnSeq b -> m c
onSameSeq f (OnSeq xname x) (OnSeq yname y) 
    | xname == yname = f x y
    | otherwise      = throwError $ strMsg $ "onSameSeq: " ++ show (LBS.unpack xname) ++ " /= " ++ show (LBS.unpack yname)

-- | Data type for a collection of objects indexed by sequence name
type OnSeqs a = M.Map SeqName a

-- | Lifts a function on an underlying object to look up the sequence
-- name in a name-indexed collection.
perSeq :: (Monoid b) => (a -> b -> c) -> OnSeq a -> OnSeqs b -> c
perSeq f (OnSeq seqname x) = f x . M.findWithDefault mempty seqname

-- | Lifts a function that updates an underlying object to look up the
-- named sequence and update a named-index collection.
perSeqUpdate :: (Monoid b) => (a -> b -> b) -> OnSeq a -> OnSeqs b -> OnSeqs b
perSeqUpdate upd onseq@(OnSeq seqname _) seqmap0 = M.insert seqname (perSeq upd onseq seqmap0) seqmap0

-- | Lifts a function on underlying objects to look up a sequence in a
-- name-indexed collection
withNameAndSeq :: (Monad m) => (SeqName -> a -> b -> m c) -> OnSeq a -> OnSeqs b -> m c
withNameAndSeq f (OnSeq seqname x) = mylookup seqname >=> f seqname x
    where mylookup k = maybe nameNotFound return . M.lookup k
              where nameNotFound = fail $ "withNameAndSeq: sequence " ++ show (LBS.unpack k) ++ " not found"