import streamlit as st
import torch
from transformers import DistilBertForSequenceClassification, DistilBertTokenizer
from PIL import Image
import numpy as np
from io import BytesIO

# Load the Fastai Learner model
learn = torch.load('model.pkl', map_location=torch.device('cpu'))
learn.eval()

# Load the X-ray detection model
learn_xray = torch.load("xraydet.pkl", map_location=torch.device('cpu'))
learn_xray.eval()

# Load the DistilBERT model and tokenizer
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')

# Define language constants
UZBEK = 'uz'
ENGLISH = 'en'
RUSSIAN = 'ru'

uz_welcome = '''- Assalomu alaykum, Men dasturchi Abdrurasulov Mirsaid tomonidan tarbiyalangan sun'iy intellektman. 
- Bemorning o'pka rengin rasmiga qarab, bemorda o'pka yallig'lanishi bor yoki yo'q ekanligini aniqlab beraman.
- Model aniqligi 98% ga teng.
- Muhim: Men test rejimida ishlamoqdaman, iltimos, menga ishonib hulosa qilmang, malakalik shifokorga murojat qiling.
- Boshlash uchun, iltimos bemor o'pkasining rengin rasmini yuboring. Muhim: rasm ko'rinishida, file emas.'''

en_welcome = '''- Hello, I am a trained artificial intelligence by developer Abdurasulov Mirsaid.
- Looking at the picture of the patient's lungs, I can determine whether the patient has pneumonia or not.
- My model accuracy is 98 %.
- Important: I am working in test mode, please do not have a conclusion based on my responce, consult a qualified doctor.
- To begin, please send the XRay image of the patient's lungs. Note: as a picture, not a file. '''

rus_welcome = '''- Здравствуйте, я обученный искусственный интеллект от разработчика Мирсаидa.
- Глядя на снимок легких больного, я могу определить, есть ли у больного туберкулез или нет.
- У меня действительно высокая точность (98 %).
- Важно: я работаю в тестовом режиме, пожалуйста, не делайте вывод по моему ответу, обратитесь к квалифицированному врачу. 
- Для начала отправьте рентгеновский снимок легких пациента. Примечание: как изображение, а не файл.'''

# Define welcome messages
WELCOME_MESSAGES = {
    UZBEK: uz_welcome,
    ENGLISH: en_welcome,
    RUSSIAN: rus_welcome,
}

# Function to make predictions
def predict_pneumonia(image):
    img_array = np.array(image)
    img_fastai = Image.fromarray(img_array).convert('RGB')
    img_fastai = img_fastai.resize((224, 224))
    img_fastai = np.array(img_fastai) / 255.0
    img_fastai = torch.tensor(img_fastai).permute(2, 0, 1).unsqueeze(0).float()

    pred_xray = learn_xray.predict(img_fastai)[0]
    
    if pred_xray == 1:
        inputs = tokenizer(img_array, return_tensors="pt")
        outputs = model(**inputs)
        predicted_class_idx = torch.argmax(outputs.logits[0]).item()
        if predicted_class_idx == 1:
            return "Pneumonia Positive"
        else:
            return "Normal"
    else:
        return "Invalid X-ray image. Please upload a clear lung X-ray image."

# Streamlit app
def main():
    st.title("Pneumonia Detection")
    st.write("Upload a chest X-ray image to detect pneumonia.")

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

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)
        if st.button('Predict'):
            prediction = predict_pneumonia(image)
            st.write(f"Prediction: {prediction}")

if __name__ == '__main__':
    main()