    {- |

       This module implements the 2bit format for sequences.

       Based on: <http://genome.ucsc.edu/FAQ/FAQformat#format7>
       Note! the description is not accurate, it is missing a reserved word
             in each sequence record.

       There are also other, completely different ideas of the 2bit format, e.g.
          <http://jcomeau.freeshell.org/www/genome/2bitformat.html>
    -}

    module Bio.Sequence.TwoBit
       ( decode2Bit,
         read2Bit,
         hRead2Bit,
         encode2Bit,
         write2Bit,
         hWrite2Bit
       ) where

    import Bio.Sequence.SeqData
    import qualified Data.ByteString.Lazy as BB
    import qualified Data.ByteString.Lazy.Char8 as B

    import System.IO
    import Control.Monad

    import Data.Char
    import Data.Binary
    import Data.Int
    import Data.List
    import Data.Bits

    -- import Test.QuickCheck hiding (check)    -- QC 1.0
    import Test.QuickCheck hiding ((.&.)) -- QC 2.0

    -- constants
    default_magic, default_version :: Word32
    default_magic   = 0x1A412743
    default_version = 0

    -- binary extras
    check :: Monad m => (a -> Bool) ->  a -> m a
    check p x = if (p x) then return x else fail "check failed"

    bswap :: Integral a => Int -> a -> a
    bswap n x = let s = bytes x in unbytes . reverse $ s ++ replicate (n-length s) 0

    bytes :: Integral a => a -> [Word8]
    bytes = Data.List.unfoldr (\w -> let (q,r) = quotRem w 256
                                     in if q == 0 && r == 0 then Nothing
                                        else Just (fromIntegral r,q))
    unbytes :: Integral a => [Word8] -> a
    unbytes = Data.List.foldr (\x y -> y*256+x) 0 . map fromIntegral

    -- Conflicts with Bio.Util.TestBase
    -- instance Arbitrary Word8 where
    --    arbitrary = choose (0,255::Int) >>= return . fromIntegral

    -- prop_bswap :: Word8 -> Word8 -> Word8 -> Word8 -> Bool
    -- prop_bswap a1 a2 a3 a4 = (bswap 4 . decode . BB.pack) [a1,a2,a3,a4] == ((decode . BB.pack) [a4,a3,a2,a1] :: Word32)

    -- "in-core" representation of 2Bit data types

    data Header = Header { swap :: Bool,  version, count, reserved :: Word32 }
    instance Show Header where show (Header _ v c r) = "H "++show (v,c,r)

    instance Binary Header where
        get = do
           m <- get
           v <- get >>= check (==default_version)
           c <- get
           r <- get
           let s = if m == default_magic
                       then id
                       else if m == bswap 4 default_magic
                                then bswap 4
                                else error "2bit decode: incorrect magic number"
           return (Header (m /= default_magic) (s v) (s c) (s r))
        put (Header m v c r) = do
           put default_magic
           put default_version
           put c
           put (0 :: Word32)

    data Entry  = Entry  { name :: B.ByteString, offset :: Word32 }
                  deriving Show

    -- Byte swap an Entry's offset!
    swapEntry :: Entry -> Entry
    swapEntry entry = entry { offset = bswap 4 (offset entry) }

    instance Binary Entry where
        get = do
           len <- getWord8
           name <- replicateM (fromIntegral len) getWord8
           offset <- get
           return (Entry (BB.pack name) offset)
        put (Entry byteString offset) = do
           let len = fromIntegral $ B.length byteString :: Word8
           put len
           mapM_ put $ BB.unpack byteString
           put offset

    data Entries = Entries Header [Entry]

    instance Show Entries
       where show (Entries h es) = unlines (show h : map show es)

    instance Binary Entries where
        get = do
           h <- get
           es <- replicateM (fromIntegral $ count h) get
           return (Entries h $ if swap h then map swapEntry es else es)
        put (Entries h es) = do
           put h
           put es

    {-

       Sequence Record definition

     dnaSize         - number of bases of DNA in the sequence
     nBlockCount     - the number of blocks of Ns in the file (representing unknown sequence)
     nBlockStarts    - the starting position for each block of Ns
     nBlockSizes     - the size of each block of Ns
     maskBlockCount  - the number of masked (lower-case) blocks
     maskBlockStarts - the starting position for each masked block
     maskBlockSizes  - the size of each masked block
     packedDna       - the DNA packed to two bits per base

    -}

    data SR = SR { dnaSize         :: Word32,
                   nBlockCount     :: Word32,
                   nBlockStarts    :: [Word32],
                   nBlockSizes     :: [Word32],
                   maskBlockCount  :: Word32,
                   maskBlockStarts :: [Word32],
                   maskBlockSizes  :: [Word32],
                   packedDna       :: [Word8],
                   reserved2       :: Word32 }
            deriving Show

    -- big- and little-endian variants (what a mess)
    newtype SRBE = SRBE SR deriving Show
    newtype SRLE = SRLE SR deriving Show

    instance Binary SRBE where
        get = do
           dz <- get :: Get Word32
           nc <- get :: Get Word32
           let n = fromIntegral nc
           nbs <- replicateM n get
           nbsz <- replicateM n get
           mc <- get :: Get Word32
           let m = fromIntegral mc
           mbs <- replicateM m get
           mbsz <- replicateM m get
           _reserved <- get :: Get Word32 -- !!!! oops?
           let d = fromIntegral dz
           pdna <- replicateM ((d+3) `div` 4)  get
           return (SRBE $ SR dz nc nbs nbsz
                  mc mbs mbsz pdna _reserved)
        -- should this happen?  Why not just write default format?
        put (SRBE sr) = do
           put $ dnaSize sr
           put $ nBlockCount sr
           mapM_ put (nBlockStarts sr)
           mapM_ put (nBlockSizes sr)
           put $ maskBlockCount sr
           mapM_ put (maskBlockStarts sr)
           mapM_ put (maskBlockSizes sr)
           put (0::Word32)
           mapM_ put (packedDna sr)

    instance Binary SRLE where
        get = do
           dz <- get :: Get Word32
           nc <- get :: Get Word32
           let n = fromIntegral $ bswap 4 nc
           nbs <- replicateM n get
           nbsz <- replicateM n get
           mc <- get :: Get Word32
           let m = fromIntegral $ bswap 4 mc
           mbs <- replicateM m get
           mbsz <- replicateM m get
           _reserved <- get :: Get Word32 -- !!!! oops?
           let d = fromIntegral $ bswap 4 dz
           pdna <- replicateM ((d+3) `div` 4)  get
           return (SRLE $ SR (bswap 4 dz) (bswap 4 nc)
                  (map (bswap 4) nbs) (map (bswap 4) nbsz)
                  (bswap 4 mc) (map (bswap 4) mbs) (map (bswap 4) mbsz) pdna (bswap 4 _reserved))
        -- should this happen?  Why not just write default format?
        put (SRLE sr) = do
           put (bswap 4 $ dnaSize sr)
           put (bswap 4 $ nBlockCount sr)
           mapM_ (put . bswap 4) (nBlockStarts sr)
           mapM_ (put . bswap 4) (nBlockSizes sr)
           put (bswap 4 $ maskBlockCount sr)
           mapM_ (put . bswap 4) (maskBlockStarts sr)
           mapM_ (put . bswap 4) (maskBlockSizes sr)
           put (0::Word32)
           mapM_ put (packedDna sr)

    -- Used to convert from sequence data in the Sequence data structure to ByteString
    fromSR :: SR -> B.ByteString
    fromSR sr = B.unfoldr go (0,low,ns,take (fromIntegral $ dnaSize sr) dna)
        where
           low = combine (maskBlockStarts sr) (maskBlockSizes sr)
           ns  = combine (nBlockStarts sr) (nBlockSizes sr)

           combine :: (Num t, Enum t) => [t] -> [t] -> [t]
           combine starts lengths = concatMap (\(p,l) -> [p..p+l-1]) $ zip starts lengths

           -- Unpack a 2Bit packed DNA sequence into a list of 2 bit values 0-3
           dna = decodeDNA $ packedDna sr
           decodeDNA = concatMap (\x -> [shiftR (x .&. 0xC0) 6, shiftR (x .&. 0x30) 4, shiftR (x .&. 0x0C) 2, x .&. 0x03])

           -- Map 2Bit nucleotide encodings to their character equivalents
           dec1 :: (Num t) => t -> Char
           dec1 x = case x of
                       0 -> 'T';
                       1 -> 'C';
                       2 -> 'A';
                       3 -> 'G';
                       _ -> error ("can't decode value "++show x)

           go :: (Num a, Num t) => (a, [a], [a], [t]) -> Maybe (Char, (a, [a], [a], [t]))
           go (_,_,_,[])              = Nothing
           go (pos,(l:ls),(n:ns),(d:ds))
               | pos == l && pos == n = Just ('n',(pos+1,ls,ns,ds))
               | pos == l             = Just (toLower (dec1 d),(pos+1,ls,n:ns,ds))
               |             pos == n = Just ('N',(pos+1,l:ls,ns,ds))
               | otherwise            = Just (dec1 d, (pos+1,l:ls,n:ns,ds))
           go (pos,[],n:ns,d:ds)
               |             pos == n = Just ('N',(pos+1,[],ns,ds))
               | otherwise            = Just (dec1 d, (pos+1,[],n:ns,ds))
           go (pos,l:ls,[],d:ds)
               | pos == l             = Just (toLower (dec1 d),(pos+1,ls,[],ds))
               | otherwise            = Just (dec1 d, (pos+1,l:ls,[],ds))
           go (pos,[],[],d:ds)        = Just (dec1 d, (pos+1,[],[],ds))
    --    go x = error (show x)

    toSR :: B.ByteString -> SR
    toSR bs = undefined

    splits :: [Int64] -> B.ByteString -> [B.ByteString]
    splits [] cs = [cs]
    splits (e:es) cs = let (this,rest) = B.splitAt e cs
                       in this : splits es rest

    -- | Parse a (lazy) ByteString as sequences in the 2bit format.
    decode2Bit :: B.ByteString -> [Sequence Unknown]
    decode2Bit cs = let
                        -- decode to (Header, [Entry]) from ByteString
                        (Entries h es) = decode cs :: Entries

                        -- map to [64-bit offset] from [32-bit offset]!!!
                        ms = map (fromIntegral . offset) es

                        -- break the ByteString up into chunks of raw 2Bit "sequences"
                        (c:chunks) = zipWith (-) ms (0:ms)

                        --  build raw undecoded [Sequence]!!!
                        --  chunks :: [Int64]
                        --          "drop c cs" .. chop off all cruft in the beginning!!!
                        ss = splits chunks $ B.drop c cs

                    --                              ($ Nothing) :: (Maybe a -> b) -> b
                    --                              zipWith :: (a -> b -> c) -> [a] -> [b] -> [c]
                    --                              Seq :: SeqData -> SeqData -> Maybe QualData -> Sequence
                    --                              map fromSR :: [SR] -> [B.ByteString]
                    --                              map (unSRLE.decode) :: [B.ByteString] -> [SR]
                    --                              map (unSRBE.decode) :: [B.ByteString] -> [SR]
                    in map ($ Nothing) $ zipWith Seq (map name es) $ map fromSR $ case swap h of
                                                                                     True -> map (unSRLE.decode) ss
                                                                                     False -> map (unSRBE.decode) ss

    type SequenceLabel = SeqData
    type SequenceSize = Offset
    type SequenceData = SeqData
    type TwoBitData = (SequenceLabel, SequenceSize, SequenceData)

    -- | Marshall from neutral representation to the 2Bit ByteString rep
    encode2Bit :: [Sequence a] -> B.ByteString
    encode2Bit ss = let

                       buildHeader :: [Sequence a] -> Header
                       buildHeader ss = Header {swap = True,
                                                version = default_version,
                                                count = sequenceListLength ss,
                                                reserved = 0}

                       -- Build the list of 2Bit Entries
                       buildEntries :: [TwoBitData] -> Offset -> [Entry]
                       buildEntries [] _ = []
                       buildEntries ((label, length, _):xs) currentOffset = Entry {name=label,
                                                                                   offset=fromIntegral $ currentOffset} : buildEntries xs (currentOffset+length)

                       -- Build a 2Bit Sequence Record!!
                       buildSR :: TwoBitData -> SR
                       buildSR (_, size, dnaData) = SR {dnaSize = (fromIntegral size),
                                                        nBlockCount = 0,
                                                        nBlockStarts = [],
                                                        nBlockSizes = [],
                                                        maskBlockCount = 0,
                                                        maskBlockStarts = [],
                                                        maskBlockSizes = [],
                                                        packedDna = encodeDNA $ splitWord8 $ explode dnaData,
                                                        reserved2 = 0}

                       -- Total # of sequences present
                       sequenceListLength :: [Sequence a] -> Word32
                       sequenceListLength [] = 0
                       sequenceListLength (s:ss) = 1 + sequenceListLength ss

                       -- Build a list of vital data to do 2Bit marshalling
                       sequenceListExtract :: [Sequence a] -> [TwoBitData]
                       sequenceListExtract ss = map (\seq -> (seqlabel seq, seqlength seq, seqdata seq)) ss

                       -- Map a nucleotide to its respective 2Bit encoding
                       enc1 :: (Num t) => Char -> t
                       enc1 c = case c of
                                   'T' -> 0;
                                   'C' -> 1;
                                   'A' -> 2;
                                   'G' -> 3;

                       -- Take a ByteString of nucleotides in character encoding and map to a list of their corresponding TwoBit encodings
                       explode :: SeqData -> [Word8]
                       explode seq = map enc1 $ B.unpack seq

                       -- Split a list of Word8's into 4 character sublists (i.e. quads)
                       splitWord8 ::  [Word8] -> [[Word8]]
                       splitWord8 [] = []
                       splitWord8 cs = let (this,rest) = splitAt 4 cs
                                          in this : splitWord8 rest

                       -- Build a 2Bit packed DNA sequence
                       encodeDNA :: [[Word8]] -> [Word8]
                       encodeDNA [] = []
                       encodeDNA (w:ws) = (pack2Bit $ getQuad) : encodeDNA ws

                                         where

                                            -- Length of a sequence
                                            len = length w

                                            -- Build a 2Bit encoded Word8 from "exploded" encoding
                                            pack2Bit :: (Word8, Word8, Word8, Word8) -> Word8
                                            pack2Bit (d1, d2, d3, d4) = shiftL (d1) 6 .|. shiftL (d2) 4 .|. shiftL (d3) 2 .|. d4

                                            -- Build a quad, i.e. a 4 tuple of Word8
                                            getQuad :: (Word8, Word8, Word8, Word8)
                                            getQuad = case len of
                                                            1 -> (w !! 0, 0, 0, 0);
                                                            2 -> (w !! 0, w !! 1, 0, 0);
                                                            3 -> (w !! 0, w !! 1, w !! 2, 0);
                                                            4 -> (w !! 0, w !! 1, w !! 2, w !! 3);

                   in

                       -- Serialize/marshall into ByteString representation
                       B.append (encode (buildHeader ss))
                          (B.append
                             (B.concat (map encode (buildEntries (sequenceListExtract ss) 56)))  -- TEMP WNH
                             (B.concat (map (encode . SRBE) (map buildSR (sequenceListExtract ss)))))      -- TEMP WNH
                             -- WNH (B.concat (map (encode . SRBE) (map buildSR (sequenceListExtract ss)))))

    unSRBE :: SRBE -> SR
    unSRBE (SRBE x) = x

    unSRLE :: SRLE -> SR
    unSRLE (SRLE x) = x

    -- | Read sequences from a file in 2bit format and
    -- | unmarshall/deserialize into Sequence format.
    read2Bit  :: FilePath -> IO [Sequence Unknown]
    read2Bit f = B.readFile f >>= return . decode2Bit

    -- | Read sequences from a file handle in the 2bit format and
    -- | unmarshall/deserialze into Sequence format.
    hRead2Bit :: Handle   -> IO [Sequence Unknown]
    hRead2Bit h = B.hGetContents h >>= return . decode2Bit

    -- | Marshall/serialize [Sequence] into 2Bit format and write to a file.
    write2Bit  :: FilePath -> [Sequence a] -> IO ()
    write2Bit f seq = do
                       let byteString = encode2Bit seq
                       B.writeFile f byteString
                       return ()

    -- | Marshall/serialize [Sequence] into 2Bit format and write to a file using handle.
    hWrite2Bit :: Handle   -> [Sequence a] -> IO ()
    hWrite2Bit h seq = do
                       let byteString = encode2Bit seq
                       B.hPut h byteString
                       return ()