#!/usr/bin/env python3
"""
This script generates an FDSTools library file that can be used to analyze PCR
amplicons of bisulfite-converted DNA. The library file shall be compatible with
FDSTools version 2.1.1 and STRNaming version 1.2.0. STRNaming must be installed
before running this script.

The input file should be a tab-separated text file with 7 columns:
  * Marker name
  * Chromosome number
  * GRCh38 position of first base in first primer
  * GRCh38 position of last base in first primer
  * GRCh38 position of first base in second primer
  * GRCh38 position of last base in second primer
  * A word indicating which strand is targeted: forward or reverse

The output file may contain "[WARNING]" comments indicating manual actions
that may be required to achieve fully concordant allele calling.
"""
import argparse
import re
import sys

from strnaming.classes import ReportedRangeStore
from strnaming import refseq_cache

PAT_REPL_FW = re.compile(r"C(?!G)")
PAT_REPL_RV = re.compile(r"(?<!C)G")

RANGES_STORE = ReportedRangeStore()
REFSEQ_STORE = RANGES_STORE.get_structure_store().get_refseq_store()


def check_range(marker, chromosome, start, end):
    """
    Check whether a range would cause STRNaming to infer reference bases in the output.
    If it does, the start and/or end positions are adjusted accordingly.
    """
    adjusted = 0
    range = RANGES_STORE.add_range(marker + "__CHECK", chromosome, start, end, load_structures=True)
    if range.library and len(range.library) == 1:
        if range.preinsert:
            # The left flank extends slightly into a repeat.
            # Remove those bases from the flank and adjust the positions.
            start -= len(range.preinsert)
            adjusted |= 1
        if range.postinsert:
            # The left flank extends slightly into a repeat.
            # Remove those bases from the flank and adjust the positions.
            end += len(range.postinsert)
            adjusted |= 2
    return adjusted, start, end


def load_marker(marker, chromosome, fw_start, fw_end, rv_start, rv_end, revcomp):
    """
    * Retrieve refseq and convert:
      - If targeting forward strand: C(?!G) --> T
      - If targeting reverse strand: (?<!C)G --> A
    * Output flanks:
      - If targeting forward strand: replace CG with YG
      - If targeting reverse strand: replace CG with CR
    * Output repeat stretches normally:
      - Display a warning if repeat region contains any C, or G if reverse strand is targeted
      - Error out if STRNaming would pull in ref bases (this is not possible with explicitly-configured markers)
    * Write prefix/suffix normally (CG stays there, allelename will contain C>T or G>A for unmethylated CpG sites)
    * Fallback to use [no_repeat] and [microhaplotype_positions] sections for non-STR targets.
    """
    # Load the reference sequence and convert it.
    if revcomp:
        refseq = PAT_REPL_RV.sub("A", refseq_cache.get_refseq(chromosome, fw_start - 1, rv_end))[1:]
    else:
        refseq = PAT_REPL_FW.sub("T", refseq_cache.get_refseq(chromosome, fw_start, rv_end + 1))[:-1]

    # Store the reference sequence and get the STRNaming structures.
    REFSEQ_STORE.add_refseq(chromosome, fw_start, refseq)
    adjusted, start, end = check_range(marker, chromosome, fw_end + 1, rv_start)
    range = RANGES_STORE.add_range(marker, chromosome, start, end, load_structures=True)

    # Extract the flanks and the stretches.
    left_flank = refseq[:start-fw_start].replace("CG", "CR" if revcomp else "YG")
    right_flank = refseq[end-fw_start:].replace("CG", "CR" if revcomp else "YG")
    ref_stretches = range.get_tssv(refseq[start-fw_start : end-fw_start], as_string=False)

    # Construct output.
    result = {
        "genome_position": ", ".join(map(str, range.location)),
        "flanks": f"{left_flank}, {right_flank}"}
    if adjusted:
        adjusted_end = ("", "the 5' (left) end", "the 3' (right) end", "both ends")[adjusted]
        message = f"[WARNING] Range for {marker} was adjusted at {adjusted_end} to include primer that extends slightly into the repeat."
        result["genome_position"] += f"\n; {message}"
        print(message, file=sys.stderr)
    if range.library:
        # STR marker.
        if len(range.library) > 1:
            print(f"[*ERROR*] Range for {marker} includes multiple STRNaming reference structures, "
                   "this is not supported. Please contact the authors for guidance.", file=sys.stderr)
            return {}
        length_adjust = range.length_adjust
        if range.library[0]["prefix"]:
            result["prefix"] = range.library[0]["prefix"]
            length_adjust += len(result["prefix"])
            del ref_stretches[0]
        if range.library[-1]["suffix"]:
            result["suffix"] = range.library[-1]["suffix"]
            length_adjust += len(result["suffix"])
            del ref_stretches[-1]
        result["repeat"] = " ".join(f"{unit} 0 {count}" for unit, count, i in ref_stretches)
        if ("G" if revcomp else "C") in result["repeat"]:
            message = f"[WARNING] Repeat structure for {marker} includes CpG sites, this could require attention."
            result["repeat"] += f"\n; {message}"
            print(message, file=sys.stderr)
        result["length_adjust"] = str(length_adjust)
        result["block_length"] = str(range.block_length)
    else:
        # Non-STR marker. Let's output the CpG sites as microhaplotypes.
        result["no_repeat"] = range.refseq
        positions = ", ".join(str(m.start() + fw_end + 1) for m in re.finditer("G" if revcomp else "C", range.refseq))
        if positions:
            result["microhaplotype_positions"] = positions
    return result


if __name__ == "__main__":
    if len(sys.argv) == 1:
        sys.argv.append("--help")
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("input", help="Pre-library file")
    args = parser.parse_args()
    sections = {}
    longest_marker = 0
    with open(args.input, "rt") as f:
        for line in f:
            marker, chromosome, fw_start, fw_end, rv_start, rv_end, strand = line.strip().split("\t")
            fw_start, fw_end, rv_start, rv_end = map(int, (fw_start, fw_end, rv_start, rv_end))
            config = load_marker(marker, chromosome, fw_start, fw_end, rv_start, rv_end, strand.lower() == "reverse")
            longest_marker = max(longest_marker, len(marker))
            for key, value in config.items():
                sections.setdefault(key, {})[marker] = value
    format = f"%-{longest_marker}s = %s"
    for section, markers in sections.items():
        print(f"\n[{section}]")
        if section == "genome_position":
            print("; This section contains the chromosome number and the genomic positions of the first and last reported nucleotide.")
        if section == "flanks":
            print("; This section contains flanking (primer) sequences, which are used by FDSTools to recognize the targets.")
        if section == "prefix":
            print("; For STR loci, this section contains the reference sequence between the left flank and the start of the repeat.")
        if section == "suffix":
            print("; For STR loci, this section contains the reference sequence between the end of the repeat and the right flank.")
        if section == "repeat":
            print("; For STR loci, this section contains the reference repeat structure. Allele names will be bracketed like it.")
        if section == "length_adjust":
            print("; For STR loci, this section contains offsets to calculate the correct CE allele number from the sequence length.")
            print("; [WARNING] The values in this section may require some manual adjustments.")
        if section == "block_length":
            print("; For STR loci, this section contains repeat unit size in nucleotides.")
        if section == "no_repeat":
            print("; For non-STR loci, this section contains the reference sequence.")
        if section == "microhaplotype_positions":
            print("; For non-STR loci, this section contains the positions of CpG islands.")
        for item in markers.items():
            print(format % item)