from easy_ko.load_genomes import SharedData
from easy_ko.CRISPR_components import component

from Bio import SeqIO
from Bio.Seq import reverse_complement

import pandas as pd
import numpy as np
import re

# for designing and evaluating primers
from primer3 import design_primers
from easy_ko.primer_interactions import *

def replace_wildcards(seq):
    # [^ATCGN] is "any letter that isn't A, T, C, G, or N"
    return re.sub("[^ATCGN]", "N", seq)

# ensures a sequence appears on a specific gene, and that it appears exactly twice in the genome
def validate_primer3_primers(data:SharedData, fwd, rev, seq_B): # also do a blacklist check to make sure the blacklisted genes don't contain anything resembling the primers
    genome_assembly_dict = data.genome_assembly_dict
    fwd_RC = reverse_complement(fwd)
    rev_RC = reverse_complement(rev)

    # primers should each appear on B allele once
    if len(list(re.findall(fwd,seq_B))) + len(list(re.findall(fwd_RC,seq_B))) != 1 or len(list(re.findall(rev_RC,seq_B))) + len(list(re.findall(rev,seq_B))) != 1:
        return 0
    
    num_found_fwd = 0
    num_found_rev = 0
    
    for seqname in genome_assembly_dict.keys():
        genome_seq = str(genome_assembly_dict.get(seqname))

        num_found_fwd += len(list(re.finditer(fwd,genome_seq)))
        num_found_fwd += len(list(re.finditer(fwd_RC,genome_seq)))

        num_found_rev += len(list(re.finditer(rev_RC,genome_seq)))
        num_found_rev += len(list(re.finditer(rev,genome_seq)))
    
    # each primer should have two places to bind to in genome
    if num_found_fwd != 2 or num_found_rev != 2:
        return 0
    
    return True

def create_primers(data:SharedData, seq, seq_B, excluded_regions = [], ok_regions = 0):
    if ok_regions == 0:
        ok_regions = [0, len(seq), 0, len(seq)]
    
    num = 1
    
    # MG note: added PRIMER_PRODUCT_SIZE_RANGE for optimal primer size
    # Kai changed the ranges of the last option
    while True: # breaks when it returns primers
        primers = design_primers(
            seq_args={
                'SEQUENCE_TEMPLATE': seq,
                'SEQUENCE_EXCLUDED_REGION': excluded_regions,
                'SEQUENCE_PRIMER_PAIR_OK_REGION_LIST': ok_regions
            },
            global_args={
                'PRIMER_NUM_RETURN': num,
                'PRIMER_OPT_TM': 55.0,
                'PRIMER_MIN_TM': 51.0,
                'PRIMER_MAX_TM': 60.0,
                'PRIMER_PRODUCT_SIZE_RANGE': [[300, 500], [200, min(500,len(seq))], [max(200,len(seq)-600), len(seq)+600]]
            }
        )

        if primers['PRIMER_PAIR_NUM_RETURNED'] != num:
            print(f"not enough primers found. found {str(primers['PRIMER_PAIR_NUM_RETURNED'])} primers.")
            return 0

        primer_id = num - 1
        fwd, rev = primers['PRIMER_LEFT_'+str(primer_id)+'_SEQUENCE'], primers[f'PRIMER_RIGHT_{str(primer_id)}_SEQUENCE']

        primers['HETERODIMER_DG'] = heterodimer_dg(fwd, rev)
        
        if num < 10: 
            if primers['HETERODIMER_DG'] <= -9:
                num+=1
                #log.append('found invalid primers; trying another pair.')
                continue # failed; go to top of while loop

            if not(validate_primer3_primers(data, fwd=fwd, rev=rev, seq_B=seq_B)):
                num+=1 # this could be made more efficient by generating more than one extra primer at a time with designPrimers
                #log.append('found invalid primers; trying another pair.')
                continue # failed; go to top of while loop
        
        primers['PRIMER_LEFT_SEQUENCE'], primers['PRIMER_RIGHT_SEQUENCE'] = fwd, rev

        # start is inclusive, end is not
        left_start = primers['PRIMER_LEFT_'+str(primer_id)][0]
        left_len = primers['PRIMER_LEFT_'+str(primer_id)][1]
        primers['PRIMER_LEFT_START'] = left_start
        primers['PRIMER_LEFT_END'] = left_start + left_len
        
        right_end = primers['PRIMER_RIGHT_'+str(primer_id)][0] + 1 # because it is flipped, inclusivity/exclusivity flips
        right_len = primers['PRIMER_RIGHT_'+str(primer_id)][1]
        primers['PRIMER_RIGHT_END'] = right_end
        primers['PRIMER_RIGHT_START'] = right_end - right_len
        

        primers['PRIMER_LEFT_TM'] = primers[f'PRIMER_LEFT_{str(primer_id)}_TM']
        primers['PRIMER_RIGHT_TM'] = primers[f'PRIMER_RIGHT_{str(primer_id)}_TM']

        return primers
    
