# IMPORTANT: Import spaces FIRST before any CUDA-related packages
import spaces
import json
import torch
import torch.nn as nn
import gradio as gr
from PIL import Image
from torchvision import models, transforms
from huggingface_hub import hf_hub_download
import warnings

# Suppress NVML warning
warnings.filterwarnings("ignore", category=UserWarning, module="torch.cuda")

MODEL_REPO = "nailarais1/image-classifier-efficientnet"


# Better device detection with error handling
def get_device():
    try:
        if torch.cuda.is_available():
            return torch.device("cuda")
        else:
            return torch.device("cpu")
    except Exception:
        return torch.device("cpu")


DEVICE = get_device()
print(f"Using device: {DEVICE}")

print("Downloading checkpoint...")
checkpoint_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename="checkpoint.pth"
)

print("Downloading config...")
config_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename="config.json"
)

print("Loading config...")
with open(config_path, "r", encoding="utf-8") as f:
    config = json.load(f)

num_classes = config["num_labels"]

id2label = {}
for key, value in config["id2label"].items():
    id2label[int(key)] = value

print("Number of classes:", num_classes)

print("Building model...")
model = models.efficientnet_b0(weights=None)

model.classifier[1] = nn.Linear(
    model.classifier[1].in_features,
    num_classes
)

print("Loading checkpoint...")
checkpoint = torch.load(
    checkpoint_path,
    map_location=DEVICE,
    weights_only=False
)

model.load_state_dict(
    checkpoint["model_state_dict"]
)

model = model.to(DEVICE)
model.eval()

print("Model loaded successfully.")

transform = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ]
)


# Add the @spaces.GPU decorator to enable GPU if available
@spaces.GPU
def predict(image):
    if image is None:
        return {"No image uploaded": 1.0}

    try:
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)

        image = image.convert("RGB")

        input_tensor = transform(image)
        input_tensor = input_tensor.unsqueeze(0)
        input_tensor = input_tensor.to(DEVICE)

        with torch.inference_mode():
            outputs = model(input_tensor)

            probabilities = torch.softmax(
                outputs,
                dim=1
            )

            top_probabilities, top_indices = torch.topk(
                probabilities,
                k=5,
                dim=1
            )

        results = {}

        # Check confidence threshold to detect non-flower images
        top_confidence = float(top_probabilities[0][0].item())
        confidence_threshold = 0.65  # 65% threshold

        if top_confidence < confidence_threshold:
            # Low confidence - likely not a flower
            results = {
                "⚠️ Not a Flower (Low Confidence)": 1.0,
                "This may not be a flower": 0.0
            }

            # Still show top predictions for reference
            for probability, index in zip(
                top_probabilities[0],
                top_indices[0]
            ):
                class_index = int(index.item())
                flower_name = id2label[class_index]
                confidence = float(probability.item())
                results[f"🤔 {flower_name}"] = confidence
        else:
            # High confidence - likely a flower
            for probability, index in zip(
                top_probabilities[0],
                top_indices[0]
            ):
                class_index = int(index.item())
                flower_name = id2label[class_index]
                confidence = float(probability.item())
                results[flower_name] = confidence

        return results

    except Exception as error:
        print("Prediction error:", error)
        return {
            "❌ Prediction error": 1.0
        }


print("Creating Gradio interface...")

