import streamlit as st
from io import StringIO
from Bio import SeqIO

import numpy as np
import os
import pandas as pd
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import Counter, OrderedDict
from copy import deepcopy
from esm import Alphabet, FastaBatchedDataset, ProteinBertModel, pretrained, MSATransformer
from esm.data import *
from esm.model.esm2 import ESM2
from torch import nn
from torch.nn import Linear
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset, DataLoader
seed = 19961231
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)


st.title("IRES-LM prediction and mutation")

# Input sequence
st.subheader("Input sequence")

seq = st.text_area("FASTA format only", value=">vir_CVB3_ires_00505.1\nTTAAAACAGCCTGTGGGTTGATCCCACCCACAGGCCCATTGGGCGCTAGCACTCTGGTATCACGGTACCTTTGTGCGCCTGTTTTATACCCCCTCCCCCAACTGTAACTTAGAAGTAACACACACCGATCAACAGTCAGCGTGGCACACCAGCCACGTTTTGATCAAGCACTTCTGTTACCCCGGACTGAGTATCAATAGACTGCTCACGCGGTTGAAGGAGAAAGCGTTCGTTATCCGGCCAACTACTTCGAAAAACCTAGTAACACCGTGGAAGTTGCAGAGTGTTTCGCTCAGCACTACCCCAGTGTAGATCAGGTCGATGAGTCACCGCATTCCCCACGGGCGACCGTGGCGGTGGCTGCGTTGGCGGCCTGCCCATGGGGAAACCCATGGGACGCTCTAATACAGACATGGTGCGAAGAGTCTATTGAGCTAGTTGGTAGTCCTCCGGCCCCTGAATGCGGCTAATCCTAACTGCGGAGCACACACCCTCAAGCCAGAGGGCAGTGTGTCGTAACGGGCAACTCTGCAGCGGAACCGACTACTTTGGGTGTCCGTGTTTCATTTTATTCCTATACTGGCTGCTTATGGTGACAATTGAGAGATCGTTACCATATAGCTATTGGATTGGCCATCCGGTGACTAATAGAGCTATTATATATCCCTTTGTTGGGTTTATACCACTTAGCTTGAAAGAGGTTAAAACATTACAATTCATTGTTAAGTTGAATACAGCAAA")
st.subheader("Upload sequence file")
uploaded = st.file_uploader("Sequence file in FASTA format")

# augments
global output_filename, start_nt_position, end_nt_position, mut_by_prob, transform_type, mlm_tok_num, n_mut, n_designs_ep, n_sampling_designs_ep, n_mlm_recovery_sampling, mutate2stronger
output_filename = st.text_input("output a .csv file", value='IRES_LM_prediction_mutation')
start_nt_position = st.number_input("The start position of the mutation of this sequence, the first position is defined as 0", value=0)
end_nt_position = st.number_input("The last position of the mutation of this sequence, the last position is defined as length(sequence)-1 or -1", value=-1)
mut_by_prob = st.checkbox("Mutated by predicted Probability or Transformed Probability of the sequence", value=True)
transform_type = st.selectbox("Type of probability transformation", 
                              ['', 'sigmoid', 'logit', 'power_law', 'tanh'],
                              index=2)
mlm_tok_num = st.number_input("Number of masked tokens for each sequence per epoch", value=1)
n_mut = st.number_input("Maximum number of mutations for each sequence", value=3)
n_designs_ep = st.number_input("Number of mutations per epoch", value=10)
n_sampling_designs_ep = st.number_input("Number of sampling mutations from n_designs_ep per epoch", value=5)
n_mlm_recovery_sampling = st.number_input("Number of MLM recovery samplings (with AGCT recovery)", value=1)
mutate2stronger = st.checkbox("Mutate to stronger IRES variant, otherwise mutate to weaker IRES", value=True)
if not mut_by_prob and transform_type != '':
    st.write("--transform_type must be '' when --mut_by_prob is False")
    transform_type = ''

global idx_to_tok, prefix, epochs, layers, heads, fc_node, dropout_prob, embed_dim, batch_toks, repr_layers, evaluation, include, truncate, return_contacts, return_representation, mask_toks_id, finetune

