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

# Load the Hugging Face model for ECG classification
model_name = "adilsaleem/ecg-image-multilabel-classification"
try:
    model = ViTForImageClassification.from_pretrained(model_name)
    model.eval()  # Set the model to evaluation mode
except Exception as e:
    st.error(f"Error loading model: {str(e)}")
    st.stop()

# CSS for styling
st.markdown(
    """
    <style>
        body {
            background-color: #F5F7FA;  /* Light background color */
            font-family: 'Arial', sans-serif;
            display: flex;
            flex-direction: column;
            height: 100vh;
            margin: 0;
        }
        .title {
            font-size: 40px;
            text-align: center;
            color: #2980B9;  /* Beautiful blue color */
            font-weight: bold;  /* Bold text */
            margin: 20px 0;
        }
        .subtitle {
            font-size: 20px;
            text-align: center;
            color: #34495E;  /* Darker shade for subtitle */
            margin-bottom: 20px;
        }
        .footer {
            background-color: #2C3E50;  /* Darker blue for footer */
            color: white;  /* Text color in footer */
            text-align: center;
            padding: 15px;  /* Padding for footer */
            position: sticky;
            bottom: 0;
            width: 100%;
            font-size: 14px;  /* Font size for footer */
        }
        .about {
            margin: 20px 0;
            text-align: center;
            color: #34495E;  /* Darker shade for about section */
        }
    </style>
    """,
    unsafe_allow_html=True,
)

# App Title
st.markdown('<div class="title">ECG READER</div>', unsafe_allow_html=True)

# About Us Section in Expander
with st.expander("About the App", expanded=True):
    st.write("""
        This ECG Reader application uses AI to analyze ECG images and provide insights into heart health.
        It helps users determine whether their ECG results are clear or unclear.
        Developed by Hamaad Ayub Khan.
    """)
    st.write("Connect with me:")
    st.markdown('- [Instagram](https://www.instagram.com/hamaadayubkhan)', unsafe_allow_html=True)
    st.markdown('- [GitHub](https://github.com/hakgs1234)', unsafe_allow_html=True)
    st.markdown('- [My Email](mailto:hakgss1234@gmail.com)', unsafe_allow_html=True)
    st.markdown('- [Facebook](https://www.facebook.com/hamaadayubkhan)', unsafe_allow_html=True)

def preprocess_image(image):
    try:
        image = image.convert("RGB")
        image = image.resize((224, 224))
        image_array = np.array(image) / 255.0
        image_array = np.transpose(image_array, (2, 0, 1))
        return torch.tensor(image_array).unsqueeze(0)
    except Exception as e:
        st.error(f"Error processing image: {str(e)}")
        return None

def classify_ecg(image):
    processed_image = preprocess_image(image)
    if processed_image is None:
        return "Error in image processing."

    try:
        with torch.no_grad():
            outputs = model(processed_image)
        logits = outputs.logits
        predicted_class = torch.argmax(logits, dim=1).item()

        return "Clear" if predicted_class == 1 else "Unclear"
    except Exception as e:
        st.error(f"Error during model inference: {str(e)}")
        return "Error in classification."

def get_ecg_explanation(status):
    explanations = {
        "Clear": "A clear ECG indicates normal heart activity without significant issues.",
        "Unclear": "An unclear ECG may indicate possible abnormalities; further evaluation is recommended."
    }
    return explanations.get(status, "No explanation available.")

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

if uploaded_file:
    try:
        image = Image.open(uploaded_file)
        st.image(image, caption="Uploaded ECG Image.", use_column_width=True)

        # Make predictions using the Hugging Face model
        heart_status = classify_ecg(image)
        st.write(f"Predicted heart status: {heart_status}")

        # Get explanation based on the status
        explanation = get_ecg_explanation(heart_status)
        st.write(f"Explanation: {explanation}")

    except Exception as e:
        st.error(f"Error uploading or displaying image: {str(e)}")

# Footer
st.markdown(
    """
    <div class="footer">
        <p>AI can make mistakes; please consult your doctor.</p>
        <p>ECG READER All rights reserved.</p>
    </div>
    """,
    unsafe_allow_html=True,
)