# this file takes the amino acid sequences generated by matlab as input and generates dna 
# oligo sequences for synthesis and subsequent cloning
from Bio import SeqIO 
from Bio.Seq import Seq 
from Bio.Data import CodonTable
from Bio.Restriction import BsaI, HindIII, XbaI 
import sys, random, re, pdb
import numpy as np

# make the list of sequences to avoid within the oligos
re_list = []
for x in [BsaI, HindIII, XbaI]:
    re_list.append(x.site)
    rc = str(Seq(x.site).reverse_complement())
    if rc != x.site:
        re_list.append(rc)

for nt in ['A', 'C', 'G', 'T']:
    for n in range(5,10):
        site = ''.join([nt]*n)
        re_list.append(site)
print re_list

# setup global variables
serotype = ''        
WT_nt = ''
WT_aa = ''
LHS = ''
RHS = ''
MAX_OLIGO_LEN = 230
MAX_CHIP_OLIGOS = 244000

#location of PCR primer files
fwd_file = 'input/skpp15-forward.fasta'
rev_file = 'input/skpp15-reverse.fasta'
hF = open(fwd_file, "rU")
hR = open(rev_file, "rU")
primersF = list(SeqIO.parse(hF, "fasta"))
primersR = list(SeqIO.parse(hR, "fasta"))
hF.close()
hR.close()

def set_serotype(name):
    global serotype
    global WT_nt
    global WT_aa
    global LHS
    global RHS
    if name is 'AAV2':
        serotype = name        
        WT_nt = 'GACGAAGAGGAAATCAGGACAACCAATCCCGTGGCTACGGAGCAGTATGGTTCTGTATCTACCAACCTCCAGAGAGGCAACAGA'
        WT_aa = 'DEEEIRTTNPVATEQYGSVSTNLQRGNR'
        LHS = 'GGTCTCTTACA'
        RHS = 'CAAGCGAGACC'
    else:
        print 'serotype not recognized:', serotype
        exit()
    print 'set serotype to: ' + name, WT_aa

def create_backtranslation_table(include_stops = False):
    codon_aa = CodonTable.unambiguous_dna_by_name["Standard"].forward_table
    if include_stops:
        codon_aa['TAG']='*'
        codon_aa['TAA']='*'
        codon_aa['TGA']='*'
    aa_codon = {}
    for codon in codon_aa:
        # print codon, codon_aa[codon]
        aa = codon_aa[codon]
        if aa in aa_codon:
            aa_codon[aa].append(codon)
        else:
            aa_codon[aa] = [codon]
    return aa_codon

def add_outer_primers(sequences, primer_set_nums, add_cut_sites = True, mutate_sites = True):    
    # if given only a single primer set number, put it into an array
    if not isinstance(primer_set_nums, list): 
        primer_set_nums = [primer_set_nums]
    if not isinstance(sequences, list): 
        sequences = [sequences]        
    new_sequences = []
    new_lib_region = []
    # start with no cut sites
    RE5 = ''
    RE3 = ''
    if add_cut_sites:
        RE5 = LHS # these contain BsaI sites
        RE3 = RHS # these contain BsaI sites
    aa_codon = create_backtranslation_table()
    for seq in sequences:      
        oligo = RE5 + seq + RE3
        codon_start = len(RE5)
        codon_end = len(oligo) - len(RE3)
        assert(oligo[codon_start:codon_end] == seq)

        if mutate_sites:
            if check_re_sites(oligo, re_list)>2: # we add 2 sites at the ends so only fix if more than 2
                new_oligo = mutate_away_sites(oligo, re_list, aa_codon, np.arange(codon_start, codon_end, 3))
                assert(str(Seq(oligo[codon_start:codon_end]).translate()) == str(Seq(new_oligo[codon_start:codon_end]).translate()))
                oligo = new_oligo
            if check_re_sites(oligo, re_list)>2: # repeat a second time if we couldn't do it first try
                print 'trying second time'
                new_oligo = mutate_away_sites(oligo, re_list, aa_codon, np.arange(codon_start, codon_end, 3))
                assert(str(Seq(oligo[codon_start:codon_end]).translate()) == str(Seq(new_oligo[codon_start:codon_end]).translate()))
                oligo = new_oligo

            if check_re_sites(oligo, re_list)>2: # print if failed
                print 'could not mutate away sites in oligo!'
                print re_list
                print check_re_sites(oligo, re_list), [re for re in re_list if re in oligo]
                print np.arange(codon_start, codon_end, 3)
                print oligo
                aa_seq = translate_keep_case(oligo[codon_start:codon_end])
                print aa_seq, mask_wt(WT_aa, aa_seq)
                sys.exit()
        # sequentially add the primers sets on the the ends of the oligo
        full_oligo = oligo
        lib_seq = oligo[codon_start:codon_end]
        for index in primer_set_nums:
            pF = primersF[index]
            pR = primersR[index]
            full_oligo = str(pF.seq + full_oligo + pR.reverse_complement().seq)
        new_sequences.append(full_oligo)
        new_lib_region.append(lib_seq)
    return {'nt':new_sequences, 'lib':new_lib_region}

