""" MultiQC module to parse output from STAR """


import logging
import os
import re

from multiqc import config
from multiqc.modules.base_module import BaseMultiqcModule, ModuleNoSamplesFound
from multiqc.plots import bargraph

# Initialise the logger
log = logging.getLogger(__name__)


class MultiqcModule(BaseMultiqcModule):
    def __init__(self):
        # Initialise the parent object
        super(MultiqcModule, self).__init__(
            name="STAR",
            anchor="star",
            href="https://github.com/alexdobin/STAR",
            info="is an ultrafast universal RNA-seq aligner.",
            doi="10.1093/bioinformatics/bts635",
        )

        # Find and load any STAR reports
        self.star_data = dict()
        for f in self.find_log_files("star"):
            parsed_data = self.parse_star_report(f["f"])
            if parsed_data is not None:
                s_name = f["s_name"]
                if s_name == "" or s_name == "Log.final.out":
                    s_name = self.clean_s_name(os.path.basename(f["root"]), f, root=os.path.dirname(f["root"]))
                if s_name in self.star_data:
                    log.debug(f"Duplicate sample name found! Overwriting: {s_name}")
                self.add_data_source(f, section="SummaryLog")
                self.star_data[s_name] = parsed_data

        # Find and load any STAR gene count tables
        self.star_genecounts_unstranded = dict()
        self.star_genecounts_first_strand = dict()
        self.star_genecounts_second_strand = dict()
        for f in self.find_log_files("star/genecounts", filehandles=True):
            parsed_data = self.parse_star_genecount_report(f)
            if parsed_data is not None:
                s_name = f["s_name"]
                if s_name == "" or s_name == "ReadsPerGene.out.tab":
                    s_name = self.clean_s_name(os.path.basename(f["root"]), f, root=os.path.dirname(f["root"]))
                if s_name in self.star_data:
                    log.debug(f"Duplicate ReadsPerGene sample name found! Overwriting: {s_name}")
                self.add_data_source(f, section="ReadsPerGene")
                self.star_genecounts_unstranded[s_name] = parsed_data["unstranded"]
                self.star_genecounts_first_strand[s_name] = parsed_data["first_strand"]
                self.star_genecounts_second_strand[s_name] = parsed_data["second_strand"]

        # Filter to strip out ignored sample names
        self.star_data = self.ignore_samples(self.star_data)
        self.star_genecounts_unstranded = self.ignore_samples(self.star_genecounts_unstranded)
        self.star_genecounts_first_strand = self.ignore_samples(self.star_genecounts_first_strand)
        self.star_genecounts_second_strand = self.ignore_samples(self.star_genecounts_second_strand)

        if len(self.star_data) == 0 and len(self.star_genecounts_unstranded) == 0:
            raise ModuleNoSamplesFound

        # Superfluous function call to confirm that it is used in this module
        # Replace None with actual version if it is available
        self.add_software_version(None)

        if len(self.star_data) > 0:
            if len(self.star_genecounts_unstranded) > 0:
                log.info(
                    "Found {} reports and {} gene count files".format(
                        len(self.star_data), len(self.star_genecounts_unstranded)
                    )
                )
            else:
                log.info(f"Found {len(self.star_data)} reports")
        else:
            log.info(f"Found {len(self.star_genecounts_unstranded)} gene count files")

        if len(self.star_data) > 0:
            # Write parsed report data to a file
            self.write_data_file(self.star_data, "multiqc_star")

            # Basic Stats Table
            self.star_stats_table()

            # Alignment bar plot
            self.add_section(name="Alignment Scores", anchor="star_alignments", plot=self.star_alignment_chart())

        if len(self.star_genecounts_unstranded) > 0:
            self.add_section(
                name="Gene Counts",
                anchor="star_geneCounts",
                description="Statistics from results generated using <code>--quantMode GeneCounts</code>. "
                + "The three tabs show counts for unstranded RNA-seq, counts for the 1st read strand "
                + "aligned with RNA and counts for the 2nd read strand aligned with RNA.",
                plot=self.star_genecount_chart(),
            )

    def parse_star_report(self, raw_data):
        """Parse the final STAR log file."""

        regexes = {
            "total_reads": r"Number of input reads \|\s+(\d+)",
            "avg_input_read_length": r"Average input read length \|\s+([\d\.]+)",
            "uniquely_mapped": r"Uniquely mapped reads number \|\s+(\d+)",
            "uniquely_mapped_percent": r"Uniquely mapped reads % \|\s+([\d\.]+)",
            "avg_mapped_read_length": r"Average mapped length \|\s+([\d\.]+)",
            "num_splices": r"Number of splices: Total \|\s+(\d+)",
            "num_annotated_splices": r"Number of splices: Annotated \(sjdb\) \|\s+(\d+)",
            "num_GTAG_splices": r"Number of splices: GT/AG \|\s+(\d+)",
            "num_GCAG_splices": r"Number of splices: GC/AG \|\s+(\d+)",
            "num_ATAC_splices": r"Number of splices: AT/AC \|\s+(\d+)",
            "num_noncanonical_splices": r"Number of splices: Non-canonical \|\s+(\d+)",
            "mismatch_rate": r"Mismatch rate per base, % \|\s+([\d\.]+)",
            "deletion_rate": r"Deletion rate per base \|\s+([\d\.]+)",
            "deletion_length": r"Deletion average length \|\s+([\d\.]+)",
            "insertion_rate": r"Insertion rate per base \|\s+([\d\.]+)",
            "insertion_length": r"Insertion average length \|\s+([\d\.]+)",
            "multimapped": r"Number of reads mapped to multiple loci \|\s+(\d+)",
            "multimapped_percent": r"% of reads mapped to multiple loci \|\s+([\d\.]+)",
            "multimapped_toomany": r"Number of reads mapped to too many loci \|\s+(\d+)",
            "multimapped_toomany_percent": r"% of reads mapped to too many loci \|\s+([\d\.]+)",
            "unmapped_mismatches_percent": r"% of reads unmapped: too many mismatches \|\s+([\d\.]+)",
            "unmapped_tooshort_percent": r"% of reads unmapped: too short \|\s+([\d\.]+)",
            "unmapped_other_percent": r"% of reads unmapped: other \|\s+([\d\.]+)",
        }
        parsed_data = {}
        for k, r in regexes.items():
            r_search = re.search(r, raw_data, re.MULTILINE)
            if r_search:
                parsed_data[k] = float(r_search.group(1))
        # Figure out the numbers for unmapped as for some reason only the percentages are given
        try:
            total_mapped = (
                parsed_data["uniquely_mapped"] + parsed_data["multimapped"] + parsed_data["multimapped_toomany"]
            )
            unmapped_count = parsed_data["total_reads"] - total_mapped
            total_unmapped_percent = (
                parsed_data["unmapped_mismatches_percent"]
                + parsed_data["unmapped_tooshort_percent"]
                + parsed_data["unmapped_other_percent"]
            )
            try:
                parsed_data["unmapped_mismatches"] = int(
                    round(unmapped_count * (parsed_data["unmapped_mismatches_percent"] / total_unmapped_percent), 0)
                )
                parsed_data["unmapped_tooshort"] = int(
                    round(unmapped_count * (parsed_data["unmapped_tooshort_percent"] / total_unmapped_percent), 0)
                )
                parsed_data["unmapped_other"] = int(
                    round(unmapped_count * (parsed_data["unmapped_other_percent"] / total_unmapped_percent), 0)
                )
            except ZeroDivisionError:
                parsed_data["unmapped_mismatches"] = 0
                parsed_data["unmapped_tooshort"] = 0
                parsed_data["unmapped_other"] = 0
        except KeyError:
            pass

        if len(parsed_data) == 0:
            return None
        return parsed_data

    def parse_star_genecount_report(self, f):
        """Parse a STAR gene counts output file"""
        # Three numeric columns: unstranded, stranded/first-strand, stranded/second-strand
        keys = ["N_unmapped", "N_multimapping", "N_noFeature", "N_ambiguous"]
        unstranded = {"N_genes": 0}
        first_strand = {"N_genes": 0}
        second_strand = {"N_genes": 0}
        num_errors = 0
        num_genes = 0
        for line in f["f"]:
            s = line.split("\t")
            try:
                for i in [1, 2, 3]:
                    s[i] = float(s[i])
                if s[0] in keys:
                    unstranded[s[0]] = s[1]
                    first_strand[s[0]] = s[2]
                    second_strand[s[0]] = s[3]
                else:
                    unstranded["N_genes"] += s[1]
                    first_strand["N_genes"] += s[2]
                    second_strand["N_genes"] += s[3]
                    num_genes += 1
            except IndexError:
                # Tolerate a few errors in case there is something random added at the top of the file
                num_errors += 1
                if num_errors > 10 and num_genes == 0:
                    log.warning(f"Error parsing {f['fn']}")
                    return None
        if num_genes > 0:
            return {"unstranded": unstranded, "first_strand": first_strand, "second_strand": second_strand}
        else:
            return None

    def star_stats_table(self):
        """Take the parsed stats from the STAR report and add them to the
        basic stats table at the top of the report"""

        headers = {
            "uniquely_mapped_percent": {
                "title": "% Aligned",
                "description": "% Uniquely mapped reads",
                "max": 100,
                "min": 0,
                "suffix": "%",
                "scale": "YlGn",
            },
            "uniquely_mapped": {
                "title": f"{config.read_count_prefix} Aligned",
                "description": f"Uniquely mapped reads ({config.read_count_desc})",
                "min": 0,
                "scale": "PuRd",
                "modify": lambda x: x * config.read_count_multiplier,
                "shared_key": "read_count",
            },
        }
        self.general_stats_addcols(self.star_data, headers)

    def star_alignment_chart(self):
        """Make the plot showing alignment rates"""

        # Specify the order of the different possible categories
        keys = {
            "uniquely_mapped": {"color": "#437bb1", "name": "Uniquely mapped"},
            "multimapped": {"color": "#7cb5ec", "name": "Mapped to multiple loci"},
            "multimapped_toomany": {"color": "#f7a35c", "name": "Mapped to too many loci"},
            "unmapped_mismatches": {"color": "#e63491", "name": "Unmapped: too many mismatches"},
            "unmapped_tooshort": {"color": "#b1084c", "name": "Unmapped: too short"},
            "unmapped_other": {"color": "#7f0000", "name": "Unmapped: other"},
        }

        # Config for the plot
        pconfig = {
            "id": "star_alignment_plot",
            "title": "STAR: Alignment Scores",
            "ylab": "# Reads",
            "cpswitch_counts_label": "Number of Reads",
        }

        return bargraph.plot(self.star_data, keys, pconfig)

    def star_genecount_chart(self):
        """Make a plot for the ReadsPerGene output"""

        # Specify the order of the different possible categories
        keys = {
            "N_genes": {"color": "#2f7ed8", "name": "Overlapping Genes"},
            "N_noFeature": {"color": "#0d233a", "name": "No Feature"},
            "N_ambiguous": {"color": "#492970", "name": "Ambiguous Features"},
            "N_multimapping": {"color": "#f28f43", "name": "Multimapping"},
            "N_unmapped": {"color": "#7f0000", "name": "Unmapped"},
        }

        # Config for the plot
        pconfig = {
            "id": "star_gene_counts",
            "title": "STAR: Gene Counts",
            "ylab": "# Reads",
            "cpswitch_counts_label": "Number of Reads",
            "data_labels": ["Unstranded", "Same Stranded", "Reverse Stranded"],
        }
        datasets = [
            self.star_genecounts_unstranded,
            self.star_genecounts_first_strand,
            self.star_genecounts_second_strand,
        ]
        return bargraph.plot(datasets, [keys, keys, keys, keys], pconfig)
