#!/usr/bin/env python
# copied from nf-core/viralrecon and adjusted

import os
import sys
import errno
import argparse
import yaml


def parse_args(args=None):
    Description = (
        "Create custom spreadsheet for pertinent MultiQC bowtie 2 metrics generated by the nf-core/mag pipeline."
    )
    Epilog = "Example usage: python multiqc_to_custom_tsv.py"
    parser = argparse.ArgumentParser(description=Description, epilog=Epilog)
    parser.add_argument(
        "-md",
        "--multiqc_data_dir",
        type=str,
        dest="MULTIQC_DATA_DIR",
        default="multiqc_data",
        help="Full path to directory containing YAML files for each module, as generated by MultiQC. (default: 'multiqc_data').",
    )
    parser.add_argument(
        "-se",
        "--single_end",
        dest="SINGLE_END",
        action="store_true",
        help="Specifies that the input is single-end reads.",
    )
    return parser.parse_args(args)


def make_dir(path):
    if not len(path) == 0:
        try:
            os.makedirs(path)
        except OSError as exception:
            if exception.errno != errno.EEXIST:
                raise


# Find key in dictionary created from YAML file recursively
# From https://stackoverflow.com/a/37626981
def find_tag(d, tag):
    if tag in d:
        yield d[tag]
    for k, v in d.items():
        if isinstance(v, dict):
            for i in find_tag(v, tag):
                yield i


def yaml_fields_to_dict(YAMLFile, AppendDict={}, FieldMappingList=[]):
    with open(YAMLFile) as f:
        yaml_dict = yaml.safe_load(f)
        for k in yaml_dict.keys():
            key = k
            if key not in AppendDict:
                AppendDict[key] = {}
            if FieldMappingList != []:
                for i, j in FieldMappingList:
                    val = list(find_tag(yaml_dict[k], j[0]))
                    if len(val) != 0:
                        val = val[0]
                        if len(j) == 2:
                            val = list(find_tag(val, j[1]))[0]
                        if i not in AppendDict[key]:
                            AppendDict[key][i] = val
                        else:
                            print(
                                "WARNING: {} key already exists in dictionary so will be overwritten. YAML file {}.".format(
                                    i, YAMLFile
                                )
                            )
            else:
                AppendDict[key] = yaml_dict[k]
    return AppendDict


# customized
def metrics_dict_to_file(FileFieldList, MultiQCDataDir, OutFile, se):
    MetricsDict = {}
    FieldList = []
    for yamlFile, mappingList in FileFieldList:
        yamlFile = os.path.join(MultiQCDataDir, yamlFile)
        if os.path.exists(yamlFile):
            MetricsDict = yaml_fields_to_dict(YAMLFile=yamlFile, AppendDict=MetricsDict, FieldMappingList=mappingList)
            FieldList += [x[0] for x in mappingList]
        else:
            print("WARNING: File does not exist: {}".format(yamlFile))

    if MetricsDict != {}:
        make_dir(os.path.dirname(OutFile))
        with open(OutFile, "w") as fout:
            if se:
                fout.write(
                    "{}\n".format("\t".join(["Sample", "SE reads not mapped (kept)", "SE reads mapped (discarded)"]))
                )
            else:
                fout.write(
                    "{}\n".format(
                        "\t".join(
                            [
                                "Sample",
                                "PE reads not mapped concordantly (kept)",
                                "PE reads mapped concordantly (discarded)",
                            ]
                        )
                    )
                )
            for k in sorted(MetricsDict.keys()):
                # write out # not mapped reads and # mapped reads (uniquely + multi mapping reads)
                fout.write(
                    "{}\n".format(
                        "\t".join(
                            [
                                k,
                                str(MetricsDict[k][FieldList[0]]),
                                str(MetricsDict[k][FieldList[1]] + MetricsDict[k][FieldList[2]]),
                            ]
                        )
                    )
                )
    return MetricsDict


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

    ## File names for MultiQC YAML along with fields to fetch from each file
    Bowtie2FileFieldList = []
    if args.SINGLE_END:
        Bowtie2FileFieldList = [
            (
                "multiqc_bowtie2.yaml",
                [
                    ("# Not mapped reads", ["unpaired_aligned_none"]),
                    ("# Mapped reads 1", ["unpaired_aligned_one"]),
                    ("# Mapped reads multi", ["unpaired_aligned_multi"]),
                ],
            ),
        ]
    else:
        Bowtie2FileFieldList = [
            (
                "multiqc_bowtie2.yaml",
                [
                    ("# Not mapped reads", ["paired_aligned_none"]),
                    ("# Mapped reads 1", ["paired_aligned_one"]),
                    ("# Mapped reads multi", ["paired_aligned_multi"]),
                ],
            ),
        ]

    ## Write Bowtie 2 metrics to file
    metrics_dict_to_file(
        FileFieldList=Bowtie2FileFieldList,
        MultiQCDataDir=args.MULTIQC_DATA_DIR,
        OutFile="host_removal_metrics.tsv",
        se=args.SINGLE_END,
    )


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