def translate_keep_case(seq):
    if len(seq)%3 != 0 and len(seq)>0:
        sys.exit(['length of sequence not divisible by 3', seq, len(seq)])
    aa = str(Seq(seq).translate())
    codons = [seq[max(i-3,0):i] for i in range(len(seq), 0, -3)][::-1]
    print aa
    print codons
    temp_aa = []
    for i in range(len(aa)):
        codon = codons[i]
        all_lower = all([nt.islower() for nt in codon])
        if all_lower:
            temp_aa.append(aa[i].lower())
        else:
            temp_aa.append(aa[i])
    return ''.join(temp_aa)

def mask_wt(wt, seq):
    # compare wt to seq, return seq with positions that match WT set to '_'
    if len(wt) != len(seq):
        # print 'wt', wt
        # print 'seq', seq
        index = 0
        assert(len(seq)>len(wt))
        # insert gaps into wt, remove gaps from seq
        temp_wt = []
        temp_seq = []
        for i in range(len(seq)):
            if seq[i] != '-':
                temp_seq.append(seq[i])

            if seq[i].islower():
                temp_wt.append('-')
            else:
                temp_wt.append(wt[index])
                index += 1
        wt = ''.join(temp_wt)
        seq = ''.join(temp_seq)
    diff = list(seq)
    count = 0
    for i in range(len(wt)):
        if wt[i] == seq[i]:
            diff[i] = '_'
        else:
            count+=1
    return ''.join(diff), count

def summarize_skpp_primers(f, primer_nums = range(30)): 
    if 370 not in primer_nums: primer_nums.append(370) # always include last skpp primer for spike in control
    counts = np.zeros( (len(primer_nums), 4) )
    with open(f) as fh:
        for line in fh:
            # output+=line + '\n'
            # print line
            for i in range(len(primer_nums)):
                fwd = str(primersF[primer_nums[i]].seq)
                rev = str(primersR[primer_nums[i]].seq)                
                fwd_rc = str(Seq(fwd).reverse_complement())
                rev_rc = str(Seq(rev).reverse_complement())
                # print i, fwd, rev, fwd_rc, rev_rc
                primers = [fwd, rev_rc, fwd_rc, rev]
                for pi in range(len(primers)):
                    primer = primers[pi]
                    if primer in line: counts[i, pi]+=1
    # print counts
    indices = np.where(np.sum(counts,1))[0]
    txt = []
    for i in indices:
        txt.append('\tskp{num}, {counts}\n'.format(num=primer_nums[i]+1, counts=counts[i]))
    return ''.join(txt)

def collate_seqs(output_seq, output_log, files, pad_T = 10):
    total = 0
    with open(output_log, 'w') as lh:
        with open(output_seq, 'w') as oh:
            for f in files:
                lh.write(f+'\n') 
                lh.write(summarize_skpp_primers(f) + '\n') 
                count=0
                print f
                with open(f) as fh:
                    for line in fh:
                        # if count >= 10: break
                        line_id = line.split(',')[0]
                        seq = line.split(',')[1].strip()
                        if len(seq) < MAX_OLIGO_LEN and pad_T > 0:
                            seq = seq + ''.join(['T']*min(pad_T, MAX_OLIGO_LEN-len(seq)))
                            assert len(seq)<=MAX_OLIGO_LEN, ['oligo too long', len(seq), f]
                        oh.write(seq+'\n')           
                        count+=1
                        total+=1
                print '{count} sequences in file: {file}'.format(count=count, file=f)
    print '{total} total sequences'.format(total=total)
    if total > MAX_CHIP_OLIGOS:
        print '{num} too many sequences'.format(num=total-MAX_CHIP_OLIGOS)
    else:
        print '{num} available sequences'.format(num=MAX_CHIP_OLIGOS-total)

