"""
JANUS - DNA sequence optimizer for heterologous protein expression

This script generates an optimized DNA sequence to express a given protein in Escherichia coli
or another organism, using a codon usage table (CUT) to avoid rare codons and structural motifs
that may reduce expression efficiency.
When applying a more stringent limit, to block codons with a frequency of under 0.1 for example,
the program may still introduce a small number of these to prevent mRNA loops or remove repeats.
Very restrictive codon selection may be counterproductive. The -lim parameter is optional.
Also optional is the -max parameter which sets the loop counter for the sequence generator.
More trials can give a better result, but values much above 250 (the default) are unlikely to
help in most cases. Use low values (eg. -max 5) for quick tests.
Run the program with a variety of lim and max settings before choosing the DNA sequence you
prefer. The code is NOT written for speed, but with a usual length of protein sequence the output
should appear in a matter of seconds.

Features:
- Codon usage control (with frequency cutoff)
- Randomized codon selection with reproducible bias
- Avoids homopolymer runs, codon repeats, and reverse-complementary sequences
- Outputs usage summary and detailed diagnostics

Usage:
    python janus.py -i input.fasta -o output.dna -lim 0.07 -max 250
"""

# To make DNA sequences for expression of desired proteins
#
import argparse, sys
import random as rand
import re
import numpy as np
from random import choice

###############################################################
# Add your protein sequence here if you like
P='EFFLYDIFLKFCLKYIDGEICHDLFLLLGKYNILPYDTSNDSIYACTNIKHLDFINPFGVAAGFDKNGVCIDSILKLGFSFIEIGTITPRGQTGNAKPRIFRDVESRSIINSCGFNNMGCDKVTENLILFRKRQEEDKLLSKHIVGVSIGKNKDTVNIVDDLKYCINKIGRYADYIAINVSSPNTPGLRDNQEAGKLKNIILSVKEEIDNLEKNNIMNDESTYNEDNKIVEKKNNFNKNNSHMMKDAKDNFLWFNTTKKKPLVFVKLAPDLNQEQKKEIADVLLETNIDGMIISNTTTQINDIKSFENKKGGVSGAKLKDISTKFICEMYNYTNKQIPIIASGGIFSGLDALEKIEAGASVCQLYSCLVFNGMKSAVQIKRELNHLLYQRGYYNLKEAIGRKHSKS'

################################################################
WC = {'A':'T', 'T':'A', 'C':'G', 'G':'C', ' ':''}

def fullstitch(protein_seq, codon_table, usage_weights, max_attempts=250):
    """
    Generate multiple candidate DNA sequences using the 'stitch' method
    and return the one with the best (lowest) score.

    Parameters:
    protein_seq (str): The amino acid sequence to translate.
    codon_table (dict): Codons allowed per amino acid.
    usage_weights (dict): Normalized usage weights for codons.
    max_attempts (int): Number of randomized trials to evaluate.

    Returns:
    str: The best-scoring DNA sequence found.
    """
    bestscore = float("inf")
    bestseq = ''

    for attempt in range(max_attempts):
        candidate = stitch(protein_seq, codon_table, usage_weights)
        score = scoreseq(candidate)
        if score < bestscore:
            bestscore = score
            bestseq = candidate

    return bestseq

def fullstitch_old(P, x, y):
    """
    This function creates DNA sequences using the stitch routine
    and scores them, returning the sequence with the best score.
    """
    i = 0
    bestscore = 10000
    bestseq = ''
    while (i<250):
        s = stitch(P, x, y)
        i = i+1
        score = scoreseq(s)
        if score < bestscore:
            bestseq = s
            bestscore = score
    return bestseq

def scoreseq(seq):
    """
    This routine scores a given sequence according to the presence
    of repeated codons, runs of a given nucleotide and repeats.
    Modify this routine according to the situation, whether doubled
    codons are considered a problem and so on.
    """
    out1 = finddoubles(seq)
    out2 = findruns(seq)
    out3 = findreps(seq)
    out4 = findloops(seq)
    score = len(out1*2+out2+out3+out4*10)
    return score

def revc(seq):
    """
    Returns a reverse compliment sequence.
    """
    return "".join([WC[base] for base in seq[::-1]])

