"""
LIFT - UCSC iGEM 2024

Module: CodonChoice

Author: Aurko Mahesh, Vibhitha Nandakumar
"""
from stealthHits import StealthHits
from codonUsage import CodonUsage
from tsvReader import TSVReader

codon_to_aa = {
    # codon to amino acid code
    # convert T to U because they are translated...
    # T
    "TTT": "F",
    "TCT": "S",
    "TAT": "Y",
    "TGT": "C",  # UxU
    "TTC": "F",
    "TCC": "S",
    "TAC": "Y",
    "TGC": "C",  # UxC
    "TTA": "L",
    "TCA": "S",
    "TAA": "-",
    "TGA": "-",  # UxA
    "TTG": "L",
    "TCG": "S",
    "TAG": "-",
    "TGG": "W",  # UxG
    # C
    "CTT": "L",
    "CCT": "P",
    "CAT": "H",
    "CGT": "R",  # CxU
    "CTC": "L",
    "CCC": "P",
    "CAC": "H",
    "CGC": "R",  # CxC
    "CTA": "L",
    "CCA": "P",
    "CAA": "Q",
    "CGA": "R",  # CxA
    "CTG": "L",
    "CCG": "P",
    "CAG": "Q",
    "CGG": "R",  # CxG
    # A
    "ATT": "I",
    "ACT": "T",
    "AAT": "N",
    "AGT": "S",  # AxU
    "ATC": "I",
    "ACC": "T",
    "AAC": "N",
    "AGC": "S",  # AxC
    "ATA": "I",
    "ACA": "T",
    "AAA": "K",
    "AGA": "R",  # AxA
    "ATG": "M",
    "ACG": "T",
    "AAG": "K",
    "AGG": "R",  # AxG
    # G
    "GTT": "V",
    "GCT": "A",
    "GAT": "D",
    "GGT": "G",  # GxU
    "GTC": "V",
    "GCC": "A",
    "GAC": "D",
    "GGC": "G",  # GxC
    "GTA": "V",
    "GCA": "A",
    "GAA": "E",
    "GGA": "G",  # GxA
    "GTG": "V",
    "GCG": "A",
    "GAG": "E",
    "GGG": "G",  # GxG
}


class CodonChoice:

    """
    Class that creates a sequence with the optimal codon choices. Is recursively called to edit a target site until it exits the site

    Input: Host and Target organisms complete genome files, Stealth site list, and target sequence being editted

    Output: a sub-sequence that does not contain a stealth hit after correction

    """

    def __init__(self, hostIn, targetIn):

        self.hostIn = hostIn
        self.targetIn = targetIn

    def hostUsage(self):
        tsvReader = TSVReader(self.hostIn)
        hostUsage_dict = tsvReader.format_usage()

        return hostUsage_dict

    def host(codon):
        hostList = CodonChoice.hostUsage()

        return hostList[codon_to_aa[codon]]

    def targetUsage(self):
        targetCodonUsage = CodonUsage(self.targetIn)
        targetUsage_dict = targetCodonUsage.returnDict()

        return targetUsage_dict

    def target(codon):
        targetList = CodonChoice.targetUsage()

        return targetList[codon_to_aa[codon]]

    def if_same(new_index):
        try:
            new_index -= 1
            return new_index
        except IndexError:
            new_index += 1
            return new_index

    def replaceCodon(codon):

        hostList = CodonChoice.host(codon)
        targetList = CodonChoice.target(codon)

        hostIndex = next(
            (i for i, sublist in enumerate(hostList) if codon in sublist), None
        )
        targetIndex = next(
            (i for i, sublist in enumerate(targetList) if codon in sublist), None
        )

        newIndex = 0

        if targetList[targetIndex][0] != hostList[hostIndex][0]:
            newIndex = hostIndex
        if targetList[targetIndex][0] in ["ATG", "TGG"]:
            newIndex = hostIndex
        else:
            newIndex = CodonChoice.if_same(hostIndex)
        if targetList[newIndex][0] == hostList[hostIndex][0]:
            newIndex = CodonChoice.if_same(newIndex)

        try:
            newCodon = targetList[newIndex][0]

        except IndexError:
            newIndex = 0
            newCodon = targetList[newIndex][0]

        return newCodon

    def altSeqMaker(self, input_dna, stealthHits, start=None, altList=None):

        if start == None:
            start = 6
        if altList == None or altList == []:
            altList = list()

        if len(input_dna[start : start + 3]) == 3:
            hitStart = input_dna[start : start + 3]

        else:
            return altList[0][1]

        aa = codon_to_aa[hitStart]

        targets = self.targetUsage()
        targetList = targets[aa]

        for target in targetList:
            newSeq = input_dna[:start] + target[0] + input_dna[start + 3 :]
            stealthHit = StealthHits(stealthHits, newSeq)
            hits = stealthHit.find_hits()  # run stealth on this here

            altList.append((len(hits), newSeq))
            altList.sort(key=lambda x: x[0])

        for alt in altList:

            if alt[0] == 0:
                alt1 = alt[1]
                altList = []
                return alt1
        print(start, altList)
        if all(alt[0] > 0 for alt in altList):
            return self.altSeqMaker(input_dna, stealthHits, start + 3, altList)
