from __future__ import absolute_import
from falcon_kit.pype import Dist
from falcon_kit import pype_tasks
from falcon_kit.pype import (wrap_gen_task as gen_task, gen_parallel_tasks, Dist)
from .. import io
import logging
import os
import sys

LOG = logging.getLogger(__name__)
EXE = os.path.basename(sys.executable)

TASK_PLACEMENT_AS_SINGLE_TASK = """\
nproc={params.pypeflow_nproc}
min_aln_len={params.min_aln_len}

# Define intermediate files.
local_p=$(basename {input.cns_p_ctg_fasta})
local_h=$(basename {input.cns_h_ctg_fasta})
clean_cns_p_ctg_fasta=${{local_p%%.*}}.clean.fasta
clean_cns_h_ctg_fasta=${{local_h%%.*}}.clean.fasta
name_mapping=name_mapping.txt
coords_all=ctg.all.coords
ctg_hp=ctg.hp.txt

# The fc_scrub_names.pl will create files in the source directory. To prevent
# this, we symlink for now.
ln -sf {input.cns_p_ctg_fasta} $(basename {input.cns_p_ctg_fasta})
ln -sf {input.cns_h_ctg_fasta} $(basename {input.cns_h_ctg_fasta})

# Generate the name_mapping file and cleaned FASTA files.
fc_scrub_names.pl ${{local_p}} ${{local_h}} > ${{name_mapping}}

# Get a list of all primary contigs which will be processed individually.
primary=$(awk '{{ print $1 }}' ${{name_mapping}} | sort | uniq)

for p_ctg_id in ${{primary}}
do
    # Define intermediate files and folder.
    ctg_dir=${{p_ctg_id}}
    delta=${{ctg_dir}}/ctg.${{p_ctg_id}}.delta
    delta_filt=${{delta}}.filt
    coords=${{ctg_dir}}/ctg.${{p_ctg_id}}.coords

    mkdir -p ${{ctg_dir}}

    # Get a list of haplotig IDs to extract from the general haplotig file.
    h_ctg_ids=$(grep "^${{p_ctg_id}}" ${{name_mapping}} | awk '{{ print $2 }}' | sort | uniq)

    # Extract the primary contig and the corresponding haplotigs for alignment.
    samtools faidx ${{clean_cns_p_ctg_fasta}} ${{p_ctg_id}} > ${{ctg_dir}}/ctg.${{p_ctg_id}}.p.fasta
    samtools faidx ${{clean_cns_h_ctg_fasta}} ${{h_ctg_ids}} > ${{ctg_dir}}/ctg.${{p_ctg_id}}.h.fasta

    # Align.
    nucmer --maxmatch -l 100 -c 500 --delta=${{delta}} -t ${{nproc}} ${{ctg_dir}}/ctg.${{p_ctg_id}}.p.fasta ${{ctg_dir}}/ctg.${{p_ctg_id}}.h.fasta

    # Filter and get coords.
    delta-filter -g ${{delta}} > ${{delta_filt}}
    show-coords -qTlHL ${{min_aln_len}} ${{delta_filt}} > ${{coords}}
    cat ${{coords}} | sort -k11,11 -k1,1n  >> ${{coords_all}}

    # These have abspaths on first line, which complicates diff.
    rm -f ${{delta}} ${{delta_filt}}
done

# Yield the final haplotig placement output.
%(python)s -m falcon_phase.mains.coords2hp ${{coords_all}} | sort -k6,6 -k8,8n > ${{ctg_hp}}
%(python)s -m falcon_phase.mains.filt_hp ${{ctg_hp}} > {output.placement_file}
"""%dict(python=EXE)  # and note %% above

