    {- |
       Data structures for manipulating (biological) sequences.

       Generally supports both nucleotide and protein sequences, some functions,
       like @revcompl@, only makes sense for nucleotides.
    -}

    {-# LANGUAGE EmptyDataDecls #-}

    module Bio.Sequence.SeqData
        (
          -- * Data structure
          -- | A sequence is a header, sequence data itself, and optional quality data.
          --   Sequences are type-tagged to identify them as nucleotide, amino acids,
          --   or unknown type.
          --   All items are lazy bytestrings.  The Offset type can be used for indexing.
          Sequence(..), Offset, SeqData,

          -- | Quality data is normally associated with nucleotide sequences
          Qual, QualData

          -- * Accessor functions
        , (!), seqlength,seqlabel,seqheader,seqdata
        , (?), hasqual, seqqual

          -- * Adding information to header
        , appendHeader, setHeader

          -- * Converting to and from [Char]
        , fromStr, toStr

          -- * Sequence utilities
        , defragSeq

          -- * Nucleotide functionality
          -- | Nucleotide sequences contain the alphabet [A,C,G,T].
          -- IUPAC specifies an extended nucleotide alphabet with wildcards, but
          -- it is not supported at this point.
        , compl, revcompl, revcompl', Nuc, castToNuc

          -- * Protein functionality
          -- | Proteins are chains of amino acids, represented by the IUPAC alphabet.
        , Amino(..), translate, fromIUPAC, toIUPAC -- amino acids
        , castToAmino

          -- * Display a nicely formated sequence.
        , putSeqLn, seqToStr

          -- * Default type for sequences
        , Unknown
        ) where

    import Data.List (unfoldr, intercalate, isPrefixOf)
    import Data.Char (toUpper, isNumber)
    import Data.Int
    import Data.Word
    import Data.Maybe (fromJust, isJust)
    import qualified Data.ByteString.Lazy.Char8 as B
    import qualified Data.ByteString.Lazy as BB
    import qualified Data.ByteString.Char8 as BS
    import qualified Data.ByteString as BBS
    import Text.Printf (printf)

    -- | An offset, index, or length of a 'SeqData'
    type Offset  = Int64

    -- | The basic data type used in 'Sequence's
    type SeqData = B.ByteString

    -- | Basic type for quality data.  Range 0..255.  Typical Phred output is in
    --   the range 6..50, with 20 as the line in the sand separating good from bad.
    type Qual = Word8

    -- | Quality data is a 'Qual' vector, currently implemented as a 'ByteString'.
    type QualData = BB.ByteString -- mild abuse

    -- | A sequence consists of a header, the sequence data itself, and optional quality data.
    --   The type parameter is a phantom type to separate nucleotide and amino acid sequences
    data Sequence t = Seq !SeqData !SeqData !(Maybe QualData) -- ^ header and actual sequence
                      deriving Eq

    -- | For type tagging sequences (protein sequences use 'Amino' below)
    data Nuc
    data Unknown

    -- | Phantom type functionality
    castToNuc :: Sequence a -> Sequence Nuc
    castToNuc (Seq h d mq) = Seq h d mq

    castToAmino :: Sequence a -> Sequence Amino
    castToAmino (Seq h d mq) = Seq h d mq

    -- | A more arranged show instance for Sequences reassembling the display of
    --   the fasta-format
    instance Show (Sequence a) where
        show s = seqToStr s 6 10 []

    -- | Returns a properly formatted and probably highlighted string
    -- | representation of a sequence. Highlighting is done using ANSI-Escape
    -- | sequences.
    seqToStr :: Sequence a -> Int -> Int -> [(Int, Int)] -> [Char]
    seqToStr (Seq header sd _) chunks len parts =
        let strSeq     = toStr sd
            strHeader  = toStr header
            identifier = head (words strHeader)
            comment    = intercalate " " $ tail (words strHeader)

            numLength  = (length . show . length) strSeq
            width      = chunks * len + numLength + chunks
            format     = intercalate "\n" . splits (width)
        in
        (line "ID" width)      ++ (format identifier) ++ "\n\n" ++
        (line "COMMENT" width) ++ (format comment)    ++ "\n\n" ++
        (line "DATA" width)    ++ (showDNA strSeq len chunks parts)

    -- | A simple function to display a sequence: we generate the sequence string and
    -- | call putStrLn
    putSeqLn :: Sequence a -> Int -> Int -> [(Int, Int)] -> IO ()
    putSeqLn s chunks len parts = putStrLn $ seqToStr s chunks len parts

    -- Creates a header line of length n
    line :: [Char] -> Int -> [Char]
    line header width =
      header ++ " " ++ (take (width - (length header) - 1) (repeat '-')) ++ "\n"

    -- | Splits a string into parts of size width. The last element can be shorter.
    splits :: Int -> [t] -> [[t]]
    splits _     [] = []
    splits width xs =
        let (a, b) = splitAt width xs in
        [a] ++ (splits width b)

    -- Highlighting and Dehighlighting ANSI-Escape sequences
    beginHighlight, endHighlight :: String
    beginHighlight = "\ESC[7m"
    endHighlight   = "\ESC[0m"

    -- Creates a string representation of the dna-part of a sequence. We can either
    -- highlight specific areas, or, if this is empty, numberize the rows. At the
    -- moment, both at the same time is not possible (see comment in the numberize-
    -- function). If you want numberized output, take a look at the showDNANumber-
    -- function.
    showDNA :: [Char] -> Int -> Int -> [(Int, Int)] -> [Char]
    showDNA xs len chunks highs =
        let hs  = highlight xs highs
            ks  = splitsSkip hs len [beginHighlight, endHighlight]
            ch  = chunkify ks chunks
            ns  = numberize ch chunks len
            ns' = clean ns False
        in ns'

    -- Add orientation numbers to each line.
    numberize :: [[Char]] -> Int -> Int -> [Char]
    numberize ls chunks len =
        let -- The initial numbering of the single elements in lines
            numbered = zip [0, (chunks * len)..] ls

            -- Create the format string with the right alignment, depending on the
            -- number of elements in lines
            numLength   = (show . length . show ) (chunks * len * (length ls))
            formatStr   = "%" ++ numLength ++ "d "
            toNum (a,b) = (printf formatStr a) ++ b
        in intercalate "\n" (map toNum numbered)

    -- Concatenates each n elements of a list.
    chunkify :: [[Char]] -> Int -> [[Char]]
    chunkify [] _n = []
    chunkify xs  n =
        [intercalate " " (take n xs)] ++ (chunkify (drop n xs) n)

    -- Highlighting should not occur on spaces or over newline borders. We walk
    -- through a generated string and add appropiate stop-highlight and start-
    -- highlight markers.
    --
    -- This is one of the functions which can be definitely optimized by a seasoned
    -- Haskell programmer.
    clean :: [Char] -> Bool -> [Char]
    clean []         _      = []

    -- Handling spaces
    clean (' ':xs)   True   = endHighlight ++ " " ++ beginHighlight ++ clean xs True
    clean (' ':xs)   False  = " " ++ clean xs False

    -- Newlines and numbering
    clean ('\n':xs)   True   =
        let (num, dna) = span (\c -> isNumber c) xs
        in endHighlight ++ "\n" ++ num ++ beginHighlight ++ clean dna True
    clean ('\n':xs)   False  = "\n" ++ clean xs False

    -- Checking for highlighting at all
    clean str@(x:xs) isHigh =
        -- Special cases: check for the beginning or ending of highlighting
        let a = isPrefixOf beginHighlight str
            b = isPrefixOf endHighlight str
        in if a then x : clean xs True
                else if b then x : clean xs False
                          else x : clean xs isHigh

    -- Highlights a range of strings. Each start and end to be highlighted is
    -- described by a position-tuple (start, end).
    highlight :: [Char] -> [(Int, Int)] -> [Char]
    highlight ss [] = ss
    highlight ss ((start, end):rest) =
        let sl = [beginHighlight, endHighlight]
            a  = insert ss start   beginHighlight sl
            b  = insert a  (end+1) endHighlight   sl
        in highlight b rest

    -- Generates lists of length len from xs, skipping sublist in skips while
    -- counting list elements.
    splitsSkip :: (Eq t) => [t] -> Int -> [[t]] -> [[t]]
    splitsSkip [] _   _     = []
    splitsSkip xs len skips =
        (takeSkip xs len skips) : splitsSkip (dropSkip xs len skips) len skips

    -- Emulates the normal take behaviour, skipping over sublists in skips.
    takeSkip :: (Eq a) => [a] -> Int -> [[a]] -> [a]
    takeSkip []         _ _     = []
    takeSkip _          0 _     = []
    takeSkip str@(x:xs) n skips =
        let n' = case (testPrefixes str skips) of
                   Nothing  -> n-1
                   Just len -> n + len - 1
        in x : takeSkip xs n' skips

    -- Emulates the normal drop behaviour, skipping over sublists in skips.
    dropSkip :: (Eq a) => [a] -> Int -> [[a]] -> [a]
    dropSkip []         _ _     = []
    dropSkip xs         0 _     = xs
    dropSkip str@(x:xs) n skips =
        let n' = case (testPrefixes str skips) of
                   Nothing  -> n-1
                   Just len -> n + len - 1
        in dropSkip xs n' skips

    -- Another definition of insert, which provides support for handling skipping
    -- of strings. Inserts elem *before* pos into xs, but skips all strings in
    -- skips while calculating the specified position
    insert :: (Eq a) => [a] -> Int -> [a] -> [[a]] -> [a]
    insert []         _   _    _     = []
    insert xs         0   elem _     = elem ++ xs
    insert str@(x:xs) pos elem skips =
        let pos' = case (testPrefixes str skips) of
                   Nothing  -> pos-1
                   Just len -> pos + len - 1
        in x : insert xs pos' elem skips

    -- Test, if one of the prefixes is the prefix of xs. If yes, return it's length
    -- else Nothing
    testPrefixes :: (Eq a) => [a] -> [[a]] -> Maybe Int
    testPrefixes _  []     = Nothing
    testPrefixes xs (p:ps) =
        if (isPrefixOf p xs)
            then return (length p)
            else testPrefixes xs ps

    -- | Convert a String to 'SeqData'
    fromStr :: String -> SeqData
    fromStr = B.pack

    -- | Convert a 'SeqData' to a String
    toStr :: SeqData -> String
    toStr   = B.unpack

    -- | Read the character at the specified position in the sequence.
    {-# INLINE (!) #-}
    (!) :: Sequence a -> Offset -> Char
    (!) (Seq _ bs _) = B.index bs

    {-# INLINE (?) #-}
    (?) :: Sequence a -> Offset -> Qual
    (?) (Seq _ _ (Just qs)) = BB.index qs
    (?) (Seq _ _ Nothing) =
        error "Attempting to use 'seqqual' on sequence without associated quality data"

    -- | Return sequence length.
    seqlength :: Sequence a -> Offset
    seqlength (Seq _ bs _) = B.length bs

    -- | Return sequence label (first word of header)
    seqlabel :: Sequence a -> SeqData
    seqlabel (Seq l _ _) = head (B.words l)

    -- | Return full header.
    seqheader :: Sequence a -> SeqData
    seqheader (Seq l _ _) = l

    -- | Return the sequence data.
    seqdata :: Sequence a -> SeqData
    seqdata (Seq _ d _) = d

    -- | Return the quality data, or error if none exist.  Use hasqual if in doubt.
    seqqual :: Sequence a -> QualData
    seqqual (Seq _ _ (Just q)) = q
    seqqual (Seq _ _ Nothing) =
        error "Attempting to use 'seqqual' on sequence without associated quality data"

    -- | Check whether the sequence has associated quality data.
    hasqual :: Sequence a -> Bool
    hasqual (Seq _ _ q) = isJust q

    -- | Modify the header by appending text, or by replacing
    --   all but the sequence label (i.e. first word).
    appendHeader, setHeader :: Sequence a -> String -> Sequence a
    appendHeader (Seq h d q) t = (Seq (B.append h (B.pack (" "++t))) d q)
    setHeader (Seq h d q) t = (Seq (B.append (head $ B.words h) (B.pack (" "++t))) d q)

    -- | Returns a sequence with all internal storage freshly copied and
    -- with sequence and quality data present as a single chunk.
    --
    -- By freshly copying internal storage, 'defragSeq' allows garbage
    -- collection of the original data source whence the sequence was
    -- read; otherwise, use of just a short sequence name can cause an
    -- entire sequence file buffer to be retained.
    --
    -- By compacting sequence data into a single chunk, 'defragSeq' avoids
    -- linear-time traversal of sequence chunks during random access into
    -- sequence data.
    defragSeq :: Sequence t -> Sequence t
    defragSeq (Seq name sequ qual) = Seq (defragB name) (defragB sequ) (fmap defragBB qual)
        where defragB = B.fromChunks . (: []) . BS.concat . B.toChunks
              defragBB = BB.fromChunks . (: []) . BBS.concat . BB.toChunks

    ------------------------------------------------------------
    -- Nucleotide (DNA, RNA) specific stuff
    ------------------------------------------------------------

    -- | Calculate the reverse complement.
    --   This is only relevant for the nucleotide alphabet,
    --   and it leaves other characters unmodified.
    revcompl :: Sequence Nuc -> Sequence Nuc
    revcompl (Seq l bs qs) = Seq l (revcompl' bs)
                            (maybe Nothing (Just . BB.reverse) qs)

    -- | Calculate the reverse complent for SeqData only.
    revcompl' :: SeqData -> SeqData
    revcompl' = B.reverse . B.map compl

    -- | Complement a single character.  I.e. identify the nucleotide it
    --   can hybridize with.  Note that for multiple nucleotides, you usually
    --   want the reverse complement (see 'revcompl' for that).
    compl :: Char -> Char
    compl 'A' = 'T'
    compl 'T' = 'A'
    compl 'C' = 'G'
    compl 'G' = 'C'
    compl 'a' = 't'
    compl 't' = 'a'
    compl 'c' = 'g'
    compl 'g' = 'c'
    compl  x  =  x

    ------------------------------------------------------------
    -- Amino acid (protein) stuff
    ------------------------------------------------------------

    -- | Translate a nucleotide sequence into the corresponding protein
    --   sequence.  This works rather blindly, with no attempt to identify ORFs
    --   or otherwise QA the result.
    translate :: Sequence Nuc -> Offset -> [Amino]
    translate s' o' = unfoldr triples (s',o')
       where triples (s,o) =
                 if o > seqlength s - 3 then Nothing
                 else Just (trans1 (map (s!) [o,o+1,o+2]),(s,o+3))

    -- prop_trans s = tail (translate s 0) == translate s 3

    trans1 :: String -> Amino
    trans1 = maybe Xaa id . flip lookup trans_tbl . map (repUT . toUpper)
        where repUT x = if x == 'U' then 'T' else x -- RNA uses U for T

    data Amino = Ala | Arg | Asn | Asp | Cys | Gln | Glu | Gly
               | His | Ile | Leu | Lys | Met | Phe | Pro | Ser
               | Thr | Tyr | Trp | Val
               | STP | Asx | Glx | Xle | Xaa -- unknowns
         deriving (Show,Eq)

    trans_tbl :: [(String,Amino)]
    trans_tbl = [("GCT",Ala),("GCC",Ala),("GCA",Ala),("GCG",Ala),
                 ("CGT",Arg),("CGC",Arg),("CGA",Arg),("CGG",Arg),("AGA",Arg),("AGG",Arg),
                 ("AAT",Asn),("AAC",Asn),
                 ("GAT",Asp),("GAC",Asp),
    --             ("RAT",Asx),("RAC",Asx), -- IUPAC: R is purine (A or G)
                 ("TGT",Cys),("TGC",Cys),
                 ("CAA",Gln),("CAG",Gln),
                 ("GAA",Glu),("GAG",Glu),
    --             ("SAA",Glx),("SAG",Glx), -- IUPAC: S is C or G
                 ("GGT",Gly),("GGC",Gly),("GGA",Gly),("GGG",Gly),
                 ("CAT",His),("CAC",His),
                 ("ATT",Ile),("ATC",Ile),("ATA",Ile),
                 ("TTA",Leu),("TTG",Leu),("CTT",Leu),("CTC",Leu),("CTA",Leu),("CTG",Leu),
                 ("AAA",Lys),("AAG",Lys),
                 ("ATG",Met),
                 ("TTT",Phe),("TTC",Phe),
                 ("CCT",Pro),("CCC",Pro),("CCA",Pro),("CCG",Pro),
                 ("TCT",Ser),("TCC",Ser),("TCA",Ser),("TCG",Ser),("AGT",Ser),("AGC",Ser),
                 ("ACT",Thr),("ACC",Thr),("ACA",Thr),("ACG",Thr),
                 ("TAT",Tyr),("TAC",Tyr),
                 ("TGG",Trp),
                 ("GTT",Val),("GTC",Val),("GTA",Val),("GTG",Val),
                 ("TAG",STP),("TGA",STP),("TAA",STP)
                ]
    -- todo: extend with more IUPAC nucleotide wildcards?

    -- | Convert a list of amino acids to a sequence in IUPAC format.
    toIUPAC :: [Amino] -> SeqData
    toIUPAC = B.pack . map (fromJust . flip lookup iupac)

    -- | Convert a sequence in IUPAC format to a list of amino acids.
    fromIUPAC :: SeqData -> [Amino]
    fromIUPAC = map (maybe Xaa id . flip lookup iupac') . B.unpack

    iupac :: [(Amino,Char)]
    iupac = [(Ala,'A')        ,(Arg,'R')        ,(Asn,'N')
            ,(Asp,'D')        ,(Cys,'C')        ,(Gln,'Q')
            ,(Glu,'E')        ,(Gly,'G')        ,(His,'H')
            ,(Ile,'I')        ,(Leu,'L')        ,(Lys,'K')
            ,(Met,'M')        ,(Phe,'F')        ,(Pro,'P')
            ,(Ser,'S')        ,(Thr,'T')        ,(Tyr,'Y')
            ,(Trp,'W')        ,(Val,'V')
            ,(Asx,'B') -- Asn or Asp
            ,(Glx,'Z') -- Gln or Glu
            ,(Xle,'J') -- Ile or Leu
            ,(STP,'X')        ,(Xaa,'X')
            ]
    iupac' :: [(Char,Amino)]
    iupac' = map (\(a,b)->(b,a)) iupac