import os
import re
import sys
import json
import yaml
import argparse
from pathlib import Path

import numpy as np
import pandas as pd

import vcf
from BCBio import GFF
from Bio import SeqIO
from Bio.Seq import Seq


def convert_vcf(fname):
    """Convert VCF to JSON."""
    output = []

    if os.path.getsize(fname) == 0:
        print(f'Empty VCF: "{fname}"')
        return output

    print(f'Parsing VCF: "{fname}"')
    with open(fname) as fd:
        vcf_reader = vcf.Reader(fd)

        # check caller
        caller_source = vcf_reader.metadata["source"][0].lower()
        if caller_source.startswith("lofreq"):
            mode = "lofreq"
        elif caller_source.startswith("shorah"):
            mode = "shorah"
        else:
            raise RuntimeError(f"Invalid variant caller: {caller_source}")

        # parse records
        for record in vcf_reader:
            if mode == "lofreq":
                freq = round(record.INFO["AF"], 3)
            elif mode == "shorah":
                freq = round(
                    np.mean(
                        [v for k, v in record.INFO.items() if k.startswith("Freq")]
                    ),
                    3,
                )

            output.append(
                {
                    "position": record.POS,
                    "reference": record.REF,
                    "variant": [v.sequence for v in record.ALT],
                    "frequency": freq,
                    "posterior": round(1 - 10 ** (-record.QUAL / 10), 3),
                }
            )

    return json.dumps(output)


def get_metainfo(metainfo_yaml):
    """load metainformation for the GFF and primers from YAML file"""

    if metainfo_yaml:
        print(f'Parsing metainformation: "{metainfo_yaml}"')
        with open(metainfo_yaml) as fd:
            metainfo = yaml.load(fd.read(), Loader=yaml.SafeLoader)
        assert (
            type(metainfo) is dict
        ), f"Probable syntax error in {metainfo_yaml} - need a dictionnary at top level, got {type(metainfo)} instead."
        return metainfo
    else:
        print("No metainformation YAML provided, skipping.")
        return {}


def parse_gff(fname):
    """Convert GFF to map."""
    features = []

    print(f'Parsing GFF: "{fname}"')
    with open(fname) as fd:
        for record in GFF.parse(fd):
            for feature in record.features:
                features.append(
                    {
                        "id": record.id,
                        "type": feature.type,
                        "name": feature.qualifiers.get("Name", [feature.id])[0],
                        "start": int(feature.location.start),
                        "end": int(feature.location.end),
                    }
                )
    return features


def arrange_gff_data(features):
    """Add row number to each feature."""
    features.sort(key=lambda f: f["start"])

    rows = []
    for feature in features:
        if not rows:
            feature["row_cnt"] = 0
            rows.append([feature])
        else:
            found = False
            for idx, row in enumerate(rows):
                if row[-1]["end"] <= feature["start"]:
                    feature["row_cnt"] = idx
                    row.append(feature)
                    found = True
                    break
            if not found:
                feature["row_cnt"] = len(rows)
                rows.append([feature])

    return [item for row in rows for item in row]


def get_gff_data(gff_dir, gff_metainfo={}):
    """Returns a map with filename key and gff json data."""
    if gff_metainfo == None:
        gff_metainfo = {}
    assert (
        type(gff_metainfo) is dict
    ), f"Probable syntax error in metainfo YAML - need a dictionnary at [gff], got {type(gff_metainfo)} instead."

    gff_map = {}
    if not gff_dir:
        print("No gff directory provided, skipping.")
        return gff_map
    for path in os.listdir(gff_dir):
        full_path = os.path.join(gff_dir, path)
        description = os.path.splitext(path)[0]
        if description in gff_metainfo:
            description = gff_metainfo[description]
        gff_map[description] = arrange_gff_data(parse_gff(full_path))
    return gff_map