TASK_PLACEMENT_SCRUB_NAMES = """\
# Define intermediate files.
local_p=$(basename {input.cns_p_ctg_fasta})
local_h=$(basename {input.cns_h_ctg_fasta})

# The fc_scrub_names.pl will create files in the source directory. To prevent
# this, we symlink for now.
ln -sf {input.cns_p_ctg_fasta} $(basename {input.cns_p_ctg_fasta})
ln -sf {input.cns_h_ctg_fasta} $(basename {input.cns_h_ctg_fasta})

# Generate the name_mapping file and cleaned FASTA files.
fc_scrub_names.pl ${{local_p}} ${{local_h}} > {output.name_mapping}

samtools faidx {output.clean_p_ctg_fasta}
samtools faidx {output.clean_h_ctg_fasta}

# Get a list of all primary contigs which will be processed individually.
# primary=$(awk '{{ print $1 }}' {output.name_mapping} | sort | uniq)
"""

TASK_PLACEMENT_SPLIT = """\
%(python)s -m falcon_phase.mains.placement_split --name-mapping-fn={input.name_mapping} --p-ctg-fasta-fn={input.clean_p_ctg_fasta} --h-ctg-fasta-fn={input.clean_h_ctg_fasta} --min-aln-len={params.min_aln_len} --split-fn={output.split} --bash-template-fn={output.bash_template}
"""%dict(python=EXE)

TASK_PLACEMENT_RUN_SCRIPT = """\
nproc={params.pypeflow_nproc}
min_aln_len={params.min_aln_len}
delta=out.delta
delta_filt=out.filt.delta

# Align.
nucmer --maxmatch -l 100 -c 500 --delta=${{delta}} -t ${{nproc}} {input.p_ctg_fasta} {input.h_ctg_fasta}

# Filter and get coords.
delta-filter -g ${{delta}} > ${{delta_filt}}
show-coords -qTlHL ${{min_aln_len}} ${{delta_filt}} > {output.coords}

# These have abspaths on first line, which complicates diff.
rm -f ${{delta}} ${{delta_filt}}
"""

TASK_PLACEMENT_GATHER_SCRIPT = """\
%(python)s -m falcon_phase.mains.placement_gather --gathered-fn={input.gathered} --gathered-key=coords --output-fn={output.coords_all}

# Yield the final haplotig placement output.
ctg_hp=unfiltered.placement
%(python)s -m falcon_phase.mains.coords2hp {output.coords_all} | sort -k6,6 -k8,8n > ${{ctg_hp}}
%(python)s -m falcon_phase.mains.filt_hp ${{ctg_hp}} > {output.placement_file}

# For placement-file format, see https://github.com/lh3/miniasm/blob/master/PAF.md
"""%dict(python=EXE)

TASK_MAKE_MINCED_BED = """\
awk '{{print $1 ":" 0 "-" $4 "\t" $6 ":" $8 "-" $9}}' {input.DAT} > {output.ABP}

cat {output.ABP} | cut -f1 | perl -lane '@Z = split /-|:/, $_; print join "\t", @Z ' | sort -k1,1 -k2,2n > {output.A}
cat {output.ABP} | cut -f2 | perl -lane '@Z = split /-|:/, $_; print join "\t", @Z ' | sort -k1,1 -k2,2n > {output.B}

mkdir -p mince
bedtools merge -i {output.B} > mince/B_haplotigs_merged.bed

samtools faidx {input.PCTG} # TODO: Should be part of placement-gather.

cut -f1,2 {input.PCTG}.fai | sort -k1,1 > tmp.genome

# Make BED file for collapsed region, exclude haplotigs
bedtools complement -i mince/B_haplotigs_merged.bed -g tmp.genome > {output.C}

rm tmp.genome
"""

TASK_MINCE_FASTA = """\
# GENERATE FASTA FILES
bedtools getfasta -fi {input.PCTG} -bed {input.BB} > {output.BF}
bedtools getfasta -fi {input.PCTG} -bed {input.CB} > {output.CF}
bedtools getfasta -fi {input.HCTG} -bed {input.AB} > {output.AF}
cat {output.AF} {output.BF} {output.CF} > {output.minced_fasta}
samtools faidx {output.minced_fasta}
"""