# HF-Safe CSS: uses theme variables, no transparent body hacks
custom_css = """
.gradio-container {
   max-width: 1180px !important;
   margin: 0 auto !important;
   padding: 28px 22px 24px !important;
}

/* Header */
.main-header {
   background: var(--block-background-fill) !important;
   border: 1px solid var(--block-border-color) !important;
   border-radius: 18px !important;
   padding: 30px 28px !important;
   text-align: center;
   margin-bottom: 18px !important;
   box-shadow: 0 14px 40px rgba(0,0,0,.18) !important;
}

.main-header h1 {
   margin: 0 !important;
   font-size: clamp(1.9rem, 4vw, 2.65rem) !important;
   line-height: 1.15 !important;
   font-weight: 750 !important;
   color: var(--body-text-color) !important;
   letter-spacing: -0.02em;
}

.main-header p {
   margin: 11px 0 0 !important;
   font-size: 1.05rem !important;
   color: var(--body-text-color-subdued) !important;
}

/* Information card */
.info-card {
   background: var(--block-background-fill) !important;
   border: 1px solid var(--block-border-color) !important;
   border-radius: 14px !important;
   padding: 18px 22px !important;
   margin-bottom: 18px !important;
   box-shadow: 0 8px 25px rgba(0,0,0,.10) !important;
}

.info-card h3 {
   margin: 0 0 10px !important;
   color: var(--body-text-color) !important;
   font-size: 1.05rem !important;
}

.info-card ul {
   margin: 0 !important;
   padding-left: 20px !important;
}

.info-card li {
   margin: 5px 0 !important;
   color: var(--body-text-color) !important;
}

.info-card strong {
   color: var(--body-text-color) !important;
}

/* Upload component */
.upload-area {
   min-height: 350px !important;
   border: 1px dashed var(--border-color-primary) !important;
   border-radius: 14px !important;
   background: var(--block-background-fill) !important;
   overflow: hidden !important;
   transition: border-color .2s ease, box-shadow .2s ease, transform .2s ease;
}

.upload-area:hover {
   border-color: var(--primary-500) !important;
   box-shadow: 0 0 0 3px rgba(255,255,255,.04) !important;
}

/* Main classify button */
.classify-btn {
   width: 100% !important;
   margin-top: 10px !important;
   min-height: 52px !important;
   border: 0 !important;
   border-radius: 12px !important;
   font-size: 1rem !important;
   font-weight: 700 !important;
   transition: transform .18s ease, box-shadow .18s ease !important;
   box-shadow: 0 8px 22px rgba(0,0,0,.18) !important;
}

.classify-btn:hover {
   transform: translateY(-1px);
   box-shadow: 0 12px 28px rgba(0,0,0,.25) !important;
}

.classify-btn:active {
   transform: translateY(0);
}

/* Prediction panel */
.prediction-box {
   min-height: 350px !important;
   background: var(--block-background-fill) !important;
   border: 1px solid var(--block-border-color) !important;
   border-radius: 14px !important;
   padding: 18px !important;
   box-shadow: 0 8px 25px rgba(0,0,0,.10) !important;
}

.prediction-box label {
   color: var(--body-text-color) !important;
   font-weight: 650 !important;
}

/* Gradio label/progress text */
.prediction-box .label-wrap,
.prediction-box .wrap {
   color: var(--body-text-color) !important;
}

/* Tip */
.tip-box {
   margin: 16px 0 !important;
   padding: 14px 18px !important;
   border-radius: 12px !important;
   background: rgba(59,130,246,.08) !important;
   border: 1px solid rgba(59,130,246,.20) !important;
}

.tip-box p {
   margin: 0 !important;
   color: var(--body-text-color) !important;
}

.tip-box strong {
   color: var(--body-text-color) !important;
}

/* Warning */
.warning-box {
   margin-top: 16px !important;
   padding: 16px 18px !important;
   border-radius: 12px !important;
   background: rgba(245,158,11,.08) !important;
   border: 1px solid rgba(245,158,11,.22) !important;
   border-left: 4px solid #f59e0b !important;
}

.warning-box h4 {
   margin: 0 0 7px !important;
   color: var(--body-text-color) !important;
}

.warning-box p {
   margin: 0 !important;
   color: var(--body-text-color-subdued) !important;
   line-height: 1.55 !important;
}

.warning-box strong {
   color: var(--body-text-color) !important;
}

/* Footer */
.footer {
   margin-top: 22px !important;
   padding: 18px 0 4px !important;
   text-align: center !important;
   color: var(--body-text-color-subdued) !important;
   font-size: .88rem !important;
   border-top: 1px solid var(--block-border-color) !important;
}

.footer * {
   color: var(--body-text-color-subdued) !important;
}

/* Mobile */
@media (max-width: 768px) {
   .gradio-container {
       padding: 16px 12px 20px !important;
   }

   .main-header {
       padding: 24px 18px !important;
   }

   .main-header h1 {
       font-size: 1.9rem !important;
   }

   .main-header p {
       font-size: .95rem !important;
   }

   .upload-area,
   .prediction-box {
       min-height: 300px !important;
   }
}
"""

# Create the interface
demo = gr.Blocks(
    title="🌸 Flower Classifier",
    theme=gr.themes.Base(
        primary_hue="indigo",
        secondary_hue="purple",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter")
    )
)

with demo:
    # Simple header
    gr.HTML("""
    <div class="main-header">
        <h1>🌸 102-Flower Image Classifier</h1>
        <p>Upload a photo and let AI identify the flower species</p>
    </div>
    """)

    # Model info
    gr.HTML("""
    <div class="info-card">
        <h3>📊 Model Information</h3>
        <ul>
            <li><strong>Model:</strong> EfficientNet-B0</li>
            <li><strong>Classes:</strong> 102 flower species</li>
            <li><strong>Accuracy:</strong> 94.38%</li>
            <li><strong>Input:</strong> 224 × 224 pixels</li>
        </ul>
    </div>
    """)

    # Main layout
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil",
                label="📷 Upload Image",
                height=350,
                elem_classes="upload-area"
            )

            classify_btn = gr.Button(
                "🔍 Classify Flower",
                variant="primary",
                size="lg",
                elem_classes="classify-btn"
            )

        with gr.Column(scale=1):
            with gr.Group(elem_classes="prediction-box"):
                output = gr.Label(
                    num_top_classes=5,
                    label="🎯 Top Predictions",
                    show_label=True
                )

    # Tip
    gr.HTML("""
    <div class="tip-box">
        <p>💡 <strong>Tip:</strong> The model shows the top 5 predictions.
        Confidence above <strong>65%</strong> = reliable prediction.</p>
    </div>
    """)

    # Warning
    gr.HTML("""
    <div class="warning-box">
        <h4>⚠️ Important Note</h4>
        <p>
            This model was trained <strong>only on 102 flower species</strong>.
            Images of <strong>cats, dogs, people, or other objects</strong>
            will still be classified as flowers with <strong>low confidence</strong>.
        </p>
    </div>
    """)

    # Footer
    gr.HTML("""
    <div class="footer">
        Made with ❤️ using Gradio • Powered by EfficientNet-B0
    </div>
    """)

    # Wire up the button
    classify_btn.click(
        fn=predict,
        inputs=image_input,
        outputs=output
    )

print("Launching app...")

# HF Spaces compatible launch — no blocking loops
demo.queue()
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    show_error=True,
    debug=False,
    css=custom_css
)