import argparse
import logging
import os
import sys
import collections
from falcon_unzip import io
from ..tasks.tasks_falcon_phase import TASK_PLACEMENT_RUN_SCRIPT

LOG = logging.getLogger()

def load_name_mapping(name_mapping_fn):
    ret = collections.defaultdict(list)
    with open(name_mapping_fn, 'r') as fp_in:
        for line in fp_in:
            sl = line.strip().split()
            ret[sl[0]].append(sl[1])
    return ret

def run(name_mapping_fn, p_ctg_fasta_fn, h_ctg_fasta_fn, min_aln_len, split_fn, bash_template_fn):
    LOG.info('Splitting placement tasks into {!r}.'.format(
        split_fn))
    with open(bash_template_fn, 'w') as stream:
        stream.write(TASK_PLACEMENT_RUN_SCRIPT)

    # Load the name mapping produced by the scrub_names script.
    name_mapping = load_name_mapping(name_mapping_fn)

    jobs = list()
    for ctg_id, h_ctg_ids in name_mapping.items(): # was iteritems() in py2
        out_p_ctg_fn = os.path.join('./ctgs', ctg_id, 'p_ctg.fasta')
        out_h_ctg_fn = os.path.join('./ctgs', ctg_id, 'h_ctg.fasta')

        if not os.path.exists(os.path.dirname(out_p_ctg_fn)):
            io.mkdirs(os.path.dirname(out_p_ctg_fn))

        io.syscall('samtools faidx {fasta_fn} {ctg_ids} > {out_fn}'.format(fasta_fn=p_ctg_fasta_fn, ctg_ids=ctg_id, out_fn=out_p_ctg_fn))
        io.syscall('samtools faidx {fasta_fn} {ctg_ids} > {out_fn}'.format(fasta_fn=h_ctg_fasta_fn, ctg_ids=' '.join(h_ctg_ids), out_fn=out_h_ctg_fn))

        new_job = {}
        new_job['params'] = dict(
                ctg_id=ctg_id,
                min_aln_len=min_aln_len,
        )
        new_job['input'] = dict(
                p_ctg_fasta=out_p_ctg_fn,
                h_ctg_fasta=out_h_ctg_fn,
        )
        new_job['output'] = dict(
                coords='out.coords',
        )
        new_job['wildcards'] = {
                'ctg_id': ctg_id,
        }
        jobs.append(new_job)

    io.serialize(split_fn, jobs)

class HelpF(argparse.RawTextHelpFormatter, argparse.ArgumentDefaultsHelpFormatter):
    pass


def parse_args(argv):
    description = 'Scatter to prepare for parallel Quiver jobs'
    epilog = ''
    parser = argparse.ArgumentParser(
        description=description,
        epilog=epilog,
        formatter_class=HelpF,
    )
    parser.add_argument(
        '--name-mapping-fn',
        type=str,
        help='The name_mapping.txt file produced by scrub_names.pl script.',
        required=True,
    )
    parser.add_argument(
        '--p-ctg-fasta-fn',
        type=str,
        help='Typically from 4-quiver/cns-output/cns_p_ctg.fasta',
        required=True,
    )
    parser.add_argument(
        '--h-ctg-fasta-fn',
        type=str,
        help='Typically from 4-quiver/cns-output/cns_h_ctg.fasta',
        required=True,
    )
    parser.add_argument(
        '--min-aln-len',
        type=int,
        default=5000,
        help='Minimum alignment length for filtering Nucmer alignments.',
    )
    parser.add_argument(
        '--split-fn',
        type=str,
        help='Output. JSON list of units of work.',
        required=True,
    )
    parser.add_argument(
        '--bash-template-fn',
        type=str,
        help='Output. Copy of known bash template, for use later.',
        required=True,
    )
    args = parser.parse_args(argv[1:])
    return args

def main(argv=sys.argv):
    args = parse_args(argv)
    logging.basicConfig(level=logging.INFO)
    run(**vars(args))

if __name__ == '__main__':  # pragma: no cover
    main()
