rule assemblGtf:
  input:
    config['bamDir'] + "{sample}.bam"
  output:
    fullgtf = config["outputDir"] + "denovo/{sample}.scallop.gtf",
    posGtf = config["outputDir"] + "denovo/{sample}.scallop.pos.gtf",
    negGtf = config["outputDir"] + "denovo/{sample}.scallop.neg.gtf"
  conda: 
    "../envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      scallop --min_transcript_coverage 10 --min_num_hits_in_bundle 100 --min_splice_bundary_hits 10 \
        --library_type first  -i {input}  -o {output.fullgtf} --verbose 0  > /dev/null
      mawk 'BEGIN{{FS="\\t";OFS="\\t"}} {{if($7=="-") print; }}' {output.fullgtf} > {output.negGtf}
      mawk 'BEGIN{{FS="\\t";OFS="\\t"}} {{if($7=="+") print; }}' {output.fullgtf} > {output.posGtf}
    """

rule mergeDevoGtf:
  input:
    inGtfs = expand(config["outputDir"] + "denovo/{sample}.scallop.{{strand}}.gtf", sample = IDS),
    refGtf = config["outputDir"] + "reference.{strand}.gtf"
  output:
    mergedDenovoGtf = config["outputDir"] + "denovo/merged.{strand}.gtf"
  conda: 
    "../envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      stringtie --merge -g 250 {input.refGtf} {input.inGtfs} -l DENO | \
        grep -v '^MT' > {output}
    """