def output_seqs(output_file, log_file, primer_set_nums, input_types, num_variants=2, num_wt=100, fill_replicas_until=0):
    primer_set_nums = [p-1 for p in primer_set_nums] # convert to zero indexing
    aa_mutants = []
    add_single_codon=False
    for i in input_types:
        print i
        if i == 'wt':
            aa_mutants.append(WT_aa)
        elif i == 'single': 
            add_single_codon=True # set flag and add later
        else:
            input_file = files[i]
            with open(input_file) as fh:
                lines = 0
                for line in fh:
                    aa = line.strip()
                    # print aa, WT_aa
                    # print aa, mask_wt(WT_aa, aa)
                    aa_mutants.append(aa)
                    lines += 1
                    # if lines > 10: break
    unique_aa_mutants = list(set((aa_mutants)))
    seed = lambda: .20
    random.shuffle(unique_aa_mutants,seed) 

    mutants_created = 0
    if 'single' in input_types:
        num_singles = 3734
    else:
        num_singles = 0
    num_replica_fillers = fill_replicas_until - num_singles - num_wt - num_variants*len(unique_aa_mutants)
    print 'number of unique mutants', len(unique_aa_mutants)
    print 'number of replica fillers', num_replica_fillers
    
    nt_mutants = []
    aa_codon = create_backtranslation_table()
    for aa_seq in unique_aa_mutants:      
        replicas = []
        if mutants_created < num_replica_fillers:
            extra_replicas = 1
        else:
            extra_replicas = 0
        for i in range(num_variants + extra_replicas):
            if i == 0:
                nt_seq = backtranslate(aa_seq, aa_codon, WT_nt, default_wt=True)
                mutants_created += 1
            else:
                nt_seq = backtranslate(aa_seq, aa_codon, WT_nt, default_wt=False)            
            nt_mutants.append({'aa':aa_seq, 'nt':nt_seq, 'variant':i+1, 'repeat':sum([nt_seq not in replicas])})
            replicas.append(nt_seq)
    replicas = []

    if add_single_codon:
        single_nt, variant_name, variant_wt_aa = single_codon_mutants(WT_nt, num_variants, max_wt=num_variants)
        for i in range(len(single_nt)):
            nt = single_nt[i]
            wt_aa = variant_wt_aa[i]
            nt_mutants.append({'aa':str(Seq(nt).translate()), 'wt_aa':wt_aa, 'nt':nt, 'variant':variant_name[i], 'repeat':sum([nt not in replicas])})
            replicas.append(nt)     
    for i in range(num_wt-num_variants):
        wt_syn_nt_seq = backtranslate(WT_aa, aa_codon, WT_nt, False)    
        nt_mutants.append({'aa':WT_aa, 'nt':wt_syn_nt_seq, 'variant':i+1+num_variants, 'repeat':sum([wt_syn_nt_seq not in replicas])})
        replicas.append(wt_syn_nt_seq)
    print 'number of mutants:', len(nt_mutants)
    oligo_id = 1
    with open(output_file, 'w') as fh:
        with open(log_file, 'w') as lh:
            lh.write('{num} mutants\n'.format(num=len(nt_mutants)))
            for mut in nt_mutants:
                nt = mut['nt']
                seq = add_outer_primers(nt, primer_set_nums, add_cut_sites=True, mutate_sites=True)
                oligo = seq['nt'][0] # returns as array, take first element
                lib = seq['lib'][0] # returns as array, take first element
                
                if isinstance(mut['variant'], str) and 'ins' in mut['variant']:
                    mask = mask_wt(mut['wt_aa'], mut['aa'])[0]
                else:
                    mask = mask_wt(WT_aa, mut['aa'])[0]

                # check the length    
                assert len(oligo)<=MAX_OLIGO_LEN, ['oligo length > ' + str(MAX_OLIGO_LEN), len(oligo), mask, oligo,]

                fh.write('{id},{oligo}\n'.format(id=oligo_id, oligo=oligo))
                lh.write('{skp},{aa},{mask},{lib_region},{variant},{repeat},{cuts}\n'.format(
                    skp=[i+1 for i in primer_set_nums],
                    aa=mut['aa'],
                    mask=mask,
                    lib_region=lib,
                    nt=oligo, 
                    variant=mut['variant'], 
                    repeat=mut['repeat'], 
                    cuts=check_re_sites(oligo, re_list)))
                oligo_id+=1
                # if oligo_id > 100: break
    print 'outputted mutants:', oligo_id-1

