version = "0.26.0"
container_url = "lbmc/kb:${version}"

params.index_fasta = ""
params.index_fasta_out = ""

workflow index_fasta {
  take:
    fasta
    cdna
    gtf
    transcript_to_gene

  main:
    transcript_to_gene 
      .ifEmpty(["NO T2G", ""])
      .set{ transcript_to_gene_optional }
    tr2g(gtf, transcript_to_gene_optional)
    index_default(fasta, cdna, gtf, tr2g.out.t2g)

  emit:
    index = index_default.out.index
    t2g = index_default.out.t2g
    report = index_default.out.report
}

process tr2g {
  // create transcript to gene table from gtf if no transcript to gene file is provided
  container = "${container_url}"
  label "big_mem_mono_cpus"
  tag "$file_id"
  if (params.index_fasta_out != "") {
    publishDir "results/${params.index_fasta_out}", mode: 'copy'
  }

  input:
    tuple val(file_id), path(gtf)
    tuple val(t2g_id), file(transcript_to_gene)

  output:
    tuple val(file_id), path("t2g.txt"), emit: t2g

  script:

  if (t2g_id == "NO T2G") 
  """
  t2g.py --gtf ${gtf}
  """
  else
  """
  mv ${transcript_to_gene} t2g.txt
  """
}

process index_default {
  container = "${container_url}"
  label "big_mem_mono_cpus"
  tag "$file_id"
  if (params.index_fasta_out != "") {
    publishDir "results/${params.index_fasta_out}", mode: 'copy'
  }

  input:
    tuple val(file_id), path(fasta)
    tuple val(cdna_id), path(cdna)
    tuple val(gtf_id), path(gtf)
    tuple val(t2g_id), path(transcript_to_gene)

  output:
    tuple val(file_id), path("*.idx"), emit: index
    tuple val(t2g_id), path("${transcript_to_gene}"), emit: t2g
    tuple val(file_id), path("*_report.txt"), emit: report

  script:
"""
kb ref \
  -i ${fasta.simpleName}.idx \
  -g ${transcript_to_gene} \
  ${params.index_fasta} \
  -f1 ${cdna} ${fasta} ${gtf} > ${fasta.simpleName}_kb_index_report.txt
"""
}


include { split } from "./../flexi_splitter/main.nf"

params.kb_protocol = "10x_v3"
params.count = ""
params.count_out = ""
workflow count {
  take:
    index
    fastq
    transcript_to_gene
    whitelist
    config

  main:
  whitelist
    .ifEmpty(["NO WHITELIST", 0])
    .set{ whitelist_optional }
  switch(params.kb_protocol) {
    case "marsseq":
      split(fastq, config)
      kb_marseq(index.collect(), split.out.fastq.view(), transcript_to_gene, whitelist_optional)
      kb_marseq.out.counts.set{res_counts}
      kb_marseq.out.report.set{res_report}
    break;
    default:
      kb_default(index, fastq, transcript_to_gene, whitelist_optional)
      kb_default.out.counts.set{res_counts}
      kb_default.out.report.set{res_report}
    break;
  }

  emit:
    counts = res_counts
    report = res_report
}

process kb_default {
  container = "${container_url}"
  label "big_mem_multi_cpus"
  tag "$file_prefix"
  if (params.count_out != "") {
    publishDir "results/${params.count_out}", mode: 'copy'
  }

  input:
  tuple val(index_id), path(index)
  tuple val(file_id), path(reads)
  tuple val(t2g_id), path(transcript_to_gene)
  tuple val(whitelist_id), path(whitelist)

  output:
  tuple val(file_id), path("${file_prefix}"), emit: counts
  tuple val(file_id), path("*_report.txt"), emit: report

  script:
  def kb_memory = "${task.memory}" - ~/GB/
  if (file_id instanceof List){
    file_prefix = file_id[0]
  } else {
    file_prefix = file_id
  }
  def whitelist_param = ""
  if (whitelist_id != "NO WHITELIST"){
    whitelist_param = "-w ${whitelist}"
  }

  if (reads.size() == 2)
  """
  mkdir ${file_prefix}
  kb count  -t ${task.cpus} \
    -m ${kb_memory} \
    -i ${index} \
    -g ${transcript_to_gene} \
    ${whitelist_param} \
    -x 10XV3
    ${params.count} \
    -o ${file_prefix} \
    ${reads[0]} ${reads[1]} &> ${file_prefix}_kb_mapping_report.txt
  """
}

process kb_marseq {
  // With the MARS-Seq protocol, we have:
  // on the read 1: 4 nt of bc plate
  // on the read 2: 6 nt of bc cell, and 8 nt of UMI
  // this process expect that the bc plate is removed from the read 1
  container = "${container_url}"
  label "big_mem_multi_cpus"
  tag "$file_prefix"
  if (params.count_out != "") {
    publishDir "results/${params.count_out}", mode: 'copy'
  }

  input:
  tuple val(index_id), path(index)
  tuple val(file_id), path(reads)
  tuple val(t2g_id), path(transcript_to_gene)
  tuple val(whitelist_id), path(whitelist)

  output:
  tuple val(file_id), path("${file_prefix}"), emit: counts
  tuple val(file_id), path("*_report.txt"), emit: report

  script:
  def kb_memory = "${task.memory}" - ~/GB/
  if (file_id instanceof List){
    file_prefix = file_id[0]
  } else {
    file_prefix = file_id
  }
  def whitelist_param = ""
  if (whitelist_id != "NO WHITELIST"){
    whitelist_param = "-w ${whitelist}"
  }

  if (reads.size() == 2)
  """
  mkdir ${file_prefix}
  kb count  -t ${task.cpus} \
    -m ${kb_memory} \
    -i ${index} \
    -g ${transcript_to_gene} \
    ${whitelist_param} \
    ${params.count} \
    -x 1,0,6:1,6,14:1,14,0
    -o ${file_prefix} \
    ${reads[0]} ${reads[1]} &> ${file_prefix}_kb_mapping_report.txt
  """
}