import streamlit as st
from transformers import AutoModel
from PIL import Image
import torch
import numpy as np

@st.cache_resource
def load_model():
    model = AutoModel.from_pretrained("ragavsachdeva/magi", trust_remote_code=True)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    return model

@st.cache_data
def read_image_as_np_array(image_path):
    with open(image_path, "rb") as file:
        image = Image.open(file).convert("L").convert("RGB")
        image = np.array(image)
    return image

@st.cache_data
def predict_detections_and_associations(
        image_path,
        character_detection_threshold,
        panel_detection_threshold,
        text_detection_threshold,
        character_character_matching_threshold,
        text_character_matching_threshold,
):
    image = read_image_as_np_array(image_path)
    with torch.no_grad():
        result = model.predict_detections_and_associations(
            [image],
            character_detection_threshold=character_detection_threshold,
            panel_detection_threshold=panel_detection_threshold,
            text_detection_threshold=text_detection_threshold,
            character_character_matching_threshold=character_character_matching_threshold,
            text_character_matching_threshold=text_character_matching_threshold,
            )[0]
    return result

@st.cache_data
def predict_ocr(
    image_path,
    character_detection_threshold,
    panel_detection_threshold,
    text_detection_threshold,
    character_character_matching_threshold,
    text_character_matching_threshold,
):
    if not generate_transcript:
        return
    image = read_image_as_np_array(image_path)
    result = predict_detections_and_associations(
        path_to_image,
        character_detection_threshold,
        panel_detection_threshold,
        text_detection_threshold,
        character_character_matching_threshold,
        text_character_matching_threshold,
    )
    text_bboxes_for_all_images = [result["texts"]]
    with torch.no_grad():
        ocr_results = model.predict_ocr([image], text_bboxes_for_all_images)
    return ocr_results

model = load_model()

path_to_image = "/scratch/shared/beegfs/rs/comics/mangas/bakuman/1.0/p_00009.png"
st.markdown("<style>.title{font-size:2em;text-align:center;color:#fff;font-family:'Comic Sans MS',cursive;text-transform:uppercase;letter-spacing:.1em;padding:.5em 0 .2em;background:0 0}.title span{background:-webkit-linear-gradient(45deg,#6495ed,#4169e1);-webkit-background-clip:text;-webkit-text-fill-color:transparent}.subheading{font-size:1.5em;text-align:center;color:#ddd;font-family:'Comic Sans MS',cursive}.affil,.authors{font-size:1em;text-align:center;color:#ddd;font-family:'Comic Sans MS',cursive}.authors{padding-top:1em}</style><div class='title-container'> <div class='title'> The <span>Ma</span>n<span>g</span>a Wh<span>i</span>sperer </div> <div class='subheading'> Automatically Generating Transcriptions for Comics </div> <div class='authors'> Ragav Sachdeva and Andrew Zisserman </div> <div class='affil'> University of Oxford </div></div>", unsafe_allow_html=True)
path_to_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

st.sidebar.markdown("**Mode**")
generate_detections_and_associations = st.sidebar.toggle("Generate detections and associations", True)
generate_transcript = st.sidebar.toggle("Generate transcript (slower)", False)
st.sidebar.markdown("**Hyperparameters**")
input_character_detection_threshold = st.sidebar.slider('Character detection threshold', 0.0, 1.0, 0.30, step=0.01)
input_panel_detection_threshold = st.sidebar.slider('Panel detection threshold', 0.0, 1.0, 0.2, step=0.01)
input_text_detection_threshold = st.sidebar.slider('Text detection threshold', 0.0, 1.0, 0.25, step=0.01)
input_character_character_matching_threshold = st.sidebar.slider('Character-character matching threshold', 0.0, 1.0, 0.7, step=0.01)
input_text_character_matching_threshold = st.sidebar.slider('Text-character matching threshold', 0.0, 1.0, 0.4, step=0.01)


if path_to_image is None:
    st.stop()

image = read_image_as_np_array(path_to_image)

st.markdown("**Prediction**")
if generate_detections_and_associations or generate_transcript:
    result = predict_detections_and_associations(
    path_to_image,
    input_character_detection_threshold,
    input_panel_detection_threshold,
    input_text_detection_threshold,
    input_character_character_matching_threshold,
    input_text_character_matching_threshold,
)

if generate_transcript:
    ocr_results = predict_ocr(
        path_to_image,
        input_character_detection_threshold,
        input_panel_detection_threshold,
        input_text_detection_threshold,
        input_character_character_matching_threshold,
        input_text_character_matching_threshold,
    )

if generate_detections_and_associations and generate_transcript:
    col1, col2 = st.columns(2)
    output = model.visualise_single_image_prediction(image, result)
    col1.image(output)
    text_bboxes_for_all_images = [result["texts"]]
    ocr_results = model.predict_ocr([image], text_bboxes_for_all_images)
    transcript = model.generate_transcript_for_single_image(result, ocr_results[0])
    col2.text(transcript)

elif generate_detections_and_associations:
    output = model.visualise_single_image_prediction(image, result)
    st.image(output)

elif generate_transcript:
    transcript = model.generate_transcript_for_single_image(result, ocr_results[0])
    st.text(transcript)