"""
DimensioDepth - Add Dimension to Everything
Advanced AI Depth Estimation with 3D Visualization

Powered by Depth-Anything V2 | Runs on Hugging Face Spaces
"""

import streamlit as st
import numpy as np
import cv2
from PIL import Image
from pathlib import Path
import sys

# Page config
st.set_page_config(
    page_title="DimensioDepth - AI Depth Estimation",
    page_icon="🎨",
    layout="wide"
)

# Add backend to path
sys.path.append(str(Path(__file__).parent / "backend"))

# Import backend utilities
from backend.utils.image_processing import (
    depth_to_colormap,
    create_side_by_side
)

# Try to import REAL AI model
@st.cache_resource
def load_model():
    try:
        print("[*] Attempting to import TransformersDepthEstimator...")
        from backend.utils.transformers_depth import TransformersDepthEstimator
        print("[*] Import successful! Loading REAL AI Depth-Anything V2 BASE model...")
        print("[*] This will download ~372MB on first run (one-time download)")
        depth_estimator = TransformersDepthEstimator(model_size="base")
        print("[+] REAL AI MODE ACTIVE - BASE MODEL!")
        print("[+] Quality: SUPERB (best available)")
        return depth_estimator, True, "BASE (372MB)"
    except Exception as e:
        print(f"[!] FULL ERROR TRACEBACK:")
        import traceback
        traceback.print_exc()
        print(f"[!] Error type: {type(e).__name__}")
        print(f"[!] Error message: {str(e)}")
        print("[*] Falling back to DEMO MODE")
        return None, False, "Demo Mode"

depth_estimator, USE_REAL_AI, MODEL_SIZE = load_model()


def estimate_depth(image, colormap_style):
    """Estimate depth from an input image using REAL AI or DEMO MODE"""
    try:
        # Convert PIL to numpy if needed
        if isinstance(image, Image.Image):
            image = np.array(image)

        # Generate depth map
        if USE_REAL_AI:
            depth = depth_estimator.predict(image)
            mode_text = "REAL AI (Depth-Anything V2)"
        else:
            from backend.utils.demo_depth import generate_smart_depth
            depth = generate_smart_depth(image)
            mode_text = "DEMO MODE (Synthetic)"

        # Convert colormap style to cv2 constant
        colormap_dict = {
            "Inferno": cv2.COLORMAP_INFERNO,
            "Viridis": cv2.COLORMAP_VIRIDIS,
            "Plasma": cv2.COLORMAP_PLASMA,
            "Turbo": cv2.COLORMAP_TURBO,
            "Magma": cv2.COLORMAP_MAGMA,
            "Hot": cv2.COLORMAP_HOT,
            "Ocean": cv2.COLORMAP_OCEAN,
            "Rainbow": cv2.COLORMAP_RAINBOW
        }

        # Create colored depth map
        depth_colored = depth_to_colormap(depth, colormap_dict[colormap_style])

        # Create grayscale depth map
        depth_gray = (depth * 255).astype(np.uint8)
        depth_gray = cv2.cvtColor(depth_gray, cv2.COLOR_GRAY2RGB)

        return depth_colored, depth_gray, mode_text, image.shape, depth.shape

    except Exception as e:
        st.error(f"Error during depth estimation: {str(e)}")
        import traceback
        traceback.print_exc()
        return None, None, None, None, None


# Header
st.title("🎨 DimensioDepth - Add Dimension to Everything")
st.markdown("### Transform 2D images into stunning 3D depth visualizations")

# Status banner
if USE_REAL_AI:
    st.success(f"🚀 REAL AI MODE ACTIVE! - Powered by Depth-Anything V2 {MODEL_SIZE} - SUPERB Quality!")
else:
    st.info("Running in DEMO MODE - Ultra-fast synthetic depth estimation")

st.markdown("---")

# Main interface
col1, col2 = st.columns(2)

with col1:
    st.subheader("Input")
    uploaded_file = st.file_uploader("Upload Your Image", type=['png', 'jpg', 'jpeg'])

    colormap_style = st.selectbox(
        "Colormap Style",
        ["Inferno", "Viridis", "Plasma", "Turbo", "Magma", "Hot", "Ocean", "Rainbow"]
    )

    process_btn = st.button("🚀 Generate Depth Map", type="primary")

with col2:
    st.subheader("Output")
    depth_placeholder = st.empty()

# Processing
if uploaded_file is not None and process_btn:
    # Load image
    image = Image.open(uploaded_file)

    with col1:
        st.image(image, caption="Original Image", use_column_width=True)

    with st.spinner("Generating depth map..."):
        depth_colored, depth_gray, mode_text, input_shape, output_shape = estimate_depth(image, colormap_style)

    if depth_colored is not None:
        with col2:
            tab1, tab2 = st.tabs(["Colored", "Grayscale"])

            with tab1:
                st.image(depth_colored, caption="Depth Map (Colored)", use_column_width=True)

            with tab2:
                st.image(depth_gray, caption="Depth Map (Grayscale)", use_column_width=True)

        # Info
        st.success(f"✅ Depth Estimation Complete!")
        st.info(f"""
**Mode**: {mode_text}
**Input Size**: {input_shape[1]}x{input_shape[0]}
**Output Size**: {output_shape[1]}x{output_shape[0]}
**Colormap**: {colormap_style}
{f'**Powered by**: Depth-Anything V2 {MODEL_SIZE}' if USE_REAL_AI else '**Processing**: Ultra-fast (<50ms) synthetic depth'}
        """)

# Info section
st.markdown("---")
st.markdown("""
## 💡 About DimensioDepth

### Features:
- ✅ Real AI depth estimation with Depth-Anything V2
- ✅ Multiple colormap styles for visualization
- ✅ Fast processing (~800ms on CPU, ~200ms on GPU)
- ✅ SUPERB quality depth maps

### Use Cases:
- 🎨 **Creative & Artistic**: Depth-enhanced photos, 3D effects
- 🎬 **VFX & Film**: Depth map generation for compositing
- 🔬 **Research**: Computer vision, depth perception studies
- 📱 **Content Creation**: Engaging 3D effects for social media

Made with ❤️ for the AI community
""")