epochs = 5
layers = 6
heads = 16
embed_dim = 128
batch_toks = 4096
fc_node = 64
dropout_prob = 0.5
folds = 10
repr_layers = [-1]
include = ["mean"]
truncate = True
finetune = False
return_contacts = False
return_representation = False

global tok_to_idx, idx_to_tok, mask_toks_id
alphabet = Alphabet(mask_prob = 0.15, standard_toks = 'AGCT')
assert alphabet.tok_to_idx == {'<pad>': 0, '<eos>': 1, '<unk>': 2, 'A': 3, 'G': 4, 'C': 5, 'T': 6, '<cls>': 7, '<mask>': 8, '<sep>': 9}

# tok_to_idx = {'<pad>': 0, '<eos>': 1, '<unk>': 2, 'A': 3, 'G': 4, 'C': 5, 'T': 6, '<cls>': 7, '<mask>': 8, '<sep>': 9}
tok_to_idx = {'-': 0, '&': 1, '?': 2, 'A': 3, 'G': 4, 'C': 5, 'T': 6, '!': 7, '*': 8, '|': 9}
idx_to_tok = {idx: tok for tok, idx in tok_to_idx.items()}
# st.write(tok_to_idx)
mask_toks_id = 8

global w1, w2, w3
w1, w2, w3 = 1, 1, 100

class CNN_linear(nn.Module):
    def __init__(self):
        super(CNN_linear, self).__init__()
        
        self.esm2 = ESM2(num_layers = layers,
                         embed_dim = embed_dim,
                         attention_heads = heads,
                         alphabet = alphabet)
        
        self.dropout = nn.Dropout(dropout_prob)
        self.relu = nn.ReLU()
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(in_features = embed_dim, out_features = fc_node)
        self.output = nn.Linear(in_features = fc_node, out_features = 2)        
            
    def predict(self, tokens):
        
        x = self.esm2(tokens, [layers], need_head_weights=False, return_contacts=False, return_representation = True)
        x_cls = x["representations"][layers][:, 0]
        
        o = self.fc(x_cls)
        o = self.relu(o)
        o = self.dropout(o)
        o = self.output(o)
        
        y_prob = torch.softmax(o, dim = 1)
        y_pred = torch.argmax(y_prob, dim = 1)
        
        if transform_type:
            y_prob_transformed = prob_transform(y_prob[:,1])
            return y_prob[:,1], y_pred, x['logits'], y_prob_transformed
        else:
            return y_prob[:,1], y_pred, x['logits'], o[:,1]
    
    def forward(self, x1, x2):
        logit_1, repr_1 = self.predict(x1)
        logit_2, repr_2 = self.predict(x2)
        return (logit_1, logit_2), (repr_1, repr_2)


def prob_transform(prob, **kwargs): # Logits
    """
    Transforms probability values based on the specified method.

    :param prob: torch.Tensor, the input probabilities to be transformed
    :param transform_type: str, the type of transformation to be applied
    :param kwargs: additional parameters for transformations
    :return: torch.Tensor, transformed probabilities
    """

    if transform_type == 'sigmoid':
        x0 = kwget('x0', 0.5)
        k = kwget('k', 10.0)
        prob_transformed = 1 / (1 + torch.exp(-k * (prob - x0)))
    
    elif transform_type == 'logit':
        # Adding a small value to avoid log(0) and log(1)
        prob_transformed = torch.log(prob + 1e-6) - torch.log(1 - prob + 1e-6)
    
    elif transform_type == 'power_law':
        gamma = kwget('gamma', 2.0)
        prob_transformed = torch.pow(prob, gamma)
    
    elif transform_type == 'tanh':
        k = kwget('k', 2.0)
        prob_transformed = torch.tanh(k * prob)
    
    return prob_transformed

def random_replace(sequence, continuous_replace=False):
    global start_nt_position, end_nt_position
    if end_nt_position == -1: end_nt_position = len(sequence)-1
    if start_nt_position < 0 or end_nt_position > len(sequence)-1 or start_nt_position > end_nt_position:
