import gradio as gr
import torch
from transformers import NougatProcessor, VisionEncoderDecoderModel
from PIL import Image
import tempfile
import os
import re
import fitz  # PyMuPDF
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
from reportlab.lib.styles import getSampleStyleSheet
from reportlab.lib.units import inch

# -----------------------------
# Load Model
# -----------------------------

MODEL_NAME = "facebook/nougat-base"

processor = NougatProcessor.from_pretrained(MODEL_NAME)
model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

# -----------------------------
# Utilities
# -----------------------------

def clean_output(text):
    text = re.sub(r"<.*?>", "", text)
    return text.strip()

def image_to_latex(image):
    pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)
    outputs = model.generate(pixel_values, max_new_tokens=1024)
    latex = processor.batch_decode(outputs, skip_special_tokens=True)[0]
    return clean_output(latex)

def pdf_to_images(pdf_path):
    images = []
    pdf = fitz.open(pdf_path)

    for page in pdf:
        pix = page.get_pixmap()
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
        images.append(img)

    return images

def generate_pdf_from_text(text, output_path):
    doc = SimpleDocTemplate(output_path)
    elements = []
    styles = getSampleStyleSheet()

    for line in text.split("\n"):
        elements.append(Paragraph(line, styles["Normal"]))
        elements.append(Spacer(1, 0.2 * inch))

    doc.build(elements)

# -----------------------------
# Main Function
# -----------------------------

def process_file(file):

    if file is None:
        return "Upload file first", None, None

    try:
        temp_dir = tempfile.mkdtemp()
        full_latex = ""

        if file.name.endswith(".pdf"):
            images = pdf_to_images(file.name)

            for img in images:
                latex = image_to_latex(img)
                full_latex += latex + "\n\n"

        else:
            image = Image.open(file.name).convert("RGB")
            full_latex = image_to_latex(image)

        # Save .tex file
        tex_path = os.path.join(temp_dir, "output.tex")
        with open(tex_path, "w", encoding="utf-8") as f:
            f.write(full_latex)

        # Generate readable PDF
        pdf_path = os.path.join(temp_dir, "output.pdf")
        generate_pdf_from_text(full_latex, pdf_path)

        return full_latex, tex_path, pdf_path

    except Exception as e:
        return f"Error occurred: {str(e)}", None, None

# -----------------------------
# Gradio UI
# -----------------------------

with gr.Blocks(title="Handwritten Assignment Digitizer") as demo:

    gr.Markdown("""
    # Handwritten Assignment Digitizer

   
    """)

    file_input = gr.File(label="Upload Image or PDF")

    latex_output = gr.Textbox(label="Generated LaTeX Code", lines=15)

    tex_download = gr.File(label="Download .tex file")

    pdf_download = gr.File(label="Download Readable PDF")

    process_button = gr.Button("Convert Document")

    process_button.click(
        fn=process_file,
        inputs=file_input,
        outputs=[latex_output, tex_download, pdf_download]
    )

demo.launch()