utils.codon

from bio import *

@inline
def get_qual_sorted_idx(qual):
    """
    Sort quality scores from low to high and return the index positions of the sorted
    array. (sorted is a in-built funciton of bio)

    """
    qual_sorted_idx = list(
        map(
            lambda x: x[0],
            sorted(
                enumerate(qual),
                key=lambda x: x[1],
                algorithm="insertion",
                reverse=False
            ) 
        )
    )

    return qual_sorted_idx


def hamming_dist1_neighbors_by_qual_order(kmer, qual):
    """
    Yield all kmers with hamming distance of 1, but start mutating original kmer at
    postition with lowest base quality score to position with highest base quality
    score.

    """
    assert len(kmer) == len(qual)

    qual_sorted_idx = get_qual_sorted_idx(qual)

    for i in range(len(kmer)):
        for b in (k"A", k"C", k"G", k"T"):
            if kmer[qual_sorted_idx[i]] != b:
                #base(kmer, idx, b) returns a new kmer equal to K but with the base at index idx set to b
                yield kmer |> base(qual_sorted_idx[i], b)


def hamming_dist2_neighbors_by_qual_order(kmer, qual):
    """
    Yield all kmers with hamming distance of 2, but start mutating original kmer at
    postition with lowest base quality score to position with highest base quality
    score.
    """

    assert len(kmer) == len(qual)

    qual_sorted_idx = get_qual_sorted_idx(qual)

    for i in range(0, len(kmer) - 1):
        for j in range(i + 1, len(kmer)):
            if i < j:
                for b1 in (k"A", k"C", k"G", k"T"):
                    if kmer[qual_sorted_idx[i]] != b1:
                        for b2 in (k"A", k"C", k"G", k"T"):
                           if kmer[qual_sorted_idx[j]] != b2:
                               yield kmer |> base(qual_sorted_idx[i], b1) |> base(qual_sorted_idx[j], b2)


def hamming_dist3_neighbors_by_qual_order(kmer, qual):
    """
    Yield all kmers with hamming distance of 3, but start mutating original kmer at
    postition with lowest base quality score to position with highest base quality
    score.

    """
    assert len(kmer) == len(qual)

    qual_sorted_idx = get_qual_sorted_idx(qual)

    for i in range(0, len(kmer) - 2):
        for j in range(i + 1, len(kmer) - 1):
            for k in range(i + 2, len(kmer)):
                if i < j:
                    for b1 in (k"A", k"C", k"G", k"T"):
                        if kmer[qual_sorted_idx[i]] != b1:
                            for b2 in (k"A", k"C", k"G", k"T"):
                                if kmer[qual_sorted_idx[j]] != b2:
                                    for b3 in (k"A", k"C", k"G", k"T"):
                                        if kmer[qual_sorted_idx[k]] != b3:
                                            yield kmer |> base(qual_sorted_idx[i], b1) |> base(qual_sorted_idx[j], b2) |> base(qual_sorted_idx[k], b3)

@tuple
class CorrectedBc:
    hamming_dist: int
    corrected_bc: str

    def __new__(hamming_dist: int, corrected_bc: str) -> CorrectedBc:
        return (hamming_dist, corrected_bc)