#         raise ValueError("Invalid start/end positions")
        st.write("Invalid start/end positions")
        start_nt_position, end_nt_position = 0, len(sequence)-1
    
    # 将序列切片成三部分：替换区域前、替换区域、替换区域后
    pre_segment = sequence[:start_nt_position]
    target_segment = list(sequence[start_nt_position:end_nt_position + 1])  # +1因为Python的切片是右开区间
    post_segment = sequence[end_nt_position + 1:]
    
    if not continuous_replace:
        # 随机替换目标片段的mlm_tok_num个位置
        indices = random.sample(range(len(target_segment)), mlm_tok_num)
        for idx in indices:
            target_segment[idx] = '*'
    else:
        # 在目标片段连续替换mlm_tok_num个位置
        max_start_idx = len(target_segment) - mlm_tok_num  # 确保从i开始的n_mut个元素不会超出目标片段的长度
        if max_start_idx < 1:  # 如果目标片段长度小于mlm_tok_num，返回原始序列
            return target_segment
        start_idx = random.randint(0, max_start_idx)
        for idx in range(start_idx, start_idx + mlm_tok_num):
            target_segment[idx] = '*'
    
    # 合并并返回最终的序列
    return ''.join([pre_segment] + target_segment + [post_segment])


def mlm_seq(seq):
    seq_token, masked_sequence_token = [7],[7]
    seq_token += [tok_to_idx[token] for token in seq]
    
    masked_seq = random_replace(seq, n_mut) # 随机替换n_mut个元素为'*'
    masked_seq_token += [tok_to_idx[token] for token in masked_seq]
    
    return seq, masked_seq, torch.LongTensor(seq_token), torch.LongTensor(masked_seq_token)

def batch_mlm_seq(seq_list, continuous_replace = False):
    batch_seq = []
    batch_masked_seq = []
    batch_seq_token_list = []
    batch_masked_seq_token_list = []

    for i, seq in enumerate(seq_list):
        seq_token, masked_seq_token = [7], [7]
        seq_token += [tok_to_idx[token] for token in seq]

        masked_seq = random_replace(seq, continuous_replace)  # 随机替换n_mut个元素为'*'
        masked_seq_token += [tok_to_idx[token] for token in masked_seq]

        batch_seq.append(seq)
        batch_masked_seq.append(masked_seq)

        batch_seq_token_list.append(seq_token)
        batch_masked_seq_token_list.append(masked_seq_token)

    return batch_seq, batch_masked_seq, torch.LongTensor(batch_seq_token_list), torch.LongTensor(batch_masked_seq_token_list)

def recovered_mlm_tokens(masked_seqs, masked_toks, esm_logits, exclude_low_prob = False):
    # Only remain the AGCT logits
    esm_logits = esm_logits[:,:,3:7]
    # Get the predicted tokens using argmax
    predicted_toks = (esm_logits.argmax(dim=-1)+3).tolist()
    
    batch_size, seq_len, vocab_size = esm_logits.size()
    if exclude_low_prob: min_prob = 1 / vocab_size
    # Initialize an empty list to store the recovered sequences
    recovered_sequences, recovered_toks = [], []
    
    for i in range(batch_size):
        recovered_sequence_i, recovered_tok_i = [], []
        for j in range(seq_len):
            if masked_toks[i][j] == 8:
                st.write(i,j)
                ### Sample M recovery sequences using the logits
                recovery_probs = torch.softmax(esm_logits[i, j], dim=-1)
                recovery_probs[predicted_toks[i][j]-3] = 0  # Exclude the most probable token
                if exclude_low_prob: recovery_probs[recovery_probs < min_prob] = 0  # Exclude tokens with low probs < min_prob
                recovery_probs /= recovery_probs.sum()  # Normalize the probabilities
                
                ### 有放回抽样
                max_retries = 5
                retries = 0
                success = False

                while retries < max_retries and not success:
                    try:
                        recovery_indices = list(np.random.choice(vocab_size, size=n_mlm_recovery_sampling, p=recovery_probs.cpu().detach().numpy(), replace=False))
                        success = True  # 设置成功标志
                    except ValueError as e:
                        retries += 1
                        st.write(f"Attempt {retries} failed with error: {e}")
                        if retries >= max_retries:
                            st.write("Max retries reached. Skipping this iteration.")
                
                ### recovery to sequence
                if retries < max_retries:
                    for idx in [predicted_toks[i][j]] + [3+i for i in recovery_indices]:
                        recovery_seq = deepcopy(list(masked_seqs[i]))
                        recovery_tok = deepcopy(masked_toks[i])
                        
                        recovery_tok[j] = idx
                        recovery_seq[j-1] = idx_to_tok[idx]
                        
                        recovered_tok_i.append(recovery_tok)
                        recovered_sequence_i.append(''.join(recovery_seq))
                        
        recovered_sequences.extend(recovered_sequence_i)
        recovered_toks.extend(recovered_tok_i)
    return recovered_sequences, torch.LongTensor(torch.stack(recovered_toks))

