import os
import yaml
import gdown 
import gradio as gr
from predict import PredictTri
from gradio import blocks 

output_path = "tashkeela-d2.pt"
gdrive_templ = "https://drive.google.com/file/d/{}/view?usp=sharing"
if not os.path.exists(output_path):
    model_gdrive_id = "1FGelqImFkESbTyRsx_elkKIOZ9VbhRuo"
    gdown.download(gdrive_templ.format(model_gdrive_id), output=output_path, quiet=False)

output_path = "vocab.vec"
if not os.path.exists(output_path):
    vocab_gdrive_id = "1-0muGvcSYEf8RAVRcwXay4MRex6kmCii"
    gdown.download(gdrive_templ.format(vocab_gdrive_id), output=output_path, quiet=False)

with open("config.yaml", 'r', encoding="utf-8") as file:
    config = yaml.load(file, Loader=yaml.FullLoader)

config["train"]["max-sent-len"] = config["predictor"]["window"]
config["train"]["max-token-count"] = config["predictor"]["window"] * 3

def diacritze(text, do_partial):
    predictor = PredictTri(config, text)
    diacritized_lines = predictor.predict_partial(do_partial=do_partial)
    return diacritized_lines

with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Partial Diacritization: A Context-Contrastive Inference Approach
    ## Authors: Muhammad ElNokrashy, Badr AlKhamissi
    """)

    with gr.Row():
        check_box = gr.Checkbox(label="Partial", info="Apply Partial Diacritics or Full Diacritics")
        threshold_txt = gr.Textbox("")

    input_txt = gr.Textbox(
        placeholder="اكتب هنا", 
        lines=5, 
        label="Input",
        type='text',
        rtl=True,
        text_align='right',
    )

    output_txt = gr.Textbox(
        lines=5, 
        label="Output",
        type='text',
        rtl=True,  
        text_align='right',
    )

    btn = gr.Button(value="Shakkel")
    btn.click(diacritze, inputs=[input_txt, check_box], outputs=[output_txt])

if __name__ == "__main__":
    demo.queue().launch(
        # share=False,
        # debug=False,
        # server_port=7860,
        # server_name="0.0.0.0",
        # ssl_verify=False,
        # ssl_certfile="cert.pem",
        # ssl_keyfile="key.pem"
    )