rule annoDevoGtf:
  input:
    inGtf = rules.mergeDevoGtf.output.mergedDenovoGtf,
    refGtf = config["outputDir"] + "reference.{strand}.gtf"
  output:
    config["outputDir"] + "denovo/anno.{strand}.annotated.gtf"
  params:
    outprefix = lambda wildcards: config["outputDir"] + "denovo/anno." + wildcards.strand
  conda: 
    "../envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      gffcompare -o {params.outprefix} -r {input.refGtf} {input.inGtf}
    """

from snakemake.utils import R
from itertools import repeat
from os.path import join
import os
import shutil
import pandas
import jsonschema
import math

# this container defines the underlying OS for each job when using the workflow
# with --use-conda --use-singularity
singularity: "docker://continuumio/miniconda3"

##### load config and sample sheets #####

configfile: "config/config.yaml"
validate(config, schema="../schemas/config.schema.yaml")

samples = pd.read_csv(config["samples"], sep="\t").set_index("sample", drop=False)
samples.index.names = ["sample_id"]
validate(samples, schema="../schemas/samples.schema.yaml")

rule mapFastq:
  input:
    R1 = config["fastqDir"] + "{fastq}_R1.fastq.gz",
    R2 = config["fastqDir"] + "{fastq}_R2.fastq.gz"
  output:
    countMtx = config["outputDir"] + "starSolo/{fastq}/{fastq}.Solo.out/Gene/raw/matrix.mtx",
    sortBam = config["outputDir"] + "starSolo/{fastq}/{fastq}.Aligned.sortedByCoord.out.bam"
  params:
    outputPrefix = lambda wildcards: config["outputDir"] + "starSolo/" + wildcards.fastq + "/" + wildcards.fastq + ".",
    whitelist = config["whitelist"],soloType = config["STARArgs"].get("soloType","Droplet"),
    barcodeLen = 0, CBLen = config["CBLen"], UMIStart = config["UMIStart"], UMILen = config["UMILen"],
    genome = config["StarsoloGenome"],
    readFilesCommand = config["STARArgs"].get("readFilesCommand","zcat"),
    outSAMtype = config["STARArgs"].get("outSAMtype","BAM SortedByCoordinate"),
    soloStrand = config["STARArgs"].get("soloStrand","Forward"),
    winAnchorMultimapNmax = config["STARArgs"].get("winAnchorMultimapNmax",2000),
    outFilterMultimapNmax = config["STARArgs"].get("outFilterMultimapNmax",2000),
    outSAMprimaryFlag = config["STARArgs"].get("outSAMprimaryFlag","AllBestScore"),
    outSAMmultNmax = config["STARArgs"].get("outSAMmultNmax",1),
    limitBAMsortRAM = config["STARArgs"].get("limitBAMsortRAM",40000000000),
    limitOutSJoneRead = config["STARArgs"].get("limitOutSJoneRead",10000),
    limitOutSJcollapsed = config["STARArgs"].get("limitOutSJcollapsed",3000000),
    limitIObufferSize = config["STARArgs"].get("limitIObufferSize",300000000),
    outSAMattributes = config["STARArgs"].get("outSAMattributes"),
    additionalArguments = config["STARArgs"].get("additionalArguments", ""),
    starOptions = starsoloParameters
  threads: math.ceil(config["STARArgs"].get("threads",16) * scaleDownThreads)
  conda:
    "../envs/generateH5.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * config["STARArgs"].get("memory",40000)
  shell:
    """
      STAR \
        --genomeDir {params.genome} \
        --runThreadN {threads} \
        --readFilesIn {input.R2} {input.R1}  \
        --readFilesCommand {params.readFilesCommand} \
        --outSAMtype {params.outSAMtype} \
	      --winAnchorMultimapNmax {params.winAnchorMultimapNmax} \
	      --outFilterMultimapNmax {params.outFilterMultimapNmax} \
	      --outSAMprimaryFlag {params.outSAMprimaryFlag} \
	      --outSAMmultNmax {params.outSAMmultNmax} \
	      --limitBAMsortRAM {params.limitBAMsortRAM} \
	      --limitOutSJoneRead {params.limitOutSJoneRead} \
        --limitOutSJcollapsed {params.limitOutSJcollapsed} \
        --limitIObufferSize {params.limitIObufferSize} \
        --outFileNamePrefix {params.outputPrefix} \
        --outSAMattributes {params.outSAMattributes} \
	      {starsoloParameters} \
        {params.additionalArguments}
    """

rule featureCount:
  input:
    bamfile = rules.mapFastq.output.sortBam,
    gtf = config["featureCountGTF"]
  output:
    temp(config["outputDir"] + "featureCount/" + config["countBy"] + "/{fastq}.Aligned.sortedByCoord.out.bam.featureCounts.bam")
  params:
    outDir = config["outputDir"] + "featureCount/" + config["countBy"] + "/", strand = config["featureCountArgs"].get("strand",1)
  threads: math.ceil(config["featureCountArgs"].get("threads",16) * scaleDownThreads)
  conda:
    "../envs/generateH5.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * config["featureCountArgs"].get("memory",10000)
  shell:
    """
      featureCounts --fraction -M -s {params.strand} -T {threads} -t exon -g gene_name -a {input.gtf} -o {params.outDir}{wildcards.fastq}.counts.txt {input.bamfile} -R BAM --Rpath {params.outDir}
    """


rule featureCountToDT:
  input:
    rules.featureCount.output
  output:
    config["outputDir"] + "featureCount/" + config["countBy"] + "/" + "{fastq}.tsv.gz"
  params:
    countScript = srcdir("../scripts/countBamFromSTARSolo.sh"),
    outDir = config["outputDir"] + "featureCount/",
    numHit = config['numHit']
  threads: math.ceil(config["featureCountArgs"].get("threads",8) * scaleDownThreads)
  conda:
    "../envs/generateH5.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * config["featureCountArgs"].get("memory",4000)
  shell:
    """
      samtools view {input} | NUMHIT={params.numHit} bash {params.countScript} | pigz -p {threads} -c > {output}
    """

rule DTToH5:
  input:
    rules.featureCountToDT.output
  output:
    countMtx = config["outputDir"] + "featureCount/" +  config["countBy"] + "/{fastq}.h5",
    molInfo = config["outputDir"] + "featureCount/" +  config["countBy"] + "/{fastq}.molInfo.h5"
  params:
    writeToh5 = srcdir("../scripts/writeHDF5.R"), whitelist = config["whitelist"]
  threads: math.ceil(config["DTToH5Args"].get("threads",16) * scaleDownThreads)
  conda:
    "../envs/generateH5.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * config["DTToH5Args"].get("memory",4000)
  shell:
    """
      Rscript {params.writeToh5} {input} {params.whitelist} {output.countMtx} {threads} {wildcards.fastq}
    """

rule extractUTR:
  input: 
    config["goodGTF"]
  output: 
    temp(config["outputDir"] + "ortholog/" + "threeUTR.gtf")
  conda: 
    "../envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      mawk 'BEGIN{{FS="\\t";OFS="\\t"}} {{if($3=="three_prime_utr") {{$3 = "exon"; print "chr"$0;}} }}' {input} > {output}
    """

