{- |
   Module: Bio.Sequence.Fasta

   This module incorporates functionality for reading and writing
   sequence data in the Fasta format.
   Each sequence consists of a header (with a '>' prefix)
   and a set of lines containing the sequence data.

   As Fasta is used for both amino acids and nucleotides, the
   resulting 'Sequence's are type-tagged with 'Unknown'.  If you know the 
   type of sequence you are reading, use 'castToAmino' or 'castToNuc'.
-}

module Bio.Sequence.Fasta
    (
    -- * Reading and writing plain FASTA files
    readFasta, writeFasta, hReadFasta, hWriteFasta
    -- * Reading and writing quality files
    , readQual, writeQual, hWriteQual
    -- * Combining FASTA and quality files
    , readFastaQual, hWriteFastaQual, writeFastaQual
    -- * Counting sequences in a FASTA file
    , countSeqs
    -- * Helper function for reading your own sequences
    , mkSeqs
    -- * Data structures
    , Qual
) where


import Data.Char (chr, isSpace)
import Data.List (groupBy,intersperse)
import System.IO

import qualified Data.ByteString.Lazy.Char8 as B
import qualified Data.ByteString.Lazy as BB
import Data.ByteString.Lazy.Char8 (ByteString)

import Bio.Sequence.SeqData

splitsAt :: Offset -> ByteString -> [ByteString]
splitsAt n s = let (s1,s2) = B.splitAt n s
               in if B.null s2 then [s1] else s1 : splitsAt n s2

{-
-- ugly?
class SeqType sd where
    toSeq :: sd -> sd -> Sequence
    fromSeq :: Sequence -> (sd,sd)

instance SeqType B.ByteString where
    toSeq = Seq
    fromSeq (Seq x y) = (x,y)

instance SeqType BS.ByteString where
    toSeq h s = Seq (B.fromChunks [h]) (B.fromChunks [s])
    fromSeq (Seq x y) = (c x, c y) where c = BS.concat . B.toChunks
-}

-- | Lazily read sequences from a FASTA-formatted file
readFasta :: FilePath -> IO [Sequence Unknown]
readFasta f = (mkSeqs . B.lines) `fmap` B.readFile f

-- | Write sequences to a FASTA-formatted file.
--   Line length is 60.
writeFasta :: FilePath -> [Sequence a] -> IO ()
writeFasta f ss = do
  h <- openFile f WriteMode
  hWriteFasta h ss
  hClose h

-- | Read quality data for sequences to a file.
readQual :: FilePath -> IO [Sequence Unknown]
readQual f = (mkQual . B.lines) `fmap` B.readFile f

-- | Write quality data for sequences to a file.
writeQual :: FilePath -> [Sequence a] -> IO ()
writeQual f ss = do
  h <- openFile f WriteMode
  hWriteQual h ss
  hClose h

-- | Read sequence and associated quality.  Will error if
--   the sequences and qualites do not match one-to-one in sequence.
readFastaQual :: FilePath -> FilePath -> IO [Sequence Unknown]
readFastaQual  s q = do
  ss <- readFasta s
  qs <- readQual q
  -- warning: assumes correct qual file!
  let mkseq s1@(Seq x y _) s2@(Seq _ _ (Just z))
          | seqlabel s1 /= seqlabel s2 = error ("mismatching sequence and quality: "
                                                ++toStr (seqlabel s1)++","++toStr (seqlabel s2))
          | B.length y /= B.length z   = error ("mismatching sequence and quality lengths:"
                                                ++ toStr (seqlabel s1)++","++show (B.length y,B.length z))
          | otherwise   = Seq x y (Just z)
      mkseq _ _ = error "readFastaQual: could not combine Fasta and Qual information"
      zipWith' f (x:xs) (y:ys) = f x y : zipWith' f xs ys
      zipWith' _ [] [] = []
      zipWith' _ _ _ = error "readFastaQual: Unbalanced sequence and quality"
  return $ zipWith' mkseq ss qs

