-- require deeptools installation: https://github.com/fidelram/deepTools
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE FlexibleInstances #-}

module Bio.Foreign.DeepTools
    ( CorMethod(..)
    , BamCorOpt
    , plot
    , bed
    , binSize
    , method
    , fragLen
    , bamCorrelate

    , BamCoverageOpt
    , Normalization(..)
    , FileFormat(..)
    , centerReads
    , outFileFormat
    , smooth
    , nCores
    , normalization
    , bamCoverage

    , ComputeMatrixOpt
    , ComputeMatrixMode(..)
    , ReferencePoint(..)
    , mode
    , cmd
    , refPoint
    , regionLen
    , upstream
    , downstream
    , scale
    , computeMatrix
    ) where

import Control.Lens (makeFields, (^.))
import qualified Data.ByteString.Char8 as B
import Data.Default.Class (Default(..))
import qualified Data.Text as T
import qualified Data.Text.IO as TIO
import Shelly (shelly, run_, run, liftIO, withTmpDir, toTextIgnore, fromText)

data CorMethod = Spearman
               | Pearson
    deriving (Show, Read)

data BamCorOpt = BamCorOpt
    { bamCorOptPlot :: !(Maybe FilePath)
    , bamCorOptBed :: !(Maybe FilePath)
    , bamCorOptBinSize :: !Int
    , bamCorOptMethod :: !CorMethod
    , bamCorOptFragLen ::  !Int
    } deriving (Show, Read)

makeFields ''BamCorOpt

instance Default BamCorOpt where
    def = BamCorOpt
        { bamCorOptPlot = Nothing
        , bamCorOptBed = Nothing
        , bamCorOptBinSize = 10000
        , bamCorOptMethod = Spearman
        , bamCorOptFragLen = 100
        }

bamCorrelate :: [(String, FilePath)] -> BamCorOpt -> IO [[Double]]
bamCorrelate dat config = shelly $ withTmpDir $ \dir -> do
    let tmpDir = T.unpack $ toTextIgnore dir
        plotoutput = case config^.plot of
                Nothing -> tmpDir ++ "/tmp.png"
                Just x -> x

    run_ "bamCorrelate" $ map T.pack $
        mode ++ ("--bamfiles" : fls) ++ cor ++ ["-o", plotoutput] ++
        ("--labels" : labels) ++ ["--binSize", show $ config^.binSize] ++
        ["--fragmentLength", show $ config^.fragLen, "--outFileCorMatrix", tmpDir ++ "/tmp.tsv"]

    r <- liftIO $ B.readFile $ tmpDir ++ "/tmp.tsv"
    return $ map (map (read . B.unpack) . tail . B.split '\t') $ tail $ B.lines r
  where
    (labels, fls) = unzip dat
    mode = case config^.bed of
        Nothing -> ["bins"]
        Just x -> ["BED-file", "--BED", x]
    cor = case config^.method of
        Spearman -> ["--corMethod", "spearman"]
        Pearson -> ["--corMethod", "pearson"]
{-# INLINE bamCorrelate #-}

data BamCoverageOpt = BamCoverageOpt
    { bamCoverageOptBinSize :: !Int
    , bamCoverageOptFragLen :: !Int
    , bamCoverageOptCenterReads :: !Bool
    , bamCoverageOptOutFileFormat :: !FileFormat
    , bamCoverageOptSmooth :: !(Maybe Int)
    , bamCoverageOptNCores :: !Int
    , bamCoverageOptNormalization :: !Normalization
    } deriving (Show, Read)

data FileFormat = Bedgraph
                | Bigwig
    deriving (Read, Show)

data Normalization = Scaling Double
                   | RPKM
                   | RPGC Int  -- ^ Report read coverage normalized to 1x sequencing depth
                   | None
    deriving (Read, Show)

makeFields ''BamCoverageOpt

instance Default BamCoverageOpt where
    def = BamCoverageOpt
        { bamCoverageOptBinSize = 50
        , bamCoverageOptFragLen = 200
        , bamCoverageOptCenterReads = False
        , bamCoverageOptOutFileFormat = Bedgraph
        , bamCoverageOptSmooth = Nothing
        , bamCoverageOptNCores = 1
        , bamCoverageOptNormalization = None
        }

bamCoverage :: FilePath -> FilePath -> FilePath -> BamCoverageOpt -> IO ()
bamCoverage bam bai output opt = shelly $ do
    let params = [ "--bam", bam, "--outFileName", output, "--bamIndex", bai
                 , "--binSize", show $ opt^.binSize, "--fragmentLength"
                 , show $ opt^.fragLen, "--numberOfProcessors"
                 , show $ opt^.nCores ] ++ format ++ center ++ smoothLength ++ normal
    run_ "bamCoverage" $ map T.pack params
  where
    format = case opt^.outFileFormat of
        Bigwig -> ["--outFileFormat", "bigwig"]
        Bedgraph -> ["--outFileFormat", "bedgraph"]
    center | opt^.centerReads = ["--centerReads"]
           | otherwise = []
    smoothLength = case opt^.smooth of
        Nothing -> []
        Just x -> ["--smoothLength", show x]
    normal = case opt^.normalization of
        Scaling x -> ["--scaleFactor", show x]
        RPKM -> ["--normalizeUsingRPKM"]
        RPGC x -> ["--normalizeTo1x", show x]
        None -> []


data ComputeMatrixMode = Reference
                       | Scale
    deriving (Read, Show)

data ReferencePoint = Start
                    | End
                    | Center
    deriving (Read, Show)

data ComputeMatrixOpt = ComputeMatrixOpt
    { computeMatrixOptMode :: !ComputeMatrixMode
    , computeMatrixOptBinSize :: !Int
    , computeMatrixOptCmd :: !FilePath
    , computeMatrixOptRegionLen :: !Int
    , computeMatrixOptRefPoint :: !ReferencePoint
    , computeMatrixOptUpstream :: !Int
    , computeMatrixOptDownstream :: !Int
    , computeMatrixOptScale :: !(Maybe Double)
    } deriving (Show, Read)

makeFields ''ComputeMatrixOpt

instance Default ComputeMatrixOpt where
    def = ComputeMatrixOpt
        { computeMatrixOptMode = Scale
        , computeMatrixOptBinSize = 10
        , computeMatrixOptCmd = "computeMatrix"
        , computeMatrixOptRegionLen = 1000
        , computeMatrixOptRefPoint = Start
        , computeMatrixOptUpstream = 500
        , computeMatrixOptDownstream = 1500
        , computeMatrixOptScale = Nothing
        }

computeMatrix :: FilePath
              -> FilePath
              -> FilePath
              -> ComputeMatrixOpt
              -> IO ()
computeMatrix regions bw output opt = shelly $
    run_ (fromText $ T.pack $ opt^.cmd) $ map T.pack $ md ++
        [ "--regionsFileName", regions, "--scoreFileName", bw, "--binSize"
        , show $ opt^.binSize,  "--sortRegions", "no", "--outFileName", output ++ ".gz"
        , "--missingDataAsZero"
        , "--beforeRegionStartLength", show $ opt^.upstream
        , "--afterRegionStartLength", show $ opt^.downstream
        ] ++ sc
  where
    rp = case opt^.refPoint of
        Start -> "TSS"
        End -> "TES"
        Center -> "center"
    md = case opt^.mode of
        Reference -> ["reference-point", "--referencePoint", rp]
        Scale -> ["scale-regions", "--regionBodyLength", show $ opt^.regionLen]
    sc = case opt^.scale of
        Just x -> ["--scale", show x]
        Nothing -> []
