import os
import sys
from copy import deepcopy

import traceback
import functools
from collections import defaultdict

import yaml

from argparse import ArgumentParser,\
    RawDescriptionHelpFormatter, ArgumentDefaultsHelpFormatter

from sgains.configuration.parser import SgainsValidator, Config
from sgains.configuration.schema import sgains_schema

from sgains.executor import Executor

from sgains.pipelines.mappableregions_pipeline import MappableRegionsPipeline
from sgains.pipelines.genomeindex_pipeline import GenomeIndexPipeline
from sgains.pipelines.bins_pipeline import BinsPipeline
from sgains.pipelines.mapping_pipeline import MappingPipeline
from sgains.pipelines.extract_10x_pipeline import Extract10xPipeline
from sgains.pipelines.varbin_10x_pipeline import Varbin10xPipeline

from sgains.pipelines.varbin_pipeline import VarbinPipeline
from sgains.pipelines.r_pipeline import Rpipeline
from sgains.pipelines.composite_pipeline import CompositePipeline


SGAINS_COMMANDS = {
    "genomeindex": {
        "config_groups": ["aligner", "genome"],
        "help": "builds appropriate hisat2 or bowtie index for the "
        "reference genome",
    },
    "mappable_regions": {
        "config_groups": ["aligner", "genome", "mappable_regions", "sge"],
        "help": "finds all mappable regions in specified genome",
    },
    "bins": {
        "config_groups": ["genome", "mappable_regions", "bins", "sge"],
        "help": "calculates all bins boundaries for specified bins count "
        "and read length",
    },
    "prepare": {
        "config_groups": [
            "aligner", "genome", "mappable_regions", "bins", "sge"],
        "help": "combines all preparation steps ('genome', 'mappable-regions' "
        "and 'bins') into single command",
    },
    "mapping": {
        "config_groups": ["aligner", "genome", "reads", "mapping", "sge"],
        "help": "performs mapping of cells reads to the reference genome",
    },
    "extract_10x": {
        "config_groups": [
            "data_10x", "reads", "sge"],
        "help": "extracts cells reads from 10x Genomics datasets",
    },
    "varbin": {
        "config_groups": ["bins", "mapping", "varbin", "sge"],
        "help": "applies varbin algorithm to count read mappings in each bin",
    },
    "varbin_10x": {
        "config_groups": [
            "data_10x", "bins", "varbin", "sge"],
        "help": "applies varbin algorithm to count read mappings in each bin "
        "to 10x Genomics datasets without realigning",
    },
    "scclust": {
        "config_groups": ["bins", "varbin", "scclust"],
        "help": "segmentation and clustering based bin counts and "
        "preparation of the SCGV input data"
    },
    "process": {
        "config_groups": [
            "aligner", "genome", "reads", "mapping", "bins", "varbin",
            "scclust", "sge"],
        "help": "combines all process steps ('mapping', 'varbin' "
        "and 'scclust') into single command"
    },
}


def build_common_options(parser):
    parser.add_argument(
        "-v", "--verbose",
        dest="verbose",
        action="count",
        help="set verbosity level [default: %(default)s]",
        default=0
    )
    parser.add_argument(
        "-c", "--config",
        dest="config",
        help="configuration file",
        metavar="path"
    )

    parser.add_argument(
        "-n", "--dry-run",
        dest="dry_run",
        action="store_true",
        help="perform a trial run with no changes made",
        default=False
    )

    parser.add_argument(
        "--force", "-F",
        dest="force",
        action="store_true",
        help="allows overwriting nonempty results directory",
        default=False
    )

    parser.add_argument(
        "--parallel", "-p",
        dest="parallel",
        help="number of task to run in parallel",
        type=int,
        default=1
    )

    parser.add_argument(
        "--sge",
        dest="sge",
        action="store_true",
        help="parallelilizes commands using SGE cluster manager",
        default=False
    )


def _get_config_value(config, group_name, name):

    if config is None:
        return None
    group = config.config.get(group_name)
    if group is None:
        return None
    result = getattr(group, name)
    return result


def build_cli_options(argparser, command=None, config=None, sge_flag=False):
    work_dirname = os.getcwd()
    if config is not None:
        work_dirname = config.work_dirname

    validator = SgainsValidator(
        deepcopy(sgains_schema), work_dirname=work_dirname)

    if command is None:
        config_groups = list(validator.schema.keys())
    else:
        assert command in SGAINS_COMMANDS

        command = SGAINS_COMMANDS[command]
        config_groups = command["config_groups"]

    for group_name in config_groups:
        if group_name == "sge" and not sge_flag:
            continue

        group = validator.schema.get(group_name)
        group_parser = argparser.add_argument_group(f"{group_name} group:")
        assert group["type"] == "dict", (group_name, group)

        group_schema = group["schema"]

        for arg_name, arg_spec in group_schema.items():
            name = f"--{arg_name.replace('_', '-')}"
            arg_type = str
            arg_type = arg_spec.get("type", "string")
            if arg_type == "string":
                arg_type = str
            elif arg_type == "integer":
                arg_type = int
            elif arg_type == "float":
                arg_type = float
            elif arg_type == "list":
                arg_type = list
            else:
                raise ValueError(f"wrong argument type {arg_type}")

            help_data = None
            meta_data = arg_spec.get("meta")
            if meta_data is not None:
                help_data = meta_data.get("help")

            arg_default = _get_config_value(config, group_name, arg_name)
            if arg_default is None:
                arg_default = arg_spec.get("default")

            group_parser.add_argument(
                name,
                help=help_data,
                dest=arg_name,
                type=arg_type,
                default=arg_default)

    return argparser


