from easy_ko.load_genomes import SharedData

from easy_ko.helper_functions import get_gene_info
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 calculating entropy
from collections import Counter
from skbio.diversity.alpha import shannon

log = [] # this should be defined before running these functions

def get_A_and_B_allele(gene_id_A):
    # return A and B alleles
    if gene_id_A[-2:] == "_A":
        gene_id_B = gene_id_A[:-1] + "B"
        return gene_id_A, gene_id_B
    elif gene_id_A[-2:] == "_B":
        # no A allele for this gene
        gene_id_B = gene_id_A
        gene_id_A = 0
        return gene_id_A, gene_id_B
    else:
        # no B allele for this gene
        gene_id_B = 0
        return gene_id_A, gene_id_B

def get_allelic_info(data:SharedData, gene_id):
    if gene_id:
        # create a dictionary for allele
        allele = {}

        try:
            allele['start'], allele['end'], allele['seqname'], allele['strand'] = get_gene_info(data.genome_GTF_df, gene_id,['seqname','strand'])
        except Exception as e:
            # print(e)
            # print(f'gene id not recognized:  {gene_id}')
            return 0

        #allele['start'], allele['start'] - 1

        # shouldn't need abs, but just in case.
        allele['len'] = abs(allele['end'] - allele['start'])

        return allele
    else:
        return 0

# ensures a sequence appears on a specific gene, and that it appears exactly twice in the genome
def validate_seq(data: SharedData, seq, gene_seq_B):
    seq_RC = reverse_complement(seq)

    # must appear once on B allele once (in any form)
    if len(list(re.findall(seq,gene_seq_B))) != 1 and len(list(re.findall(seq_RC,gene_seq_B))) != 1:
        return 0
    
    num_sites_found = 0
    for seqname in data.genome_assembly_dict.keys():
        genome_seq = str(data.genome_assembly_dict.get(seqname))
        num_sites_found += len(list(re.finditer(seq,genome_seq)))
        num_sites_found += len(list(re.finditer(seq_RC,genome_seq)))
    
    # seq should appear no more than twice in genome
    if num_sites_found > 2:
        return 0
    
    return True

# this function was generated by ChatGPT. The closer to 6, the higher the entropy.
def entropy(seq, k=3):
    kmers = [seq[i:i+k] for i in range(len(seq)-k+1)]
    vector = list(Counter(kmers).values())
    return shannon(vector, base=2)

# returns either 0 or the site formatted to ...NGG, but without the PAM
def evaluate_site(data: SharedData, evaluated, gene_seq_B, blacklisted = 0):
    # site_seq contains the PAM seq; seq does not
    site_seq = evaluated['seq']
    site_RC = reverse_complement(site_seq)

    if evaluated['A_PAM'] == 'NGG':
        # originally ...NGG
        # remove PAM
        seq = site_seq[:-3]
    else:
        # RC should end in NGG
        seq = site_RC[:-3]
    
    # should not have a "wild card" outside of PAM
    if set(seq) > set('ATCG'):
        log.append('removed a seq with an ambiguous base')
        return 0
    
    # C or G : A or T should be 40-60%
    gc_count = seq.count('G') + seq.count('C')
    gc_percent = gc_count / len(seq)
    if not(0.40 <= gc_percent <= 0.60):
        return 0
    
    if not(validate_seq(data, site_seq, gene_seq_B)):
        return 0
    
    evaluated['entropy'] = entropy(seq)

    evaluated['priority'] = 1 #biopython blast and for primers

    # remove priority status from sequences that appear in other places in the genome but with a broken PAM site
    if blacklisted:
        seq_RC = reverse_complement(seq)
        for gene_id in blacklisted:
            start, end, seqname = get_gene_info(data.genome_GTF_df, gene_id,['seqname'])
            start -= 1
            end += 3
            gene_seq = str(data.genome_assembly_dict[seqname][start:end])
            if seq in gene_seq or seq_RC in gene_seq:
                # another gene contains the sequence with a broken PAM site
                evaluated['priority'] = 0

    # seq is standardized to ...NGG, but with PAM removed
    evaluated['seq'] = seq
    
    return evaluated


def list_guide_sites(data:SharedData, gene_seq_A, gene_seq_B, blacklisted, start_A, end_A, start_B, end_B):

    # . is any character except newline
    PAM = ".GG"
    # reverse complement:
    PAM_RC = "CC."

    # find PAMs in A allele
    found_PAMs = list(re.finditer(PAM, gene_seq_A))
    found_PAMs_RC = list(re.finditer(PAM_RC, gene_seq_A))

    # includes PAM
    potential_guide_sites = []

    # look in A allele for NGG and CCN
    for found_PAM in found_PAMs:
        if found_PAM.start() - 20 >= 0:
            start, end = found_PAM.start() - 20, found_PAM.start()
            potential_guide_sites.append({
                'seq': gene_seq_A[start:end] + PAM,
                'A_PAM': 'NGG',
                'PAM_loc': found_PAM.start(),
                'crRNA_start': start + start_A
                })
        else:
            log.append("there's a potential guide site that's out of range")

    for found_PAM_RC in found_PAMs_RC:
        if found_PAM_RC.end() + 20 <= len(gene_seq_A):
            start, end = found_PAM_RC.end(),  found_PAM_RC.end() + 20
            potential_guide_sites.append({
                'seq': PAM_RC + gene_seq_A[start:end],
                'A_PAM': 'CCN',
                'PAM_loc': found_PAM_RC.start(),
                'crRNA_start': start + start_A
                })
        else:
            log.append("there's a potential guide site that's out of range")

    # doesn't include PAMs
    guide_sites = []

    # evaluate potential guide sites. if they are valid, add them to a list
    for site in potential_guide_sites:
        # evaluate_site standardizes seq to ...NGG and removes PAM
        evaluated = evaluate_site(data, site, gene_seq_B, blacklisted)
        if evaluated:
            guide_sites.append(evaluated)
        
    return guide_sites