def recovered_mlm_multi_tokens(masked_seqs, masked_toks, esm_logits, exclude_low_prob = False):
    # Only remain the AGCT logits
    esm_logits = esm_logits[:,:,3:7]
    # Get the predicted tokens using argmax
    predicted_toks = (esm_logits.argmax(dim=-1)+3).tolist()
    
    batch_size, seq_len, vocab_size = esm_logits.size()
    if exclude_low_prob: min_prob = 1 / vocab_size
    # Initialize an empty list to store the recovered sequences
    recovered_sequences, recovered_toks = [], []
    
    for i in range(batch_size):
        recovered_sequence_i, recovered_tok_i = [], []
        recovered_masked_num = 0
        for j in range(seq_len):
            if masked_toks[i][j] == 8:
                ### Sample M recovery sequences using the logits
                recovery_probs = torch.softmax(esm_logits[i, j], dim=-1)
                recovery_probs[predicted_toks[i][j]-3] = 0  # Exclude the most probable token
                if exclude_low_prob: recovery_probs[recovery_probs < min_prob] = 0  # Exclude tokens with low probs < min_prob
                recovery_probs /= recovery_probs.sum()  # Normalize the probabilities
                
                ### 有放回抽样
                max_retries = 5
                retries = 0
                success = False

                while retries < max_retries and not success:
                    try:
                        recovery_indices = list(np.random.choice(vocab_size, size=n_mlm_recovery_sampling, p=recovery_probs.cpu().detach().numpy(), replace=False))
                        success = True  # 设置成功标志
                    except ValueError as e:
                        retries += 1
                        st.write(f"Attempt {retries} failed with error: {e}")
                        if retries >= max_retries:
                            st.write("Max retries reached. Skipping this iteration.")
                
                ### recovery to sequence 
                        
                if recovered_masked_num == 0:
                    if retries < max_retries:
                        for idx in [predicted_toks[i][j]] + [3+i for i in recovery_indices]:
                            recovery_seq = deepcopy(list(masked_seqs[i]))
                            recovery_tok = deepcopy(masked_toks[i])

                            recovery_tok[j] = idx
                            recovery_seq[j-1] = idx_to_tok[idx]

                            recovered_tok_i.append(recovery_tok)
                            recovered_sequence_i.append(''.join(recovery_seq))

                elif recovered_masked_num > 0:
                    if retries < max_retries:
                        for idx in [predicted_toks[i][j]] + [3+i for i in recovery_indices]:
                            for recovery_seq, recovery_tok in zip(list(recovered_sequence_i), list(recovered_tok_i)): # 要在循环开始之前获取列表的副本来进行迭代。这样，在循环中即使我们修改了原始的列表，也不会影响迭代的行为。

                                recovery_seq_temp = list(recovery_seq)
                                recovery_tok[j] = idx
                                recovery_seq_temp[j-1] = idx_to_tok[idx]
                                
                                recovered_tok_i.append(recovery_tok)
                                recovered_sequence_i.append(''.join(recovery_seq_temp))
                
                recovered_masked_num += 1                
        recovered_indices = [i for i, s in enumerate(recovered_sequence_i) if '*' not in s]
        recovered_tok_i = [recovered_tok_i[i] for i in recovered_indices]
        recovered_sequence_i = [recovered_sequence_i[i] for i in recovered_indices]

        recovered_sequences.extend(recovered_sequence_i)
        recovered_toks.extend(recovered_tok_i)
        
        recovered_sequences, recovered_toks = remove_duplicates_double(recovered_sequences, recovered_toks)

    return recovered_sequences, torch.LongTensor(torch.stack(recovered_toks))

def mismatched_positions(s1, s2):
    # 这个函数假定两个字符串的长度相同。
    """Return the number of positions where two strings differ."""
    
    # The number of mismatches will be the sum of positions where characters are not the same
    return sum(1 for c1, c2 in zip(s1, s2) if c1 != c2)

