{- TooManyCells.Motifs.FindMotif
Gregory W. Schwartz

Collects functions pertaining to finding motifs per cluster.
-}

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TupleSections #-}

module TooManyCells.Motifs.FindMotif
    ( getMotif
    , getMotifGenome
    , getNodes
    ) where

-- Remote
import Control.Monad (mfilter)
import Data.List (foldl')
import Data.Maybe (catMaybes, fromMaybe, isJust)
import System.Directory (getTemporaryDirectory)
import TextShow (showt)
import qualified Control.Foldl as Fold
import qualified Data.ByteString.Lazy.Char8 as B
import qualified Data.Csv as CSV
import qualified Data.Map.Strict as Map
import qualified Data.Text as T
import qualified Data.Text.Read as T
import qualified Text.Printf as TP
import qualified Turtle as TU
import qualified Turtle.Line as TU

-- Local
import TooManyCells.File.Types
import TooManyCells.Motifs.Types

newtype ColMatch = ColMatch T.Text
newtype Match = Match T.Text

-- | Get the temporary directory.
getTmpDir :: IO TU.FilePath
getTmpDir = fmap (TU.fromText . T.pack) getTemporaryDirectory

-- | Read a CSV to a shell.
readCsv :: TU.FilePath -> TU.Shell (Map.Map T.Text T.Text)
readCsv file = do
  contents <- TU.liftIO . B.readFile . T.unpack . TU.format TU.fp $ file

  TU.select
    . either error snd
    . CSV.decodeByName
    $ contents

-- | Cut a shell csv.
csvCut :: T.Text -> TU.Shell (Map.Map T.Text T.Text) -> TU.Shell (Maybe TU.Line)
csvCut c = fmap (fmap TU.unsafeTextToLine . Map.lookup c)

-- | Sort a csv.
csvSort :: T.Text
        -> TU.Shell (Map.Map T.Text T.Text)
        -> TU.Shell (Map.Map T.Text T.Text)
csvSort c = (=<<) TU.select . TU.sortOn (Map.lookup c)

-- | Sort a csv by a numeric column.
csvNumSort :: T.Text
        -> TU.Shell (Map.Map T.Text T.Text)
        -> TU.Shell (Map.Map T.Text T.Text)
csvNumSort c =
  (=<<) TU.select . TU.sortOn (fmap readNum . Map.lookup c)
  where
    readNum "Infinity" = 1 / 0 :: Double
    readNum "-Infinity" = -1 / 0 :: Double
    readNum x = either (error . (<> ": " <> T.unpack x)) fst . T.double $ x

-- | Match a column for a shell csv.
csvMatch :: ColMatch
         -> Match
         -> TU.Shell (Map.Map T.Text T.Text)
         -> TU.Shell (Map.Map T.Text T.Text)
csvMatch (ColMatch c) (Match m) =
  mfilter ((== m) . Map.findWithDefault m c)

-- | Get the nodes from a differential file
getNodes :: DiffFile -> IO [Node]
getNodes (DiffFile df)
      = fmap catMaybes
      . TU.reduce Fold.nub
      . (fmap . fmap) (either error (Node . fst) . T.decimal . TU.lineToText)
      . csvCut "node"
      . readCsv
      $ df

-- | Make a temporary fasta file for input into the motif command.
mkTmpFasta :: TU.FilePath
           -> DiffFile
           -> GenomeFile
           -> TopN
           -> Maybe Node
           -> TU.Shell ()
mkTmpFasta tmpFasta diffFile gf (TopN topN) node = TU.sh $ do
  tmpDir <- TU.liftIO getTmpDir
  tmpBed <- TU.mktempfile tmpDir "motif_input.bed"

  (=<<) (TU.output tmpBed . TU.select)
    . TU.reduce (Fold.lastN topN)
    . TU.sed ("-" *> pure "\t")
    . TU.sed (":" *> pure "\t")
    . fmap (fromMaybe (error "feature column not found."))
    . csvCut "feature"
    . csvNumSort "log2FC"
    . csvMatch (ColMatch "node") (Match $ maybe "0" (showt . unNode) node)
    . readCsv
    . unDiffFile
    $ diffFile

  TU.output tmpFasta
    . TU.inproc "bedtools" [ "getfasta"
                           , "-fi", unGenomeFile gf
                           , "-bed", TU.format TU.fp tmpBed
                           ]
    $ mempty

-- | Make a temporary bed file for input into the genome motif command.
mkTmpBed :: TU.FilePath
           -> DiffFile
           -> TopN
           -> Maybe Node
           -> TU.Shell ()
mkTmpBed tmpBed diffFile (TopN topN) node = TU.sh $ do
  tmpDir <- TU.liftIO getTmpDir

  let formatRow xs =  -- Format recommended by findMotifsGenome.pl
        TU.unsafeTextToLine
          $ foldl' (\acc x -> T.replace x "\t" acc) (TU.lineToText xs) [":", "-"]
         <> "\t"
         <> TU.lineToText xs
         <> "\t\t+"

  (=<<) (TU.output tmpBed . TU.select)
    . TU.reduce (Fold.lastN topN)
    . fmap (formatRow . fromMaybe (error "feature column not found."))
    . csvCut "feature"
    . csvNumSort "log2FC"
    . csvMatch (ColMatch "node") (Match $ maybe "0" (showt . unNode) node)
    . readCsv
    . unDiffFile
    $ diffFile

getMotif :: DiffFile
         -> Maybe BackgroundDiffFile
         -> OutputPath
         -> MotifCommand
         -> GenomeFile
         -> TopN
         -> Maybe Node
         -> IO ()
getMotif diffFile bgDiffFile outPath mc gf topN node = TU.sh $ do
  tmpDir <- TU.liftIO getTmpDir
  tmpFasta <- TU.mktempfile tmpDir "motif_input.fasta"
  mkTmpFasta tmpFasta diffFile gf topN node

  tmpBgFasta <- TU.mktempfile tmpDir "motif_bg_input.fasta"
  maybe
    (return ())
    (\x -> mkTmpFasta tmpBgFasta (DiffFile . unBackgroundDiffFile $ x) gf topN node)
    bgDiffFile

  let cmd = if isJust bgDiffFile
              then
                T.pack $
                  TP.printf
                    (unMotifCommand mc)
                    (TU.format TU.fp tmpFasta)
                    (TU.format TU.fp . unOutputPath $ outPath)
                    (TU.format TU.fp tmpBgFasta)
              else
                T.pack $
                  TP.printf
                    (unMotifCommand mc)
                    (TU.format TU.fp tmpFasta)
                    (TU.format TU.fp . unOutputPath $ outPath)

  TU.stderr . TU.inshell cmd $ mempty

getMotifGenome :: DiffFile
               -> Maybe BackgroundDiffFile
               -> OutputPath
               -> MotifGenomeCommand
               -> GenomeFile
               -> TopN
               -> Maybe Node
               -> IO ()
getMotifGenome diffFile bgDiffFile outPath mc gf topN node = TU.sh $ do
  tmpDir <- TU.liftIO getTmpDir
  tmpBed <- TU.mktempfile tmpDir "motif_input.bed"
  mkTmpBed tmpBed diffFile topN node

  tmpBgBed <- TU.mktempfile tmpDir "motif_bg_input.bed"
  maybe
    (return ())
    (\x -> mkTmpBed tmpBgBed (DiffFile . unBackgroundDiffFile $ x) topN node)
    bgDiffFile

  let cmd = if isJust bgDiffFile
              then
                T.pack $
                  TP.printf
                    (unMotifGenomeCommand mc)
                    (TU.format TU.fp tmpBed)
                    (unGenomeFile gf)
                    (TU.format TU.fp . unOutputPath $ outPath)
                    (TU.format TU.fp tmpBgBed)
              else
                T.pack $
                  TP.printf
                    (unMotifGenomeCommand mc)
                    (TU.format TU.fp tmpBed)
                    (unGenomeFile gf)
                    (TU.format TU.fp . unOutputPath $ outPath)

  TU.stderr . TU.inshell cmd $ mempty