def parse_cli_options(args):
    config_dict = defaultdict(dict)
    work_dirname = os.getcwd()

    if args.config is not None:
        assert os.path.exists(args.config), args.config

        with open(args.config, "r") as infile:
            config_dict = yaml.safe_load(infile)
        work_dirname = os.path.dirname(args.config)

    validator = SgainsValidator(
        deepcopy(sgains_schema), work_dirname=work_dirname)

    result = defaultdict(dict)

    config_groups = list(validator.schema.keys())

    for group_name in config_groups:
        if group_name == "sge" and not args.sge:
            continue

        group = validator.schema.get(group_name)
        group_schema = group.get("schema")
        if group_schema is None:
            continue

        group_result = {}
        for arg_name in group_schema.keys():
            arg_value = getattr(args, arg_name, None)
            if arg_value is not None:
                group_result[arg_name] = arg_value
            else:
                config_value = config_dict.get(group_name, None)
                if config_value is not None:
                    config_value = config_value.get(arg_name, None)
                if config_value is not None:
                    group_result[arg_name] = config_value
        if group_result:
            result[group_name] = group_result

    config = Config.from_dict(result, work_dirname)
    config.verbose = args.verbose
    config.config_file = args.config
    config.dry_run = args.dry_run
    config.force = args.force
    config.parallel = args.parallel
    config.sge = args.sge

    return config


def main(argv=sys.argv[1:]):
    program_name = os.path.basename(sys.argv[0])
    program_shortdesc = \
        'sgains - sparse genomic analysis of individual nuclei by ' \
        'sequencing pipeline'
    program_description = '''%s

USAGE
''' % (program_shortdesc, )

    try:

        config = Config.parse_argv(argv)
        sge_flag = Config.check_sge_argv(argv)

        argparser = ArgumentParser(
            description=program_description,
            formatter_class=ArgumentDefaultsHelpFormatter)
        
        build_common_options(argparser)
        subparsers = argparser.add_subparsers(
            title="sGAINS subcommands"
        )

        for command in SGAINS_COMMANDS:
            command_name = command.replace("_", "-")
            command_help = SGAINS_COMMANDS[command].get("help", "")
            subparser = subparsers.add_parser(
                name=command_name,
                help=command_help,
                formatter_class=ArgumentDefaultsHelpFormatter
            )

            build_cli_options(subparser, command, config, sge_flag=sge_flag)
            subparser.set_defaults(func=functools.partial(execute, command))

        args = argparser.parse_args(argv)
        args.func(args)

    except KeyboardInterrupt:
        traceback.print_exc()
        return 0
    except Exception as e:
        traceback.print_exc()

        indent = len(program_name) * " "
        sys.stderr.write(program_name + ": " + repr(e) + "\n")
        sys.stderr.write(indent + "  for help use --help")
        sys.stderr.write('\n')
        return 2

def create_pipeline(command, config):
    if command == "genomeindex":
        return GenomeIndexPipeline(config)
    elif command == "mappable_regions":
        return MappableRegionsPipeline(config)
    elif command == "bins":
        return BinsPipeline(config)
    elif command == "mapping":
        return MappingPipeline(config)
    elif command == "varbin":
        return VarbinPipeline(config)
    elif command == "scclust":
        return Rpipeline(config)
    elif command == "extract_10x":
        return Extract10xPipeline(config)
    elif command == "varbin_10x":
        return Varbin10xPipeline(config)
    elif command == "prepare":
        pipelines = [
            GenomeIndexPipeline(config),
            MappableRegionsPipeline(config),
            BinsPipeline(config),
        ]
        return CompositePipeline(config, pipelines)
    elif command == "process":
        pipelines = [
            MappingPipeline(config),
            VarbinPipeline(config),
            Rpipeline(config),
        ]
        return CompositePipeline(config, pipelines)

    raise ValueError(f"Unexpected command: {command}")

def execute(command, args):
    config = parse_cli_options(args)

    pipeline = create_pipeline(command, config)
    assert pipeline is not None, command

    executor = Executor(config)
    executor.run_pipeline(pipeline)


if __name__ == "__main__":
    sys.exit(main())

