import gradio as gr

import os
import glob
import cv2
import numpy as np
import torch
from rxnscribe import RxnScribe

from huggingface_hub import hf_hub_download

REPO_ID = "yujieq/RxnScribe"
FILENAME = "pix2seq_reaction_full.ckpt"
ckpt_path = hf_hub_download(REPO_ID, FILENAME)

device = torch.device('cpu')
model = RxnScribe(ckpt_path, device)


def get_markdown(reaction):
    output = []
    for x in ['reactants', 'conditions', 'products']:
        s = ''
        for ent in reaction[x]:
            if 'smiles' in ent:
                s += "\n```\n" + ent['smiles'] + "\n```\n"
            elif 'text' in ent:
                s += ' '.join(ent['text']) + '<br>'
            else:
                s += ent['category']
        output.append(s)
    return output


def predict(image, molscribe, ocr):
    predictions = model.predict_image(image, molscribe=molscribe, ocr=ocr)
    pred_image = model.draw_predictions_combined(predictions, image=image)
    markdown = [[i] + get_markdown(reaction) for i, reaction in enumerate(predictions)]
    return pred_image, markdown


with gr.Blocks() as demo:
    gr.Markdown("""
    <center> <h1>RxnScribe</h1> </center>

    Extract chemical reactions from a diagram. Please upload a reaction diagram, RxnScribe will predict the reaction structures in the diagram. 

    The predicted reactions are visualized in separate images. 
    <b style="color:red">Red boxes are <i><u style="color:red">reactants</u></i>.</b>
    <b style="color:green">Green boxes are <i><u style="color:green">reaction conditions</u></i>.</b>
    <b style="color:blue">Blue boxes are <i><u style="color:blue">products</u></i>.</b>

    It usually takes 5-10 seconds to process a diagram with this demo.
    Check the options to run [MolScribe](https://huggingface.co/spaces/yujieq/MolScribe) and [OCR](https://huggingface.co/spaces/tomofi/EasyOCR) (it will take a longer time, of course).

    Paper: [RxnScribe: A Sequence Generation Model for Reaction Diagram Parsing](https://pubs.acs.org/doi/10.1021/acs.jcim.3c00439)
    
    Code: [https://github.com/thomas0809/RxnScribe](https://github.com/thomas0809/RxnScribe)

    Authors: [Yujie Qian](mailto:yujieq@csail.mit.edu), Jiang Guo, Zhengkai Tu, Connor W. Coley, Regina Barzilay. _MIT CSAIL_.
    """)
    with gr.Column():
        with gr.Row():
            image = gr.Image(label="Upload reaction diagram", show_label=False, type='pil').style(height=256)
        with gr.Row():
            molscribe = gr.Checkbox(label="Run MolScribe to recognize molecule structures")
            ocr = gr.Checkbox(label="Run OCR to recognize text")
            btn = gr.Button("Submit").style(full_width=False)
        with gr.Row():
            gallery = gr.Image(label='Predicted reactions', show_label=True).style(height="auto")
            markdown = gr.Dataframe(
                headers=['#', 'reactant', 'condition', 'product'],
                datatype=['number'] + ['markdown'] * 3,
                wrap=False
            )

    btn.click(predict, inputs=[image, molscribe, ocr], outputs=[gallery, markdown])

    gr.Examples(
        examples=sorted(glob.glob('examples/*.png')),
        inputs=[image],
        outputs=[gallery, markdown],
        fn=predict,
    )

demo.launch()