import streamlit as st 
import tensorflow as tf
import numpy as np
import os
import pandas as pd
import altair as alt
from PIL import Image
from tensorflow.keras.preprocessing.image import ImageDataGenerator
import zipfile
import io
import shutil

# 🎨 App Title
st.title("🐶🐱 Cat vs Dog Classifier")

# 📖 About the App
st.write(
    """
    ## About This App
    This is a machine learning application that classifies images into two categories:
    **Cats 🐱** and **Dogs 🐶**. The model is trained using a deep learning architecture
    called Convolutional Neural Networks (CNNs) and is able to distinguish between images 
    of cats and dogs with high accuracy.
    ### Features:
    - **Dataset Overview**: View the number of images in the dataset, categorized by "Cats" and "Dogs".
    - **Model Evaluation**: Check the model's performance on the validation set, including accuracy and loss.
    - **Image Classification**: Upload an image, and the model will predict whether it's a cat or a dog, along with the confidence level.
    - **Download Test Folder**: Download a ZIP file containing the test images.
    The app is powered by **Streamlit** for an interactive user interface and **TensorFlow** for image classification.
    """
)

# ✅ Detect Environment & Set Dataset Path
BASE_DIR = "dataset"  # In Hugging Face Spaces, the dataset folder should be at the root of the Space
ZIP_PATH = "dataset.zip"  # If dataset is uploaded as a ZIP (make sure it's in the same directory as app.py)

TRAIN_DIR = os.path.join(BASE_DIR, "train")
TEST_DIR = os.path.join(BASE_DIR, "test")

# ✅ Extract Dataset if Needed (Hugging Face)
if ZIP_PATH and os.path.exists(ZIP_PATH):
    if not os.path.exists(BASE_DIR):  # Avoid re-extracting
        with zipfile.ZipFile(ZIP_PATH, "r") as zip_ref:
            zip_ref.extractall(BASE_DIR)  # Extract into dataset folder
    st.success("✅ Dataset extracted!")

# 📌 Check if dataset exists
if not os.path.exists(TRAIN_DIR):
    st.error(f"❌ Dataset folder 'train' not found at {TRAIN_DIR}. Please upload the dataset.")
    st.stop()

# 📌 Verify Cats & Dogs Folders
cat_dir = os.path.join(TRAIN_DIR, "cats")
dog_dir = os.path.join(TRAIN_DIR, "dogs")

if not os.path.exists(cat_dir) or not os.path.exists(dog_dir):
    st.error("❌ Missing 'cats' or 'dogs' folders inside 'train'. Please check your dataset.")
    st.stop()

# 📌 Constants
IMG_SIZE = (150, 150)
BATCH_SIZE = 32
MODEL_PATH = "cats_dogs_model.h5"  # Ensure the model is uploaded to Hugging Face Space

# 🎯 Load Model
if os.path.exists(MODEL_PATH):
    model = tf.keras.models.load_model(MODEL_PATH)
else:
    st.error("⚠ No trained model found. Please upload 'cats_dogs_model.h5' to your Hugging Face repository.")
    st.stop()

# 📷 Image Preprocessing
def preprocess_image(image):
    image = image.convert('RGB').resize(IMG_SIZE)
    img_array = np.array(image, dtype=np.float32) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# 🐶🐱 Classify Image
def classify_image(image):
    processed_img = preprocess_image(image)
    prediction = model.predict(processed_img)[0][0]
    label = "Dog 🐶" if prediction > 0.5 else "Cat 🐱"
    confidence = prediction if label == "Dog 🐶" else 1 - prediction
    return label, confidence

# 📊 Model Evaluation
def evaluate_model():
    datagen = ImageDataGenerator(rescale=1.0 / 255, validation_split=0.2)
    
    test_data = datagen.flow_from_directory(
        TRAIN_DIR,
        target_size=IMG_SIZE,
        batch_size=BATCH_SIZE,
        class_mode='binary',
        subset='validation'
    )

    loss, accuracy = model.evaluate(test_data, verbose=0)

    return accuracy, loss

# 🎛 Streamlit Tabs
tab1, tab2, tab3 = st.tabs(["📁 Dataset Preview", "📊 Model Performance", "🐶🐱 Image Classification"])

# 📁 Tab 1: Dataset Preview
with tab1:
    st.write("### Dataset Overview")

    dataset_info = {
        "Total Images": len(os.listdir(cat_dir)) + len(os.listdir(dog_dir)),
        "Cat Images": len(os.listdir(cat_dir)),
        "Dog Images": len(os.listdir(dog_dir))
    }
    
    df_info = pd.DataFrame(list(dataset_info.items()), columns=["Category", "Count"])
    st.dataframe(df_info)

    # Visualization
    st.write("### Image Distribution")
    chart = alt.Chart(df_info).mark_bar().encode(
        x="Category", y="Count", color="Category"
    )
    st.altair_chart(chart, use_container_width=True)

# 📊 Tab 2: Model Performance
with tab2:
    st.write("### Model Evaluation")
    
    accuracy, loss = evaluate_model()
    
    st.write(f"✅ **Validation Accuracy:** {accuracy*100:.2f}%")
    st.write(f"✅ **Validation Loss:** {loss:.4f}")

# 🐶🐱 Tab 3: Image Classification
with tab3:
    st.write("### Upload an Image for Classification")
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])

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

        with st.spinner("Classifying..."):
            label, confidence = classify_image(image)

        st.subheader("Prediction:")
        st.write(f"This is a **{label}** with **{confidence*100:.2f}%** confidence.")

    # **New Feature: Download the 'test' folder as a ZIP**
    def zip_folder(folder_path):
        # Create an in-memory zip file
        zip_buffer = io.BytesIO()
        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            for root, dirs, files in os.walk(folder_path):
                for file in files:
                    zip_file.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file), folder_path))
        zip_buffer.seek(0)  # Go to the beginning of the file
        return zip_buffer

    # Button to download 'test' folder
    if os.path.exists(TEST_DIR):
        st.write("### Download Test Folder")
        zip_buffer = zip_folder(TEST_DIR)
        st.download_button(
            label="Download Test Folder as ZIP",
            data=zip_buffer,
            file_name="test_folder.zip",
            mime="application/zip"
        )
    else:
        st.warning(f"❌ Test folder not found at {TEST_DIR}")