from multiqc import config
from multiqc.plots import table, bargraph
from multiqc.base_module import BaseMultiqcModule
from multiqc.base_module import ModuleNoSamplesFound
import logging
from csv import DictReader

log = logging.getLogger(__name__)


class MultiqcModule(BaseMultiqcModule):
    """
    Possible mgikit output files are:

    1. **Sample stats file** ('*.L{1,2,3,4}.mgikit.sample_stats'): Sample statistics for each lane like yield, quality scores, cluster count
    2. **General info file** ('*.L{1,2,3,4}.mgikit.general'): Rounded up sample stats, but also includes lane-level stats
    3. **General info file** ('*.L{1,2,3,4}.mgikit.info'): Matching indexes within the data generated by a specific lane
    4. **Undetermined barcodes file** ('*.L{1,2,3,4}.mgikit.undetermined_barcode'): Barcodes that did not match with any sample.
    5. **Ambiguous barcodes file** ('*.L{1,2,3,4}.mgikit.ambiguous_barcode'): Barcodes that match with multiple samples. This situation can happen when setting a high mismatch threshold.

    Configuration options:

    ```yaml
    mgikit:
      # ignore undetermined and ambiguous cases in the report
      keep_core_samples: false
      # number of undetermined barcodes to be presented in the report. It takes any positive value less than or equal to the number of barcodes in the demultiplexer reports which is usually 50
      undetermined_barcode_threshold: 25
      # generate a brief version of the report. Ignores the reports for cluster per sample per lane.
      brief_report: false
      # number of decimal positions to be used for counts in the tables
      decimal_positions: 2
    ```
    """

    def __init__(self):
        super(MultiqcModule, self).__init__(
            name="mgikit",
            anchor="mgikit",
            href="https://github.com/sagc-bioinformatics/mgikit",
            info="Demultiplexes FASTQ files from an MGI sequencing instrument",
            doi="https://doi.org/10.1101/2024.01.09.574938",
        )

        for f in self.find_log_files("mgikit"):
            self.add_data_source(f)

        self.add_software_version(None)
        module_name = "MGIKIT"
        # Colours from Set3 of ColorBrewer R package, with a few additions for perfect matches (#1f78b4),
        # undetermined reads (#000000) and ambiguous reads (#a6761d)
        mgikit_colors = [
            "#1f78b4",
            "#8dd3c7",
            "#ffffb3",
            "#bebada",
            "#fb8072",
            "#80b1d3",
            "#fdb462",
            "#b3de69",
            "#fccde5",
            "#d9d9d9",
            "#bc80bd",
            "#ccebc5",
            "#ffed6f",
            "#000000",
            "#a6761d",
        ]
        undetermined_label = "Undetermined"
        ambiguous_label = "Ambiguous"

        visualisation_threshold = getattr(config, "mgikit", {}).get("undetermined_barcode_threshold", 25)
        brief_report = getattr(config, "mgikit", {}).get("brief_report", False)
        if brief_report:
            log.warning("The brief version version is activated!")

        show_all_samples = not getattr(config, "mgikit", {}).get("keep_core_samples", False)
        if not show_all_samples:
            log.warning("Undetermined and Ambiguous cases will be eliminated in this report!")

        decimal_positions = str(getattr(config, "mgikit", {}).get("decimal_positions", 2))
        fmt = "{:,." + decimal_positions + "f}"
        columns_headers = {}
        columns_headers["Clusters"] = {
            "title": "Clusters",
            "description": f"Total number of clusters for this sample ({config.read_count_desc})",
            "scale": "Blues",
            "shared_key": "read_count",
            "suffix": config.read_count_prefix,
            "format": fmt,
        }
        columns_headers["Yield"] = {
            "title": "Yield",
            "description": f"Total number of bases for this sample ({config.base_count_desc})",
            "scale": "Greens",
            "shared_key": "base_count",
            "suffix": config.base_count_prefix,
            "format": fmt,
        }
        columns_headers["Yield ≥ Q30"] = {
            "title": "Yield ≥ Q30",
            "description": f"Number of bases with a Phred score of 30 or higher, passing filter ({config.base_count_desc})",
            "scale": "Greens",
            "shared_key": "base_count",
            "suffix": config.base_count_prefix,
            "format": fmt,
            "hidden": True,
        }
        columns_headers["% R1 Yield ≥ Q30"] = {
            "title": "R1 Yield ≥ Q30",
            "description": "Percent of bases in R1 with phred qualty score ≥ Q30",
            "scale": "Greens",
            "max": 100,
            "min": 0,
            "suffix": "%",
            "hidden": True,
        }
        columns_headers["% R2 Yield ≥ Q30"] = {
            "title": "R2 Yield ≥ Q30",
            "description": "Percent of bases in R2 with phred qualty score ≥ Q30",
            "scale": "Greens",
            "max": 100,
            "min": 0,
            "suffix": "%",
            "hidden": True,
        }
        columns_headers["% R3 Yield ≥ Q30"] = {
            "title": "R3 Yield ≥ Q30",
            "description": "Percent of bases in R3 with phred qualty score ≥ Q30",
            "scale": "Greens",
            "max": 100,
            "min": 0,
            "suffix": "%",
            "hidden": True,
        }
        columns_headers["% Perfect Index"] = {
            "title": "Perfect Index",
            "description": "Percent of reads with perfect index (0 mismatches)",
            "scale": "Blues",
            "max": 100,
            "min": 0,
            "suffix": "%",
        }

        mgi_general_statistics = {}
        file_cnt = 0
        for f in self.find_log_files("mgikit/mgi_sample_stats"):
            file_cnt += 1

            for row in DictReader(f["f"].splitlines(), delimiter="\t"):
                sample = row["sample_id"]

                # Calculate total bases and Q30 bases
                total_bases = float(row["r1_bases"]) + float(row["r2_bases"]) + float(row["r3_bases"])
                q30_bases = float(row["r1_qc_30"]) + float(row["r2_qc_30"]) + float(row["r3_qc_30"])

                # Calculate percentages
                r1_q30_pct = (float(row["r1_qc_30"]) / float(row["r1_bases"])) * 100
                r2_q30_pct = (float(row["r2_qc_30"]) / float(row["r2_bases"])) * 100
                r3_q30_pct = (float(row["r3_qc_30"]) / float(row["r3_bases"])) * 100
                perfect_index_pct = (float(row["0-mismatches"]) / float(row["all_reads"])) * 100

                mgi_general_statistics[sample] = {
                    "Clusters": int(row["all_reads"]),
                    "Yield": total_bases,
                    "Yield ≥ Q30": q30_bases,
                    "% R1 Yield ≥ Q30": r1_q30_pct,
                    "% R2 Yield ≥ Q30": r2_q30_pct,
                    "% R3 Yield ≥ Q30": r3_q30_pct,
                    "% Perfect Index": perfect_index_pct,
                }

        if len(mgi_general_statistics.keys()) == 0:
            raise ModuleNoSamplesFound
        if file_cnt > 0:
            log.info(f"{file_cnt} general information log files (*.mgikit.sample_stats) were loaded")
        self.general_stats_addcols(mgi_general_statistics, columns_headers)
        self.write_data_file(mgi_general_statistics, "multiqc_mgikit_general")

        # Lane statistics
        mgi_lane_statistics = {}
        columns_headers = {}
        columns_headers["Mb Total Yield"] = {
            "title": "Yield",
            "description": "Number of megabases",
            "scale": "Greens",
            "suffix": "Mb",
            "format": fmt,
            "min": 0,
        }

        columns_headers["M Total Clusters"] = {
            "title": "Clusters",
            "description": "Total number of clusters for this lane (millions)",
            "scale": "Blues",
            "suffix": "M",
            "format": fmt,
            "min": 0,
        }

        columns_headers["% bases ≥ Q30"] = {
            "title": "Bases ≥ Q30",
            "description": "Percent of bases with a Phred score of 30 or higher, passing filter",
            "scale": "Greens",
            "max": 100,
            "min": 0,
            "suffix": "%",
        }

        columns_headers["Mean Quality"] = {
            "title": "Mean Quality",
            "description": "Average phred qualty score",
            "scale": "Oranges",
            "max": 40,
            "min": 0,
        }

        columns_headers["% Perfect Index"] = {
            "title": "Perfect Index",
            "description": "Percent of reads with perfect index (0 mismatches)",
            "scale": "Blues",
            "max": 100,
            "min": 0,
            "suffix": "%",
        }

        for f in self.find_log_files("mgikit/mgi_general_info"):
            lines = f["f"].splitlines()
            for line_itr in range(len(lines)):
                line = lines[line_itr]
                if line.startswith("#Lane statistics"):
                    # found it. Information is in the next 2 lines
                    header = lines[line_itr + 1].split("\t")
                    vals = lines[line_itr + 2].split("\t")
                    mgi_lane_statistics[vals[0]] = {}
                    for i in range(1, len(header)):
                        mgi_lane_statistics[vals[0]][header[i]] = vals[i]
                    break
        if len(mgi_lane_statistics.keys()) > 0:
            self.add_section(
                name="Lane Statistics",
                anchor="mgikit-lane-statistics",
                description="Statistics about each lane for each flowcell.",
                helptext="This longer string (can be **markdown**) helps explain how to interpret the plot",
                plot=table.plot(
                    mgi_lane_statistics,
                    columns_headers,
                    {
                        "id": "mgi_lane_stats_plot",
                        "title": "Lane Statistics",
                        "sort_rows": False,
                        "col1_header": "Run ID - Lane",
                    },
                ),
            )

        # Statistics about matching reads
        mgi_sample_reads_data = dict()
        mgi_lane_read_data = {}
        mgi_lane_sample_read_data = {}
        file_cnt = 0
        for f in self.find_log_files("mgikit/mgi_sample_reads"):
            file_cnt += 1
            log_details = f["fn"].split(".")
            if len(log_details[-4]) == 0:
                curr_label = log_details[-3]
            else:
                curr_label = log_details[-4] + "-" + log_details[-3]

            mgi_lane_read_data[curr_label] = {}
            mgi_lane_sample_read_data[curr_label] = {}
            lines = f["f"].splitlines()
            header = [x.strip() for x in lines[0].split("\t")]
            for i in range(1, len(header)):
                mgi_lane_read_data[curr_label][header[i]] = 0

            for line in lines[1:]:
                vals = line.split("\t")
                if len(vals) < 2 or vals[0] == "Total":
                    continue

                if vals[0] in [undetermined_label, ambiguous_label]:
                    mgi_lane_read_data[curr_label][vals[0]] = int(vals[1])

                if not show_all_samples and vals[0] in [undetermined_label, ambiguous_label]:
                    continue
                if vals[0] in mgi_lane_sample_read_data[curr_label].keys():
                    log.warning("Same sample ({}) appears twice in the same log file!".format(vals[0]))

                mgi_lane_sample_read_data[curr_label][vals[0]] = {}

                if vals[0] not in mgi_sample_reads_data.keys():
                    mgi_sample_reads_data[vals[0]] = {}

                for i in range(1, len(header)):
                    if header[i] not in mgi_sample_reads_data[vals[0]].keys():
                        mgi_sample_reads_data[vals[0]][header[i]] = 0

                    mgi_lane_sample_read_data[curr_label][vals[0]][header[i]] = 0
                    mgi_sample_reads_data[vals[0]][header[i]] += int(vals[i])
                    if vals[0] not in [undetermined_label, ambiguous_label]:
                        mgi_lane_read_data[curr_label][header[i]] += int(vals[i])
                    mgi_lane_sample_read_data[curr_label][vals[0]][header[i]] += int(vals[i])
        if len(mgi_sample_reads_data) == 0:
            raise ModuleNoSamplesFound

        mgi_sample_reads_data_filtered = None
        if file_cnt > 0:
            log.info(f"{file_cnt} read information log files (*.mgikit.info) were loaded")

            # Filter out samples matching ignored sample names
            mgi_sample_reads_data_filtered = self.ignore_samples(mgi_sample_reads_data)
            cats = {}
            for i in range(1, len(header)):
                cats[header[i]] = {"name": header[i], "color": mgikit_colors[i - 1]}

            cats[undetermined_label] = {"name": undetermined_label, "color": mgikit_colors[len(header) - 1]}
            cats[ambiguous_label] = {"name": ambiguous_label, "color": mgikit_colors[len(header)]}
            pconfig = {
                "id": "mgi_lane_read_plot",
                "title": module_name + ": Clusters by lane",
                "xlab": "Lane",
                "ylab": "# Reads",
                "hide_zero_cats": True,
            }

            lane_read_plt = bargraph.plot(mgi_lane_read_data, cats, pconfig=pconfig)

            # Add a report section with the line plot
            self.add_section(
                name="Clusters by lane",
                anchor="mgikit-first",
                description="Number of reads per lane.",
                helptext="The number of reads with different amounts of mismatches in their indexes, up to the maximum allowed during demultiplexing.",
                plot=lane_read_plt,
            )

            # Nothing found - raise a UserWarning to tell MultiQC
            if not mgi_sample_reads_data_filtered or len(mgi_sample_reads_data_filtered.keys()) == 0:
                log.info("No sample left after filteration!")
                raise ModuleNoSamplesFound

            pconfig = {
                "id": "mgi_sample_read_plot",
                "title": module_name + ": Clusters by sample",
                "xlab": "Sample",
                "ylab": "Matching reads",
                "hide_zero_cats": True,
                "tt_decimals": 0,
            }

            sample_read_plt = bargraph.plot(mgi_sample_reads_data_filtered, cats, pconfig=pconfig)

            # Add a report section with the line plot
            self.add_section(
                name="Clusters by sample",
                anchor="mgikit-second",
                description="Number of reads per sample.",
                helptext="Perfect index reads are those that do not have a single mismatch. All samples are aggregated across lanes combinned. Undetermined reads are treated as a separate sample.",
                plot=sample_read_plt,
            )

            if not brief_report:
                cntr = 0
                for lane in sorted(mgi_lane_read_data.keys()):
                    cntr += 1
                    pconfig = {
                        "id": "mgi_sample_read_plot" + str(cntr),
                        "title": module_name + ": Clusters by sample for lane (" + lane + ")",
                        "xlab": "Sample",
                        "ylab": "Matching reads",
                        "hide_zero_cats": True,
                        "tt_decimals": 0,
                    }

                    sample_read_plt = bargraph.plot(mgi_lane_sample_read_data[lane], cats, pconfig=pconfig)

                    # Add a report section with the line plot
                    self.add_section(
                        name="Clusters by sample for lane: " + lane,
                        anchor="mgikit-lane-sample_reads-" + str(cntr),
                        description="This plot shows the number of reads per sample for Lane: " + lane,
                        helptext="This longer string (can be **markdown**) helps explain how to interpret the plot",
                        plot=sample_read_plt,
                    )

        # Check undetermined barcodes
        mgi_undetermined_barcode_data = {}
        mgi_undetermined_barcode_cnt = {}

        file_cnt = 0
        cat_lane = {}
        for f in self.find_log_files("mgikit/mgi_undetermined_barcode"):
            file_cnt += 1
            log_details = f["fn"].split(".")
            if len(log_details[-4]) == 0:
                curr_label = log_details[-3]
            else:
                curr_label = log_details[-4] + "-" + log_details[-3]

            lines = f["f"].splitlines()
            cat_lane[curr_label] = {
                "name": curr_label,
                "color": mgikit_colors[file_cnt - 1],
            }
            for line in lines:
                vals = [x.strip() for x in line.split()]
                if len(vals) < 2:
                    continue
                if vals[0] not in mgi_undetermined_barcode_data.keys():
                    mgi_undetermined_barcode_data[vals[0]] = {}
                mgi_undetermined_barcode_data[vals[0]][curr_label] = int(vals[1])

                if vals[0] not in mgi_undetermined_barcode_cnt.keys():
                    mgi_undetermined_barcode_cnt[vals[0]] = int(vals[1])
                else:
                    mgi_undetermined_barcode_cnt[vals[0]] += int(vals[1])

        if file_cnt > 0:
            log.info(f"{file_cnt} undetermined barcode log files (*.mgikit.undetermined_barcode) were loaded")

            pconfig = {
                "id": "mgi_lane_undetermined_plot",
                "title": module_name + ": Undetermined barcodes by lane",
                "xlab": "Barcode",
                "ylab": "# Reads",
                "hide_zero_cats": True,
                "sort_samples": False,
                "tt_decimals": 0,
            }

            sorted_dic = {
                k[0]: mgi_undetermined_barcode_data[k[0]]
                for k in list(
                    sorted(
                        mgi_undetermined_barcode_cnt.items(),
                        key=lambda item: -int(item[1]),
                    )
                )[:visualisation_threshold]
            }
            lane_barcode_plt = bargraph.plot(sorted_dic, cat_lane, pconfig=pconfig)

            # Add a report section with the line plot
            self.add_section(
                name="Undetermined barcodes by lane",
                anchor="mgikit-undetermined-barcode",
                description="Count of the top twenty five most abundant undetermined barcodes by lanes.",
                plot=lane_barcode_plt,
            )

        # Check ambiguous barcodes
        mgi_ambiguous_barcode_data = {}
        file_cnt = 0
        cat_lane = {}
        for f in self.find_log_files("mgikit/mgi_ambiguous_barcode"):
            file_cnt += 1
            log_details = f["fn"].split(".")
            if len(log_details[-4]) == 0:
                curr_label = log_details[-3]
            else:
                curr_label = log_details[-4] + "-" + log_details[-3]
            lines = f["f"].splitlines()
            for line in lines:
                vals = [x.strip() for x in line.split()]
                if len(vals) < 2:
                    continue
                mgi_ambiguous_barcode_data[curr_label + ":" + vals[0]] = {
                    "Frequency": int(vals[1]),
                }

        if file_cnt > 0:
            log.info(f"{file_cnt} ambiguous barcode log files (*.mgikit.ambiguous_barcode) were loaded")
            sorted_dic = {
                k: v
                for k, v in list(
                    sorted(
                        mgi_ambiguous_barcode_data.items(),
                        key=lambda item: -int(item[1]["Frequency"]),
                    )
                )[:visualisation_threshold]
            }

            # Add a report section with the line plot
            columns_headers = {}
            columns_headers["Frequency"] = {
                "title": "Frequency",
                "description": "The frequency of this barcode in the dataset",
                "min": 0,
            }
            self.add_section(
                name="Ambiguous barcodes by lane",
                anchor="mgikit-ambiguous-barcode",
                description="Ambiguous barcodes that could match to multiple samples.",
                helptext="This longer string (can be **markdown**) helps explain how to interpret the plot",
                plot=table.plot(
                    sorted_dic,
                    columns_headers,
                    {
                        "id": "mgi_lane_ambiguous_plot",
                        "title": module_name + ": Ambiguous barcodes by lane",
                        "sort_rows": False,
                        "col1_header": "Barcode",
                    },
                ),
            )
