import streamlit as st
import os
import cv2
import yaml
import numpy as np
import torch
from PIL import Image

from src.pipeline import EcoPulsePipeline
from src.cnn_model import load_model
from src.transforms import EUROSAT_TRANSFORM
from src.visualization import apply_grad_cam, create_greenery_overlay

# --- Configuration & Caching ---
st.set_page_config(page_title="EcoPulse Dashboard", page_icon="🌿", layout="wide")

@st.cache_resource(show_spinner="Loading deep learning models (this may take a moment)...")
def load_pipeline(version=1):
    """Load the full pipeline once and cache it in GPU memory."""
    # Ensure config exists
    config_path = "config/config.yaml"
    return EcoPulsePipeline(config_path)

@st.cache_resource(show_spinner="Loading Grad-CAM resources...")
def load_grad_cam_resources():
    """Load the CNN specifically for Grad-CAM."""
    config_path = "config/config.yaml"
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_model(
        weights_path=os.path.join(config['paths']['output_models'], 'resnet50_eurosat.pth'),
        num_classes=config['model']['num_classes'],
        device=device
    )
    model.eval()

    transform = EUROSAT_TRANSFORM

    return model, transform, config['classes'], config['greenery_classes']

# --- Main App ---
def main():
    st.title("EcoPulse Satellite Analysis")

    # --- Sidebar Controls ---
    with st.sidebar:
        st.markdown("### System Controls")
        st.markdown("---")

        st.markdown("**Hardware Status**")
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            vram_used = torch.cuda.memory_allocated(0) / (1024**3)
            vram_total = torch.cuda.get_device_properties(0).total_memory / (1024**3)

            st.markdown(f"`GPU:` {gpu_name}")
            st.progress(vram_used / vram_total, text=f"VRAM Allocation: {vram_used:.1f}GB / {vram_total:.1f}GB")
        else:
            st.warning("Running on CPU (No CUDA detected)")

        st.markdown("\n**Active Pipeline**\n`SAM (ViT-B) + ResNet-50`")

        st.markdown("<br><br>", unsafe_allow_html=True)
        st.markdown("**Maintenance**")

        if st.button("Clear Model Cache", use_container_width=True):
            st.cache_resource.clear()
            st.toast("Cache Cleared. Models will reload on next execution.")

        if st.button("Terminate Session", help="Stop the Streamlit process securely", type="primary", use_container_width=True):
            st.warning("Terminating server process...")
            st.stop()

        st.markdown("---")
        st.caption("EcoPulse v1.0.0 | Environmental Auditing Platform")

    st.markdown("""
    Welcome to the **EcoPulse Dashboard**. Use the tabs below to either analyze a single image or compare two different regions.
    """)

    # Initialize models (cached — runs once, then reuses)
    with st.spinner("Loading deep learning models (this may take a moment)..."):
        pipeline = load_pipeline(version=1)
        cnn_model, transform, all_classes, greenery_classes = load_grad_cam_resources()

    tab1, tab2 = st.tabs(["Single Image Analysis", "Region Comparison"])

    with tab1:
        st.header("Single Image Analysis")
        uploaded_file = st.file_uploader("Upload a Satellite Image (.jpg, .png)", type=["jpg", "png", "jpeg"], key="single")

        if uploaded_file is not None:
            # Save temp file
            temp_dir = "data/temp"
            os.makedirs(temp_dir, exist_ok=True)
            temp_path = os.path.join(temp_dir, uploaded_file.name)
            with open(temp_path, "wb") as f:
                f.write(uploaded_file.getbuffer())

            st.info("Image uploaded successfully. Running EcoPulse pipeline...")

            # Process Image
            with st.spinner("Segmenting and Classifying..."):
                image_np, results = pipeline.process_image(temp_path)

            # --- Display Metrics ---
            st.header("Analysis Results")

            green_pct = results['greenery_percentage']
            total_px = results['total_pixels']
            green_px = results['green_pixels']

            col1, col2, col3 = st.columns(3)
            col1.metric("Greenery Coverage", f"{green_pct:.1f}%", delta=None)
            col2.metric("Green Pixels", f"{green_px:,}")
            col3.metric("Total Pixels", f"{total_px:,}")

            # --- Visualizations ---
            st.subheader("Visual Overlays")

            # Build composite greenery overlay
            composite, green_masks = create_greenery_overlay(image_np, results['mask_classifications'])

            v_col1, v_col2 = st.columns(2)
            with v_col1:
                st.image(image_np, caption="Original Satellite Image", width='stretch')
            with v_col2:
                st.image(composite, caption="Greenery Segmentation Overlay", width='stretch')


            # --- Grad-CAM Interpretability ---
            st.divider()
            st.header("Model Interpretability (Grad-CAM)")
            st.markdown("Select a detected greenery region below to see exactly which features the CNN focused on to make its classification.")

            if len(green_masks) > 0:
                # Sort masks by size (pixel count) descending
                green_masks = sorted(green_masks, key=lambda x: x['pixels'], reverse=True)

                # Create dropdown options
                options = {f"Region {i+1} (Class: {m['class']}, Size: {m['pixels']:,} px)": m for i, m in enumerate(green_masks)}

                selected_option = st.selectbox("Select a Greenery Region to Analyze:", list(options.keys()))
                selected_mask_data = options[selected_option]

                # Generate Grad-CAM for the selected mask
                bbox = selected_mask_data['bbox'] # [x, y, w, h]
                x, y, w_box, h_box = [int(v) for v in bbox]
                h_img, w_img = image_np.shape[:2]

                # Clamp bounding box coordinates to image boundaries
                x = max(0, min(x, w_img - 1))
                y = max(0, min(y, h_img - 1))
                w_box = min(w_box, w_img - x)
                h_box = min(h_box, h_img - y)

                if w_box > 0 and h_box > 0:
                    crop = image_np[y:y+h_box, x:x+w_box]
                    crop_pil = Image.fromarray(crop)
                    input_tensor = transform(crop_pil).unsqueeze(0)

                    with st.spinner("Generating Grad-CAM Heatmap..."):
                        heatmap, pred_idx = apply_grad_cam(cnn_model, input_tensor, target_class=None)

                        # Create overlay
                        heatmap_resized = cv2.resize(heatmap, (crop.shape[1], crop.shape[0]))
                        heatmap_colored = cv2.applyColorMap(np.uint8(255 * heatmap_resized), cv2.COLORMAP_JET)
                        heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)

                        alpha = 0.5
                        gradcam_overlay = np.uint8(crop * (1 - alpha) + heatmap_colored * alpha)

                        g_col1, g_col2 = st.columns(2)
                        with g_col1:
                            st.image(crop, caption=f"Cropped Region (Original)", width='stretch')
                        with g_col2:
                            st.image(gradcam_overlay, caption=f"Grad-CAM Heatmap (Class: {all_classes[pred_idx]})", width='stretch')
                else:
                    st.warning("Selected region is too small to analyze.")
            else:
                st.info("No greenery regions detected in this image.")

    with tab2:
        st.header("Region Comparison")
        st.markdown("Upload two satellite images to compare their greenery coverage side-by-side.")

        c_col1, c_col2 = st.columns(2)

        with c_col1:
            file_a = st.file_uploader("Upload Area A", type=["jpg", "png", "jpeg"], key="area_a")
        with c_col2:
            file_b = st.file_uploader("Upload Area B", type=["jpg", "png", "jpeg"], key="area_b")

        if file_a and file_b:
            if st.button("Run Comparison Analysis"):
                # Save temp files
                temp_dir = "data/temp"
                os.makedirs(temp_dir, exist_ok=True)
                path_a = os.path.join(temp_dir, "compare_a_" + file_a.name)
                path_b = os.path.join(temp_dir, "compare_b_" + file_b.name)

                with open(path_a, "wb") as f: f.write(file_a.getbuffer())
                with open(path_b, "wb") as f: f.write(file_b.getbuffer())

                with st.spinner("Analyzing both regions (this may take a minute)..."):
                    img_a_np, results_a = pipeline.process_image(path_a)
                    img_b_np, results_b = pipeline.process_image(path_b)

                pct_a = results_a['greenery_percentage']
                pct_b = results_b['greenery_percentage']

                st.divider()
                st.subheader("Comparison Result")

                diff = pct_a - pct_b
                if abs(diff) < 1:
                    st.success("Both regions have nearly identical greenery coverage.")
                else:
                    winner = "Area A" if diff > 0 else "Area B"
                    st.info(f"**{winner}** is more vegetated by **{abs(diff):.1f}%**.")

                st.markdown("### Detailed Metrics")
                res_col1, res_col2 = st.columns(2)

                with res_col1:
                    st.markdown("**Area A**")
                    st.metric("Greenery Coverage", f"{pct_a:.1f}%")
                    st.caption(f"Green Pixels: {results_a['green_pixels']:,} / Total: {results_a['total_pixels']:,}")

                with res_col2:
                    st.markdown("**Area B**")
                    st.metric("Greenery Coverage", f"{pct_b:.1f}%")
                    st.caption(f"Green Pixels: {results_b['green_pixels']:,} / Total: {results_b['total_pixels']:,}")

                # Visual comparison
                st.markdown("### Visual Side-by-Side Analysis")

                # Generate overlays
                comp_a, _ = create_greenery_overlay(img_a_np, results_a['mask_classifications'])
                comp_b, _ = create_greenery_overlay(img_b_np, results_b['mask_classifications'])

                v_res_col1, v_res_col2 = st.columns(2)
                v_res_col1.image(comp_a, caption=f"Area A Overlay ({pct_a:.1f}% Green)", width='stretch')
                v_res_col2.image(comp_b, caption=f"Area B Overlay ({pct_b:.1f}% Green)", width='stretch')

if __name__ == "__main__":
    main()