from huggingface_hub import HfApi, ModelFilter
from transformers import AutoTokenizer, AutoModelForMaskedLM
import pandas as pd
import re
from tqdm import tqdm
import torch
import gradio as gr
import warnings
warnings.filterwarnings('ignore')

MODEL, MODEL_NAME, BATCH_CONVERTER, ALPHABET = None, None, None, None
OFFSET = 1
MODELS = [m.modelId for m in HfApi().list_models(filter=ModelFilter(author="facebook", model_name="esm", task="fill-mask"), sort="lastModified", direction=-1)]
SCORING = ["masked-marginals (more accurate)", "wt-marginals (faster)"]

def label_row(row, sequence, token_probs):
    wt, idx, mt = row[0], int(row[1:-1]) - OFFSET, row[-1]
    assert sequence[idx] == wt, "The listed wildtype does not match the provided sequence"

    wt_encoded, mt_encoded = ALPHABET[wt], ALPHABET[mt]

    score = token_probs[0, 1 + idx, mt_encoded] - token_probs[0, 1 + idx, wt_encoded]
    return score.item()

def initialise_model(model_name):
    global MODEL, MODEL_NAME, BATCH_CONVERTER, ALPHABET
    MODEL_NAME = model_name
    MODEL = AutoModelForMaskedLM.from_pretrained(model_name)
    BATCH_CONVERTER = AutoTokenizer.from_pretrained(model_name)
    ALPHABET = BATCH_CONVERTER.get_vocab()
    if torch.cuda.is_available():
        MODEL = MODEL.cuda()

def parse_input(seq, sub):
    assert seq.isalpha(), "Sequence must be alphabetic"
    substitutions, mode = list(), None

    if len(sub.split()) == 1 and len(sub.split()[0]) == len(seq):
        mode = 'seq vs seq'
        for resi,(src,trg) in enumerate(zip(seq,sub), OFFSET):
            if src != trg:
                substitutions.append(f"{src}{resi}{trg}")
    elif len(targets := sub.split()) > 1:
        if all(re.match(r'\d+', x) for x in targets):
            mode = 'deep mutational scan'
            for resi in map(int, sub.split()):
                src = seq[resi-OFFSET]
                for trg in "ACDEFGHIKLMNPQRSTVWY".replace(src,''):
                    substitutions.append(f"{src}{resi}{trg}")
        elif all(re.match(r'[A-Z]\d+[A-Z]', x) for x in targets):
            mode = 'aa substitutions'
            substitutions = targets
    
    if not mode:
        raise RuntimeError("Unrecognised running mode")
    
    return mode, pd.DataFrame(substitutions, columns=['0'])

def run_model(sequence, substitutions, batch_tokens, scoring_strategy):
    if scoring_strategy.startswith("wt-marginals"):
        with torch.no_grad():
            token_probs = torch.log_softmax(MODEL(batch_tokens)["logits"], dim=-1)
        substitutions[MODEL_NAME] = substitutions.apply(
            lambda row: label_row(
                row['0'],
                sequence,
                token_probs,
            ),
            axis=1,
        )
    elif scoring_strategy.startswith("masked-marginals"):
        all_token_probs = []
        for i in tqdm(range(batch_tokens.size()[1])):
            batch_tokens_masked = batch_tokens.clone()
            batch_tokens_masked[0, i] = ALPHABET['<mask>']
            with torch.no_grad():
                token_probs = torch.log_softmax(
                    MODEL(batch_tokens_masked)["logits"], dim=-1
                )
            all_token_probs.append(token_probs[:, i])
        token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0)
        substitutions[MODEL_NAME] = substitutions.apply(
            lambda row: label_row(
                row['0'],
                sequence,
                token_probs,
            ),
            axis=1,
        )

    return substitutions

