#!/usr/bin/env nextflow

nextflow.enable.dsl = 2

// import modules
include { COPY_CONTAINER_CONFIG; PREPARE_SRA; DIACHROMATIC_TRUNCATE } from './modules/prepare_sra'
include { BOWTIE2_BUILD; ALIGN_READS } from './modules/align_reads'
include { GET_GENOME_LENGHT } from './modules/genome_lenght.nf'
include { RUN_TEG_NESTER; ANNOTATE_FRAGMENTS; SHUFFLE_FRAGMENTS } from './modules/annotate_fragments'
include { INTERSECT_BED; INTERSECT_BED as INTERSECT_BED_SHUFFLE } from './modules/intersect_bed'
include { GENERATE_TABLE; GENERATE_TABLE as GENERATE_TABLE_SHUFFLED } from './modules/generate_table'

include { FORMAT_READS; SEQTK_SEQ; RUN_RE } from './modules/repeatexplorer'
include { BLAST_CONTIGS; ASSIGN_READS_TO_CLUSTERS; MERGE_TABLES; REFORMAT_MATRIX; SAVE_RE_TABLE; RE_HEATMAP_NORM2 } from './modules/re_to_heatmap'
include { MAP_CONTIGS; RE_MODIFY_SAM; BAM_TO_BED; RE_BED_ANNOTATION; BED_TO_GFF; FILTER_GFF } from './modules/annotate_with_re'
include { MAP_CONTIGS as MAP_CONTIGS_PL; BAM_TO_BED as BAM_TO_BED_PL; BED_TO_GFF as BED_TO_GFF_PL; FILTER_GFF as FILTER_GFF_PL } from './modules/annotate_with_re'

include { ANNOTATE_EXONS } from './modules/annotate_exons'
include { CAT_ANNOTATIONS } from './modules/cat_annotations'

include { PREPARE_TABLE; PLOT_CHR; PLOT_BARPLOT; TE_HEATMAP_NORM1; TE_HEATMAP_NORM2; TE_HEATMAP_NORM3 } from './modules/te_to_heatmap'
include { PLOT_CIRCOS } from './modules/plot_circos'


workflow prepare_sra_wf {
    main:
    if (!params.skip_diachromatic) {
        if (!params.skip_convert_fastq) {
            if (workflow.containerEngine == 'singularity' || workflow.containerEngine == 'docker') {
                msg = COPY_CONTAINER_CONFIG()
            }
            else {
                msg = "using local user-settings.mkfg"
            }
            PREPARE_SRA(msg)
            fastq_data = PREPARE_SRA.out.fastq_ch
        }
        else {
            fastq_data = Channel.fromFilePairs("${params.fastq_dir}/${params.sra_run}*_{1,2}.fastq")
        }

        sra_truncated = DIACHROMATIC_TRUNCATE(fastq_data)
    }
    else {
        sra_truncated = Channel.fromFilePairs("${params.input_diachromatic}/${params.sra_run}*-HiC.truncated_R{1,2}.fastq.gz")
    }

    //sra_truncated.view()

    emit:
    sra_truncated_ch = sra_truncated
}

workflow align_reads_wf {
    take:
    sra_truncated

    main:
    if (!params.skip_align) {
        if (!params.skip_bow_index) {
            BOWTIE2_BUILD(params.reference)
            bowtie_index = BOWTIE2_BUILD.out.bowtie2_index_ch
        }
        else {
            bowtie_index = params.input_bow_index
        }
        bam_file = ALIGN_READS(sra_truncated, bowtie_index)
    }
    else {
        bam_file = params.input_align_bam
    }

    emit:
    aligned_bam_ch = bam_file
}

workflow nester_wf {
    main:
    nester_out = ( !params.skip_nester ? RUN_TEG_NESTER("${params.reference}") : "${params.input_nester}" )
    annotated_fragments = ANNOTATE_FRAGMENTS(nester_out)

    emit:
    nester_gff_ch = annotated_fragments
}

workflow re_wf {
    take:
    sra_truncated

    main:
    if (!params.skip_re) {
        RUN_RE(SEQTK_SEQ(FORMAT_READS(sra_truncated)))
        clTABLE = RUN_RE.out.clTABLE
        contigs_fasta = RUN_RE.out.contigs_fasta
        hitsort_cls = RUN_RE.out.hitsort_cls
    }
    else {
        clTABLE = "${params.input_re}/CLUSTER_TABLE.csv"
        contigs_fasta = "${params.input_re}/contigs.fasta"
        hitsort_cls = "${params.input_re}/seqclust/clustering/hitsort.cls"
    }

    emit:
    cluster_table_ch = clTABLE
    hitsort_ch = hitsort_cls
    re_contigs_ch = contigs_fasta
}

