__author__ = "Christopher Schröder, Patrik Smeds"
__copyright__ = "Copyright 2020, Christopher Schröder, Patrik Smeds"
__email__ = "[email protected], [email protected]"
__license__ = "MIT"

from os import path
from snakemake.shell import shell

log = snakemake.log_fmt_shell(stdout=True, stderr=True)

# Check inputs/arguments.
if len(snakemake.input) == 0:
    raise ValueError("A reference genome has to be provided.")
elif len(snakemake.input) > 1:
    raise ValueError("Please provide exactly one reference genome as input.")

valid_suffixes = {".0123", ".amb", ".ann", ".bwt.2bit.64", ".pac"}


def get_valid_suffix(path):
    for suffix in valid_suffixes:
        if path.endswith(suffix):
            return suffix


prefixes = set()
for s in snakemake.output:
    suffix = get_valid_suffix(s)
    if suffix is None:
        raise ValueError(f"{s} cannot be generated by bwa-mem2 index (invalid suffix).")
    prefixes.add(s[: -len(suffix)])

if len(prefixes) != 1:
    raise ValueError("Output files must share common prefix up to their file endings.")
(prefix,) = prefixes

shell("bwa-mem2 index -p {prefix} {snakemake.input[0]} {log}")

__author__ = "Christopher Schröder, Johannes Köster, Julian de Ruiter"
__copyright__ = (
    "Copyright 2020, Christopher Schröder, Johannes Köster and Julian de Ruiter"
)
__email__ = "[email protected] [email protected], [email protected]"
__license__ = "MIT"


import tempfile
from os import path
from snakemake.shell import shell
from snakemake_wrapper_utils.java import get_java_opts
from snakemake_wrapper_utils.samtools import get_samtools_opts


# Extract arguments.
extra = snakemake.params.get("extra", "")
log = snakemake.log_fmt_shell(stdout=False, stderr=True)
sort = snakemake.params.get("sort", "none")
sort_order = snakemake.params.get("sort_order", "coordinate")
sort_extra = snakemake.params.get("sort_extra", "")
samtools_opts = get_samtools_opts(snakemake)
java_opts = get_java_opts(snakemake)


index = snakemake.input.get("index", "")
if isinstance(index, str):
    index = path.splitext(snakemake.input.idx)[0]
else:
    index = path.splitext(snakemake.input.idx[0])[0]


# Check inputs/arguments.
if not isinstance(snakemake.input.reads, str) and len(snakemake.input.reads) not in {
    1,
    2,
}:
    raise ValueError("input must have 1 (single-end) or 2 (paired-end) elements")

if sort_order not in {"coordinate", "queryname"}:
    raise ValueError(f"Unexpected value for sort_order ({sort_order})")

# Determine which pipe command to use for converting to bam or sorting.
if sort == "none":
    # Simply convert to bam using samtools view.
    pipe_cmd = "samtools view {samtools_opts}"

elif sort == "samtools":
    # Sort alignments using samtools sort.
    pipe_cmd = "samtools sort {samtools_opts} {sort_extra} -T {tmpdir}"

    # Add name flag if needed.
    if sort_order == "queryname":
        sort_extra += " -n"

elif sort == "picard":
    # Sort alignments using picard SortSam.
    pipe_cmd = "picard SortSam {java_opts} {sort_extra} --INPUT /dev/stdin --TMP_DIR {tmpdir} --SORT_ORDER {sort_order} --OUTPUT {snakemake.output[0]}"

else:
    raise ValueError(f"Unexpected value for params.sort ({sort})")


with tempfile.TemporaryDirectory() as tmpdir:
    shell(
        "(bwa-mem2 mem"
        " -t {snakemake.threads}"
        " {extra}"
        " {index}"
        " {snakemake.input.reads}"
        " | " + pipe_cmd + ") {log}"
    )

from bisect import bisect_right
from collections.abc import Sequence
import pandas as pd
from typing import List, Union, Dict

import geffa

# Give the snakemake object a type...
snakemake: object

# Read the pilon changes file
df = pd.read_table(snakemake.input.changes, sep=' ', header=None, names=['old_coordinate', 'new_coordinate', 'old', 'new'])
# Split the coordinates into contig name and position
split = df.old_coordinate.str.split(':', expand=True)
df.loc[:, 'contig'] = split.iloc[:, 0]
df.loc[:, 'old_coordinate'] = split.iloc[:, 1]
split = df.new_coordinate.str.split(':', expand=True)
df.loc[:, 'new_coordinate'] = split.iloc[:, 1]