def findloops(dnaseq, replen=12):
    """
    Returns complementary regions replen bases long within a given sequence.
    Such regions may form loops in mRNA that inhibit protein expression.
    """
    n = len(dnaseq)
    out = ""
    start = 1
    while (start < (n-replen)):
        target = dnaseq[start:(start+replen)]
        revtarget = revc(target)
        m = re.search(revtarget, dnaseq)
        if (m):
            #print(start, target, str(m.start()), revtarget)
            out += f"{start}: {target}. Complement {m.span()[0]}-{m.span()[1]}\n "
        start = start +1
    return out

def findreps(dna_sequence, repeat_length=12):
    """
    Returns repeated regions of the given length within a DNA sequence.
    """
    min_seq_len = 2 * repeat_length
    if len(dna_sequence) < min_seq_len:
        return ""
    
    results = []
    for start in range(len(dna_sequence) - min_seq_len + 1):
        target = dna_sequence[start:start+repeat_length]
        match = re.search(target, dna_sequence[start+repeat_length:])
        if match:
            results.append(f"{start}: {target} repeats at {match.start() + start + repeat_length}")
    
    return "\n".join(results)


def findruns(seq, min_length=5):
    """
    Returns runs of repeated bases in a given DNA sequence.
    
    Parameters:
    seq (str): The DNA sequence to search for repeats.
    min_length (int): The minimum length of repeated sequences to consider as 
a run.
    
    Returns:
    str: A string containing details of all runs found, or "No runs." 
if none are detected.
    """
    n = len(seq)
    out = ""
    
    for start in range(n - min_length + 1):
        i = 1
        while (start + i < n and seq[start] == seq[start + i]):
            i += 1
        if i >= min_length:
            out += f"run at {start}-{start+i}: {seq[start:start+i]}\n"
    
    return "No runs." if not out else out

def reportruns(seq):
    """
    Returns positions of repeated bases in a given DNA sequence.
    """
    n = len(seq)
    out = []
    start = 0
    while (start < (n-5)):
        i=1
        while (seq[start]==seq[start+i]):
            i = i+1
        if (i>5):
            out.append((start,start+i))
        start += i
    return out

def finddoubles(seq):
    """
    Reports repeated codons in a DNA sequence, assuming the phase
    with the first three bases forming the first codon.
    """
    n = len(seq)
    out = ""
    start = 0
    while (start < (n-5)):
        teststr = seq[start+3:start+6]
        if seq[start:start+3] == teststr:
            out += f"double at {start}: {teststr}\n "
        start += 3
    return out

CODE = {'G':['GGG', 'GGC', 'GGA', 'GGT'], \
'E':['GAG', 'GAA'], \
'D':['GAT', 'GAC'], \
'V':['GTG', 'GTA', 'GTT', 'GTC'], \
'A':['GCG', 'GCA', 'GCT', 'GCC'], \
'R':['AGG', 'AGA', 'CGG', 'CGA', 'CGT', 'CGC'], \
'K':['AAG', 'AAA'], \
'N':['AAT', 'AAC'], \
'M':['ATG'], \
'I':['ATA', 'ATT', 'ATC'], \
'T':['ACG', 'ACA', 'ACT', 'ACC'], \
'W':['TGG'], \
'C':['TGT', 'TGC'], \
'Y':['TAT', 'TAC'], \
'S':['AGC', 'AGT', 'TCG', 'TCA', 'TCT', 'TCC'], \
'Q':['CAG', 'CAA'], \
'H':['CAT', 'CAC'], \
'F':['TTT', 'TTC'], \
'L':['CTG', 'TTA', 'TTG', 'CTA', 'CTT', 'CTC'], \
'P':['CCG', 'CCA', 'CCT', 'CCC']}

