# app.py
from PIL import Image
import streamlit as st
from llama_cpp import Llama
from image_generation import generate_image

st.set_page_config(page_title="GreenThumber", layout="centered")
st.title("🌱 GreenThumber")

@st.cache_resource(hash_funcs={Llama: lambda _: None})
def load_mistral_model():
    llm = Llama(
        model_path="./models/mistral-7b-instruct-v0.1.Q4_K_M.gguf",
        n_ctx=2048,
        n_threads=4,
        n_batch=512,
        verbose=False
    )
    return llm

llm = load_mistral_model()
           

def generate_growth_description(plant_type, plant_age, soil_type, sunlight_hours, water_frequency,
                                fertilizer_type, temperature, humidity, days, additional_info):
    prompt = (
        f" Instruction:\n"
        f"You are a botanical expert. Describe how a {plant_type} plant will likely look in {days} days "
        f"based on these conditions:\n"        
        f"Important environment conditions: {additional_info}\n"
        f"- Plant age: {plant_age} days\n"
        f"- Soil Type: {soil_type}\n"
        f"- Sunlight: {sunlight_hours} hours per day\n"
        f"- Water Frequency: {water_frequency} times per week\n"
        f"- Fertilizer Type: {fertilizer_type}\n"
        f"- Temperature: {temperature}°C\n"
        f"- Humidity: {humidity}%\n"
        f"### Response:\n"
    )
    output = llm(prompt, max_tokens=100, stop=["###"])
    return output["choices"][0]["text"].strip()


st.header("Plant Info")
plant_input_mode = st.radio("How would you like to provide plant info?", ("Type name", "Upload image"))
plant_type = None
uploaded_image = None

if plant_input_mode == "Type name":
    plant_type = st.selectbox("Select Plant Type", ["Basil", "Tomato", "Lettuce"])
    plant_age = st.number_input("Enter Plant Age (in days)", min_value=1, max_value=365, value=25)
    
    

elif plant_input_mode == "Upload image":
    plant_type = st.selectbox("Select Plant Type", ["Basil", "Tomato", "Lettuce"])
    plant_age = st.number_input("Enter Plant Age (in days)", min_value=1, max_value=365, value=25)
    image_file = st.file_uploader("Upload an image of your plant", type=["jpg", "jpeg", "png"])
    if image_file:
        uploaded_image = Image.open(image_file)
        st.image(uploaded_image, caption="Uploaded Plant Image", use_container_width =True)


col1, col2 = st.columns(2)

with col1:
    st.markdown("<h3 style='text-align: center;'>Environmental Parameters</h3>", unsafe_allow_html=True)
    soil_options = ["Sandy", "Clay", "Loamy", "Peaty", "Chalky", "Silty"]

    soil_type = st.selectbox("Soil Type", soil_options)
    sunlight_hours = st.slider("Sunlight Hours per day", 0, 24, 6)
    water_frequency = st.slider("Water Frequency (times per week)", 0, 14, 3)

with col2:
    st.markdown("---")
    fertilizer_options = ["Organic", "Chemical", "None"]

    fertilizer_type = st.selectbox("Fertilizer Type", fertilizer_options)
    temperature = st.slider("Temperature (°C)", -10, 50, 22)
    humidity = st.slider("Humidity (%)", 0, 100, 60)
    
days = st.slider("Prediction Interval (in days)", min_value=1, max_value=30, value=7)

additional_info = st.text_area("Feel free to include any additional detail", height=100)


if st.button("Start Prediction"): # Comment out this line to enable automatic refresh on input change
    if plant_type and plant_type.strip() != "":
        with st.spinner("Analyzing data and generating description..."):
            description = generate_growth_description(
                plant_type, plant_age, soil_type, sunlight_hours, water_frequency,
                fertilizer_type, temperature, humidity, days, additional_info
            )
        st.subheader(f"📝 Predicted Plant Condition in {days} Days:")
        st.write(description)

        manipulated_img = generate_image(plant_type, description, plant_age)
        st.image(manipulated_img, caption="Predicted Plant Condition Image")


st.markdown("---")
st.caption("Made with ❤️ by Sandwich Craftz.")