def single_codon_mutants(WT_nt, num_variants=1, max_wt=10):
    nt = ['A', 'C', 'G', 'T']
    codons = [''.join([x, y, z]) for x in nt for y in nt for z in nt]
    codon_pos = np.arange(0, len(WT_nt), 3)
    ins_codon_pos = np.arange(0, len(WT_nt)+3, 3)
    # print WT_nt
    mutant_nt = []
    variant_names = []
    variant_wt_aa = []
    wt_count = 0
    aa_codon_stops = create_backtranslation_table(include_stops = True)
    wt_aa = str(Seq(WT_nt).translate())
    for i in range(len(codon_pos)):
        p = codon_pos[i]
        for c in codons:
            for v in range(num_variants):
                sub_seq = WT_nt[:p] + c + WT_nt[p+3:]
                if v>0:
                    aa_seq = str(Seq(sub_seq).translate())
                    sub_seq = backtranslate(aa_seq, aa_codon_stops)
                    assert(str(Seq(sub_seq).translate()) == aa_seq)
                sub_seq = mutate_away_sites(sub_seq, re_list, aa_codon_stops, codon_pos[codon_pos!=p])
                name = 'sub_{i}_{aa}_{nt}_{var}'.format(i=i, nt=c, aa=str(Seq(c).translate()), var=v+1)
                if sub_seq == WT_nt:
                    wt_count += 1
                if sub_seq != WT_nt or wt_count <= max_wt:
                    num_times = 2 if c in ['ATG', 'TTG'] else 1
                    for n in range(num_times):
                        mutant_nt.append(sub_seq)                
                        variant_names.append(name)     
                        variant_wt_aa.append(wt_aa)
                        # print sub_seq[:p] + " " + c + " " + sub_seq[p+3:], name
    for i in range(len(ins_codon_pos)):
        p = ins_codon_pos[i]
        ins_wt_aa = wt_aa[:i] + "-" + wt_aa[i:]    
        for c in codons:
            for v in range(num_variants):
                ins_seq = WT_nt[:p] + c + WT_nt[p:]
                if v>0:
                    aa_seq = str(Seq(ins_seq).translate())
                    ins_seq = backtranslate(aa_seq, aa_codon_stops)
                    assert(str(Seq(ins_seq).translate()) == aa_seq)
                ins_seq = mutate_away_sites(ins_seq, re_list, aa_codon_stops, ins_codon_pos[ins_codon_pos!=p])
                name = 'ins_{i}_{aa}_{nt}_{var}'.format(i=i, nt=c, aa=str(Seq(c).translate()), var=v+1)
                
                num_times = 2 if c in ['ATG', 'TTG'] else 1
                for n in range(num_times):
                    mutant_nt.append(ins_seq)                
                    variant_names.append(name)
                    variant_wt_aa.append(ins_wt_aa)
                    # if i < len(ins_codon_pos)-1:
                    #     print ins_seq[:p] + " " + c + " " + ins_seq[p+3:], name
                    # else:
                    #     print ins_seq[:p] + " " + c, name
    return mutant_nt, variant_names, variant_wt_aa  

def check_re_sites(nt_seq, re_list):
    count = 0
    for site in re_list:
        site_loc = re.finditer(site, nt_seq.upper())
        for loc in site_loc:
            count+=1
    return count