def make_orf_primers(data:SharedData, gene_id, gene_name, chr_seq_A, chr_seq_B, A_allele, B_allele):
    # designing primers for within the gene

    # 10% padding
    A_padding = A_allele['len'] * 0.10
    B_padding = B_allele['len'] * 0.10

    A_start, A_end = round(A_allele['start'] + A_padding), round(A_allele['end'] - A_padding)
    B_start, B_end = round(B_allele['start'] + B_padding), round(B_allele['end'] - B_padding)

    gene_seq_A = chr_seq_A[A_start:A_end]
    gene_seq_B = chr_seq_B[B_start:B_end]

    gene_seq_A = replace_wildcards(gene_seq_A)
    if A_allele['strand'] == "-":
        gene_seq_A = reverse_complement(gene_seq_A)


    orf_primers = create_primers(data, gene_seq_A, gene_seq_B)

    F_orf_seq, R_orf_seq = orf_primers['PRIMER_LEFT_SEQUENCE'], orf_primers['PRIMER_RIGHT_SEQUENCE']

    F_orf_notes = f"Tm = {orf_primers['PRIMER_LEFT_TM']} "
    R_orf_notes = f"Tm = {orf_primers['PRIMER_RIGHT_TM']} "

    orf_heterodimer_dg = orf_primers['HETERODIMER_DG']

    F_orf_notes += f"dg with R_orf = {orf_heterodimer_dg}"
    R_orf_notes += f"dg with F_orf = {orf_heterodimer_dg}"

    # note: A_start acccounts for padding. It is not A_allele['start]

    F_orf_start = orf_primers['PRIMER_LEFT_START'] + A_start
    F_orf_end = orf_primers['PRIMER_LEFT_END'] + A_start
    R_orf_start = orf_primers['PRIMER_RIGHT_START'] + A_start
    R_orf_end = orf_primers['PRIMER_RIGHT_END'] + A_start

    """
    df_orf_primers = pd.DataFrame({
        'name': [f'F_orf_{gene_name}',f'R_orf_{gene_name}'],
        'seq': [F_orf_seq, R_orf_seq],
        'notes': [f"Tm = {orf_primers['PRIMER_LEFT_TM']}", f"Tm = {orf_primers['PRIMER_RIGHT_TM']}"],
        'pair with': ['',''],
    })
    """

    output_F_orf, output_R_orf = ({
        'systematic name': gene_id,
        'gene name': gene_name,
        'part name': f'F_orf_{gene_name}',
        'seq': F_orf_seq,
        'notes': F_orf_notes,
        'pair with': '',
        'start': (F_orf_start),
        'end': (F_orf_end)
    },
    {
        'systematic name': gene_id,
        'gene name': gene_name,
        'part name': f'R_orf_{gene_name}',
        'seq': R_orf_seq,
        'notes': R_orf_notes,
        'pair with': '',
        'start': (R_orf_start),
        'end': (R_orf_end)
    })

    return output_F_orf, output_R_orf

