import os
import torch
import config
import streamlit as st
from utils import (
    load_dataset,
    get_model_instance,
    load_checkpoint,
    can_load_checkpoint,
    normalize_text,
)
from PIL import Image
import torchvision.transforms as transforms

# Define device
DEVICE = 'cpu'

# Define image transformations
TRANSFORMS = transforms.Compose([
    transforms.Resize((224, 224)),  # Replace with your model's expected input size
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

def load_model():
    """
    Loads the model with the vocabulary and checkpoint.
    """
    st.write("Loading dataset and vocabulary...")
    dataset = load_dataset()  # Load dataset to access vocabulary
    vocabulary = dataset.vocab  # Assuming 'vocab' is an attribute of the dataset

    st.write("Initializing the model...")
    model = get_model_instance(vocabulary)  # Initialize the model

    if can_load_checkpoint():
        st.write("Loading checkpoint...")
        load_checkpoint(model)
    else:
        st.write("No checkpoint found, starting with untrained model.")

    model.eval()  # Set the model to evaluation mode
    st.write("Model is ready for inference.")
    return model

def preprocess_image(image_path):
    """
    Preprocess the input image for the model.
    """
    st.write(f"Preprocessing image: {image_path}")
    image = Image.open(image_path).convert("RGB")  # Ensure RGB format
    image = TRANSFORMS(image).unsqueeze(0)  # Add batch dimension
    return image.to(DEVICE)

def generate_report(model, image_path):
    """
    Generates a report for a given image using the model.
    """
    image = preprocess_image(image_path)

    st.write("Generating report...")
    with torch.no_grad():
        # Assuming the model has a 'generate_caption' method
        output = model.generate_caption(image, max_length=25)
        report = " ".join(output)

    st.write(f"Generated report: {report}")
    return report

# Streamlit app
def main():
    st.title("Chest X-Ray Report Generator")
    st.write("Upload a Chest X-Ray image to generate a medical report.")

    # Upload image
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
        st.write("")

        # Save the uploaded file temporarily
        image_path = "./temp_image.png"
        with open(image_path, "wb") as f:
            f.write(uploaded_file.getbuffer())
        st.write("Image uploaded successfully.")

        # Load the model
        model = load_model()

        # Generate report
        report = generate_report(model, image_path)
        st.write("### Generated Report:")
        st.write(report)

        # Clean up temporary file
        os.remove(image_path)

if __name__ == "__main__":
    main()