# This CUT is for E. coli
CUT = {
    'GGG' : 0.15,
'GGA': 0.11,
'GGT' : 0.34,
'GGC' : 0.40,
'GAG' : 0.31,
'GAA' : 0.69,
'GAT' : 0.63,
'GAC' : 0.37,
'GTG' : 0.37,
'GTA' : 0.15,
'GTT' : 0.26,
'GTC' : 0.22,
'GCG' : 0.36,
'GCA' : 0.21,
'GCT' : 0.16,
'GCC' : 0.27,
'AGG' : 0.02,
'AGA' : 0.04,
'CGG' : 0.08,
'CGA' : 0.06,
'CGT' : 0.40,
'CGC' : 0.40,
'AAG' : 0.23,
'AAA' : 0.77,
'AAT' : 0.45,
'AAC' : 0.55,
'ATG' : 1.00,
'ATA' : 0.07,
'ATT' : 0.51,
'ATC' : 0.42,
'ACG' : 0.27,
'ACA' : 0.13,
'ACT' : 0.17,
'ACC' : 0.44,
'TGG' : 1.00,
'TGT' : 0.45,
'TGC' : 0.55,
'TAG' : 0.07,
'TAA' : 0.64,
'TGA' : 0.29,
'TAT' : 0.57,
'TAC' : 0.43,
'TTT' : 0.57,
'TTC' : 0.43,
'AGT' : 0.15,
'AGC' : 0.28,
'TCG' : 0.15,
'TCA' : 0.12,
'TCT' : 0.15,
'TCC' : 0.15,
'CAG' : 0.65,
'CAA' : 0.35,
'CAT' : 0.57,
'CAC' : 0.43,
'TTG' : 0.13,
'TTA' : 0.13,
'CTG' : 0.50,
'CTA' : 0.04,
'CTT' : 0.10,
'CTC' : 0.10,
'CCG' : 0.52,
'CCA' : 0.19,
'CCT' : 0.16,
'CCC' : 0.12
}

def make_table(lim = 0.07, offset = 0.04):
    """
    Creates a codon usage table after excluding the unwanted codons.
    Codons with a usage below lim are excluded.
    """
    NEWCODE = {}
    NEWCUT = {}
    for i in CODE.keys():
        codons = CODE[i]
        freq = []
        list = []
        for j in range(len(codons)):
            cutval = CUT[codons[j]]
            if (cutval > lim):
                list.append(codons[j])
                freq.append(cutval - offset)
        sum = 0.0
        for j in freq:
            sum += j
        for j in range(len(freq)):
            freq[j] = freq[j]/sum
        for j in range(len(freq)):
            NEWCUT[list[j]] = freq[j]
        NEWCODE[i] = list
    return(NEWCODE, NEWCUT)

def stitch(seq, code, cut):
    """
    Creates a deterministic codon sequence for a given protein sequence
    by assigning codons in proportion to their frequency in the codon usage table (CUT),
    then randomizing their order before assigning them to residues.

    Parameters:
    seq (str): Protein sequence in one-letter amino acid code.
    code (dict): Codon table with allowed codons per amino acid.
    cut (dict): Frequency table of codon usage.

    Returns:
    str: Codon-optimized DNA sequence.
    """
    codon_pool = {}

    for aa in code:
        count = seq.count(aa)
        if count == 0:
            continue

        codons = code[aa]
        weights = [cut[codon] for codon in codons]
        total = sum(weights)
        if total == 0:
            raise ValueError(f"No usable codons for amino acid {aa} in CUT.")

        # Normalize weights and create proportional counts
        proportions = [w / total for w in weights]
        counts = [int(p * count) for p in proportions]

        # Adjust the most frequent codon to match total residue count
        shortfall = count - sum(counts)
        if shortfall > 0:
            max_index = proportions.index(max(proportions))
            counts[max_index] += shortfall

        # Build the codon list
        codon_list = []
        for i, c in enumerate(codons):
            codon_list.extend([c] * counts[i])
        rand.shuffle(codon_list)
        codon_pool[aa] = codon_list

    # Build the final sequence
    dna_sequence = ""
    for aa in seq:
        try:
            dna_sequence += codon_pool[aa].pop()
        except KeyError:
            raise KeyError(f"Residue '{aa}' is not in the codon table.")
        except IndexError:
            raise IndexError(f"Ran out of codons for residue '{aa}'. This should not happen.")

    return dna_sequence

def list_codons(seq):
    """ 
    Counts codon usage.
    Returns a dictionary. 
    """
    usage = dict()
    for i in range(0, len(seq), 3):
        codon = seq[i:i+3]
        try:
            usage[codon] += 1
        except KeyError:
            usage[codon] = 1
    return usage

def report(d):
    """
    Prints the codon usage for each codon in a dictionary.
    """
    for k in d.keys():
        print(k, CUT.get(k, 0.0), d[k])
    return 0    