def print_results(priority, non_priority, gene_id_A, log):
    #print(log)

    print(f"viable sgRNA sites targeting gene with id \'{gene_id_A}\', \narranged by bp entropy (top=most and bottom=least): ")

    if len(priority) == 0 and len(non_priority) == 0:
        non_priority.append({'seq': 'no viable gRNA site found'})

    if len(priority) > 0:
        print("\npriority sites (these sequences are gene-specific):")
        for site in priority:
            print(site['seq'])
    else:
        print ("\nno priority sites")

    if len(non_priority) > 0:
        print("\nnon-priority sites (other genes share these sequences, but the other genes have broken PAM sites):")
        for site in non_priority:
            print(site['seq']) # + "NGG"
    else:
        print("\nno non-priority sites")


def get_seq_info(data:SharedData, gene_id, allele, padding_fraction):
    output = {}
    padding = allele['len'] * padding_fraction
    start, end = round(allele['start'] + padding), round(allele['end'] - padding)
    output['seq'] = str(data.genome_assembly_dict[allele['seqname']][start:end])
    output['start'], output['end'] = start, end
    return output


# ----- orchestrator function -----
def find_guides(data: SharedData, gene_id, gene_name, gene_id_A, gene_id_B, A_allele, B_allele):
    
    #print(gene_id_A, gene_id_B)

    # what portion of base pairs to exclude on edges on each side
    padding_fraction = 0.10

    blacklisted = 0

    try:
        filter_df = data.black_list_df[data.black_list_df["genes"].str.contains(gene_id_A)]
        cluster_size = filter_df(['cluster size']).iloc[0]
    except:
        filter_df, cluster_size = 0, 0
        

    if  cluster_size != 2 and filter_df:
        blacklisted = filter_df['genes'].to_string(index=False).split(",")
        blacklisted.remove(gene_id_A)
        if gene_id_B:
            blacklisted.remove(gene_id_B)


    # first, search excluding the padding

    A_info = get_seq_info(data, gene_id_A, allele=A_allele, padding_fraction=padding_fraction)
    B_info = get_seq_info(data, gene_id_B, allele=B_allele, padding_fraction=padding_fraction)

    guide_sites = list_guide_sites(data, gene_seq_A=B_info['seq'], gene_seq_B=B_info['seq'], blacklisted=blacklisted, start_A=A_allele['start'], end_A=A_allele['end'], start_B=B_allele['start'], end_B=B_allele['end'])

    if guide_sites == []:
        # if no sequences are found, search for them in the padding

        # for now, search without any padding
        padding_fraction = 0
        A_info = get_seq_info(data, gene_id_A, allele=A_allele, padding_fraction=padding_fraction)
        B_info = get_seq_info(data, gene_id_B, allele=B_allele, padding_fraction=padding_fraction)

        guide_sites = list_guide_sites(data, gene_seq_A=B_info['seq'], gene_seq_B=B_info['seq'], blacklisted=blacklisted, start_A=A_allele['start'], end_A=A_allele['end'], start_B=B_allele['start'], end_B=B_allele['end'])
        if guide_sites == []:
            print(f'No guide sites found for {gene_name}. ')
            return ({
                'systematic name': gene_id,
                'gene name': gene_name,
                'part name': f'gRNA_{gene_name}',
                'seq': '', 
                'notes': 'no sites found',
                'pair with': '',
                'start': pd.NA,
                'end': pd.NA
            })
        

    guide_sites = sorted(guide_sites, key=lambda site: 0 - site['entropy'])

    # priority:
    # 1. the 20 nucleotides do not appear in any blacklisted genes
    # 2. the 20 nucleotides appear in the blacklisted genes, but with a broken PAM site
    # REMOVE: the 20 nucleotide sequence WITH the PAM site appears anywhere besides on the A and B allele


    priority, non_priority = [], []

    for site in guide_sites:
        if site['priority']:
            priority.append(site)
        else:
            non_priority.append(site)

    #print_results(priority, non_priority, gene_id_A, log)

    ordered_sites = priority + non_priority
    guide = ordered_sites[0]
    
    # add the best gRNA site to a df

    guide_seq = guide['seq']

    num_sites = len(ordered_sites)

    notes = f'guide: {guide_seq}'

    backup_guides = []

    if num_sites >= 2:
        backup_guide1 = ordered_sites[1]
        notes += f', backup_guide1: {backup_guide1["seq"]}'
        backup_guides.append(backup_guide1)
        if num_sites >= 3:
            backup_guide2 = ordered_sites[2]
            notes += f', backup_guide2: {backup_guide2["seq"]}'
            backup_guides.append(backup_guide2)


    output_guide = {
        'systematic name': gene_id,
        'gene name': gene_name,
        'part name': f'gRNA_{gene_name}',
        'seq': component(data,'pSNR52')+guide_seq+component(data,'conserved gRNA'), 
        'notes': notes,
        'pair with': '',
        'start': (guide['crRNA_start']),
        'end': (guide['crRNA_start'] + 20)
        }
    
    return output_guide, backup_guides