workflow map_re_contigs_wf {
    take:
    contigs_fasta
    clTABLE

    main:
    if (!params.skip_re_contigs_map) {
        annotated_re = FILTER_GFF(BED_TO_GFF(RE_BED_ANNOTATION(BAM_TO_BED(RE_MODIFY_SAM(MAP_CONTIGS(params.reference, contigs_fasta))), clTABLE)))
    }
    else {
        annotated_re = params.input_re_contigs_gff
    }

    emit:
    re_gff_ch = annotated_re
}

workflow plot_re_wf {
    take:
    sra_truncated
    hitsort_cls
    contigs_fasta
    clTABLE
    table_repeats

    main:
    formatted_reads = FORMAT_READS(sra_truncated)

    //inClusters clTABLE blastResults
    inClusters = MERGE_TABLES(formatted_reads, ASSIGN_READS_TO_CLUSTERS(hitsort_cls))
    blastResults = BLAST_CONTIGS(contigs_fasta)
    reformatted_table = REFORMAT_MATRIX(inClusters, clTABLE, blastResults)

    RE_HEATMAP_NORM2(SAVE_RE_TABLE(reformatted_table),table_repeats)
}

workflow map_plantsat_wf {
    main:
    FILTER_GFF_PL(BED_TO_GFF_PL(BAM_TO_BED_PL(MAP_CONTIGS_PL(params.reference, params.plantsat))))

    emit:
    plantsat_gff_ch = FILTER_GFF_PL.out.re_gff_filtered_ch
}

workflow table_wf {
    take:
    annotations_gff
    aligned_bam

    main:
    table_file = GENERATE_TABLE(INTERSECT_BED(aligned_bam, annotations_gff))
    PREPARE_TABLE(table_file)

    emit:
    table_ch = table_file
    table_rdata_ch = PREPARE_TABLE.out.table_ch
}

workflow shuffled_wf {
    take:
    annotations_gff
    genome_lenght
    aligned_bam

    main:
    shuffled_data = SHUFFLE_FRAGMENTS(annotations_gff, genome_lenght)
    bed_shuffled = INTERSECT_BED_SHUFFLE(aligned_bam, shuffled_data)
    GENERATE_TABLE_SHUFFLED(bed_shuffled)

    emit:
    table_shuffled_ch = GENERATE_TABLE_SHUFFLED.out.table_ch
}

workflow visualization_wf {
    take:
    genome_lenght
    aligned_bam
    table
    table_shuffled
    table_repeats

    main:
    PLOT_CHR(genome_lenght, aligned_bam)
    PLOT_BARPLOT(table)

    heatmap_types = ['"ALL"', '"TAD"', '"DIST"']
    TE_HEATMAP_NORM1(table, table_repeats, heatmap_types)
    TE_HEATMAP_NORM2(table, table_repeats, heatmap_types)
    TE_HEATMAP_NORM3(table, table_shuffled, table_repeats, heatmap_types)
}

/*
 * main script flow
 */

workflow {
    println("Process SRA data:" + params.sra_run)

    GET_GENOME_LENGHT(params.reference)

    sra_truncated = prepare_sra_wf()

    bam_file = align_reads_wf(sra_truncated)

    annotated_fragments = nester_wf()

    re_wf(sra_truncated)
    contigs_fasta = re_wf.out.re_contigs_ch

    clTABLE = re_wf.out.cluster_table_ch
    hitsort_cls = re_wf.out.hitsort_ch

    annotated_re = map_re_contigs_wf(contigs_fasta, clTABLE)

    annotated_plantsat = map_plantsat_wf()

    annotated_exons = ANNOTATE_EXONS(params.exons_gff3)

    annotated_all = CAT_ANNOTATIONS(annotated_fragments, annotated_re, annotated_plantsat, annotated_exons)

    table_wf(annotated_all, bam_file)
    table_file = table_wf.out.table_ch
    table_rdata = table_wf.out.table_rdata_ch

    table_shuffled = shuffled_wf(annotated_all, GET_GENOME_LENGHT.out.genome_len_ch, bam_file)
    table_repeats = "${baseDir}/data/repeats.tab"

    visualization_wf(GET_GENOME_LENGHT.out.genome_len_ch, bam_file, table_rdata, table_shuffled, table_repeats)
    plot_re_wf(sra_truncated, hitsort_cls, contigs_fasta, clTABLE, table_repeats)

    PLOT_CIRCOS(table_file, GET_GENOME_LENGHT.out.circos_genome_len_ch)
}