def correct_bc_by_qual_order_with_whitelist(
    bc_whitelist,
    bc,
    qual,
    max_hamming_dist: int
) -> Optional[CorrectedBc]:
    """
    Check if provided barcode matches the whitelist with up to max_hamming_dist
    (0, 1, 2 or 3) mismatches by starting mutating original kmer at postition
    with lowest base quality score to position with highest base quality score.

    """

    if bc in bc_whitelist:
        # Exact match.
        return CorrectedBc(0, str(bc))

    if max_hamming_dist >= 1:
        for hamming_dist1_neighbor in hamming_dist1_neighbors_by_qual_order(bc, qual):
            if hamming_dist1_neighbor in bc_whitelist:
                # Hamming neighbor in whitelist (distance = 1).
                return CorrectedBc(1, str(hamming_dist1_neighbor))

    if max_hamming_dist >= 2:
        for hamming_dist2_neighbor in hamming_dist2_neighbors_by_qual_order(bc, qual):
            if hamming_dist2_neighbor in bc_whitelist:
                # Hamming neighbor in whitelist (distance = 2).
                return CorrectedBc(2, str(hamming_dist2_neighbor))

    if max_hamming_dist >= 3:
        for hamming_dist3_neighbor in hamming_dist3_neighbors_by_qual_order(bc, qual):
            if hamming_dist3_neighbor in bc_whitelist:
                # Hamming neighbor in whitelist (distance = 3).
                return CorrectedBc(3, str(hamming_dist3_neighbor))

    # Hamming distance greater than max_hamming_dist or greater than 3.
    return None


def correct_bc_with_Ns_with_whitelist(
    bc_whitelist,
    bc_length: Static[int],
    bc_with_Ns: str,
    max_hamming_dist: int
) -> Optional[CorrectedBc]:
    """
    Check if provided barcode with Ns matches the whitelist with up to
    max_hamming_dist (0, 1, 2 or 3) mismatches (only Ns will be corrected).

    """

    n_count = 0
    n_positions = List[int](capacity=len(bc_with_Ns))

    for idx, nuc in enumerate(bc_with_Ns):
        if nuc == "N":
            n_count += 1
            n_positions.append(idx)

    if n_count > max_hamming_dist:
        # Do not try to correct barcode, if there are more Ns than max allowed hamming
        # distance.
        return None

    # Convert barcode with Ns to a kmer.
    bc = Kmer[bc_length](bc_with_Ns.replace("N", "A"))

    if n_count == 1:
        for b in (k"A", k"C", k"G", k"T"):
            bc = bc |> base(n_positions[0], b)

            if bc in bc_whitelist:
                return CorrectedBc(1, str(bc))
    elif n_count == 2:
        for b1 in (k"A", k"C", k"G", k"T"):
            for b2 in (k"A", k"C", k"G", k"T"):
                bc = bc |> base(n_positions[0], b1) |> base(n_positions[1], b2)

                if bc in bc_whitelist:
                    return CorrectedBc(2, str(bc))
    elif n_count == 3:
        for b1 in (k"A", k"C", k"G", k"T"):
            for b2 in (k"A", k"C", k"G", k"T"):
                for b3 in (k"A", k"C", k"G", k"T"):
                    bc = bc |> base(n_positions[0], b1) |> base(n_positions[1], b2) |> base(n_positions[2], b3)

                    if bc in bc_whitelist:
                        return CorrectedBc(3, str(bc))

    # Ns in barcode could not be corrected.
    return None

def correct_bc_by_qual_order_with_whitelist_or_correct_bc_with_Ns_with_whitelist(
    bc_whitelist,
    bc_length: Static[int],
    bc_seq: seq,
    bc_qual: str,
    max_hamming_dist: int
) -> Optional[CorrectedBc]:
    """
    Check if provided barcode matches the whitelist with up to max_hamming_dist
    (0, 1, 2 or 3) mismatches by starting mutating original kmer at postition
    with lowest base quality score to position with highest base quality score
    or check if provided barcode with Ns matches the whitelist with up to
    max_hamming_dist (0, 1, 2 or 3) mismatches (only Ns will be corrected).

    """
    assert len(str(bc_seq)) == bc_length, "The barcodes are not the expected length..."
    
    if "N" in str(bc_seq):
        corrected_bc = correct_bc_with_Ns_with_whitelist(
                bc_whitelist,
                bc_length,
                str(bc_seq),
                max_hamming_dist)
    else:
        corrected_bc = correct_bc_by_qual_order_with_whitelist(
            bc_whitelist,
            Kmer[bc_length](bc_seq),
            bc_qual,
            max_hamming_dist)

    return corrected_bc