def pretty(d):
    """
    Produces a codon usage table from a DNA sequence, organized by base order (T, C, A, G),
    and labelled with the corresponding amino acid in three-letter code.

    Parameters:
    d (dict): Dictionary mapping codons to usage counts.

    Output:
    Prints formatted codon usage per codon, including amino acid.
    Example line: (Phe) TTT  12  0.57
    """
    aa_map = {
        'F': 'Phe', 'L': 'Leu', 'I': 'Ile', 'M': 'Met', 'V': 'Val',
        'S': 'Ser', 'P': 'Pro', 'T': 'Thr', 'A': 'Ala', 'Y': 'Tyr',
        'H': 'His', 'Q': 'Gln', 'N': 'Asn', 'K': 'Lys', 'D': 'Asp',
        'E': 'Glu', 'C': 'Cys', 'W': 'Trp', 'R': 'Arg', 'G': 'Gly'
    }

    codon_to_aa = {}
    for aa, codons in CODE.items():
        for codon in codons:
            codon_to_aa[codon] = aa_map.get(aa, ' X ')

    base = ['T','C','A','G']
    print()
    for a in base:
        for c in base:
            for b in base:
                codon = a + b + c
                aa = codon_to_aa.get(codon, ' X ')
                count = d.get(codon, 0)
                freq = CUT.get(codon, 0.0)
                print(f"({aa}) {codon} {count:3d} {freq:5.2f}  ", end='')
            print()

def pretty_basic(d):
    """
    Produces a codon usage table from a dictionary.
    """
    a = 0
    b = 0
    c = 0
    codon = []
    row = 0
    base = ['T','C','A','G']
    print("\n", end=" ")
    while (a<4):
        while (c<4):
            while (b<4):
                codon = base[a] + base[b] + base[c]
                try:
                    k = d[codon]
                except KeyError:
                    k = 0
                print("%5s%3d%6.2f"%(codon, k, CUT.get(codon, 0.0)), end=" ") 
                b = b + 1
            print("\n", end=" ")
            b = 0
            c = c + 1
        c = 0
        a = a + 1
    return 0

def final(seq):
    """
    This attempts to remove repeated codons by adding an extra
    randomization step for amino acids with more than one codon.
    """
    seqlen = len(seq)
    newcode, newcut = make_table(0.08)
    if (seqlen%3) != 0:
        print("sequence error: incorrect length")
        exit(0)
    codonlist = [seq[i:i+3] for i in range(0,len(seq),3)]
    outstr = codonlist[0]
    for i in range(len(codonlist)-1):
        codon = codonlist[i+1]
        teststr = codonlist[i]
        if codon == teststr:
            for a in newcode.keys():
                if codon in newcode[a]:
                    if len(newcode[a]) > 1:
                        while (codon == teststr):
                            codon = choice(newcode[a])
                    outstr += codon
        else:
            outstr += codon

    if (len(seq) != len(outstr)):
        print(f"Error in final: {len(outstr), len(seq)}")
        return seq
    return outstr

def removerun(seq, start, end):
    """
    Tries to correct problems by replacing codons in the
    region just downstream from the indicated start position.
    """
    #seqlen = len(seq)
    newcode, _ = make_table(0.08)
    while (start%3) != 0:
        start -=1
    jump = end - start
    while (jump%3) != 0:
        jump -=1

    codonlist = [seq[i:i+3] for i in range(0,len(seq),3)]

    codon_idx = int(start/3)
    jump = int(jump/3)
    for i in range(codon_idx, codon_idx+jump):
        codon = codonlist[i]
        for a in newcode.keys():
            if codon in newcode[a]:
                if len(newcode[a]) > 1:
                    while True:
                        newcodon = choice(newcode[a])
                        if (newcodon != codon):
                            break
                    codonlist[i] = newcodon
                    
    outstr = ''.join(codonlist)
    return outstr

def gc_content(dna):
    g = dna.count('G')
    c = dna.count('C')
    return 100 * (g + c) / len(dna)

