import os

import gdown as gdown
import nltk
import streamlit as st
import torch
from transformers import AutoTokenizer

from mt5 import MT5


def download_models(ids):
    """
    Download all models.
    :param ids: name and links of models
    :return:
    """

    # Download sentence tokenizer
    nltk.download('punkt')

    # Download model from drive if not stored locally
    for key in ids:
        if not os.path.isfile(f"model/{key}.ckpt"):
            url = f"https://drive.google.com/u/0/uc?id={ids[key]}"
            gdown.download(url=url, output=f"model/{key}.ckpt")


@st.cache(allow_output_mutation=True)
def load_model(model_path):
    """
    Load model and cache it.
    :param model_path: path to model
    :return:
    """

    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # Loading model and tokenizer
    model = MT5.load_from_checkpoint(model_path).eval().to(device)
    model.tokenizer = AutoTokenizer.from_pretrained('tokenizer')

    return model

# elif task == 'Question Answering':
#
#     # Input area
#     inputs = st.text_area('Context:', value="A few years after the First Crusade, in 1107, the Normans under "
#                                             "the command of Bohemond, Robert\'s son, landed in Valona and "
#                                             "besieged Dyrrachium using the most sophisticated military "
#                                             "equipment of the time, but to no avail. Meanwhile, they occupied "
#                                             "Petrela, the citadel of Mili at the banks of the river Deabolis, "
#                                             "Gllavenica (Ballsh), Kanina and Jericho. This time, "
#                                             "the Albanians sided with the Normans, dissatisfied by the heavy "
#                                             "taxes the Byzantines had imposed upon them. With their help, "
#                                             "the Normans secured the Arbanon passes and opened their way to "
#                                             "Dibra. The lack of supplies, disease and Byzantine resistance "
#                                             "forced Bohemond to retreat from his campaign and sign a peace "
#                                             "treaty with the Byzantines in the city of Deabolis. ", max_chars=2048,
#                           height=250)
#     question = st.text_input('Question:', value="What forced Bohemond to retreat from his campaign? ")
#
#     # Prediction
#     with st.spinner('Please wait while the inputs are being processed...'):
#         predictions = model.qa([{'question': question, 'context': inputs}], max_length=512)
#         answer = {question: predictions[0]}
#
#     # Answer area
#     st.write(answer)
#
# elif task == 'Question Generation':
#
#     # Input area
#     inputs = st.text_area('Context (highlight answers with <hl> tokens): ',
#                           value="A few years after the First Crusade, in <hl> 1107 <hl>, the <hl> Normans <hl> under "
#                                 "the command of <hl> Bohemond <hl>, Robert\'s son, landed in Valona and "
#                                 "besieged Dyrrachium using the most sophisticated military "
#                                 "equipment of the time, but to no avail. Meanwhile, they occupied "
#                                 "Petrela, <hl> the citadel of Mili <hl> at the banks of the river Deabolis, "
#                                 "Gllavenica (Ballsh), Kanina and Jericho. This time, "
#                                 "the Albanians sided with the Normans, dissatisfied by the heavy "
#                                 "taxes the Byzantines had imposed upon them. With their help, "
#                                 "the Normans secured the Arbanon passes and opened their way to "
#                                 "Dibra. The <hl> lack of supplies, disease and Byzantine resistance <hl> "
#                                 "forced Bohemond to retreat from his campaign and sign a peace "
#                                 "treaty with the Byzantines in the city of Deabolis. ", max_chars=2048,
#                           height=250)
#
#     # Split by highlights
#     hl_index = [i for i in range(len(inputs)) if inputs.startswith('<hl>', i)]
#     contexts = []
#     answers = []
#
#     # Build a context for each highlight pair
#     for i in range(0, len(hl_index), 2):
#         contexts.append(inputs[:hl_index[i]].replace('<hl>', '') +
#                         inputs[hl_index[i]: hl_index[i + 1] + 4] +
#                         inputs[hl_index[i + 1] + 4:].replace('<hl>', ''))
#         answers.append(inputs[hl_index[i]: hl_index[i + 1] + 4].replace('<hl>', '').strip())
#
#     # Prediction
#     with st.spinner('Please wait while the inputs are being processed...'):
#         predictions = model.qg(contexts, max_length=512)
#
#     # Answer area
#     content = {}
#     for pred, ans in zip(predictions, answers):
#         content[pred] = ans
#     st.write(content)