rule mapUTR:
  input:
    utrgtf = rules.extractUTR.output,
    liftChain = config["liftChain"]
  output:
    tmpMapGtf = temp(config["outputDir"] + "ortholog/" + "mapUTR.tmp.gtf"),
    mapGtfPos = config["outputDir"] + "ortholog/" + "mapUTR.pos.gtf",
    mapGtfNeg = config["outputDir"] + "ortholog/" + "mapUTR.neg.gtf"
  params:
    minMatch = config['liftMinMatch'],
    rtracklayer = srcdir("../scripts/rtracklayer.R")
  threads: 16
  conda: 
    "../envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 30000
  shell:
    """
      # liftOver -gff -minMatch={params.minMatch} {input.utrgtf} {input.liftChain} {output.tmpMapGtf} unmap.bed
      Rscript {params.rtracklayer} {input.utrgtf} {input.liftChain} {output.tmpMapGtf}
      # fix short exons
      mawk 'BEGIN{{FS="\\t";OFS="\\t"}}{{if($5-$4 > 50 && $7 == "+") {{if($1 ~ /^chr/) $1 = substr($1, 4, length($1)); print $0; $3 = "transcript"; print $0;}}}}' {output.tmpMapGtf} >  {output.mapGtfPos}
      mawk 'BEGIN{{FS="\\t";OFS="\\t"}}{{if($5-$4 > 50 && $7 == "-") {{if($1 ~ /^chr/) $1 = substr($1, 4, length($1)); print $0; $3 = "transcript"; print $0;}}}}' {output.tmpMapGtf} >  {output.mapGtfNeg}
    """

rule mergeOrthologGtf:
  input:
    config["outputDir"] + "ortholog/" + "mapUTR.{strand}.gtf",
    config["outputDir"] + "reference.{strand}.gtf"
  output:
    mergedGtf = config["outputDir"] + "ortholog/" + "merged.{strand}.gtf"
  conda: 
    "../envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      stringtie -f 0 -T 0 -F 0 -l ORTH --merge -g 250 -o {output.mergedGtf} {input}
    """

rule annoOrthologGtf:
  input:
    mergedGtf = rules.mergeOrthologGtf.output.mergedGtf,
    refGtf = config["outputDir"] + "reference.{strand}.gtf"
  output:
    annoGtf = config["outputDir"] + "ortholog/" + "anno.{strand}.annotated.gtf"
  params:
    outprefix = lambda wildcards: config["outputDir"] + "ortholog/" + "anno." + wildcards.strand
  conda: 
    "../envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      gffcompare -o {params.outprefix} -r {input.refGtf} {input.mergedGtf} -K
    """

from snakemake.utils import R
from itertools import repeat
from os.path import join
import os
import shutil
import pandas
import jsonschema
import math

# report: "report/workflow.rst"

# this container defines the underlying OS for each job when using the workflow
# with --use-conda --use-singularity
# singularity: "docker://512303535801.dkr.ecr.eu-west-1.amazonaws.com/scrnax_conda_env:v1.1.1"

container: "docker://continuumio/miniconda3:4.7.12"

##### load config and sample sheets #####


# validate(config, schema="../schemas/config.schema.yaml")

# samples = pd.read_csv(config["samples"], sep="\t").set_index("sample", drop=False)
# samples.index.names = ["sample_id"]
# validate(samples, schema="../schemas/samples.schema.yaml")

# Globals ---------------------------------------------------------------------

scaleDownThreads = config.get("scaleDownThreads", 1)

# Selftest -------------------------------------------------------------------
if config.get('example', "") == "true":
  configfile: srcdir("../config/conf.selftest.json")
else:
  configfile: srcdir("../config/conf.default.json")

# Initial optional parameters -------------------------------------------------
# optionalKeys = {
#   "numHit": 10000,
#   "featureCountArgs": {},
#   "featureCountToDTArgs": {},
#   "STARArgs": {},
#   "DTToH5Args": {}
# }
# for tmpKey, tmpValue in optionalKeys.items():
#   if tmpKey not in config:
#     config[tmpKey] = tmpValue

