import streamlit as st
import os
import cv2
import numpy as np
import pandas as pd
import random
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier

# ------------------------ PAGE CONFIG & STYLING ------------------------
st.set_page_config(page_title="Potato Leaf Disease Classifier", layout="centered")

# Custom background styling (you must place 'Plant.jpg' in the root folder)
st.markdown("""
    <style>
    .stApp {
        background-image: url('Plant.jpg');
        background-size: cover;
        background-attachment: fixed;
        background-position: center;
        backdrop-filter: blur(2px);
    }
    </style>
""", unsafe_allow_html=True)

# ------------------------ TITLE & DESCRIPTION ------------------------
st.title("🥔 Potato Leaf Disease Classifier")
st.markdown("""
<div style='font-size: 17px; line-height: 1.6;'>
This project aims to automatically classify the health condition of potato leaves using grayscale images from different disease categories.
The model helps farmers and agricultural experts identify potential diseases early and take appropriate actions to ensure healthy crop production.
</div>
""", unsafe_allow_html=True)

# ------------------------ CONFIG ------------------------
BASE_PATH = r"PotatoPlants"  # Ensure this folder exists in your app's directory
folders = [f for f in os.listdir(BASE_PATH) if os.path.isdir(os.path.join(BASE_PATH, f))]

# ------------------------ MODEL SELECTION ------------------------
st.markdown("<h4 style='margin-top: 30px;'>🧠 <b>Select Classifier</b></h4>", unsafe_allow_html=True)

model_choice = st.selectbox(
    "Choose a machine learning model to classify potato diseases",
    ["Logistic Regression", "KNN", "Decision Tree"],
    index=0
)

# ------------------------ LOAD DATA ------------------------
@st.cache_data
def load_data():
    images, labels = [], []
    for folder in folders:
        full_path = os.path.join(BASE_PATH, folder)
        for img_name in os.listdir(full_path):
            img_path = os.path.join(full_path, img_name)
            img = cv2.imread(img_path, 0)  # Read in grayscale
            if img is not None:
                resized = cv2.resize(img, (100, 100))
                images.append(resized.flatten())
                labels.append(folder.split("___")[-1])
    return pd.DataFrame(images), pd.Series(labels)

X, y = load_data()
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=27)

# ------------------------ MODEL TRAINING ------------------------
st.markdown("### 🔧 Training the model...")

with st.spinner("Training in progress..."):
    if model_choice == "Logistic Regression":
        model = LogisticRegression(max_iter=200)
    elif model_choice == "KNN":
        model = KNeighborsClassifier()
    else:
        model = DecisionTreeClassifier()

    model.fit(X_train, y_train)
    accuracy = model.score(X_test, y_test)

st.success(f"✅ {model_choice} trained with accuracy: **{accuracy:.2f}**")

# ------------------------ RANDOM IMAGE TESTING ------------------------
st.markdown("### 🔍 Test on a Random Image")

if st.button("Test Random Image"):
    random_folder = random.choice(folders)
    folder_path = os.path.join(BASE_PATH, random_folder)
    random_image_file = random.choice(os.listdir(folder_path))
    image_path = os.path.join(folder_path, random_image_file)

    img = cv2.imread(image_path, 0)
    resized = cv2.resize(img, (100, 100))
    flat = resized.flatten().reshape(1, -1)
    prediction = model.predict(flat)[0]
    actual = random_folder.split("___")[-1]

    # Convert grayscale to RGB for better browser compatibility
    color_img = cv2.cvtColor(resized, cv2.COLOR_GRAY2RGB)

    st.image(color_img, caption="📷 Randomly selected potato leaf", use_container_width=True)
    st.markdown(f"<h3 style='color:green;'>✅ Actual: <b>{actual}</b></h3>", unsafe_allow_html=True)
    st.markdown(f"<h3 style='color:blue;'>🤖 Predicted: <b>{prediction}</b></h3>", unsafe_allow_html=True)

# ------------------------ FOOTER ------------------------
st.markdown("---")
st.markdown("<center>Made with ❤️ for smart agriculture 🌾</center>", unsafe_allow_html=True)