import Bio.TwoBit
import Bio.TwoBit.Tool
import Codec.Compression.GZip ( decompress )
import Control.Exception
import Control.Monad
import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Lazy.Char8 as L
import Options.Applicative
import System.Directory ( renameFile )
import System.IO
import Text.Printf
options :: Parser (IO ())
options = hsubparser $
command "info" (info tbinfo_options (progDesc "List reference sequences"))
<> command "tofa" (info tbtofa_options (progDesc "Extract Fasta from 2bit"))
<> command "fromfa" (info fatotb_options (progDesc "Convert Fasta to 2bit"))
<> command "fromvcf" (info vcftotb_options (progDesc "Extract reference from dense VCF to 2bit"))
<> command "transcriptome" (info tbxome_options (progDesc "Transform GTF annotation into a transcriptome"))
tbinfo_options :: Parser (IO ())
tbinfo_options = mapM_ go <$> some (strArgument (metavar "2BIT-FILE"))
where
go f = do ref <- openTwoBit f
mapM_ (\TBC{..} -> printf "%s\t%d\n" (C.unpack tbc_name) tbc_dna_size) (tbf_chroms ref)
tbtofa_options :: Parser (IO ())
tbtofa_options = go <$> strArgument (metavar "2BIT-FILE") <*> many (argument rng (metavar "RANGE"))
where
rng = do s0 <- str
if all (/=':') s0
then pure (s0, Nothing)
else do (ch, ':':s1) <- pure $ break (':' ==) s0
(start,'-':s2) : _ <- pure $ reads s1
if null s2
then pure (ch, Just (start, Nothing))
else do (end, "") : _ <- pure $ reads s2
pure (ch, Just (start,Just end))
go fp [ ] = do
ref <- openTwoBit fp
forM_ (tbf_chroms ref) $ \TBC{..} ->
putStrLn ('>' : C.unpack tbc_name) >> twoBitToFa maxBound (tbc_fwd_seq 0)
go fp rns = do
ref <- openTwoBit fp
forM_ rns $ \(ch,se) ->
case (findChrom (C.pack ch) ref, se) of
(Just tbs, Nothing) -> do printf ">%s\n" ch
twoBitToFa maxBound $ tbc_fwd_seq tbs 0
(Just tbs, Just (start,Nothing)) -> do printf ">%s:%d-\n" ch start
twoBitToFa maxBound $ tbc_fwd_seq tbs start
(Just tbs, Just (start,Just end)) -> do printf ">%s:%d-%d\n" ch start end
if start <= end
then twoBitToFa (end-start) $ tbc_fwd_seq tbs start
else twoBitToFa (start-end) $ tbc_rev_seq tbs start
(Nothing, _) -> fail $ "Unknown target: " ++ ch
tbxome_options :: Parser (IO ())
tbxome_options = go <$> strOption (short 'o' <> long "output" <> metavar "FILE" <> value "-" <> help "Write output to FILE")
<*> strArgument (metavar "2BIT-FILE")
<*> many (strArgument (metavar "GFF-File"))
where
go :: FilePath -> FilePath -> [FilePath] -> IO ()
go ofile tbf gffs = do
ref <- openTwoBit tbf
withOutputFile ofile $ \h ->
mapM_ (either (hPutStrLn stderr . displayException) (B.hPutBuilder h . formatCdna ref)) .
concatMap (uncurry parseAnno)
=<< readInputs gffs
walkEncodeProgress :: FilePath -> EncodeProgress -> IO ()
walkEncodeProgress fp (Encoded b) = withOutputFile fp $ flip B.hPutBuilder b
walkEncodeProgress fp (EncodeProgress{..}) = do hPrintf stderr "%s: %d kbases (%d Ns), %dkB encoded\n"
(show ep_seqname) (ep_position `div` 1000) ep_hardmasked (ep_enclength `div` 1024)
walkEncodeProgress fp ep_tail
gunzip :: L.ByteString -> L.ByteString
gunzip s = if "\x1f\x8b" `L.isPrefixOf` s then decompress s else s
fatotb_options :: Parser (IO ())
fatotb_options = go
<$> many (strArgument (metavar "FASTA-FILE"))
<*> strOption (short 'o' <> long "output" <> metavar "FILE" <> value "-" <> help "Write output to FILE")
where
go fs fp = walkEncodeProgress fp . faToTwoBit . L.concat . map gunzip . map snd =<< readInputs fs
vcftotb_options :: Parser (IO ())
vcftotb_options = go
<$> many (strArgument (metavar "VCF-FILE"))
<*> strOption (short 'o' <> long "output" <> metavar "FILE" <> value "-" <> help "Write output to FILE")
where
go fs fp = walkEncodeProgress fp . vcfToTwoBit .
map L.toStrict . concatMap L.lines . map gunzip . map snd =<<
readInputs fs
withOutputFile :: FilePath -> (Handle -> IO a) -> IO a
withOutputFile "-" k = k stdout
withOutputFile f k = bracket (openBinaryFile (f++".#~#") WriteMode) hClose $ \hdl ->
k hdl >>= \r -> renameFile (f++".#~#") f >> return r
readInputs :: [FilePath] -> IO [( String, L.ByteString )]
readInputs [] = pure . (,) "stdin" <$> L.getContents
readInputs fs = mapM go fs
where
go "-" = (,) "stdin" <$> L.getContents
go f = (,) (show f) <$> L.readFile f
main :: IO ()
main = id <=< execParser $
info (options <**> helper)
(progDesc "Stores genomes compactly in the 2bit format and allows fast extraction of sequences from them." <>
header "Compact genome storage" <> fullDesc)