def remove_duplicates_triple(filtered_mut_seqs, filtered_mut_probs, filtered_mut_logits):
    seen = {}
    unique_seqs = []
    unique_probs = []
    unique_logits = []
    
    for seq, prob, logit in zip(filtered_mut_seqs, filtered_mut_probs, filtered_mut_logits):
        if seq not in seen:
            unique_seqs.append(seq)
            unique_probs.append(prob)
            unique_logits.append(logit)
            seen[seq] = True

    return unique_seqs, unique_probs, unique_logits

def remove_duplicates_double(filtered_mut_seqs, filtered_mut_probs):
    seen = {}
    unique_seqs = []
    unique_probs = []
    
    for seq, prob in zip(filtered_mut_seqs, filtered_mut_probs):
        if seq not in seen:
            unique_seqs.append(seq)
            unique_probs.append(prob)
            seen[seq] = True

    return unique_seqs, unique_probs

def mutated_seq(wt_seq, wt_label):
    wt_seq = '!'+ wt_seq
    wt_tok = torch.LongTensor([[tok_to_idx[token] for token in wt_seq]]).to(device)
    wt_prob, wt_pred, _, wt_logit = model.predict(wt_tok)
    
    st.write(f'Wild Type: Length = {len(wt_seq)} \n{wt_seq}')
    st.write(f'Wild Type: Label = {wt_label}, Y_pred = {wt_pred.item()}, Y_prob = {wt_prob.item():.2%}')
    
#     st.write(n_mut, mlm_tok_num, n_designs_ep, n_sampling_designs_ep, n_mlm_recovery_sampling, mutate2stronger)
    # pbar = tqdm(total=n_mut)
    mutated_seqs = []
    i = 1
    # pbar = st.progress(i, text="mutated number of sequence")
    while i <= n_mut:
        if i == 1: seeds_ep = [wt_seq[1:]]
        seeds_next_ep, seeds_probs_next_ep, seeds_logits_next_ep = [], [], []
        for seed in seeds_ep:
            seed_seq, masked_seed_seq, seed_seq_token, masked_seed_seq_token = batch_mlm_seq([seed] * n_designs_ep, continuous_replace = True) ### mask seed with 1 site to "*"

            seed_prob, seed_pred, _, seed_logit = model.predict(seed_seq_token[0].unsqueeze_(0).to(device))
            _, _, seed_esm_logit, _ = model.predict(masked_seed_seq_token.to(device))
            mut_seqs, mut_toks = recovered_mlm_multi_tokens(masked_seed_seq, masked_seed_seq_token, seed_esm_logit)
            mut_probs, mut_preds, mut_esm_logits, mut_logits = model.predict(mut_toks.to(device))

            ### Filter mut_seqs that mut_prob < seed_prob and mut_prob < wild_prob
            filtered_mut_seqs = []
            filtered_mut_probs = []
            filtered_mut_logits = []
            if mut_by_prob:
                for z in range(len(mut_seqs)):
                    if mutate2stronger:
                        if mut_probs[z] >= seed_prob and mut_probs[z] >= wt_prob:
                            filtered_mut_seqs.append(mut_seqs[z])
                            filtered_mut_probs.append(mut_probs[z].cpu().detach().numpy())
                            filtered_mut_logits.append(mut_logits[z].cpu().detach().numpy())
                    else:
                        if mut_probs[z] < seed_prob and mut_probs[z] < wt_prob:
                            filtered_mut_seqs.append(mut_seqs[z])
                            filtered_mut_probs.append(mut_probs[z].cpu().detach().numpy())
                            filtered_mut_logits.append(mut_logits[z].cpu().detach().numpy())
            else:
                for z in range(len(mut_seqs)):
                    if mutate2stronger:
                        if mut_logits[z] >= seed_logit and mut_logits[z] >= wt_logit:
                            filtered_mut_seqs.append(mut_seqs[z])
                            filtered_mut_probs.append(mut_probs[z].cpu().detach().numpy())
                            filtered_mut_logits.append(mut_logits[z].cpu().detach().numpy())
                    else:
                        if mut_logits[z] < seed_logit and mut_logits[z] < wt_logit:
                            filtered_mut_seqs.append(mut_seqs[z])
                            filtered_mut_probs.append(mut_probs[z].cpu().detach().numpy())
                            filtered_mut_logits.append(mut_logits[z].cpu().detach().numpy())
                
            

            ### Save
            seeds_next_ep.extend(filtered_mut_seqs)
            seeds_probs_next_ep.extend(filtered_mut_probs)
            seeds_logits_next_ep.extend(filtered_mut_logits)
            seeds_next_ep, seeds_probs_next_ep, seeds_logits_next_ep = remove_duplicates_triple(seeds_next_ep, seeds_probs_next_ep, seeds_logits_next_ep)

        ### Sampling based on prob
        if len(seeds_next_ep) > n_sampling_designs_ep:
            seeds_probs_next_ep_norm = seeds_probs_next_ep / sum(seeds_probs_next_ep)  # Normalize the probabilities
            seeds_index_next_ep = np.random.choice(len(seeds_next_ep), n_sampling_designs_ep, p = seeds_probs_next_ep_norm, replace = False)
            
            seeds_next_ep = np.array(seeds_next_ep)[seeds_index_next_ep]
            seeds_probs_next_ep = np.array(seeds_probs_next_ep)[seeds_index_next_ep]
            seeds_logits_next_ep = np.array(seeds_logits_next_ep)[seeds_index_next_ep]
        seeds_mutated_num_next_ep = [mismatched_positions(wt_seq[1:], s) for s in seeds_next_ep]
        
        mutated_seqs.extend(list(zip(seeds_next_ep, seeds_logits_next_ep, seeds_probs_next_ep, seeds_mutated_num_next_ep)))

        seeds_ep = seeds_next_ep
        i += 1
        # pbar.update(1)
        # pbar.progress(i/n_mut, text="Mutating")
    # pbar.close()
    # st.success('Done', icon="✅")
    mutated_seqs.extend([(wt_seq[1:], wt_logit.item(), wt_prob.item(), 0)])
    mutated_seqs = sorted(mutated_seqs, key=lambda x: x[2], reverse=True)
    mutated_seqs = pd.DataFrame(mutated_seqs, columns = ['mutated_seq', 'predicted_logit', 'predicted_probability', 'mutated_num']).drop_duplicates('mutated_seq')
    return mutated_seqs

