{-# OPTIONS -fglasgow-exts #-}

module Bio.Util.TestBase where

import Control.Monad (liftM)
import System.CPUTime
import System.Time
import Test.QuickCheck
import System.Random
-- import Data.Char (ord)
import Data.Word
import Data.ByteString.Lazy (pack)

import Bio.Sequence.SeqData

data Test = forall t . Testable t => T String t

newtype Nucleotide = N Char deriving Show
newtype Quality    = Q Word8 deriving Show

fromN :: Nucleotide -> Char
fromN (N c) = c

fromQ :: Quality -> Word8
fromQ (Q c) = c

-- | For testing, variable lengths
newtype EST = E (Sequence Nuc) deriving Show
newtype ESTq = Eq (Sequence Nuc) deriving Show
newtype Protein = P (Sequence Amino) deriving Show

-- | For benchmarking, fixed lengths
newtype EST_short = ES (Sequence Nuc) deriving Show
newtype EST_long  = EL (Sequence Nuc) deriving Show
newtype EST_set  = ESet [Sequence Nuc] deriving Show

-- | Take time (CPU and wall clock) and report it
time :: String -> IO () -> IO ()
time msg action = do
    d1 <- getClockTime
    t1 <- getCPUTime
    action
    t2 <- getCPUTime
    d2 <- getClockTime
    putStrLn $ "\n"++msg++", CPU time: " ++ showT (t2-t1) ++ ", wall clock: "
                 ++ timeDiffToString (diffClockTimes d2 d1)

-- | Print a CPUTime difference
showT :: Integral a => a -> String
showT t = show (fromIntegral t/1e12::Double)++"s"

-- | Shamelessly stolen from FPS
integralRandomR :: (Integral a, RandomGen g) => (a,a) -> g -> (a,g)
integralRandomR  (a,b) g = case randomR (fromIntegral a :: Integer,
                                         fromIntegral b :: Integer) g of
                            (x,g') -> (fromIntegral x, g')

-- | Constrained position generators

genOffset :: Gen Offset
genOffset = do isneg <- arbitrary
               nnoff <- genNonNegOffset
               return $ (if isneg then negate else id) nnoff

genNonNegOffset :: Gen Offset
genNonNegOffset = liftM (subtract 1) genPositiveOffset

genPositiveOffset :: Gen Offset
genPositiveOffset = do scale <- chooseInteger (1, 13)
                       liftM fromIntegral $ chooseInteger (1, 2^scale)
    where chooseInteger :: (Integer, Integer) -> Gen Integer
          chooseInteger = choose


instance Random Word8 where
    randomR = integralRandomR
    random = randomR (minBound,maxBound)

instance Arbitrary Word8 where
    arbitrary = choose (0,255)
    coarbitrary _ = id

instance Arbitrary Nucleotide where
    arbitrary = elements (map N "aaacccgggtttn")
    coarbitrary _ = id

instance Arbitrary Quality where
    arbitrary = do c <- choose (0,60)
                   return (Q c)
    coarbitrary _ = id

instance Arbitrary ESTq where
    arbitrary = do n <- choose (1,100)
                   s <- vector n
                   q <- vector n
                   return $ Eq $ Seq (fromStr "qctest")
                              (fromStr $ map fromN s) (Just $ pack $ map fromQ q)
    coarbitrary _ = id

instance Arbitrary EST where
    arbitrary = do n <- choose (1,100)
                   s <- vector n
                   return $ E $ Seq (fromStr "qctest")
                              (fromStr $ map fromN s) Nothing
    coarbitrary _ = id

instance Arbitrary Char where
    arbitrary = elements (['A'..'Z']++['a'..'z']++" \t\n\r")
    coarbitrary _ = id

instance Arbitrary EST_short where
    arbitrary = do let n = 200
                   s <- vector n
                   q <- vector n
                   return $ ES $ Seq (fromStr "qctest")
                              (fromStr $ map fromN s) (Just $ pack $ map fromQ q)
    coarbitrary _ = id

instance Arbitrary EST_long where
    arbitrary = do let n = 1000
                   s <- vector n
                   q <- vector n
                   return $ EL $ Seq (fromStr "qctest")
                              (fromStr $ map fromN s) (Just $ pack $ map fromQ q)
    coarbitrary _ = id

-- 1000 ESTs of length 1000
instance Arbitrary EST_set where
    arbitrary = do let n = 1000
                   s <- vector n
                   return (ESet $ map (\(EL x) -> x) s)
    coarbitrary _ = id