TASK_ALIGN_HIC = """\
nproc={params.pypeflow_nproc}
ln -sf {input.REF} minced.fasta
bwa index minced.fasta
bwa mem -5 -t ${{nproc}} minced.fasta {input.R1} {input.R2} | samtools view -S -h -b -F 2316 > {output.bam}
"""

TASK_BAM_FILT = """\
falcon-phase bamfilt -f 20 -m 10 -i {input.bam_unfiltered} -o {output.bam_filtered}
"""

TASK_BINMAT = """\
falcon-phase bam2 binmat {input.bam_filtered} {output.binmat}
"""

TASK_INDEX_PAIR = """\
fc_primary_contig_index.pl {input.AB_pairs} {input.minced_fasta}.fai > {output.ov_index}
"""

TASK_PHASE = """\
falcon-phase phase -s 1 -f {input.minced_fasta} -b {input.binmat} -m {params.enzyme} -i {input.ov_index} -p out -n {params.iterations}
mv out.results.txt {output.phased}
"""

TASK_PREP_EMIT = """\
cat {input.PCTG} {input.HCTG} > {output.p_h_ctg_fasta}
samtools faidx {output.p_h_ctg_fasta}
cat {input.B} {input.C} | sort -k1,1 -k2,2n > {output.BC_bed}
"""

TASK_EMIT_HAPLOTIGS = """\
%(python)s -m falcon_phase.mains.emit_haplotigs {input.phased} {input.BC_bed} {input.p_h_ctg_fasta} bedtools {params.output_format}

rm tmp_phase1.txt
rm tmp_phase0.txt
rm {input.p_h_ctg_fasta}
rm {input.p_h_ctg_fasta}.fai

mv f0.fasta {output.F0}
mv f1.fasta {output.F1}
mv b0.bed {output.B0}
mv b1.bed {output.B1}
"""%dict(python=EXE)