def get_primers_data(full_path, consensus, primers_metainfo={}):
    """Returns a map with filename key and primers json data."""
    if primers_metainfo == None:
        primers_metainfo = {}
    assert (
        type(primers_metainfo) is dict
    ), f"Probable syntax error in metainfo YAML - need a dictionnary at [primers], got {type(primers_metainfo)} instead."

    primers_map = {}
    if not full_path:
        print("No primers table provided, skipping.")
        return primers_map
    else:
        print(f'Parsing GFF: "{full_path}"')
    consensus_upper = consensus.upper()
    description = os.path.splitext(os.path.basename(full_path))[0]
    if description in primers_metainfo:
        description = primers_metainfo[description]
    csv = pd.read_csv(full_path, sep=",")
    primers = [
        {"name": row[0], "seq": row[1].upper()} for row in csv[["name", "seq"]].values
    ]
    primer_locations = []
    for entry in primers:
        lookup_sequence = entry["seq"]
        # If the sequence corresponds to a `right` primer, then reverse and
        # complement.
        if "RIGHT" in entry["name"].upper():
            seq = Seq(lookup_sequence)
            lookup_sequence = str(seq.reverse_complement())
        offsets = [
            m.start()
            for m in re.finditer("(?=" + lookup_sequence + ")", consensus_upper)
        ]
        for offset in offsets:
            primer_locations.append(
                {
                    "name": entry["name"],
                    "seq": entry["seq"],
                    "start": offset,
                    "end": offset + len(entry["seq"]) - 1,
                }
            )
    if primer_locations:
        primers_map[description] = arrange_gff_data(primer_locations)
    else:
        print(f'No primer was mapped from "{full_path}".')
    return primers_map


def convert_coverage(fname, sample_name=None):
    """Convert the read coverage to bp coverage."""
    print(f'Parsing coverage: "{fname}"')
    csv = pd.read_csv(
        fname, sep="\t", compression="infer", index_col=["ref", "pos"], header=0
    )
    col = None
    if len(csv.columns) == 1:
        if sample_name is not None and sample_name != csv.columns[0]:
            print(
                f'Ooops: column name "{csv.columns[0]}" in file is different from requested sample name "{sample_name}"!'
            )
        col = csv[csv.columns[0]]
    else:
        assert (
            sample_name is not None
        ), "Sample name is required when using combined coverage TSV."
        col = csv[sample_name]
    return col.values.tolist()


def assemble_html_page(html_file_in, html_file_out, placeholder_replacement_map):

    # assemble webpage
    with open(html_file_in) as fd:
        raw_html = fd.read()

    # TODO: make this more robust
    mod_html = raw_html
    for key, value in placeholder_replacement_map.items():
        mod_html = mod_html.replace(key, value)

    with open(html_file_out, "w") as fd:
        fd.write(mod_html)


def assemble_snv_calling_visualization_webpage(
    sample_name,
    consensus_file,
    coverage_file,
    vcf_file,
    gff_directory,
    primers_file,
    html_file_in,
    html_file_out,
    reference_file,
    metainfo_yaml,
):

    # parse the consensus sequence
    print(f'Parsing consensus: "{consensus_file}"')
    consensus = next(SeqIO.parse(consensus_file, "fasta")).seq.upper()

    # parse coverage file
    coverage = convert_coverage(coverage_file, sample_name.replace("/", "-"))

    # load biodata in json format
    vcf_json = convert_vcf(vcf_file)
    metainfo = get_metainfo(metainfo_yaml)
    gff_map = get_gff_data(
        gff_directory, gff_metainfo=metainfo["gff"] if "gff" in metainfo else {}
    )
    primers_map = get_primers_data(
        primers_file,
        str(consensus),
        primers_metainfo=metainfo["primers"] if "primers" in metainfo else {},
    )

    # parse the reference name
    reference_name = (
        re.search(r"(.+)\.fa.*", os.path.basename(reference_file))
        .group(1)
        .replace("_", "&nbsp;")
    )

    embed_code = f"""
        var sample_name = \"{sample_name}\"
        var consensus = \"{consensus}\"
        var coverage = {coverage}
        var vcfData = {vcf_json}
        var gffData = {gff_map}
        var primerData = {primers_map}
        var reference_name = \"{reference_name}\"
    """

    assemble_html_page(
        html_file_in, html_file_out, {"{EXTERNAL_SNAKEMAKE_CODE_MARKER}": embed_code}
    )


