import streamlit as st
import base64
import regex as re
from predictor import predict
import torch
import numpy as np
import matplotlib.pyplot as plt
import torchvision.models as models
from PIL import Image
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset

def add_bg_from_local(image_file):
    with open(image_file, "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read())
    st.markdown(
    f"""
    <style>
    .stApp {{
        background-image: url(data:image/{"png"};base64,{encoded_string.decode()});
        background-repeat: repeat;
    }}
    </style>
    """,
        unsafe_allow_html=True
    )

def header_white(text, fontsize = 40, bold = True):
    st.markdown(
        f"""
        <span style="display:inline-block; background:rgba(255, 255, 255, 1); outline: 1px solid; border-radius: 50px; padding: 10px; font-size:{fontsize}px; font-weight:{"bold" if bold else "normal"}; line-height: 2em">{text}</span>
        """,
        unsafe_allow_html=True
    )

def header_red(text, fontsize = 40, bold = True):
    st.markdown(
        f"""
        <span style="display:inline-block; background:rgba(255, 255, 255, 1); color:red; outline: 1px solid; border-radius: 30px; padding: 10px;  font-size:{fontsize}px; font-weight:{"bold" if bold else "normal"}">{text}</span>
        """,
        unsafe_allow_html=True
    )

def header_green(text, fontsize = 40, bold = True):
    st.markdown(
        f"""
        <span style="display:inline-block; background:rgba(255, 255, 255, 1); color:green; outline: 1px solid; border-radius: 30px; padding: 10px; font-size:{fontsize}px; font-weight:{"bold" if bold else "normal"}; line-height: 1.25">{text}</span>
        """,
        unsafe_allow_html=True
    )

def plant_treatment_message(predicted_string):
    if predicted_string == "Apple___Apple_scab":
        return "Remove the infected leaves and fruit and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Apple___Black_rot":
        return "Remove the infected branches and fruit and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Apple___Cedar_apple_rust":
        return "Remove the infected branches and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Cherry_(including_sour)___Powdery_mildew":
        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot":
        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Corn_(maize)___Common_rust_":
        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Corn_(maize)___Northern_Leaf_Blight":
        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Grape___Black_rot":
        return "Remove the infected branches and fruit and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Grape___Esca_(Black_Measles)":
        return "Remove the infected branches and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)":
        return "Remove the infected leaves and apply a fungicide to prevent it from spreading."
    elif predicted_string == "Orange___Haunglongbing_(Citrus_greening)":
        return "Remove the infected branches and apply a pesticide to prevent it from spreading."
    elif predicted_string == "Peach___Bacterial_spot":
        return "Remove the infected leaves and apply a copper fungicide to prevent it from spreading."
    elif predicted_string == "Squash___Powdery_mildew":
        return "This is a fungal disease that can cause white powdery spots on leaves and fruit. Consider removing infected plant parts and treating with a fungicide."
    elif predicted_string == "Strawberry___Leaf_scorch":
        return "This can be caused by drought, sunburn, or fungal diseases. Make sure your plant is getting enough water and sunlight. Consider removing infected plant parts and treating with a fungicide."
    elif predicted_string == "Tomato___Bacterial_spot":
        return "This is a bacterial disease that can cause spots on leaves and fruit. Consider removing infected plant parts and treating with a copper-based fungicide."
    elif predicted_string == "Tomato___Early_blight":
        return "This is a fungal disease that can cause dark spots on leaves and stems. Consider removing infected plant parts and treating with a fungicide."
    elif predicted_string == "Tomato___Late_blight":
        return "This is a fungal disease that can cause rapid decay of foliage and fruit. Consider removing infected plant parts and treating with a fungicide."
    elif predicted_string == "Tomato___Leaf_Mold":
        return "This is a fungal disease that can cause brown spots on leaves. Consider removing infected plant parts and treating with a fungicide."
    elif predicted_string == "Tomato___Septoria_leaf_spot":
        return "This is a fungal disease that can cause brown spots with a yellow halo on leaves. Consider removing infected plant parts and treating with a fungicide."
    elif predicted_string == "Tomato___Spider_mites Two-spotted_spider_mite":
        return "These are tiny pests that can cause yellow spots on leaves and webbing. Consider removing infected plant parts and treating with an insecticide."
    elif predicted_string == "Tomato___Target_Spot":
        return "This is a fungal disease that can cause circular spots with a bullseye pattern on leaves. Consider removing infected plant parts and treating with a fungicide."
    elif predicted_string == "Tomato___Tomato_Yellow_Leaf_Curl_Virus":
        return "This is a viral disease that can cause yellowing and curling of leaves. Consider treating with a fungicide."

def clean_prediction(prediction):
    pattern = re.compile('(.*)___(.*)')
    clean_predictions = []
    for p in prediction:
        r = pattern.search(p['predicted'])
        plant = r.groups()[0].replace('_', ' ').lower()
        diagnosis = r.groups()[1].replace('_', ' ').lower()
        treatment = plant_treatment_message(p['predicted']) if diagnosis is not 'healthy' else None
        clean_predictions.append([plant, diagnosis, "{0:.1f}%".format(float(p['probability']) * 100), treatment])

    return clean_predictions

def diagnose_health(file):
    prediction = predict(file)
    clean_predictions = clean_prediction(prediction)
    return clean_predictions

def app():

    add_bg_from_local('assets/background.png')
    header_white(f'<span style="color:green">Plant</span><span style="color:orange">Dx</span><span style="color:green">: Diagnosis in a Snap!</span> ')

    # Upload image of plant
    header_white("Upload an image of your plant below:", fontsize=32)
    header_white("For best results, remove the leaf from the plant and take the image against a dark background.", fontsize=16, bold=False)
    uploaded_file = st.file_uploader("", type=["jpg", "jpeg", "png"])

    if uploaded_file:
        header_white("Preview of the selected image:", fontsize=28, bold=False)
        st.image(uploaded_file)

    # Get diagnosis button
    if st.button("Get Diagnosis"):
        if uploaded_file is not None:
            # Diagnose plant health and display results
            results = diagnose_health(uploaded_file)

            if results[0][1] == 'healthy':
                header_green(f"We believe this is a healthy {results[0][0]} plant with {results[0][2]} confidence. Keep up the good work with proper watering, sunlight, and nutrients.", fontsize=32, bold=False)
            else:
                header_red(f"We believe this is an unhealthy {results[0][0]} plant with {results[0][1]}, with {results[0][2]} confidence. {results[0][3] if results [0][3] else ''}", fontsize=32, bold=False)

            if len(results) > 1:
                header_white("Other potential diagnoses: ", fontsize=24)

                for p in range(1, len(results)):
                    if results[p][1] == 'healthy':
                        header_white(
                            f"A healthy {results[p][0]} plant, {results[p][2]} confidence.",
                            fontsize=20, bold=False)
                    else:
                        header_white(
                            f"An unhealthy {results[p][0]} plant with {results[p][1]}, {results[p][2]} confidence. {results[p][3] if results [p][3] else ''}",
                            fontsize=20, bold=False)

        else:
            st.warning("Please upload an image of your plant first")

    # # Create user profile button
    # if st.button("Create User Profile"):
    #     st.subheader("User Profile")
    #     # Prompt user to add their name and the plant they own
    #     user_name = st.text_input("Enter your name:")
    #     plant_name = st.text_input("Enter the plant you own:")
    #     if user_name and plant_name:
    #         st.success(f"User profile created for {user_name} with plant {plant_name}")


# Run Streamlit app
if __name__ == "__main__":
    app()