def add_placement_as_parallel_tasks(wf, cns_p_ctg_fasta, cns_h_ctg_fasta,
                                       min_aln_len, default_nproc, config):

    clean_p_ctg_fasta = './5-phase/name-mapping/{bn}.clean.fasta'.format(bn = os.path.splitext(os.path.basename(cns_p_ctg_fasta))[0])
    clean_h_ctg_fasta = './5-phase/name-mapping/{bn}.clean.fasta'.format(bn = os.path.splitext(os.path.basename(cns_h_ctg_fasta))[0])
    name_mapping_fn = './5-phase/name-mapping/name_mapping.txt'

    wf.addTask(gen_task(
            script=TASK_PLACEMENT_SCRUB_NAMES,
            inputs={
                'cns_p_ctg_fasta': cns_p_ctg_fasta,
                'cns_h_ctg_fasta': cns_h_ctg_fasta,
            },
            outputs={
                'clean_p_ctg_fasta': clean_p_ctg_fasta,
                'clean_h_ctg_fasta': clean_h_ctg_fasta,
                'name_mapping': name_mapping_fn,
            },
            parameters={
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))

    placement_all_units_fn = './5-phase/placement-split/all-units-of-work.json'
    placement_run_bash_template_fn ='./5-phase/placement-split/bash-template.sh'

    wf.addTask(gen_task(
            script=TASK_PLACEMENT_SPLIT,
            inputs={
                'clean_p_ctg_fasta': clean_p_ctg_fasta,
                'clean_h_ctg_fasta': clean_h_ctg_fasta,
                'name_mapping': name_mapping_fn,
            },
            outputs={
                'split': placement_all_units_fn,
                'bash_template': placement_run_bash_template_fn,
            },
            parameters={
                'min_aln_len': min_aln_len,
            },
            dist=Dist(
                NPROC=1,
                local=True,
                job_dict=config['Phase'],
            ),
    ))

    gathered_placement_fn = './5-phase/placement-gathered/gathered.json'

    gen_parallel_tasks(
        wf,
        placement_all_units_fn, gathered_placement_fn,
        run_dict=dict(
            bash_template_fn=placement_run_bash_template_fn,
            script=TASK_PLACEMENT_RUN_SCRIPT,
            inputs={
                'units_of_work': './5-phase/placement-chunks/{ctg_id}/some-units-of-work.json',
            },
            outputs={
                'results': './5-phase/placement-run/{ctg_id}/placement-result-list.json',
            },
            parameters={},
        ),
        dist=Dist(
            NPROC=default_nproc,
            job_dict=config['job.step.phase.placement'],
        ),
    )

    coords_all_fn = './5-phase/placement-output/coords.all'
    placement_fn = './5-phase/placement-output/haplotig.placement'

    wf.addTask(gen_task(
        script=TASK_PLACEMENT_GATHER_SCRIPT,
        inputs={
            'gathered': gathered_placement_fn,
        },
        outputs={
            'coords_all': coords_all_fn,
            'placement_file': placement_fn,
        },
        parameters={},
        dist=Dist(
                NPROC=1,
                local=True,
                job_dict=config['Phase'],
        ),
    ))

    return placement_fn, clean_p_ctg_fasta, clean_h_ctg_fasta

def run_workflow(wf, config):
    general_config = config.get('General', {})
    target = general_config.get('target', 'all')
    valid_targets = ('all', 'minced')
    assert target in valid_targets, 'General.target={} not in {}'.format(target, valid_targets)

    default_njobs = int(config['job.defaults']['njobs'])
    default_nproc = int(config['job.defaults']['NPROC'])

    wf.max_jobs = default_njobs

    LOG.info('This is where FALCON-Phase workflow should be initialized and run.')

    config_phase = config['Phase']

    # There are only 2 input contig files and 2 input reads files. Make them
    # all configurable to make the workflow more flexible.
    cns_p_ctg_fasta = config_phase['cns_p_ctg_fasta'] if 'cns_p_ctg_fasta' in config_phase else './4-quiver/cns-output/cns_p_ctg.fasta'
    cns_h_ctg_fasta = config_phase['cns_h_ctg_fasta'] if 'cns_h_ctg_fasta' in config_phase else './4-quiver/cns-output/cns_h_ctg.fasta'
    reads_1 = config_phase['reads_1']
    reads_2 = config_phase['reads_2']
    iterations = config_phase['iterations']
    enzyme = config_phase['enzyme']
    min_aln_len = config_phase['min_aln_len']
    output_format = config_phase['output_format']

    placement_file, clean_p_ctg_fasta, clean_h_ctg_fasta = add_placement_as_parallel_tasks(wf,
                                    cns_p_ctg_fasta, cns_h_ctg_fasta,
                                    min_aln_len, default_nproc, config)

    AB_pairs_fn = './5-phase/minced_bed/AB_pairs.txt'
    A_haplotigs_bed_fn = './5-phase/minced_bed/A_haplotigs.bed'
    B_haplotigs_bed_fn = './5-phase/minced_bed/B_haplotigs.bed'
    collapsed_haplotypes_bed_fn = './5-phase/minced_bed/collapsed_haplotypes.bed'

    wf.addTask(gen_task(
            script=TASK_MAKE_MINCED_BED,
            inputs={
                'DAT': placement_file,
                'PCTG': clean_p_ctg_fasta,
            },
            outputs={
                'ABP': AB_pairs_fn,
                'A': A_haplotigs_bed_fn,
                'B': B_haplotigs_bed_fn,
                'C': collapsed_haplotypes_bed_fn,
            },
            parameters={
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))

    A_haplotigs_fasta_fn = './5-phase/minced_fasta/A_haplotigs.fasta'
    B_haplotigs_fasta_fn = './5-phase/minced_fasta/B_haplotigs.fasta'
    collapsed_haplotypes_fasta_fn = './5-phase/minced_fasta/collapsed_haplotypes.fasta'
    minced_fasta_fn = './5-phase/minced_fasta/minced.fasta'

    wf.addTask(gen_task(
            script=TASK_MINCE_FASTA,
            inputs={
                'AB': A_haplotigs_bed_fn,
                'BB': B_haplotigs_bed_fn,
                'CB': collapsed_haplotypes_bed_fn,
                'PCTG': clean_p_ctg_fasta,
                'HCTG': clean_h_ctg_fasta,
            },
            outputs={
                'AF': A_haplotigs_fasta_fn,
                'BF': B_haplotigs_fasta_fn,
                'CF': collapsed_haplotypes_fasta_fn,
                'minced_fasta': minced_fasta_fn,
            },
            parameters={
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))
    if target == 'minced':
        LOG.warning('Quitting early, target=minced')
        sys.exit(0)

    bam_unfiltered_fn = './5-phase/alignment/aln.unfiltered.bam'

    wf.addTask(gen_task(
            script=TASK_ALIGN_HIC,
            inputs={
                'REF': minced_fasta_fn,
                'R1': reads_1,
                'R2': reads_2,
            },
            outputs={
                'bam': bam_unfiltered_fn,
            },
            parameters={
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['job.step.phase.alignment'],
            ),
    ))

    bam_filtered_fn = './5-phase/alignment-filter/aln.filtered.bam'

    wf.addTask(gen_task(
            script=TASK_BAM_FILT,
            inputs={
                'bam_unfiltered': bam_unfiltered_fn,
            },
            outputs={
                'bam_filtered': bam_filtered_fn,
            },
            parameters={
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))

    binmat_fn = './5-phase/binmat/out.binmat'

    wf.addTask(gen_task(
            script=TASK_BINMAT,
            inputs={
                'bam_filtered': bam_filtered_fn,
            },
            outputs={
                'binmat': binmat_fn,
            },
            parameters={
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))

    ov_index_fn = './5-phase/index-pair/ov_index.txt'

    wf.addTask(gen_task(
            script=TASK_INDEX_PAIR,
            inputs={
                'minced_fasta': minced_fasta_fn,
                'AB_pairs': AB_pairs_fn,
            },
            outputs={
                'ov_index': ov_index_fn,
            },
            parameters={
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))

    phased_fn = './5-phase/phasing/phased.txt'

    wf.addTask(gen_task(
            script=TASK_PHASE,
            inputs={
                'ov_index': ov_index_fn,
                'minced_fasta': minced_fasta_fn,
                'binmat': binmat_fn,
            },
            outputs={
                'phased': phased_fn,
            },
            parameters={
                'enzyme': enzyme,
                'iterations': iterations,
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))

    p_h_ctg_fasta_fn = './5-phase/prep-emit/p_h_ctg.fasta'
    BC_bed_fn = './5-phase/prep-emit/BC.bed'

    wf.addTask(gen_task(
            script=TASK_PREP_EMIT,
            inputs={
                'PCTG': clean_p_ctg_fasta,
                'HCTG': clean_h_ctg_fasta,
                'B': B_haplotigs_bed_fn,
                'C': collapsed_haplotypes_bed_fn,
            },
            outputs={
                'p_h_ctg_fasta': p_h_ctg_fasta_fn,
                'BC_bed': BC_bed_fn,
            },
            parameters={
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))

    phased_0_fasta_fn = './5-phase/output/phased.0.fasta'
    phased_1_fasta_fn = './5-phase/output/phased.1.fasta'
    phased_0_bed_fn = './5-phase/output/phased.0.bed'
    phased_1_bed_fn = './5-phase/output/phased.1.bed'

    wf.addTask(gen_task(
            script=TASK_EMIT_HAPLOTIGS,
            inputs={
                'phased': phased_fn,
                'BC_bed': BC_bed_fn,
                'p_h_ctg_fasta': p_h_ctg_fasta_fn,
            },
            outputs={
                'F0': phased_0_fasta_fn,
                'F1': phased_1_fasta_fn,
                'B0': phased_0_bed_fn,
                'B1': phased_1_bed_fn,

            },
            parameters={
                'output_format': output_format,
            },
            dist=Dist(
                NPROC=default_nproc,
                job_dict=config['Phase'],
            ),
    ))

    wf.refreshTargets()
