#!/usr/bin/env python3

from glob import glob

import argparse

## config file parser
from samflow.workflow import Workflow, attach_back
from samflow.command import ShellCommand, PythonCommand

from gcap.funcs.config import Conf
from gcap.funcs.sequence_quality import read_quality
from gcap.funcs.mapping import read_mapping
from gcap.funcs.filter import filter_bam
from gcap.funcs.dnase_stat import sample_bam_stat
from gcap.funcs.peaks_calling import call_peaks
from gcap.funcs.json_collection import json_collect
import os

# from gcap.funcs.fragment import fragment_size
# from gcap.funcs.union_dhs_overlap import union_DHS_overlap

from configparser import SafeConfigParser
from pkg_resources import resource_filename


def parse_args(args=None):
    """
    If args is None, argparse will parse from sys.argv
    """
    description = "GCAP :  Global Chromatin Accessibility Pipeline"
    parser = argparse.ArgumentParser(description=description)

    parser.add_argument("--from", dest="start_step", default=0, type=int,
                        help="Only step after this number will be processed")
    parser.add_argument("--to", dest="end_step", default=100, type=int,
                        help="Only step before this number will be processed ")
    parser.add_argument("--skip", dest="skip_step", default="",
                        help="Steps to skip, use comma as seperator")

    parser.add_argument("-v", "--verbose-level", dest="verbose_level", type=int, default=2)
    parser.add_argument("--dry-run", dest="dry_run", action="store_true", default=False)
    parser.add_argument("--resume", dest="resume", action="store_false", default=True)
    parser.add_argument("--pe", dest="pe", action="store_true", default=False)
    parser.add_argument("--threads", dest="threads", default=8, type=int,
                        help="threads # to use when aligning using bwa (default: 4)")
    parser.add_argument("-i", "--input", dest = "input", required=True,
                        help = "DNase-seq data, comma(,) separate replicates, format(fastq) fixed by extension")

    cf = SafeConfigParser()
    cf.read(resource_filename('gcap','static/gcap.conf.filled'))
    sections_ls = ['basics', 'tool', 'bwa', 'spp', 'macs2']
    species_ls = [s for s in cf.sections() if s not in sections_ls]

    parser.add_argument("-s", "--species", dest = "species", 
                        help="species set by gcap.conf: %s" % ' '.join(species_ls))
    parser.add_argument("-r", "--readsize", dest = "readsize", default=36,
                        help="species set by gcap.conf: %s" % ' '.join(species_ls))

    parser.add_argument("-o", "--output", dest = "out", 
                        help="output directory", default="output")
    parser.add_argument("-n", "--name", dest = "name", 
                        help="prefix for your dnase seq data")

    return parser.parse_args(args)


class StepChecker:
    def __init__(self, start, end, skips):
        """control step"""
        self.start = start
        self.end = end
        self.skips = skips

    def need_run(self, step_id):
        if step_id < self.start:
            return False
        if step_id > self.end:
            return False
        if step_id in self.skips:
            return False
        return True


class GCAPBuilder:
    def __init__(self, workflow, conf, tex):
        """append to workflow"""
        self.workflow = workflow
        self.conf = conf
        self.tex = tex
        self.LaTex_fragments = []
        self.plain_fragments = []
        self.finished = set()

    def build(self, prep_func, tag=None):
        prep_func(self.workflow, self.conf, self.tex)
        if tag:
            self.finished.add(tag)

    def attach_back(self, command):
        attach_back(self.workflow, command)


def create_workflow(args, conf, tex):
    workflow = Workflow(name="Main")
    bld = GCAPBuilder(workflow, conf, tex)

    if args.skip_step:
        skipped_steps = [int(i) for i in args.skip_step.split(",")]
    else:
        skipped_steps = []

    step_checker = StepChecker(args.start_step, args.end_step, skipped_steps)

    need_run = step_checker.need_run

    bld.attach_back(ShellCommand(
        "if [ ! -d '{output}' ]; then mkdir -p {output}; fi",
        output=conf.target_dir))

    if need_run(1):
        ## funcs.sequence_quality
        bld.build(read_quality)

    if need_run(2):
        bld.build(read_mapping)

    if need_run(3):
        bld.build(filter_bam)

    if need_run(4):
        bld.build(sample_bam_stat)

    if need_run(5):
        bld.build(call_peaks)

    if need_run(6):

        bld.attach_back(PythonCommand(
            json_collect,
            input = {"dir": conf.target_dir},
            output = {"json":conf.json_prefix + "_qc_summary.json"},
            param = conf))

    return workflow

def main(args=None):
    args = parse_args(args)

    ## latex template
    tex = resource_filename("gcap", "static/gcap_template.tex")

    conf = Conf(resource_filename("gcap", "static/gcap.conf.filled"), args)
    print("""
Arguments:
pair %s
input %s
output %s
    """ % (conf.pe, conf.treatment_raws, conf.target_dir))


    workflow = create_workflow(args, conf, tex)
    workflow.set_option(
        verbose_level=args.verbose_level,
        dry_run_mode=args.dry_run,
        resume=args.resume)
    workflow.invoke()



if __name__ == "__main__":
    main()