# Prepare outputs -------------------------------------------------------------
resultFiles = [config["outputDir"] + "ortholog/fixed.gtf"]

if "bamDir" in config:
  if config['bamDir'][len(config['bamDir'])-1] != "/":
    config['bamDir'] = config['bamDir'] + "/"
  IDS, = glob_wildcards(config['bamDir'] + "{sample}.bam")
  resultFiles = [
    resultFiles,
    config["outputDir"] + "denovo/fixed.gtf",
    config["outputDir"] + "combined.fixed.gtf"
  ]
else:
  IDS = ""
  config["bamDir"] = ""

# How the fastqs shall be counted, i.e., by original GTF, or from denovo assembled GTF, or from ortholog GTF, or combined GTF.

if "countBy" in config:
  if config["countBy"] == "combined":
    config["featureCountGTF"] = config["outputDir"] + "combined.fixed.gtf"
  else:
    config["featureCountGTF"] = config["outputDir"] + config["countBy"] + "/fixed.gtf"
elif config["bamDir"] != "":
    config["countBy"] = "combined"
    config["featureCountGTF"] = config["outputDir"] + "combined.fixed.gtf"
else:
  config["countBy"] = "ortholog"
  config["featureCountGTF"] = config["outputDir"] + "ortholog/fixed.gtf"

# define the parameters for STARSolo
# Generate count matrix if fastq files are provided.

if "fastqDir" in config:
  if config['fastqDir'][len(config['fastqDir'])-1] != "/":
    config['fastqDir'] = config['fastqDir'] + "/"
  FASTQS, = glob_wildcards(config['fastqDir'] + "{fastq}_R1.fastq.gz")
  resultFiles = resultFiles + [config["outputDir"] + "featureCount/" + config["countBy"] + "/" + str(i) + ".h5" for i in FASTQS]
  jsS = open(srcdir("schemas/jsonSchemaStarSolo.json")).read()
  # And validate the json config file
else:
  jsS = open(srcdir("schemas/jsonSchemaWithoutStarSolo.json")).read()
  config["fastqDir"] = ""
# validate json config file
jsSchema = json.loads(jsS)
try:
  jsonschema.validate(config, jsSchema)
except jsonschema.exceptions.ValidationError as e:
  print("well-formed but invalid JSON:", e)
  sys.exit()
except json.decoder.JSONDecodeError as e:
  print("poorly-formed text, not JSON:", e)
  sys.exit()

# set default parameters for STARSolo
starsoloParameters = '''  --soloType {solotype} \
        --soloCBwhitelist {whitelist} \
        --soloCBlen {CBLen} \
        --soloUMIstart {UMIStart} \
        --soloUMIlen {UMILen} \
        --soloBarcodeReadLength 0 \
        --soloStrand {soloStrand} \
	      --soloFeatures {soloFeatures} \
    '''.format(solotype = config["STARArgs"].get("soloType","CB_UMI_Simple"), whitelist = config['whitelist'], CBLen = config['CBLen'], UMIStart = config["UMIStart"], UMILen = config["UMILen"], soloStrand = config["STARArgs"].get("soloStrand","Forward"), soloFeatures = config["STARArgs"].get("soloFeatures","Gene GeneFull SJ Velocyto"))

config["STARArgs"]["outSAMattributes"] = "NH HI nM AS CR UR CB UB GX GN sS sQ sM"

# print(resultFiles)
rule all:
	input: 
	  resultFiles

include: "rules/prepareOriginalGTF.py"
include: "rules/assembleGTF.py"
include: "rules/orthologGTF.py"
include: "rules/countFastq.py"


rule fixOriginGtf:
  input:
    config['refGTF']
  output:
    fixedGtf = config["outputDir"] + "origin/" + "fixed.gtf",
  params:
    fixGTFRscript = srcdir("scripts/fixGffcompareGTF.R")
  conda: 
    "envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      Rscript {params.fixGTFRscript} {input} {output.fixedGtf}
    """

rule fixOrthologGtf:
  input:
    annoGtf = rules.annoOrthologGtf.output.annoGtf,
    mtGtf = config["outputDir"] + "mt.{strand}.gtf"
  output:
    fixedGtf = config["outputDir"] + "ortholog/" + "fixed.{strand}.gtf",
    tmpGtf = temp(config["outputDir"] + "ortholog/anno.{strand}.combined.fixed.tmp.gtf")
  params:
    fixGTFRscript = srcdir("scripts/fixGffcompareGTF.R")
  conda: 
    "envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      cat {input.annoGtf} > {output.tmpGtf}
      cat {input.mtGtf} >> {output.tmpGtf}
      Rscript {params.fixGTFRscript} {output.tmpGtf} {output.fixedGtf}
    """