-- | Write sequence and quality data simulatnously
--   This may be more laziness-friendly.
writeFastaQual :: FilePath -> FilePath -> [Sequence a] -> IO ()
writeFastaQual f q ss = do
  hf <- openFile f WriteMode
  hq <- openFile q WriteMode
  hWriteFastaQual hf hq ss
  hClose hq
  hClose hf

hWriteFastaQual :: Handle -> Handle -> [Sequence a] -> IO ()
hWriteFastaQual hf hq = mapM_ wFQ
    where wFQ s = wFasta hf s >> wQual hq s

-- | Lazily read sequence from handle
hReadFasta :: Handle -> IO [Sequence Unknown]
hReadFasta h = (mkSeqs . B.lines) `fmap` B.hGetContents h

-- | Write sequences in FASTA format to a handle.
hWriteFasta :: Handle -> [Sequence a] -> IO ()
hWriteFasta h = mapM_ (wFasta h)

wHead :: Handle -> SeqData -> IO ()
wHead h l = do
   B.hPut h $ B.pack ">"
   B.hPut h l
   B.hPut h $ B.pack "\n"

wFasta :: Handle -> Sequence a -> IO ()
wFasta h (Seq l d _) = do
  wHead h l
  let ls = splitsAt 60 d
  mapM_ (B.hPut h) $ intersperse (B.pack "\n") ls
  B.hPut h $ B.pack "\n"

hWriteQual :: Handle -> [Sequence a] -> IO ()
hWriteQual h = mapM_ (wQual h)

wQual :: Handle -> Sequence a -> IO ()
wQual h (Seq l _ (Just q)) = do
  wHead h l
  let qls = splitsAt 20 q
      qs = B.pack . unwords . map show . BB.unpack
  mapM_ ((\l' -> B.hPut h l' >> B.hPut h (B.pack "\n")) . qs) qls
wQual _ (Seq _ _ Nothing) = return ()

-- | Convert a list of FASTA-formatted lines into a list of sequences.
--   Blank lines are ignored.
--   Comment lines start with "#" are allowed between sequences (and ignored).
--   Lines starting with ">" initiate a new sequence.
mkSeqs :: [ByteString] -> [Sequence Unknown]
mkSeqs = map mkSeq . blocks

mkSeq :: [ByteString] -> Sequence Unknown
mkSeq (l:ls) 
  -- maybe check this?  | B.length l < 2 || isSpace (B.head $ B.tail l) 
  --  = error "Trying to read sequence without a name...and failing."
  | otherwise = Seq (B.drop 1 l) (B.filter (not . isSpace) $ B.concat $ takeWhile isSeq ls) Nothing
    where isSeq s = (not . B.null) s && ((flip elem) (['A'..'Z']++['a'..'z']) . B.head) s
mkSeq [] = error "empty input to mkSeq"

mkQual :: [ByteString] -> [Sequence Unknown]
mkQual = map f . blocks
    where f [] = error "mkQual on empty input - this is impossible"
          f (l:ls) = Seq (B.drop 1 l) B.empty 
                     (Just $ BB.pack $ go 0 ls)
          isDigit x = x <= 58 && x >= 48 
          go i (s:ss) = case BB.uncons s of Just (c,rs) -> if isDigit c then go (c - 48 + 10*i) (rs:ss)
                                                           else let rs' = BB.dropWhile (not . isDigit) rs
                                                                in if BB.null rs' then i : go 0 ss else i : go 0 (rs':ss)
                                            Nothing -> i : go 0 ss
          go _ [] = []

-- | Split lines into blocks starting with '>' characters
--   Filter out # comments (but not semicolons?)
blocks :: [ByteString] -> [[ByteString]]
blocks = groupBy (const (('>' /=) . B.head)) . filter ((/='#') . B.head) . dropWhile (('>' /=) . B.head) . filter (not . B.null)

countSeqs :: FilePath -> IO Int
countSeqs f = do
  ss <- B.readFile f
  let hdrs = filter (('>'==).B.head) $ filter (not . B.null) $ B.lines ss
  return (length hdrs)