def mutate_away_sites(seq, enzyme_list, aa_codon, codon_pos = None):
    if codon_pos is None:
        codon_pos = np.arange(0, len(seq), 3) # all codon positions   
    # print [[p, seq[p:p+3]] for p in codon_pos]
    start_seq = seq
    num_sites = check_re_sites(seq, enzyme_list)
    if num_sites>0:
        sites_in = [site for site in enzyme_list if site in seq]
    else:
        return seq
    seq_len = len(seq)    
    current_seq = seq
    for site in sites_in: # for all enzymes
        for site_loc in re.finditer(site, seq):  # for all cut sites
            codon_iter = 0 
            site_start = site_loc.start()
            site_end = site_loc.end()
            overlapping_codon_pos = codon_pos[(codon_pos+3>site_start) & (codon_pos<site_end)] # codons that overlap the site
            for p in overlapping_codon_pos: # for all overlapping codon positions
                codon = current_seq[p:p+3]
                aa = str(Seq(codon).translate())
                syn_codons = aa_codon[aa]
                random.shuffle(syn_codons) # note that this changes the order in aa_codon
                for new_codon in syn_codons:
                    new_seq = current_seq[:p] + new_codon + current_seq[p+3:]
                    # change the sequence if we ever reduce the number of cut sites
                    if check_re_sites(new_seq, enzyme_list) < check_re_sites(current_seq, enzyme_list):
                        current_seq = new_seq
    return current_seq

def backtranslate(mut_aa_seq, aa_codon_table, wt_nt_seq = '', default_wt = False):
    if default_wt:
        if len(wt_nt_seq)%3 != 0 and len(wt_nt_seq)>0:
            sys.exit(['WT nt sequence not divisible by 3', wt_nt_seq, len(wt_nt_seq)])
        wt_aa_seq = str(Seq(wt_nt_seq).translate())
        # wt_mask = mask_wt(wt_aa_seq, mut_aa_seq)
        wt_codons = [wt_nt_seq[max(i-3,0):i] for i in range(len(wt_nt_seq), 0, -3)][::-1]
        # print wt_aa_seq, wt_mask
        # print wt_nt_seq, wt_codons
    else:
        default_wt = False
    mut_codons = [None]*len(mut_aa_seq)
    
    wt_i = 0
    insertion = False

    for i in range(len(mut_aa_seq)):
        # '-' represents a gap, treat it like an insertion so don't increment the wt counter and don't use the wt codon
        if mut_aa_seq[i] == '-': 
            mut_codons[i] = '' 
            gap = True
            print '-'
        # lowercase represents an insertion, don't increment the wt counter and don't ever use the wt codon
        elif mut_aa_seq[i].islower(): 
            insertion = True
            gap = False
        else:
            insertion = False
            gap = False

        if not gap:
            if not insertion and default_wt and mut_aa_seq[i] == wt_aa_seq[wt_i]:
                mut_codons[i] = wt_codons[wt_i]
                # print mut_aa_seq[i], wt_codons[wt_i], 'wt'
            else:
                aa = mut_aa_seq[i]
                codons = aa_codon_table[aa.upper()]
                random.shuffle(codons) # note that this changes the order in aa_codon_table
                # print aa, aa_codon_table[aa.upper()]
                mut_codons[i] = codons[0] # choose the first one by default
            if not insertion:
                wt_i += 1
        if insertion:
                mut_codons[i] = mut_codons[i].lower()
    # print mut_codons
    mut_nt = ''.join(mut_codons).upper()
    mut_aa_test = str(Seq(mut_nt).translate())
    # print mask_wt(mut_aa_seq, mut_aa_test)
    # print mut_nt
    return mut_nt

def move_chip_files(input, output):
    count=0
    with open(output, 'w') as oh:
        with open(input) as ih:            
            for line in ih:
                if count == 0: # header
                    columns = line.strip().split(',')
                    # print columns
                    seq_index  = columns.index('chip_oligo')
                else:
                    vals = line.strip().split(',')
                    oh.write('{id},{seq}\n'.format(id=vals[0], seq=vals[seq_index]))
                count += 1
                    
files = { 
    'additive':'output/additive_mutants.txt', 
    }

set_serotype('AAV2') # initialize global variables
output_seqs('output/additive_mutant_oligos.txt', 'output/log_additive_mutant_oligos.txt', [5, 4], ['single', 'wt', 'additive'], num_variants=1, num_wt=100, fill_replicas_until = 80000)