#!/usr/bin/env python3

import os
import argparse
from alignmentIntervals import read_fasta

from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Seq import Seq
import sh
import numpy as np
import pandas as pd


__author__ = "Susana Posada-Cespedes"
__license__ = "Apache2.0"
__maintainer__ = "Ivan Topolsky"
__email__ = "v-pipe@bsse.ethz.ch"


DBG = True if os.environ.get("DBG") is not None else False


def parse_args():
    """Set up the parsing of command-line arguments"""

    parser = argparse.ArgumentParser(
        description="Benchmark: test",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    requiredNamed = parser.add_argument_group("required named arguments")
    requiredNamed.add_argument(
        "-f",
        required=True,
        default=None,
        metavar="FASTA",
        dest="haplotype_seqs",
        help="Fasta file containing either the sequences of the true "
        "haplotypes or haplotypes sequences (msa) already reported using "
        "the same indexing as the reference/consensus sequence",
    )
    requiredNamed.add_argument(
        "-s",
        required=True,
        default=None,
        metavar="CSV",
        dest="snvs",
        help="File containing called SNVs",
    )
    requiredNamed.add_argument(
        "-N",
        required=False,
        default="sample",
        metavar="STR",
        dest="sampleID",
        help="Patient/sample identifiers",
    )
    parser.add_argument(
        "-m",
        required=False,
        default=None,
        metavar="FASTA",
        dest="haplotype_master",
        type=str,
        help="Fasta file containing the sequence with respect to which SNVs "
        "were called",
    )
    parser.add_argument(
        "--ref",
        required=False,
        default=None,
        metavar="FASTA",
        dest="reference",
        type=str,
        help="Fasta file containing the reference sequence with respect to "
        "which reads were aligned",
    )
    parser.add_argument(
        "-d",
        required=False,
        default="unif",
        metavar="str",
        dest="freq_dstr",
        type=str,
        choices=["unif", "geom", "dirichlet", "cust"],
        help="Distribution of haplotype frequencies",
    )
    parser.add_argument(
        "-gr",
        required=False,
        default=0.75,
        metavar="FLOAT",
        dest="ratio",
        type=float,
        help="Sucess probability for the geometric distribution",
    )
    parser.add_argument(
        "-df",
        required=False,
        default=None,
        metavar="FASTA",
        dest="dirichlet_freqs",
        type=str,
        help="File containing haplotype frequencies",
    )
    parser.add_argument(
        "-ci",
        required=False,
        default=None,
        metavar="chrm:start-end",
        dest="coverage_intervals",
        type=str,
        help="File containing coverage intervals",
    )
    parser.add_argument(
        "--no-expansion",
        required=False,
        default=False,
        action="store_true",
        dest="no_expansion",
        help="Coverage intervals do not correspond to region use to run "
        "ShoRAH, but the actual target region",
    )
    parser.add_argument(
        "--caller",
        required=False,
        default="shorah",
        metavar="str",
        dest="snv_caller",
        type=str,
        choices=["shorah", "lofreq"],
        help="Inidcate if other software different from ShoRAH was used for "
        "SNV calling",
    )
    parser.add_argument(
        "-wl",
        required=False,
        default=201,
        metavar="INT",
        dest="window_len",
        type=int,
        help="Window length used by ShoRAH to construct overlapping windows",
    )
    parser.add_argument(
        "-ws",
        required=False,
        default=3,
        metavar="INT",
        dest="window_shift",
        type=int,
        help="Number of window shifts used by ShoRAH to construct overlapping "
        "windows",
    )
    parser.add_argument(
        "-cf",
        required=False,
        default=None,
        metavar="TXT",
        dest="coverage",
        type=str,
        help="File to read coverage per window used by ShoRAH, or a "
        "tab-separated values file containing coverage per locus",
    )
    parser.add_argument(
        "-ms",
        required=False,
        default=False,
        action="store_true",
        dest="msa",
        help="Indicate if the multiple sequence alignment including "
        "reference/consensus sequence should be constructed",
    )
    parser.add_argument(
        "--only-dels",
        required=False,
        default=False,
        action="store_true",
        dest="only_deletions",
        help="Indicate if only performance based on deletions should reported",
    )
    parser.add_argument(
        "--long-dels",
        required=False,
        default=False,
        action="store_true",
        dest="long_deletions",
        help="Indicate if deletions should be parsed as multipe-base deletions",
    )
    parser.add_argument(
        "-t",
        required=False,
        default=False,
        action="store_true",
        dest="output_true",
        help="Indicate if file containing expected SNVs should be reported. "
        "Report using 1-based indexing for the position",
    )
    parser.add_argument(
        "-mafft",
        required=False,
        default="mafft",
        metavar="PATH",
        dest="mafft",
        type=str,
        help="Path to binaries for the multiple sequence aligner MAFFT",
    )
    parser.add_argument(
        "-of",
        required=False,
        default="performance.tsv",
        metavar="OUTPUT",
        dest="outfile",
        type=str,
        help="Output file - file containing expected SNVs",
    )
    parser.add_argument(
        "-od",
        required=False,
        default=None,
        metavar="DIR",
        dest="outdir",
        type=str,
        help="Output directory for intermediate files",
    )

    return parser.parse_args()


def frequencies(freq_dstr, num_haplotypes, ratio=0.75, infile=None):
    "Compute the expected haplotype frequencies"
    if freq_dstr == "unif":
        haplotype_freqs = np.repeat(1 / num_haplotypes, num_haplotypes)
    elif freq_dstr == "geom":
        haplotype_freqs = [ratio ** (i + 1) for i in range(num_haplotypes)]
        haplotype_freqs = np.asarray(haplotype_freqs)
        haplotype_freqs = haplotype_freqs / np.sum(haplotype_freqs)
    elif freq_dstr == "dirichlet":
        # Read haplotype frequencies from output file
        if infile is None:
            raise IOError("Input file containing haplotype frequencies is expected")
        ids, haplotype_freqs = read_fasta(infile)
        haplotype_freqs = np.asarray(haplotype_freqs, dtype=float)

    return haplotype_freqs


def parse_info(df, snvcaller):
    if snvcaller == "shorah":
        df_info = pd.DataFrame.from_dict(
            [
                dict([entry.strip().split("=") for entry in line.split(";")])
                for line in df["INFO"]
            ]
        ).astype("float")
        # We ignore columns with 0-counts to compute the SNV frequency. A zero
        # count means that the SNV was not found in the corresponding window.
        df_freq = df_info[["Freq1", "Freq2", "Freq3"]].copy()
        df_freq[df_freq == 0] = np.nan
        df_freq = df_freq.mean(axis=1)
    elif snvcaller == "lofreq":
        df["INFO"] = df["INFO"].str.replace("INDEL", "INDEL=1")
        df_info = pd.DataFrame.from_dict(
            [
                dict([entry.strip().split("=") for entry in line.split(";")])
                for line in df["INFO"]
            ]
        )
        df_freq = df_info["AF"]
    return df_freq


def parse_vcf(snvfile, snvcaller):
    # Read VCF file to infer how many lines to skip
    skiplines = 0
    with open(snvfile, "r") as infile:
        for line in infile:
            if not line.startswith("##"):
                break
            skiplines += 1
    try:
        df_snvs = pd.read_csv(
            snvfile, sep="\t", skiprows=skiplines, header=0, compression=None
        )
        df_snvs = df_snvs.rename(columns={"#CHROM": "CHROM"})
        df_snvs["FREQ"] = parse_info(df_snvs, snvcaller)
    except pd.errors.EmptyDataError:
        df_snvs = pd.DataFrame()
    return df_snvs


def true_snvs(
    haplotype_master_arr,
    haplotype_master,
    haplotype_seqs,
    num_haplotypes,
    haplotype_freqs,
    long_deletions,
    alphabet,
):
    """
    Extract expected SNVs using the MSA of the true haplotype sequences and
    the reference sequence
    """
    # loci = np.arange(haplotype_master_arr.size)
    haplotype_idx = np.arange(num_haplotypes)
    variants = haplotype_master_arr != haplotype_seqs

    df_snvs = pd.DataFrame(columns=("POS", "REF", "ALT", "FREQ", "HAPLOTYPES"))
    num_snvs = 0
    for locus in range(haplotype_master_arr.size):
        idxs = variants[:, locus]
        if np.any(idxs):
            var = haplotype_seqs[idxs, locus]
            snv_freq = haplotype_freqs[idxs]
            if np.sum(idxs) == 1:
                df_snvs.loc[num_snvs] = [
                    locus,
                    haplotype_master_arr[locus].decode(),
                    var[0].decode(),
                    snv_freq[0],
                    haplotype_idx[idxs].astype(str)[0],
                ]
                num_snvs += 1
            else:
                for base in alphabet:
                    idxs_base = var == base
                    if np.sum(idxs_base) > 0:
                        hap_aux = ",".join(haplotype_idx[idxs][idxs_base].astype(str))
                        df_snvs.loc[num_snvs] = [
                            locus,
                            haplotype_master_arr[locus].decode(),
                            base.decode(),
                            np.sum(snv_freq[idxs_base]),
                            hap_aux,
                        ]
                        num_snvs += 1
    df_snvs["POS"] = df_snvs["POS"].astype(int)

    if long_deletions:
        df_long_dels = pd.DataFrame(
            {
                "POS": pd.Series([], dtype="int"),
                "REF": pd.Series([], dtype="str"),
                "ALT": pd.Series([], dtype="str"),
                "FREQ": pd.Series([], dtype="float"),
                "HAPLOTYPES": pd.Series([], dtype="str"),
            }
        )
        for idx, seq in enumerate(haplotype_seqs):
            is_deletion = np.concatenate(([0], seq == b"-", [0]))
            intervals = np.where(np.abs(np.diff(is_deletion)) == 1)[0].reshape(-1, 2)
            if intervals.size > 0:
                assert (
                    intervals[:, 0] > 0
                ).all(), "Deletion reported in the first reference position"
                # Deletions are by convention reported at the preceding
                # position
                dict_dels = {
                    "POS": intervals[:, 0] - 1,
                    "REF": [haplotype_master[(x[0] - 1) : x[1]] for x in intervals],
                    "ALT": [haplotype_master[x[0] - 1] for x in intervals],
                    "FREQ": [haplotype_freqs[idx]] * intervals.shape[0],
                    "HAPLOTYPES": [str(haplotype_idx[idx])] * intervals.shape[0],
                }
                df_tmp = pd.DataFrame.from_dict(dict_dels)
                df_long_dels = pd.concat([df_long_dels, df_tmp], ignore_index=True)
        # Merge deletions found in different haplotypes together
        grpby = df_long_dels.set_index(["POS", "REF", "ALT"])[
            ["FREQ", "HAPLOTYPES"]
        ].groupby(["POS", "REF", "ALT"])

        df_long_dels = pd.concat(
            [grpby["FREQ"].sum(), grpby["HAPLOTYPES"].apply(lambda s: ",".join(s))],
            axis=1,
        )
        df_long_dels.reset_index(inplace=True)

        # Drop one-base deletions
        del_mask = df_snvs["ALT"].str.startswith("-")
        df_snvs = df_snvs[~del_mask]
        df_snvs = pd.concat([df_snvs, df_long_dels], ignore_index=True)
        df_snvs = df_snvs.set_index(["POS", "REF", "ALT"])
        df_snvs = df_snvs.sort_index()
        df_snvs.reset_index(inplace=True)

    return df_snvs


def mafft(infile, outfile, max_iter=1000, thrd=4, mafft="mafft"):
    "Use MAFFT to obtain the multiple sequence alignment"
    # --nuc sequences are nucleotide
    # --localpair pairwise alignments
    # --maxiterate number of iterative refinement
    cmd = sh.Command(mafft)
    cmd = cmd.bake("--nuc")
    cmd = cmd.bake("--preservecase")
    cmd = cmd.bake("--maxiterate", max_iter)
    cmd = cmd.bake("--localpair")
    cmd = cmd.bake("--thread", thrd)
    cmd = cmd.bake(infile)
    cmd = cmd.bake(_out=outfile)

    print(cmd)
    cmd()


def consecutive(array, stepsize=1):
    return np.split(array, np.where(np.diff(array) != stepsize)[0] + 1)


def target_snvs(start_region, end_region, start_locus, long_deletions, end_locus=None):
    if long_deletions:
        is_contained = (start_locus >= start_region) & (end_locus < end_region)
    else:
        is_contained = (start_locus >= start_region) & (start_locus < end_region)
    return is_contained


def main():

    args = parse_args()

    alphabet = ["-", "A", "C", "G", "T"]
    alphabet = np.array(alphabet, dtype="c")

    # Compute average frequency for SNVs called using ShoRAH
    df_snvs = parse_vcf(args.snvs, args.snv_caller)

    if df_snvs.empty:
        print("No called SNVs")
        with open(args.outfile, "w") as outfile:
            outfile.write("ID\tTP\tFP\tFN\tTN\n")
        return

    # Drop insertions
    ins_mask = df_snvs["ALT"].str.len() > 1
    df_snvs = df_snvs[~ins_mask]

    if args.only_deletions:
        # Only look at deletions
        # NOTE: temporary work-around while ShoRAH (v1.99.2) is modified to
        #       report indels complying to VCF format
        if args.snv_caller == "shorah":
            is_deletion = df_snvs["ALT"] == "-"
        elif args.snv_caller == "lofreq":
            is_deletion = df_snvs["REF"].str.len() > 1
        df_snvs = df_snvs[is_deletion]

    # NOTE: once ShoRAH (v1.99.2) is upgraded to report indels complying to
    #       VCF format, --long-dels can also be executed and raising an error
    #       won't be needed
    if args.long_deletions and args.snv_caller == "shorah":
        raise ValueError("No curent support for --long-dels and ShoRAH")

    if df_snvs.empty:
        print("No called SNVs")
        with open(args.outfile, "w") as outfile:
            outfile.write("ID\tTP\tFP\tFN\tTN\n")
        return

    if not args.long_deletions:
        # Unroll deletions into one-base deletions
        del_mask = df_snvs["REF"].str.len() > 1
        assert (
            df_snvs.loc[del_mask, "ALT"] == df_snvs.loc[del_mask, "REF"].str[0]
        ).all(), "Reference base preceding deletion does not match"

        del_len = df_snvs.loc[del_mask, "REF"].str.len() - 1
        df_del = pd.DataFrame(
            np.repeat(df_snvs[del_mask].values, del_len.to_list(), axis=0)
        )
        df_del.columns = df_snvs.columns
        df_del["ALT"] = "-"
        aux_idx = 0
        aux_pos = df_del.columns.get_loc("POS")
        aux_ref = df_del.columns.get_loc("REF")
        for idx, row in df_snvs[del_mask].iterrows():
            # ignore first base as it corresponds to the reference at the
            # preceding locus
            ref = list(row["REF"][1:])
            pos = [row["POS"] + x + 1 for x in range(len(ref))]
            df_del.iloc[aux_idx : (aux_idx + del_len[idx]), aux_pos] = pos
            df_del.iloc[aux_idx : (aux_idx + del_len[idx]), aux_ref] = ref
            aux_idx += del_len[idx]

        # Handle special case: reference sequence might contain a gap character
        # and a long deletion could include it. When unrolling long deletions
        # the REF and ALT fields will contain both gaps symbols
        is_gap = (df_del["REF"] == "-") & (df_del["ALT"] == "-")
        df_del = df_del[~is_gap]

        # Remove previous rows corresponding to deletions and add the one-base
        # deletions
        df_snvs = df_snvs[~del_mask]
        df_snvs = pd.concat([df_snvs, df_del], ignore_index=True)
        df_snvs = df_snvs.set_index(["POS", "ALT", "REF"])
        df_snvs = df_snvs.sort_index()

        # Merge on POS and ALT
        grpby = df_snvs.set_index("CHROM", append=True)[["INFO", "FREQ"]].groupby(
            ["POS", "ALT", "REF", "CHROM"]
        )
        df_snvs = pd.concat(
            [grpby["INFO"].apply(lambda s: ";".join(s)), grpby["FREQ"].sum()], axis=1
        )
        # grpby["REF"].first() # If not part of the index

    outdir = args.outdir if args.outdir is not None else os.getcwd()
    if args.haplotype_master is not None:
        # Parse file containing reference/consensus sequence (sequence w.r.t
        # which SNVs were called)
        header, haplotype_master = read_fasta(args.haplotype_master)
        header = header[0]
        haplotype_master = haplotype_master[0].upper()
        haplotype_master_array = np.array(list(haplotype_master))
        reference_len = haplotype_master_array.size

        if args.msa:
            # Expected if cohort consensus has gaps
            if args.reference:
                tmp, reference = read_fasta(args.reference)
                reference = reference[0].upper()
                reference = np.array(list(reference))
                assert (
                    reference.size == haplotype_master_array.size
                ), "Reference and cohort consensus have different lengths"
                idxs_gaps = haplotype_master_array == "-"
                haplotype_master_array[idxs_gaps] = reference[idxs_gaps]
                args.haplotype_master = os.path.join(outdir, "cohort_consensus.fasta")
                cohort_consensus = SeqRecord(
                    Seq("".join(haplotype_master_array)), id=header, description=""
                )
                with open(args.haplotype_master, "w") as outfile:
                    SeqIO.write(cohort_consensus, outfile, "fasta")

            haplotype_master_array = haplotype_master_array.astype("c")
            # construct msa: haplotypes + reference/consensus sequence
            infile = os.path.join(outdir, "tmp.fasta")
            sh.cat([args.haplotype_seqs, args.haplotype_master], _out=infile)
            msa_file = os.path.join(outdir, "haplotypes_re-msa.fasta")
            mafft(infile, msa_file, mafft=args.mafft)
            os.remove(infile)
            # Parse fasta file containing msa
            haplotype_ids, haplotype_seqs = read_fasta(msa_file)
            num_haplotypes = len(haplotype_ids) - 1

            haplotype_ref = haplotype_seqs[-1]
            haplotype_ref = haplotype_ref.upper()
            haplotype_ref = np.array(haplotype_ref, dtype="c")
            if haplotype_ref.size != reference_len:
                assert haplotype_ref.size > reference_len, (
                    "Length of the consensus/reference sequence after the "
                    "MSA is smaller"
                )
                # Deletions '-' were placed on the consensus/reference
                # sequence after the msa
                idx_master = 0
                idx_ref = 0
                idxs_ref = np.arange(haplotype_ref.size)
                del_idxs = np.zeros(haplotype_ref.size, dtype=bool)
                for i in range(haplotype_ref.size - reference_len):
                    left = min(
                        reference_len + i - idx_ref,
                        haplotype_master_array[idx_master:].size,
                    )
                    idxs = (
                        haplotype_ref[idx_ref : (idx_ref + left)]
                        == haplotype_master_array[idx_master:]
                    )
                    aux = idxs_ref[idx_ref : (idx_ref + left)][~idxs]
                    if aux.size == 0:
                        # gaps '-' were placed until the end of haplotype_ref
                        del_idxs[(idx_ref + left) :] = True
                        break
                    else:
                        idx_master = aux[0] - i
                        idx_ref = aux[0] + 1
                        del_idxs[aux[0]] = True

                assert np.all(
                    haplotype_ref[~del_idxs] == haplotype_master_array
                ), "After substracting gaps sequences do not agree"
                assert np.all(
                    haplotype_ref[del_idxs] == b"-"
                ), "All substracted loci do not correspond to '-'"

            # Parse sequences of the true haplotype
            haplotype_ids = haplotype_ids[0:num_haplotypes]
            haplotype_seqs = haplotype_seqs[0:num_haplotypes]
            haplotype_seqs_array = np.array(haplotype_seqs, dtype="c")
            # Remove insertions with respect to consensus/reference sequence
            if haplotype_ref.size != reference_len:
                haplotype_seqs_array = haplotype_seqs_array[:, ~del_idxs]
            # Restore gaps into the master sequence
            if args.reference:
                haplotype_master_array[idxs_gaps] = b"-"
        else:
            # Sequences of true haplotypes are already reported using the same
            # indexing as reference/consensus
            # Parse file containing true haplotype sequences
            haplotype_ids, haplotype_seqs = read_fasta(args.haplotype_seqs)
            num_haplotypes = len(haplotype_ids)
            haplotype_seqs_array = np.array(haplotype_seqs, dtype="c")
            haplotype_master_array = haplotype_master_array.astype("c")
    else:
        # if master sequence is not provided, report with respect to the
        # consensus. Note that SNVs are called with respect to the cohort
        # consensus.
        from scipy.stats import mode

        outfile = os.path.join(outdir, "true_haplotype_msa.fasta")
        mafft(args.haplotype_seqs, outfile, mafft=args.mafft)
        haplotype_ids, haplotype_seqs = read_fasta(outfile)
        num_haplotypes = len(haplotype_ids)
        haplotype_seqs_array = np.array(haplotype_seqs, dtype="c")
        if args.freq_dstr != "unif":
            haplotype_freqs = frequencies(
                args.freq_dstr, num_haplotypes, args.ratio, args.dirichlet_freqs
            )
            aux = np.repeat(
                haplotype_seqs_array,
                np.round(haplotype_freqs * 100).astype(int),
                axis=0,
            )
            consensus = mode(aux, nan_policy="omit")
        else:
            consensus = mode(haplotype_seqs_array, nan_policy="omit")
        if np.any(consensus[1] < 1):
            print("At some loci the consensus base is ambiguous")
        haplotype_master_array = consensus[0][0]

    haplotype_freqs = frequencies(
        args.freq_dstr, num_haplotypes, args.ratio, args.dirichlet_freqs
    )

    # missed = np.zeros(num_haplotypes)

    df_snvs_expected = true_snvs(
        haplotype_master_array,
        haplotype_master,
        haplotype_seqs_array,
        num_haplotypes,
        haplotype_freqs,
        args.long_deletions,
        alphabet,
    )

    if args.only_deletions:
        # Only look at deletions: drop other entries in expected SNVs dataframe
        if args.long_deletions:
            is_deletion = df_snvs_expected["REF"].str.len() > 1
        else:
            is_deletion = df_snvs_expected["ALT"].str.startswith("-")
        df_snvs_expected = df_snvs_expected[is_deletion]

    # Keep track of SNVs that fall within targeted regions
    df_snvs["IS_CONTAINED"] = False
    df_snvs_expected["IS_CONTAINED"] = False
    if args.long_deletions:
        deletion_length = df_snvs["REF"].str.len() - 1
        is_deletion = deletion_length > 0
        # Using 0-based indexing
        start_locus = df_snvs["POS"] - 1
        start_locus[is_deletion] += 1
        end_locus = start_locus + deletion_length - 1
        # Similarly for expected SNVs (Already uses 0-based indexing)
        deletion_length_exp = df_snvs_expected["REF"].str.len() - 1
        is_deletion_exp = deletion_length_exp > 0
        start_locus_exp = df_snvs_expected["POS"].copy()
        start_locus_exp[is_deletion_exp] += 1
        end_locus_exp = start_locus_exp + deletion_length_exp - 1
    else:
        # Handle SNVs and single-nucleotide deletions
        # Using 0-based indexing
        start_locus = df_snvs.index.get_level_values("POS") - 1
        end_locus = None
        # Similarly for expected SNVs (Already uses 0-based indexing)
        start_locus_exp = df_snvs_expected["POS"]
        end_locus_exp = None

    if args.coverage_intervals is not None:
        with open(args.coverage_intervals, "r") as infile:
            for line in infile:
                record = line.rstrip().split("\t")
                if record[0] == args.sampleID:
                    if len(record) == 1:
                        print("Empty target region")
                        with open(args.outfile, "w") as outfile:
                            outfile.write("ID\tTP\tFP\tFN\tTN\n")
                        return
                    regions = record[1]
                    break
        regions = regions.split(",")
        idxs = np.zeros(reference_len, dtype=bool)
        print("Reporting using 1-based indexing (and closed intervals)")
        num_loci = 0
        for r in regions:
            aux = r.split(":")
            ref_name = aux[0]
            if args.haplotype_master is not None:
                assert header == ref_name, (
                    f"Name of the reference, {ref_name}, does not agree with "
                    f"fasta file, {header}"
                )
            aux = aux[1].split("-")
            start = int(aux[0])
            end = int(aux[1])
            if args.snv_caller == "shorah" and not args.no_expansion:
                # ShoRAH was used for SNV calling
                # Assuming 3 windows were used for SNV calling, identify
                # region that is covered by at least 2 windows (below, using
                # 0-based indexing and closed intervals)
                start_ = max(0, start - args.window_len - 1)
                end_ = min(reference_len, end + args.window_len)
                num_windows = (
                    np.floor(
                        (end_ - (start_ + args.window_len - 1))
                        / (args.window_len // args.window_shift)
                    )
                    + 1
                )
                offset = (args.window_shift - 1) * args.window_len / args.window_shift

                start = max(0, start - offset - 1)
                # In order to identify the region which is covered by at least
                # two windows, add to the end of the first window the
                # increment multiply by the number of windows - 2 (i.e.,
                # discarding last window). In this case assuming half-open
                # interval [start, end)
                end = min(
                    reference_len,
                    start_
                    + args.window_len
                    + (num_windows - 2) * (args.window_len // args.window_shift),
                )
            # idxs[range(int(start), int(end))] = True
            # loci_region = loci[int(start):int(end)]

            # if DBG:
            #     print(f"DBG loci_true[i]: {loci_true[i]}")
            #     print(f"DBG loci_region[0]: {loci_region[0]}")
            # Here, loci are reported using 1-based indexing and a closed
            # interval
            num_loci += end - start
            start = int(start)
            end = int(end)
            print(f"Region with enough support: {start + 1}-{end}")

            # Mark reported and expected SNVs within the region
            is_contained = target_snvs(
                start, end, start_locus, args.long_deletions, end_locus
            )
            df_snvs["IS_CONTAINED"] = df_snvs["IS_CONTAINED"] | is_contained
            is_contained = target_snvs(
                start, end, start_locus_exp, args.long_deletions, end_locus_exp
            )
            df_snvs_expected["IS_CONTAINED"] = (
                df_snvs_expected["IS_CONTAINED"] | is_contained
            )

    else:
        loci = np.arange(reference_len)
        if args.snv_caller == "shorah":
            idxs = np.zeros(reference_len, dtype=bool)
            offset = args.window_len // args.window_shift
            # Parse coverage intervals from ShoRAH output
            with open(args.coverage, "r") as infile:
                # Look for regions at least covered by two windows
                start_w = 1
                end_w = 1
                for count, line in enumerate(infile):
                    record = line.rstrip().split("\t")
                    if count == 0:
                        start_w = int(record[2])
                        end_w = int(record[3])
                    else:
                        if int(record[2]) == start_w + offset:
                            start_w = int(record[2])
                            idxs[(start_w - 1) : end_w] = True
                        else:
                            start_w = int(record[2])
                        end_w = int(record[3])

            loci_region = np.extract(idxs, loci)

        else:
            if args.coverage is not None:
                with open(args.coverage, "r") as infile:
                    header = infile.readline().rstrip().split("\t")
                sampleID_idx = [
                    idx for idx, name in enumerate(header) if args.sampleID in name
                ]
                coverage = np.loadtxt(
                    args.coverage,
                    dtype=int,
                    delimiter="\t",
                    skiprows=1,
                    usecols=(sampleID_idx[0],),
                )
                assert coverage.size == reference_len, (
                    "Coverage file and reference file do not have the same "
                    "number of loci"
                )
                # Do not account for position with zero coverage for reporting
                # TP, FP, FN, and specially TN
                mask = coverage <= 0
                loci_region = loci[~mask]
            else:
                raise IOError(
                    "Expected coverage file as input when target region is not specified"
                )

        num_loci = loci_region.size
        regions = consecutive(loci_region)
        start = [el[0] for el in regions]
        end = [el[-1] + 1 for el in regions]
        for si, ei in zip(start, end):
            # Mark reported and expected SNVs within the region
            is_contained = target_snvs(
                si, ei, start_locus, args.long_deletions, end_locus
            )
            df_snvs["IS_CONTAINED"] = df_snvs["IS_CONTAINED"] | is_contained
            is_contained = target_snvs(
                si, ei, start_locus_exp, args.long_deletions, end_locus_exp
            )
            df_snvs_expected["IS_CONTAINED"] = (
                df_snvs_expected["IS_CONTAINED"] | is_contained
            )

    # Drop SNVs that fall outside of the targeted regions. Otherwise, these
    # rows will be counted toward false positives/negatives.
    df_snvs = df_snvs[df_snvs["IS_CONTAINED"]]
    df_snvs_expected = df_snvs_expected[df_snvs_expected["IS_CONTAINED"]]

    if args.output_true:
        output_file = os.path.join(outdir, "true_snvs.tsv")
        # Report using 1-based indexing
        df_snvs_expected["POS"] += 1
        df_snvs_expected.to_csv(
            output_file,
            sep="\t",
            columns=["POS", "REF", "ALT", "FREQ", "HAPLOTYPES"],
            header=["Loci", "Reference", "Variant", "Frequency", "Haplotypes"],
            index=False,
            compression=None,
        )

    # join on POS and ALT
    df_pairs = df_snvs_expected.merge(
        df_snvs, how="outer", on=["POS", "ALT", "REF"], suffixes=["_exp", "_rep"]
    )

    FN_mask = df_pairs["INFO"].isnull()
    FN = sum(FN_mask)

    FP_mask = df_pairs["HAPLOTYPES"].isnull()
    FP = sum(FP_mask)

    TP_mask = ~FN_mask & ~FP_mask
    TP = sum(TP_mask)

    TN = num_loci - len(df_pairs["POS"].value_counts())
    # Sensitivity
    if TP or FN:
        print("Sensitivity: {:.6f}".format(TP / (TP + FN)))

    # Precision
    if TP or FP:
        print("Precision: {:.6f}".format(TP / (TP + FP)))

    # Specificity
    if TN or FP:
        print("Specificity: {:.6f}".format(TN / (TN + FP)))

    print("TP: ", TP)
    print("FP: ", FP)
    print("FN: ", FN)
    print("TN: ", int(TN))
    # print("Number of FN per haplotype: ", missed)

    # Write to output file
    with open(args.outfile, "w") as outfile:
        outfile.write("ID\tTP\tFP\tFN\tTN\n")
        outfile.write(f"{args.sampleID}\t{TP}\t{FP}\t{FN}\t{int(TN)}\n")

    # output_file = os.path.join(outdir, 'FN_per_haplotype.tsv')
    # with open(output_file, 'w') as outfile:
    #     for idx, name in enumerate(haplotype_ids):
    #         aux = name.split(' ')[0]
    #         outfile.write(f'{aux}\t{missed[idx]}\n')

    output_file = os.path.join(outdir, "TP_frequencies.tsv")
    df_pairs[TP_mask].to_csv(
        output_file,
        sep="\t",
        columns=["POS", "REF", "ALT", "FREQ_exp", "FREQ_rep", "INFO"],
        header=[
            "Loci",
            "Reference",
            "Variant",
            "Freq (expected)",
            "Freq (reported)",
            "Info",
        ],
        index=False,
        compression=None,
    )

    output_file = os.path.join(outdir, "FP_frequencies.tsv")
    df_pairs[FP_mask].to_csv(
        output_file,
        sep="\t",
        columns=["POS", "REF", "ALT", "FREQ_rep", "INFO"],
        header=["Loci", "Reference", "Variant", "Freq", "Info"],
        index=False,
        compression=None,
    )

    output_file = os.path.join(outdir, "FN_frequencies.tsv")
    df_pairs[FN_mask].to_csv(
        output_file,
        sep="\t",
        columns=["POS", "REF", "ALT", "FREQ_exp", "HAPLOTYPES"],
        header=["Loci", "Reference", "Variant", "Freq", "Haplotypes"],
        index=False,
        compression=None,
    )


if __name__ == "__main__":
    main()