# Function to categorise the each row (deletion, insertion or snp)
# Also calculates the change in length (negative for deletion, positive for insertion, zero for snps)
def categorize_change(row):
    lold = len(row.old)
    lnew = len(row.new)
    diff = lnew - lold
    if diff == 0:
        if lold == 1:
            category = 'snp'
        elif lold == 0:
            raise ValueError('No change!?!')
        else:
            category = 'indel'
    elif diff < 0:
        category = 'deletion'
    else:
        category = 'insertion'
    return pd.Series({'difference': diff, 'category': category})

# Incorporate the categorisation into the table
df = pd.concat([df, df.replace('.', '').apply(categorize_change, axis=1)], axis=1)
df.loc[:, 'old_coordinate_int'] = df.old_coordinate.map(lambda a: int(a.split('-')[0]))

# Class that takes a list of change coordinates and associated shifts and allows us to 
# then transform old feature coordinates into new ones.
class PieceWiseConstantShift:
    def __init__(self, coordinates: Sequence[int], shifts: Sequence[int]) -> None:
        self.coordinates: list[int] = [0]
        self.shifts: list[int] = [0]
        self.cumshifts: list[int] = [0]

        # Calculate cumulative shifts
        for c, s in zip(coordinates, shifts):
            if s != 0: # Filter out snps, they don't shift anything
                self.coordinates.append(c)
                self.shifts.append(s)
                self.cumshifts.append(self.cumshifts[-1] + s)

    # This takes a coordinate and shifts it according to the changes given above
    def __call__(self, x: Union[int,List[int]]) -> Union[int,List[int]]:
        if isinstance(x, Sequence):
            return [self(y) for y in x]
        # Find the index of the shift to the left of the given coordinate
        index = bisect_right(self.coordinates, x)-1
        shift = self.shifts[index]
        cumshift = self.cumshifts[index]
        coordinate = self.coordinates[index]
        # If the old coordinate is within a deletion or insertion region, we need to shift it beyond
        if x < coordinate - shift:
            new = coordinate
        else:
            new = x
        new += cumshift
        try:
            # If we shifted beyond the next change, we need to correct for that
            if new > self.coordinates[index+1]:
                new = self.coordinates[index+1]
        except IndexError:
            # We were at the last entry, so nothing to see here.
            pass
        return new

# Build a translator for each contig
translations:Dict[str, PieceWiseConstantShift] = {
    contig: PieceWiseConstantShift(g.old_coordinate_int, g.difference)
    for contig, g in df.groupby('contig')
}

# TODO: Validate - mark incomplete CDS as pseudogenes!!!

# Load the parental GFF file and the new sequence fasta file
gff = geffa.GffFile(snakemake.input.gff, fasta_file=snakemake.input.fasta, ignore_unknown_feature_types=True)
for seqreg in gff.sequence_regions.values():
    # Iterate over contigs
    try:
        # Load the correct translator
        trans = translations[seqreg.name]
    except KeyError:
        # If we have no reads in a contig, it won't show up in the translation table
        continue

    marked_for_deletion = []
    # Iterate over all features in the contig
    for feature in seqreg.node_registry.values():
        # Translate the start and end coordinates
        start, end = feature.start, feature.end
        feature.start, feature.end = trans((feature.start, feature.end))
        try:
            # Validate the feature with the new coordinates
            feature.validate()
        except Exception as e:
            # The new feature isn't correct!
            # Print some debug info
            print(feature)
            print(feature.sequence_region.sequence[feature.start-5:feature.end+5])
            print(f"{start}->{feature.start}, {end}->{feature.end}")
            msg = str(e)
            # Delete the feature if it's a feature that has been corrupted by an indel
            if ("SLAS feature needs to be of length 2" in msg) or ("__len__() should return >= 0" in msg):
                print("Deleting this feature because it's been destroyed by an indel.")
                marked_for_deletion.append(feature)
            elif "Exon parent needs to be an RNA" in msg:
                print("Invalid exon found - skipping because it's not that severe.")
            else:
                # Fail otherwise, don't know what to do.
                print(list(zip(trans.coordinates, trans.shifts, trans.cumshifts)))
                raise
    # Delete the features that were marked for deletion
    # (couldn't do that earlier, can't change a container while iterating over it)
    for feature in marked_for_deletion:
        feature.delete()
# Save the new annotation GFF
gff.save(snakemake.output.gff, include_sequences=True)

shell: """
trim_galore --paired --phred33 --length 36 -q 5 --stringency 1 -e 0.1 -j 8 --no_report_file --dont_gzip --output_dir 'results/{wildcards.experiment}' --basename '{wildcards.sample_id}' '{input.r1}' '{input.r2}'
"""