rule mergeStrandedOrthologGtf:
  input:
    config["outputDir"] + "ortholog/" + "fixed.pos.gtf",
    config["outputDir"] + "ortholog/" + "fixed.neg.gtf"
  output:
    config["outputDir"] + "ortholog/" + "fixed.gtf"
  conda: 
    "envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      cat {input} > {output}
    """

rule fixDenovoGtf:
  input:
    annoGtf = rules.annoDevoGtf.output,
    mtGtf = config["outputDir"] + "mt.{strand}.gtf"
  output:
    fixedGtf = config["outputDir"] + "denovo/fixed.{strand}.gtf",
    tmpGtf = temp(config["outputDir"] + "denovo/anno.{strand}.combined.fixed.tmp.gtf")
  params:
    fixGTFRscript = srcdir("scripts/fixGffcompareGTF.R")
  conda: 
    "envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      cat {input.annoGtf} > {output.tmpGtf}
      cat {input.mtGtf} >> {output.tmpGtf}
      Rscript {params.fixGTFRscript} {output.tmpGtf} {output.fixedGtf}
    """

rule mergeStrandedDenovoGtf:
  input:
    config["outputDir"] + "denovo/" + "fixed.pos.gtf",
    config["outputDir"] + "denovo/" + "fixed.neg.gtf"
  output:
    config["outputDir"] + "denovo/" + "fixed.gtf"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads: 1
  shell:
    """
      cat {input} > {output}
    """

rule combineAndFixGtf:
  input:
    #rules.fixDenovoGtf.output.fixedGtf,
    #rules.fixOrthologGtf.output.fixedGtf
    denoAnnoGtf = rules.annoDevoGtf.output,
    orthAnnoGtf = rules.annoOrthologGtf.output,
    refGtf = config["outputDir"] + "reference.{strand}.gtf",
    mtGtf = config["outputDir"] + "mt.{strand}.gtf"
  output:
    combinedMergedGtf = config["outputDir"] + "merged.{strand}.gtf",
    combinedMergedGtfNR = config["outputDir"] + "merged.{strand}.nr.gtf",
    tmpGtf = config["outputDir"] + "tmp.{strand}.annotated.gtf",
    fixedGtf = config["outputDir"] + "combined.fixed.{strand}.gtf"
    # gtflist = config["outputDir"] + "gtflist.txt"
  params:
    fixGTFRscript = srcdir("scripts/fixGffcompareGTF.R"),
    # annoOrthPath = rules.annoOrthologGtf.output,
    # annoDevoPath = rules.annoDevoGtf.output,
    outprefix = lambda wildcards: config["outputDir"] + "tmp." + wildcards.strand,
    mergeGtfDir = config["outputDir"]
  conda: 
    "envs/refineUTR.yaml"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 4000
  threads:
    8
  shell:
    """
      cat {input.denoAnnoGtf} {input.orthAnnoGtf} > {output.combinedMergedGtf}
      stringtie --merge -l COMB -g -1000000000 -i -F 0 -T 0 -f 0 {output.combinedMergedGtf} > {output.combinedMergedGtfNR}
      gffcompare -o {params.outprefix} -r {input.refGtf} {output.combinedMergedGtfNR} -K
      cat {input.mtGtf} >> {output.tmpGtf}
      Rscript {params.fixGTFRscript} {output.tmpGtf} {output.fixedGtf}
   """

rule mergeStrandedCombinedGtf:
  input:
    config["outputDir"] + "combined.fixed.pos.gtf",
    config["outputDir"] + "combined.fixed.neg.gtf"
  output:
    config["outputDir"] + "combined.fixed.gtf"
  resources:
    mem_mb=lambda wildcards, attempt: attempt * 1000
  threads: 1
  shell:
    """
      cat {input} > {output}
    """