#
# The function "complete" will take the protein sequence and
# codon usage table to produce a possible gene sequence. 
# By changing the limit lim, the rarer codons can be more stringently
# excluded or more readily allowed.
# The default is 0.07 but other values (0.0-0.1) should be tested.
# Only you can decide how restrictive you want to be.
# Values higher than around 0.14 are unlikely to be helpful.
#
def complete(P, code, cut, limit, max_attempts):
    """
    Runs everything twice to produce the best DNA sequence.
    """
    stitchseq1 = fullstitch(P, code, cut, max_attempts)
    stitchseq2 = fullstitch(P, code, cut, max_attempts)

    s1 = scoreseq(stitchseq1)
    s2 = scoreseq(stitchseq2)

    seqlen=len(stitchseq1)

    if (s2 < s1):
        swapseq = stitchseq2
        stitchseq2 = stitchseq1
        stitchseq1 = swapseq
 
    swap = np.zeros(seqlen, dtype=int)

    start = 0
    while (start < (seqlen-12)):
        target = stitchseq1[start:(start+12)]
        revtarget = revc(target)
        m = re.search(revtarget, stitchseq1)
        if (m):
            n = int(start/3)*3
            swap[n:n+6] = [1,1,1,1,1,1]
        start = start +1

    seq=""
    for i in range(seqlen):
        if swap[i] == 0:
            seq += stitchseq1[i]
        else:
            seq+= stitchseq2[i]

    runs = reportruns(seq)
    for i in runs:
         newseq = removerun(seq, i[0], i[1])
         seq = newseq

    finalseq = seq
    nodouble = final(seq)

    s1 = scoreseq(finalseq)
    s2 = scoreseq(nodouble)

    if (s2 < s1):
        finalseq = nodouble
        s1 = s2

    print(f"Best score {s1}")
    report = finddoubles(finalseq)
    if len(report) > 0:
        print("Doubles:\n" + report)
    else:
        print("No doubles.")
    report = findruns(finalseq)
    if len(report) > 0:
        print("Runs:\n" + report)
    else:
        print("No runs more than 5 bases long.")
    report = findreps(finalseq)
    if len(report) > 0:
        print("Repeats:\n" + report)
    else:
        print("No repeats.")
    report = findloops(finalseq)
    if len(report) > 0:
        print("Complementary regions:\n" + report)
    else:
        print("No possible loops.")

    print(f"GC content: {round(gc_content(finalseq),1)}%\n")
    usage = list_codons(finalseq)
    pretty(usage)
    min_freq_codons = [(codon, CUT[codon]) for codon in usage if CUT[codon] < limit]
    if min_freq_codons:
        print("\n Codons below the suggested limit threshold were used:")
        for codon, freq in min_freq_codons:
            print(f"{codon}: {freq:.2f}")

    print(f"\nFinal sequence\n{finalseq}")
    return finalseq

def main():
    parser = argparse.ArgumentParser(
    description="Calculates a suitable DNA sequence for expressing a given protein.")
    parser.add_argument("-i", metavar="inputFile", dest="inputFile", help="input protein sequence in single letter code")
    parser.add_argument("-o", metavar="outputDNA", dest="outputDNA", help="output DNA sequence")
    parser.add_argument("-lim", metavar="limit", dest="limit", 
                        help="controls use of less frequent codons. Try testing values 0 to 0.12")
    parser.add_argument("-max", metavar="max_attempts", dest="max_attempts",
                    type=int, default=250,
                    help="maximum number of randomized stitch attempts (default: 250)")

    
    # parse CLI
    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)

    args = parser.parse_args()
    inputFile = args.inputFile

    with open(inputFile, "r") as in_file:
        data = in_file.readlines()

    if data[0][0] == '>':
        data[0] = ''
    inputProtein = ''.join(line.strip() for line in data if not line.startswith('>'))

    nogood = ['B','J','O','U','X','Z']
    check = True
    for i in inputProtein:
        if i.isalpha():
            if i.islower():
                check = False
                break
            if i in nogood:
                check = False
                break
        else:
            check = False
            break

    if not(check):
        print("Problem character:", i)
        print("The input protein sequence must contain only uppercase letters representing")
        print("proteinogenic amino acids.\n")
        exit(0)

    if args.limit is not None:
        limit = float(args.limit)
        if (limit < 0.) or (limit > 0.14):
            print("Set limit between zero and 0.14.")
            exit(0)
    else:
        limit = 0.07

    code,cut = make_table(limit)
    max_attempts = args.max_attempts
    seq = complete(inputProtein, code, cut, limit, max_attempts)
    outputDNA = args.outputDNA
    with open(outputDNA, "w") as out_file:
        print(seq, file=out_file)

if __name__ =='__main__':
    main()