[(146, 'os.getcwd', 'os.getcwd', False, 'import os\n'), (206, 'collections.defaultdict', 'defaultdict', False, 'from collections import defaultdict\n'), (207, 'os.getcwd', 'os.getcwd', False, 'import os\n'), (219, 'collections.defaultdict', 'defaultdict', False, 'from collections import defaultdict\n'), (246, 'sgains.configuration.parser.Config.from_dict', 'Config.from_dict', False, 'from sgains.configuration.parser import SgainsValidator, Config\n'), (258, 'os.path.basename', 'os.path.basename', False, 'import os\n'), (348, 'sgains.executor.Executor', 'Executor', False, 'from sgains.executor import Executor\n'), (151, 'copy.deepcopy', 'deepcopy', False, 'from copy import deepcopy\n'), (210, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (214, 'os.path.dirname', 'os.path.dirname', False, 'import os\n'), (217, 'copy.deepcopy', 'deepcopy', False, 'from copy import deepcopy\n'), (269, 'sgains.configuration.parser.Config.parse_argv', 'Config.parse_argv', False, 'from sgains.configuration.parser import SgainsValidator, Config\n'), (270, 'sgains.configuration.parser.Config.check_sge_argv', 'Config.check_sge_argv', False, 'from sgains.configuration.parser import SgainsValidator, Config\n'), (272, 'argparse.ArgumentParser', 'ArgumentParser', False, 'from argparse import ArgumentParser, RawDescriptionHelpFormatter, ArgumentDefaultsHelpFormatter\n'), (310, 'sgains.pipelines.genomeindex_pipeline.GenomeIndexPipeline', 'GenomeIndexPipeline', False, 'from sgains.pipelines.genomeindex_pipeline import GenomeIndexPipeline\n'), (213, 'yaml.safe_load', 'yaml.safe_load', False, 'import yaml\n'), (297, 'traceback.print_exc', 'traceback.print_exc', False, 'import traceback\n'), (300, 'traceback.print_exc', 'traceback.print_exc', False, 'import traceback\n'), (304, 'sys.stderr.write', 'sys.stderr.write', False, 'import sys\n'), (305, 'sys.stderr.write', 'sys.stderr.write', False, 'import sys\n'), (312, 'sgains.pipelines.mappableregions_pipeline.MappableRegionsPipeline', 'MappableRegionsPipeline', False, 'from sgains.pipelines.mappableregions_pipeline import MappableRegionsPipeline\n'), (314, 'sgains.pipelines.bins_pipeline.BinsPipeline', 'BinsPipeline', False, 'from sgains.pipelines.bins_pipeline import BinsPipeline\n'), (291, 'functools.partial', 'functools.partial', False, 'import functools\n'), (316, 'sgains.pipelines.mapping_pipeline.MappingPipeline', 'MappingPipeline', False, 'from sgains.pipelines.mapping_pipeline import MappingPipeline\n'), (318, 'sgains.pipelines.varbin_pipeline.VarbinPipeline', 'VarbinPipeline', False, 'from sgains.pipelines.varbin_pipeline import VarbinPipeline\n'), (320, 'sgains.pipelines.r_pipeline.Rpipeline', 'Rpipeline', False, 'from sgains.pipelines.r_pipeline import Rpipeline\n'), (322, 'sgains.pipelines.extract_10x_pipeline.Extract10xPipeline', 'Extract10xPipeline', False, 'from sgains.pipelines.extract_10x_pipeline import Extract10xPipeline\n'), (324, 'sgains.pipelines.varbin_10x_pipeline.Varbin10xPipeline', 'Varbin10xPipeline', False, 'from sgains.pipelines.varbin_10x_pipeline import Varbin10xPipeline\n'), (331, 'sgains.pipelines.composite_pipeline.CompositePipeline', 'CompositePipeline', False, 'from sgains.pipelines.composite_pipeline import CompositePipeline\n'), (327, 'sgains.pipelines.genomeindex_pipeline.GenomeIndexPipeline', 'GenomeIndexPipeline', False, 'from sgains.pipelines.genomeindex_pipeline import GenomeIndexPipeline\n'), (328, 'sgains.pipelines.mappableregions_pipeline.MappableRegionsPipeline', 'MappableRegionsPipeline', False, 'from sgains.pipelines.mappableregions_pipeline import MappableRegionsPipeline\n'), (329, 'sgains.pipelines.bins_pipeline.BinsPipeline', 'BinsPipeline', False, 'from sgains.pipelines.bins_pipeline import BinsPipeline\n'), (338, 'sgains.pipelines.composite_pipeline.CompositePipeline', 'CompositePipeline', False, 'from sgains.pipelines.composite_pipeline import CompositePipeline\n'), (334, 'sgains.pipelines.mapping_pipeline.MappingPipeline', 'MappingPipeline', False, 'from sgains.pipelines.mapping_pipeline import MappingPipeline\n'), (335, 'sgains.pipelines.varbin_pipeline.VarbinPipeline', 'VarbinPipeline', False, 'from sgains.pipelines.varbin_pipeline import VarbinPipeline\n'), (336, 'sgains.pipelines.r_pipeline.Rpipeline', 'Rpipeline', False, 'from sgains.pipelines.r_pipeline import Rpipeline\n')]