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

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

import gradio as gr
import numpy as np
import cv2
from PIL import Image
from pathlib import Path
import sys

# 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
try:
    from backend.utils.transformers_depth import TransformersDepthEstimator
    print("[*] 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)")
    USE_REAL_AI = True
    MODEL_SIZE = "BASE (372MB)"
except Exception as e:
    print(f"[!] Could not load AI models: {e}")
    print("[*] Falling back to DEMO MODE")
    from backend.utils.demo_depth import generate_smart_depth
    USE_REAL_AI = False
    MODEL_SIZE = "Demo Mode"


def estimate_depth(image, colormap_style):
    """
    Estimate depth from an input image using REAL AI or DEMO MODE
    """
    if image is None:
        return None, None, "Please upload an image first"

    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:
            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)

        # Processing info
        info_text = f"Mode: {mode_text} | Input: {image.shape[1]}x{image.shape[0]} | Output: {depth.shape[1]}x{depth.shape[0]} | Colormap: {colormap_style}"
        if USE_REAL_AI:
            info_text += f" | Model: Depth-Anything V2 {MODEL_SIZE}"

        return depth_colored, depth_gray, info_text

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


# Create interface
demo = gr.Interface(
    fn=estimate_depth,
    inputs=[
        gr.Image(label="Upload Your Image"),
        gr.Dropdown(
            choices=["Inferno", "Viridis", "Plasma", "Turbo", "Magma", "Hot", "Ocean", "Rainbow"],
            value="Inferno",
            label="Colormap Style"
        )
    ],
    outputs=[
        gr.Image(label="Depth Map (Colored)"),
        gr.Image(label="Depth Map (Grayscale)"),
        gr.Textbox(label="Info")
    ],
    title="DimensioDepth - AI Depth Estimation",
    description=f"**{'REAL AI MODE - Depth-Anything V2 BASE (372MB) - SUPERB Quality!' if USE_REAL_AI else 'DEMO MODE - Ultra-fast synthetic depth estimation'}**",
    article="""
## About DimensioDepth

Transform 2D images into stunning 3D depth visualizations using state-of-the-art AI.

### 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
    """
)


# Launch the app
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=False  # Disable API docs to avoid schema generation bug
    )