def parse_output(output, mode):
    if mode == 'aa substitutions':
        output = output.sort_values(MODEL_NAME, ascending=False)
    elif mode == 'deep mutational scan':
        output = pd.concat([(output.assign(resi=output['0'].str.extract(r'(\d+)', expand=False).astype(int))
                .sort_values(['resi', MODEL_NAME], ascending=[True,False])
                .groupby(['resi'])
                .head(19)
                .drop(['resi'], axis=1)).iloc[19*x:19*(x+1)].reset_index(drop=True) for x in range(output.shape[0]//19)]
                , axis=1).set_axis(range(output.shape[0]//19*2), axis='columns')

    return output.style.format(lambda x: f'{x:.2f}' if isinstance(x, float) else x).hide_index().hide_columns().background_gradient(cmap="RdYlGn", vmax=8, vmin=-8).to_html()

    
# mode = 'deep mutational scan'  #@param ['seq vs seq', 'deep mutational scan', 'aa substitutions']
# sequence = "MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ" #@param {type:"string"}
# target = "61 214 19 30 122 140" #@param {type:"string"}
# substitutions = list()
# scoring_strategy = "masked-marginals"

# if mode == 'seq vs seq':
#   for resi,(seq,trg) in enumerate(zip(sequence,target), OFFSET):
#     if seq != trg:
#       substitutions.append(f"{seq}{resi}{trg}")
# elif mode == 'deep mutational scan':
#   for resi in map(int, target.split()):
#     seq = sequence[resi-OFFSET]
#     for trg in "ACDEFGHIKLMNPQRSTVWY".replace(seq,''):
#       substitutions.append(f"{seq}{resi}{trg}")
# elif mode == 'aa substitutions':
#   substitutions = target.split()
# else:
#   raise RuntimeError("Unrecognised running mode")

# df = pd.DataFrame(substitutions, columns=['0'])
# mutation_col = df.columns[0]

# batch_tokens = batch_converter(sequence, return_tensors='pt')['input_ids']

# if scoring_strategy == "wt-marginals":
#     with torch.no_grad():
#         token_probs = torch.log_softmax(model(batch_tokens)["logits"], dim=-1)
#     df[model_name] = df.apply(
#         lambda row: label_row(
#             row[mutation_col],
#             sequence,
#             token_probs,
#             alphabet,
#             OFFSET,
#         ),
#         axis=1,
#     )
# elif scoring_strategy == "masked-marginals":
#     all_token_probs = []
#     for i in tqdm(range(batch_tokens.size()[1])):
#         batch_tokens_masked = batch_tokens.clone()
#         batch_tokens_masked[0, i] = alphabet['<mask>']
#         with torch.no_grad():
#             token_probs = torch.log_softmax(
#                 model(batch_tokens_masked)["logits"], dim=-1
#             )
#         all_token_probs.append(token_probs[:, i])  # vocab size
#     token_probs = torch.cat(all_token_probs, dim=0).unsqueeze(0)
#     df[model_name] = df.apply(
#         lambda row: label_row(
#             row[mutation_col],
#             sequence,
#             token_probs,
#             alphabet,
#             OFFSET,
#         ),
#         axis=1,
#     )

# if mode == 'aa substitutions':
#   df = df.sort_values(model_name, ascending=False)
# elif mode == 'deep mutational scan':
#   df = pd.concat([(df.assign(resi=df['0'].str.extract(f'(\d+)', expand=False).astype(int))
#           .sort_values(['resi', model_name], ascending=[True,False])
#           .groupby(['resi'])
#           .head(19)
#           .drop(['resi'], axis=1)).iloc[19*x:19*(x+1)].reset_index(drop=True) for x in range(df.shape[0]//19)]
#         , axis=1).set_axis(range(df.shape[0]//19*2), axis='columns')

# df.style.hide_index().hide_columns().background_gradient(cmap="RdYlGn", vmax=8, vmin=-8)

def app(*argv):
    seq, trg, model_name, scoring_strategy, *_ = argv

    mode, substitutions = parse_input(seq, trg)

    if model_name != MODEL_NAME:
        initialise_model(model_name)

    batch_tokens = BATCH_CONVERTER(seq, return_tensors='pt')['input_ids']

    df = run_model(seq, substitutions, batch_tokens, scoring_strategy)

    return parse_output(df, mode)

# demo = gr.Interface(
#     theme=gr.themes.Base(),
#     title="Protein Sequence Mutagenesis",
#     description="Predict the effect of mutations on protein stability",
#     fn=app,
#     inputs=[gr.Textbox(lines=2, label="Sequence", placeholder="Sequence here...", required=True, value='MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ'),
#             gr.Textbox(lines=2, label="Substitutions", placeholder="Substitutions here...", required=True, value="61 214 19 30 122 140"),
#             gr.Dropdown(MODELS, label="Model", value=MODELS[1]),
#             gr.Dropdown(["masked-marginals (more accurate)", "wt-marginals (faster)"], label="Scoring strategy", value="wt-marginals (faster)"),
#             ],
#     outputs=gr.HTML(formatter="html", label="Output"),
# )

with gr.Blocks() as demo:
    gr.Markdown("""Protein Sequence Mutagenesis""", name="title")
    gr.Markdown("""Predict the effect of mutations on protein stability""", name="description")
    seq = gr.Textbox(lines=2, label="Sequence", placeholder="Sequence here...", required=True, value='MVEQYLLEAIVRDARDGITISDCSRPDNPLVFVNDAFTRMTGYDAEEVIGKNCRFLQRGDINLSAVHTIKIAMLTHEPCLVTLKNYRKDGTIFWNELSLTPIINKNGLITHYLGIQKDVSAQVILNQTLHEENHLLKSNKEMLEYLVNIDALTGLHNRRFLEDQLVIQWKLASRHINTITIFMIDIDYFKAFNDTYGHTAGDEALRTIAKTLNNCFMRGSDFVARYGGEEFTILAIGMTELQAHEYSTKLVQKIENLNIHHKGSPLGHLTISLGYSQANPQYHNDQNLVIEQADRALYSAKVEGKNRAVAYREQ')
    trg = gr.Textbox(lines=1, label="Substitutions", placeholder="Substitutions here...", required=True, value="61 214 19 30 122 140")
    model_name = gr.Dropdown(MODELS, label="Model", value=MODELS[1])
    scoring_strategy = gr.Dropdown(SCORING, label="Scoring strategy", value=SCORING[1])
    btn = gr.Button(label="Submit", type="submit")
    btn.click(fn=app, inputs=[seq, trg, model_name, scoring_strategy], outputs=[gr.HTML()])

if __name__ == '__main__':
    demo.launch()
    # demo.launch(share=True, server_name="0.0.0.0", server_port=7878)