def assemble_alignment_visualization_webpage(
    sample_name, nwk_file, reference_uri_file, bam_uri_file, html_file_in, html_file_out
):

    nwk_tree = ""
    if nwk_file:
        with open(nwk_file) as fd:
            nwk_tree = fd.read().rstrip("\n")

    embed_code = f"""
        var sample_name = \"{sample_name}\"
        var phylogenyData = \"{nwk_tree}\"
    """

    # read the reference uri file
    with open(reference_uri_file) as fd:
        reference_uri = fd.read().rstrip("\n")

    # read the BAM uri file
    with open(bam_uri_file) as fd:
        bam_uri = fd.read().rstrip("\n")

    assemble_html_page(
        html_file_in,
        html_file_out,
        {
            "{EXTERNAL_SNAKEMAKE_CODE_MARKER}": embed_code,
            "{EXTERNAL_FASTA_URI}": reference_uri,
            "{EXTERNAL_BAM_URI}": bam_uri,
        },
    )


def main():
    """Parse command line, run default functions."""
    # parse command line
    parser = argparse.ArgumentParser(
        description="Generate HTML visual report from VCF variants",
        epilog="at minimum, either provide `-v` and `-c` or provide `-w`",
    )
    parser.add_argument(
        "-f",
        "--consensus",
        metavar="FASTA",
        required=False,
        type=str,
        dest="consensus_file",
        help="consensus sequence of this sample",
    )
    parser.add_argument(
        "-c",
        "--coverage",
        metavar="TSV",
        required=False,
        default="variants/coverage.tsv",
        type=str,
        dest="coverage_file",
        help="global coverage table",
    )
    parser.add_argument(
        "-v",
        "--vcf",
        metavar="VCF",
        required=False,
        type=str,
        dest="vcf_file",
        help="VCF containing the SNPs to be visualised",
    )
    parser.add_argument(
        "-g",
        "--gff",
        metavar="DIR",
        required=False,
        type=str,
        dest="gff_directory",
        help="directory containing GFF annotations",
    )
    parser.add_argument(
        "-p",
        "--primers",
        metavar="CSV",
        required=False,
        type=str,
        dest="primers_file",
        help="table with primers",
    )
    parser.add_argument(
        "-m",
        "--metainfo",
        metavar="YAML",
        required=False,
        type=str,
        dest="metainfo_yaml",
        help="metainformation for the GFF and primers",
    )
    parser.add_argument(
        "-x",
        "--snv_calling_template",
        metavar="HTML",
        required=False,
        default=f"{os.path.dirname(__file__)}/snv_calling_visualization.html",
        type=str,
        dest="html_template_snv",
        help="HTML template used to generate visual report for snv calling",
    )
    parser.add_argument(
        "-y",
        "--alignment_template",
        metavar="HTML",
        required=False,
        default=f"{os.path.dirname(__file__)}/alignment_visualization.html",
        type=str,
        dest="html_template_alignment",
        help="HTML template used to generate visual report for alignments",
    )
    parser.add_argument(
        "-s",
        "--html_out_snv_calling",
        metavar="HTML",
        required=False,
        type=str,
        dest="html_out_snv_calling",
        help="produced HTML report for SNV calling",
    )
    parser.add_argument(
        "-a",
        "--html_out_alignment",
        metavar="HTML",
        required=False,
        type=str,
        dest="html_out_alignment",
        help="produced HTML report for phylogeny and read alignment",
    )
    parser.add_argument(
        "-w",
        "--wildcards",
        metavar="SAMPLE/DATE",
        required=False,
        type=str,
        dest="wildcards_dataset",
        help="sample's two-level directory hierarchy prefix",
    )
    parser.add_argument(
        "-r",
        "--reference",
        metavar="FASTA",
        required=False,
        default="variants/cohort_consensus.fasta",
        type=str,
        dest="reference_file",
        help="reference against which SNVs were called (e.g.: cohort's consensus)",
    )
    parser.add_argument(
        "-n",
        "--nwk",
        metavar="NWK",
        required=False,
        type=str,
        dest="nwk_file",
        help="phylogenetic tree in NWK format",
    )
    parser.add_argument(
        "-u",
        "--reference_uri_file",
        metavar="FILE",
        required=False,
        type=str,
        dest="reference_uri_file",
        help="reference file uri",
    )
    parser.add_argument(
        "-b",
        "--bam_uri_file",
        metavar="FILE",
        required=False,
        type=str,
        dest="bam_uri_file",
        help="bam file uri",
    )

    args = parser.parse_args()

    # defaults which can be guess from one another
    if (
        args.vcf_file == None
    ):  # e.g.: samples/140074_395_D02/20200615_J6NRK/variants/SNVs/snvs.vcf
        assert (
            args.wildcards_dataset != None
        ), "cannot automatically find VCF without wildcards"
        args.vcf_file = os.path.join(
            args.wildcards_dataset, "variants", "SNVs", "snvs.vcf"
        )

    if (
        args.consensus_file == None
    ):  # e.g.: samples/140074_395_D02/20200615_J6NRK/references/ref_majority.fasta
        assert (
            args.wildcards_dataset != None
        ), "cannot automatically find consensus without wildcards"
        args.consensus_file = os.path.join(
            args.wildcards_dataset, "references", "ref_majority.fasta"
        )

    if args.wildcards_dataset == None:
        assert (
            args.vcf_file != None and args.consensus != None
        ), "cannot deduce wilcards without a consensus and a vcf"
        try1 = "/".join(os.path.normpath(args.vcf_file).split(os.path.sep)[-5:-3])
        try2 = "/".join(os.path.normpath(args.consensus_file).split(os.path.sep)[-4:-2])
        assert (
            try1 == try2
        ), f"cannot deduce wildcards automatically from <{args.vcf_file}> and <{args.consensus_file}>, please specify explicitly using `--wirdcards`"
        args.wildcards_dataset = try1

    if args.html_out_snv_calling == None:
        args.html_out_snv_calling = os.path.join(
            args.wildcards_dataset, "visualization", "snv_calling.html"
        )

    if args.html_out_alignment == None:
        args.html_out_alignment = os.path.join(
            args.wildcards_dataset, "visualization", "alignment.html"
        )

    # check mandatory files exist
    for n, f in {
        "vcf": args.vcf_file,
        "consensus": args.consensus_file,
        "coverage": args.coverage_file,
        "snv_calling_template": args.html_template_snv,
        "alignment_template": args.html_template_alignment,
    }.items():
        if not os.path.exists(f):
            parser.error(f"{n} file <{f}> does not exist!")

    # check optional files exist if specified
    for n, f in {
        "gff": args.gff_directory,
        "primers": args.primers_file,
        "metainfo": args.metainfo_yaml,
    }.items():
        if f and not os.path.exists(f):
            parser.error(f"{n} file <{f}> does not exist!")

    if (
        args.reference_uri_file == None
    ):  # e.g.: samples/140074_395_D02/20200615_J6NRK/visualization/reference_uri_file
        assert (
            args.wildcards_dataset != None
        ), "cannot automatically find reference_uri_file without wildcards"
        args.reference_uri_file = os.path.join(
            args.wildcards_dataset, "visualization", "reference_uri_file"
        )

    if (
        args.bam_uri_file == None
    ):  # e.g.: samples/140074_395_D02/20200615_J6NRK/visualization/bam_uri_file
        assert (
            args.wildcards_dataset != None
        ), "cannot automatically find bam_uri_file without wildcards"
        args.bam_uri_file = os.path.join(
            args.wildcards_dataset, "visualization", "bam_uri_file"
        )

    # parse the sample name
    path_components = os.path.normpath(args.wildcards_dataset).split(os.path.sep)
    sample_name = "/".join(path_components[-2:])

    # run the visual report generator for SNV calling
    assemble_snv_calling_visualization_webpage(
        sample_name,
        consensus_file=args.consensus_file,
        coverage_file=args.coverage_file,
        vcf_file=args.vcf_file,
        gff_directory=args.gff_directory,
        primers_file=args.primers_file,
        html_file_in=args.html_template_snv,
        html_file_out=args.html_out_snv_calling,
        reference_file=args.reference_file,
        metainfo_yaml=args.metainfo_yaml,
    )

    # run the visual report generator for read alignment and phylogeny

    assemble_alignment_visualization_webpage(
        sample_name,
        nwk_file=args.nwk_file,
        reference_uri_file=args.reference_uri_file,
        bam_uri_file=args.bam_uri_file,
        html_file_in=args.html_template_alignment,
        html_file_out=args.html_out_alignment,
    )


if __name__ == "__main__":
    main()