def make_flank_primers(data:SharedData, gene_id, gene_name, chr_seq_A, chr_seq_B, A_allele, B_allele, margin = 250, L = 150):
    """ 
    designing primers with margin from gene
    margin (int): bp out from either end
    L(int): bp farther out is the boundary for an acceptable primer site
     """
    
    upstream_start_A, upstream_end_A = A_allele['start'] - margin - L, A_allele['start'] - margin
    downstream_start_A, downstream_end_A = A_allele['end'] + margin, A_allele['end'] + margin + L
    upstream_A = chr_seq_A[upstream_start_A : upstream_end_A]
    downstream_A = chr_seq_A[downstream_start_A:downstream_end_A]

    A_seq = replace_wildcards(upstream_A + downstream_A)

    if len(A_seq) != 2 * L:
        print("error: cannot create outer primers because the gene is too close to the edge of the A chromosome")

    upstream_start_B, upstream_end_B = B_allele['start'] - margin - L, B_allele['start'] - margin
    downstream_start_B, downstream_end_B = B_allele['end'] + margin, B_allele['end'] + margin + L
    upstream_B = chr_seq_B[upstream_start_B : upstream_end_B]
    downstream_B = chr_seq_B[downstream_start_B:downstream_end_B]

    B_seq = replace_wildcards(upstream_B + downstream_B)

    if len(B_seq) != 2 * L:
        print("error: cannot create outer primers because the B allele is too close to the edge of the B chromosome")

    # flip the sequence if the gene is on the - strand
    """
    if A_allele['strand'] == "-":
        excluded_regions = reverse_complement(excluded_regions)
    """

    # want primers on correct sides
    ok_regions = [0, len(upstream_A), len(upstream_A), len(downstream_A)] # [left_start,left_length,right_start,right_length]

    # we don't want primers that span upstream and downstream
    excluded_regions = [[len(upstream_A) - 5, 10]]

    flank_primers = create_primers(data, A_seq, B_seq, excluded_regions, ok_regions)

    F_5flank_seq, R_3flank_seq = flank_primers['PRIMER_LEFT_SEQUENCE'], flank_primers['PRIMER_RIGHT_SEQUENCE']

    F_5flank_notes = f"Tm = {flank_primers['PRIMER_LEFT_TM']} "
    R_3flank_notes = f"Tm = {flank_primers['PRIMER_RIGHT_TM']} "

    F_3FLP_seq, R_5SAT_seq = component(data, "F_3 SAT1 (oML028)").upper(), component(data, "R_5 FLP (oML032)").upper() # the orientation of the SAT FLP part is flipped, so primer pairs need to be reversed (fixed)

    F_5flank_R_5SAT_dg = heterodimer_dg(F_5flank_seq,R_5SAT_seq)
    F_3FLP_R_3flank_dg = heterodimer_dg(F_3FLP_seq, R_3flank_seq)

    F_5flank_notes += f"dg with R_5SAT primer: {F_5flank_R_5SAT_dg} "
    R_3flank_notes += f"dg with F_3FLP primer: {F_3FLP_R_3flank_dg} "

    
    flank_heterodimer_dg = flank_primers['HETERODIMER_DG'] # dg of flank primers with each other

    F_5flank_notes += f"dg with R_3flank = {flank_heterodimer_dg} "
    R_3flank_notes += f"dg with F_5flank = {flank_heterodimer_dg} "

    # todo: maybe combine from the start of upstream to the end of downstream, and then just exclude the whole gene? this way the locs are easier
    #F_5flank_start = flank_primers['']


    """
    df_flank_primers = pd.DataFrame({
        'name': [f'F_5flank_{gene_name}',f'R_3flank_{gene_name}'],
        'seq': [F_5flank_seq, R_3flank_seq],
        'notes': [f"Tm = {flank_primers['PRIMER_LEFT_TM']}", f"Tm = {flank_primers['PRIMER_RIGHT_TM']}"],
        'pair with': [data.components['part name'].iloc[2], data.components['part name'].iloc[3]], # [rev SAT1, fwd FLP]
    })
    """

    output_F_5flank, output_R_3flank = ({
        'systematic name': gene_id,
        'gene name': gene_name,
        'part name': f'F_5flank_{gene_name}',
        'seq': F_5flank_seq,
        'notes': F_5flank_notes,
        'pair with': data.components['part name'].iloc[2],  # rev SAT1
        'start': (upstream_start_A + flank_primers['PRIMER_LEFT_START']),
        'end': (upstream_start_A + flank_primers['PRIMER_LEFT_END'])
    },
    {
        'systematic name': gene_id,
        'gene name': gene_name,
        'part name': f'R_3flank_{gene_name}',
        'seq': R_3flank_seq,
        'notes': R_3flank_notes,
        'pair with': data.components['part name'].iloc[3],  # fwd FLP
        'start': (downstream_start_A - L + flank_primers['PRIMER_LEFT_START']), # subtract the length of the upstream sequence (L) in front of the downstream sequence to get distance from downstream_start
        'end': (downstream_start_A - L + flank_primers['PRIMER_LEFT_END'])
    })

    return output_F_5flank, output_R_3flank