def read_raw(raw_input):
    ids = []
    sequences = []

    file = StringIO(raw_input)
    for record in SeqIO.parse(file, "fasta"):

        # 检查序列是否只包含A, G, C, T
        sequence = str(record.seq.back_transcribe()).upper()
        if not set(sequence).issubset(set("AGCT")):
            st.write(f"Record '{record.description}' was skipped for containing invalid characters. Only A, G, C, T(U) are allowed.")
            continue

        # 将符合条件的序列添加到列表中
        ids.append(record.id)
        sequences.append(sequence)
    
    return ids, sequences

def predict_raw(raw_input):
    model.eval()
    # st.write('====Parse Input====')
    ids, seqs = read_raw(raw_input)

    # st.write('====Predict====')
    res_pd = pd.DataFrame(columns = ['wildtype_id', 'mutated_seq', 'predicted_logit', 'predicted_probability', 'mutated_num'])
    for wt_seq, wt_id in zip(seqs, ids):
        try:
            res = mutated_seq(wt_seq, wt_id)
            res['wildtype_id'] = wt_id
            res_pd = pd.concat([res_pd,res], axis = 0)
        except:
            st.write('====Please Try Again this sequence: ', wt_id, wt_seq)
    
    # st.write(res_pd)
    return res_pd

global model, device
device = "cpu"
state_dict = torch.load('model.pt', map_location=torch.device(device))
new_state_dict = OrderedDict()

for k, v in state_dict.items():
    name = k.replace('module.','')
    new_state_dict[name] = v

model = CNN_linear().to(device)
model.load_state_dict(new_state_dict, strict = False)

# Run    
if st.button("Predict and Mutate"):
    if uploaded:
        result = predict_raw(uploaded.getvalue().decode())
    else:
        result = predict_raw(seq)
    
    result_file = result.to_csv(index=False)
    st.download_button("Download", result_file, file_name=output_filename+".csv")
    st.dataframe(result)