{-# LANGUAGE TemplateHaskell, QuasiQuotes #-}
module Tests.Count
    ( tgroup_Count
    ) where

import Test.Tasty.TH
import Test.Tasty.HUnit

import qualified Data.IntervalIntMap as IM
import qualified Data.Set as S
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as B8
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as VU
import qualified Data.Map as M

import qualified Data.Conduit.Combinators as C
import qualified Data.Conduit.Binary as CB
import qualified Data.Conduit.List as CL
import qualified Data.Conduit as C
import           Data.Conduit ((.|))
import Control.Monad.IO.Class (liftIO)
import Data.Maybe

import Interpretation.Count
import qualified Interpretation.Count.RefSeqInfoVector as RSV
import FileOrStream (FileOrStream(..))
import Tests.Utils
import Utils.Here
import NGLess


tgroup_Count = $(testGroupGenerator)


readCountFile :: FilePath -> IO (M.Map B.ByteString Double)
readCountFile fp =
    C.runConduitRes $
        C.sourceFile fp
            .| CB.lines
            .| (C.await >> (C.awaitForever C.yield)) -- skip first line
            .| CL.foldMap parseLine
    where
        parseLine line = case B8.split '\t' line of
            [h,val] -> M.singleton h (read $ B8.unpack val)
            _ -> error ("Could not parse line: " ++ show line)


runSamGffAnnotation:: B.ByteString -> B.ByteString -> CountOpts -> NGLessIO (M.Map B.ByteString Double)
runSamGffAnnotation sam_content gff_content opts = do
    sam_fp <- asTempFile sam_content "sam"
    gff_fp <- asTempFile gff_content "gff"
    ann <- loadAnnotator (AnnotateGFF gff_fp) opts
    p <- performCount (File sam_fp) "testing" ann opts
    liftIO $ readCountFile p

listNub :: (Ord a) => [a] -> [a]
listNub = S.toList . S.fromList

defCountOpts =
    CountOpts
    { optFeatures = []
    , optSubFeatures = Nothing
    , optIntersectMode = annotationRule IntersectUnion
    , optAnnotationMode = AnnotateSeqName
    , optStrandMode = SMBoth
    , optMinCount = 0.0
    , optMMMethod = MMUniqueOnly
    , optDelim = "\t"
    , optNormMode = NMRaw
    , optIncludeMinus1 = False
    }


extractIds :: [AnnotationInfo] -> [Int]
extractIds = map (\(AnnotationInfo _ ix) -> ix)

very_short_gff = "test_samples/very_short.gtf"
case_load_very_short = do
    [GFFAnnotator immap headers szmap] <- testNGLessIO
                $ loadAnnotator (AnnotateGFF very_short_gff) defCountOpts  { optFeatures = ["gene"] }
    let usedIDs = extractIds $ concatMap IM.elems $ M.elems immap
    length (listNub usedIDs) @?= V.length headers
    minimum usedIDs @?= 0
    maximum usedIDs @?= V.length headers - 1
    VU.length szmap @?= V.length headers
    let mix = do
            ix <- V.elemIndex "WBGene00010199" headers
            return $! szmap VU.! ix
    mix @?= Just (721-119+1)


short3 :: B.ByteString
short3 = [here|
V	protein_coding	gene	7322	8892	.	-	.	gene_id "WBGene00008825"; gene_name "F14H3.6"; gene_source "ensembl"; gene_biotype "protein_coding";
X	protein_coding	gene	140	218	.	+	.	gene_id "WBGene00020330"; gene_name "T07H6.1"; gene_source "ensembl"; gene_biotype "protein_coding";
X	protein_coding	gene	632	733	.	+	.	gene_id "WBGene00000526"; gene_name "clc-5"; gene_source "ensembl"; gene_biotype "protein_coding";
|]

-- this is a regression test
case_load_gff_order = do
    fp <- testNGLessIO $ asTempFile short3 "gtf"
    [GFFAnnotator immap headers _] <- testNGLessIO
                $ loadAnnotator (AnnotateGFF fp) defCountOpts  { optFeatures = ["gene"] }
    let [h] = extractIds . IM.elems  . fromJust $ M.lookup "V" immap
    (headers V.! h) @?= "WBGene00008825"

short1 :: B.ByteString
short1 = [here|
X	protein_coding	gene	610	1473	.	+	.	gene_id "WBGene00002254"; gene_name "lbp-2"; gene_source "ensembl"; gene_biotype "protein_coding";
|]

short_sam :: B.ByteString
short_sam = [here|
@SQ	SN:X	LN:18942
SRR070372.1096	0	X	1174	60	62S75M1D37M46D58M10S	*	0	0	GTTCTACAACGTCCAGATCGGAAGCAAGTTCGAAGGAGAGGGTCTTGATAACACCAAGCACGAGGTTACCTTCACTCTCAAGGACGGACACTTGTTCGAACATCACAAGCCACTTGAAGAGGGAGAATCCAAGGAAGAACCTATGAGTATTACTTTGATGGAGATTTTCTTATTCAGAAGATGAGCTTCAACAATATCGAAGGCCGCAGATTCTACAAGAGACTCCCATAAAGTTAACTATC	IIIIIIF@@@CIIIIIIIIIIIIIIIIIIIIIHHIIIB=5669CIIIIIIIIIIIIIIIIIIIIIIIIIIHHHIHIIIIIIIIIIIIIIIIIIIHIIIIIIIIIIIIIIIHIH>>>FIIGBB@E??;75444<<:62///1>?BAAAD?AE;72217<AAAA;=/1117//7AADACDDGIEEEEEGGHGD@@@GGGGD@@@@DD@@@DDEBCBEBB@:566?6333;C@@=BAA:?E9911	NM:i:47	MD:Z:75^A37^CAGGTAAAATTTGGTCAATCTATTTGACATACATTTTTGTTAATTA58	AS:i:111	XS:i:19	SA:Z:X,1053,+,7M3D59M176S,60,3;
SRR070372.1096	2048	X	1053	60	7M3D59M176H	*	0	0	GTTCTACAACGTCCAGATCGGAAGCAAGTTCGAAGGAGAGGGTCTTGATAACACCAAGCACGAGGT	IIIIIIF@@@CIIIIIIIIIIIIIIIIIIIIIHHIIIB=5669CIIIIIIIIIIIIIIIIIIIIII	NM:i:3	MD:Z:7^AAA59	AS:i:59	XS:i:0	SA:Z:X,1174,+,62S75M1D37M46D58M10S,60,47;
SRR070372.1334	0	X	1174	60	61S75M1D16M1D10M1D9M46D55M1D7M1D6M1D10M1D5M2D25M2D11M2D16M	*	0	0	GTTCTACAACGTCCAGATCGGAAGCAAGTTCGAAGGAGAGGTCTTGATAACACCAAGCACGAGGTTACCTTCACTCTCAAGGACGGACACTTGTTCGAACATCACAAGCCACTTGAAGAGGGAGAATCCAAGGAAGAACCTATGAGTATTACTTGATGGAGATTTCTTATTCAGAAGATGAGCTTCAACAATATCGAAGGCCGCAGATTCTACAAGAGACTCCCATAAAGTTTAACTTATCTATTGAAATTTCTAAATTGCAATTCAATTTCATTTCCGAAAAATAAATTATTTCAAGCAATCTTC	IIIIIII???GIIIIIIIICCCCIIIIIIIIIIIIIIB?555?IIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIIICCCIIIICCBBBBB?;;6688EEEEGEIE:::?<<?CGEDDEEBBCC44:EIIIEEEIIIIGHGGGHGGIIIIIIEGCCCCEEIIIIICC?<?EEEG?AAEB1//5=<52.---,.=../,34A?CB<<<4777/222;;@@DDDAEEEEEBGA74496:,,,,,.,,,--.221326::477<BEE	NM:i:61	MD:Z:75^A16^T10^T9^CAGGTAAAATTTGGTCAATCTATTTGACATACATTTTTGTTAATTA55^A7^A6^T10^T5^TA2T22^AA11^TA0T15	AS:i:110	XS:i:20	SA:Z:X,1053,+,7M3D32M1D26M241S,60,4;
SRR070372.1334	2048	X	1053	60	7M3D32M1D26M241H	*	0	0	GTTCTACAACGTCCAGATCGGAAGCAAGTTCGAAGGAGAGGTCTTGATAACACCAAGCACGAGGT	IIIIIII???GIIIIIIIICCCCIIIIIIIIIIIIIIB?555?IIIIIIIIIIIIIIIIIIIIII	NM:i:4	MD:Z:7^AAA32^G26	AS:i:51	XS:i:0	SA:Z:X,1174,+,61S75M1D16M1D10M1D9M46D55M1D7M1D6M1D10M1D5M2D25M2D11M2D16M,60,61;
|]

case_count_two = do
    c <- testNGLessIO $ do
        let opts = defCountOpts { optFeatures = ["gene"] }
        gff <- asTempFile short1 "gff"
        samf <- asTempFile short_sam "sam"
        ann <- loadAnnotator (AnnotateGFF gff) opts
        cfp <- performCount (File samf) "testing" ann opts
        liftIO (readCountFile cfp)
    c @?= M.fromList [("WBGene00002254", 2)]

sam1 = [here|
@SQ	SN:X	LN:10000
Read1	0	X	200	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
|]
sam1neg = [here|
@SQ	SN:X	LN:10000
Read1	16	X	200	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
|]

samPartial = [here|
@SQ	SN:X	LN:10000
Read1	0	X	80	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
|]

samAmbiguous = [here|
@SQ	SN:X	LN:10000
Read1	0	X	280	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
|]

samAmbiguous2 = [here|
@SQ	SN:X	LN:10000
Read1	0	X	280	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
Read2	0	X	100	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
|]

samAmbiguous3 = [here|
@SQ	SN:X	LN:10000
Ambiguous	0	X	280	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
Match100.1	0	X	100	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
Match100.2	0	X	100	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
Match100.3	0	X	100	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
Match300.1	0	X	420	60	35M	*	0	0	CAATTGGAGTGCATCAAGTGGTGCGATAAGGTCCTA	22777449446411100.,,,1.11..0000,,,,	NM:i:51		AS:i:430	XS:i:19
|]

gff1 = [here|
X	protein_coding	gene	100	400	.	+	.	gene_id "Gene100"; gene_name "gene1"; gene_source "ensembl"; gene_biotype "protein_coding";
X	protein_coding	gene	300	600	.	+	.	gene_id "Gene300"; gene_name "gene2"; gene_source "ensembl"; gene_biotype "protein_coding";
|]

case_gff_match = do
    c <- testNGLessIO $ runSamGffAnnotation sam1 gff1 defCountOpts { optFeatures = ["gene"] }
    c @?= M.fromList [("Gene100", 1), ("Gene300", 0)]

case_gff_strand_check = do
    c <- testNGLessIO $ runSamGffAnnotation sam1 gff1 defCountOpts { optFeatures = ["gene"], optStrandMode = SMSense }
    c @?= M.fromList [("Gene100", 1), ("Gene300", 0)]

case_gff_strand_check_negstrand = do
    c <- testNGLessIO $ runSamGffAnnotation sam1neg gff1 defCountOpts { optFeatures = ["gene"], optStrandMode = SMSense }
    c @?= M.fromList [("Gene100", 0), ("Gene300", 0)]

case_gff_feature_mismatch = do
    c <- testNGLessIO $ runSamGffAnnotation sam1 gff1 defCountOpts { optFeatures  = ["CDS"] }
    c @?= M.fromList []

case_gff_feature_partial = do
    c <- testNGLessIO $ runSamGffAnnotation samPartial gff1 defCountOpts { optFeatures  = ["gene"] }
    c @?= M.fromList [("Gene100", 1), ("Gene300", 0)]

case_gff_feature_partial_intersect = do
    c <- testNGLessIO $ runSamGffAnnotation samPartial gff1 defCountOpts { optFeatures  = ["gene"], optIntersectMode = annotationRule IntersectUnion }
    c @?= M.fromList [("Gene100", 1), ("Gene300", 0)]

case_gff_feature_ambiguous = do
    c <- testNGLessIO $ runSamGffAnnotation samAmbiguous gff1 defCountOpts { optFeatures  = ["gene"], optMMMethod = MMCountAll }
    c @?= M.fromList [("Gene100", 1), ("Gene300", 1)]

case_gff_feature_ambiguous_discard = do
    c <- testNGLessIO $ runSamGffAnnotation samAmbiguous gff1 defCountOpts { optFeatures  = ["gene"] }
    c @?= M.fromList [("Gene100", 0), ("Gene300", 0)]

case_gff_1OverN = do
    c <- testNGLessIO $ runSamGffAnnotation samAmbiguous gff1 defCountOpts { optFeatures  = ["gene"], optMMMethod = MM1OverN }
    c @?= M.fromList [("Gene100", 0.5), ("Gene300", 0.5)]

case_gff_dist1_fallback = do
    c <- testNGLessIO $ runSamGffAnnotation samAmbiguous gff1 defCountOpts { optFeatures  = ["gene"], optMMMethod = MMDist1 }
    c @?= M.fromList [("Gene100", 0.5), ("Gene300", 0.5)]

case_gff_dist1_dist = do
    c <- testNGLessIO $ runSamGffAnnotation samAmbiguous2 gff1 defCountOpts { optFeatures  = ["gene"], optMMMethod = MMDist1 }
    c @?= M.fromList [("Gene100", 2.0), ("Gene300", 0.0)]

case_gff_dist1_dist1_to_4 = do
    c <- testNGLessIO $ runSamGffAnnotation samAmbiguous3 gff1 defCountOpts { optFeatures  = ["gene"], optMMMethod = MMDist1 }
    c @?= M.fromList [("Gene100", 3.75), ("Gene300", 1.25)]

case_rsv_1 = do
    v <- RSV.newRefSeqInfoVector
    RSV.insert v (B.take 5 "hello SLICE") 1.0 
    RSV.insert v "world" 2.0 
    RSV.sort v
    fv <- RSV.unsafeFreeze v
    RSV.length fv @?= 2
    RSV.lookup fv "hello" @?= Just 0
    RSV.lookup fv (B.take 5 "hello SLICE2") @?= Just 0
    RSV.retrieveSize fv 0 @?= 1.0
    RSV.lookup fv "world" @?= Just 1


simple_map = [here|
#gene	cog	ko	module
gene1	NOG318324	NA	NA	NA
gene2	COG2813	K00564	NA
|]

case_load_map = do
    [GeneMapAnnotator tag nmap names] <- testNGLessIO $ do
        map_fp <- asTempFile simple_map "map"
        loadFunctionalMap map_fp ["ko"]
    let Just [ix] = M.lookup "gene1" nmap
    RSV.retrieveName names ix @?= "NA"
    tag @?= B.empty