from fastai.vision.all import *
import gradio as gr
import fal_client
from PIL import Image
import io
import random
import requests
from pathlib import Path

# =========================================================
# WIKIPEDIA LINKS
# =========================================================
search_terms_wikipedia = {
    "blazing star": "https://en.wikipedia.org/wiki/Mentzelia",
    "bristlecone pine": "https://en.wikipedia.org/wiki/Pinus_longaeva",
    "california bluebell": "https://en.wikipedia.org/wiki/Phacelia_minor",
    "california buckeye": "https://en.wikipedia.org/wiki/Aesculus_californica",
    "california buckwheat": "https://en.wikipedia.org/wiki/Eriogonum_fasciculatum",
    "california fuchsia": "https://en.wikipedia.org/wiki/Epilobium_canum",
    "california checkerbloom": "https://en.wikipedia.org/wiki/Sidalcea_malviflora",
    "california lilac": "https://en.wikipedia.org/wiki/Ceanothus",
    "california poppy": "https://en.wikipedia.org/wiki/Eschscholzia_californica",
    "california sagebrush": "https://en.wikipedia.org/wiki/Artemisia_californica",
    "california wild grape": "https://en.wikipedia.org/wiki/Vitis_californica",
    "california wild rose": "https://en.wikipedia.org/wiki/Rosa_californica",
    "coyote mint": "https://en.wikipedia.org/wiki/Monardella",
    "elegant clarkia": "https://en.wikipedia.org/wiki/Clarkia_unguiculata",
    "baby blue eyes": "https://en.wikipedia.org/wiki/Nemophila_menziesii",
    "hummingbird sage": "https://en.wikipedia.org/wiki/Salvia_spathacea",
    "delphinium": "https://en.wikipedia.org/wiki/Delphinium",
    "matilija poppy": "https://en.wikipedia.org/wiki/Romneya_coulteri",
    "blue-eyed grass": "https://en.wikipedia.org/wiki/Sisyrinchium_bellum",
    "penstemon spectabilis": "https://en.wikipedia.org/wiki/Penstemon_spectabilis",
    "seaside daisy": "https://en.wikipedia.org/wiki/Erigeron_glaucus",
    "sticky monkeyflower": "https://en.wikipedia.org/wiki/Diplacus_aurantiacus",
    "tidy tips": "https://en.wikipedia.org/wiki/Layia_platyglossa",
    "wild cucumber": "https://en.wikipedia.org/wiki/Marah_(plant)",
    "douglas iris": "https://en.wikipedia.org/wiki/Iris_douglasiana",
    "goldfields coreopsis": "https://en.wikipedia.org/wiki/Coreopsis"
}

# =========================================================
# AI PROMPTS
# =========================================================
prompt_templates = [

    "A dreamy watercolor painting of a {flower} in a magical forest with glowing sunlight and butterflies.",

    "A cinematic artistic interpretation of a {flower} blooming beside a mountain trail at sunrise.",

    "A fantasy botanical artwork featuring a {flower} surrounded by mist and colorful wildlife.",

    "An impressionist oil painting of a {flower} field with vibrant brush strokes and golden light.",

    "A detailed nature journal illustration of a {flower} with artistic sketches and handwritten notes."
]

# =========================================================
# EXAMPLE IMAGES
# =========================================================
example_images = [
    str(Path("example_images/example_1.jpg")),
    str(Path("example_images/example_2.jpg")),
    str(Path("example_images/example_3.jpg")),
    str(Path("example_images/example_4.jpg")),
    str(Path("example_images/example_5.jpg"))
]

# =========================================================
# LOAD MODEL
# =========================================================
learn = load_learner("resnet50_30_categories.pkl")

# =========================================================
# FAL LOGS
# =========================================================
def on_queue_update(update):

    if isinstance(update, fal_client.InProgress):

        for log in update.logs:

            print(log["message"])

# =========================================================
# MAIN FUNCTION
# =========================================================
def process_image(
    img,
    art_style,
    creativity_level,
    image_quality
):

    if img is None:

        return (
            None,
            None,
            "⚠️ Please upload a flower image.",
            {}
        )

    # =====================================================
    # CLASSIFICATION
    # =====================================================
    predicted_class, _, probs = learn.predict(img)

    classification_results = dict(
        zip(
            learn.dls.vocab,
            map(float, probs)
        )
    )

    confidence = max(classification_results.values()) * 100

    # =====================================================
    # WIKIPEDIA
    # =====================================================
    wiki_url = search_terms_wikipedia.get(
        predicted_class,
        "No article found."
    )

    # =====================================================
    # AI PROMPT
    # =====================================================
    prompt = random.choice(
        prompt_templates
    ).format(
        flower=predicted_class
    )

    final_prompt = f"""
{prompt}

Style: {art_style}

Creativity Level: {creativity_level}

Image Quality: {image_quality}
"""

    # =====================================================
    # IMAGE GENERATION
    # =====================================================
    result = fal_client.subscribe(
        "fal-ai/flux/schnell",
        arguments={
            "prompt": final_prompt,
            "image_size": "portrait_4_3"
        },
        with_logs=True,
        on_queue_update=on_queue_update,
    )

    image_url = result["images"][0]["url"]

    response = requests.get(image_url)

    generated_image = Image.open(
        io.BytesIO(response.content)
    )

    # =====================================================
    # RESULT TEXT
    # =====================================================
    result_text = f"""
# 🌸 Flower Classification Complete

### Predicted Flower:
## {predicted_class.title()}

### Confidence Score:
## {confidence:.2f}%

### AI artistic interpretation generated successfully.
"""

    # =====================================================
    # TOP PROBABILITIES
    # =====================================================
    sorted_results = dict(
        sorted(
            classification_results.items(),
            key=lambda x: x[1],
            reverse=True
        )[:5]
    )

    return (
        generated_image,
        wiki_url,
        result_text,
        sorted_results
    )

# =========================================================
# CUSTOM CSS
# =========================================================
custom_css = """
body {
    background: #f5f7fb;
    font-family: 'Segoe UI', sans-serif;
}

.gradio-container {
    max-width: 1350px !important;
    margin: auto;
}

.hero {
    background: linear-gradient(135deg,#065f46,#16a34a);
    padding: 40px;
    border-radius: 30px;
    color: white;
    margin-bottom: 20px;
}

.hero h1 {
    font-size: 52px;
    font-weight: 800;
    margin-bottom: 10px;
}

.hero p {
    font-size: 18px;
    opacity: 0.92;
}

.card {
    background: white;
    border-radius: 24px;
    padding: 22px;
    box-shadow: 0 6px 18px rgba(0,0,0,0.08);
}

button {
    height: 60px !important;
    border-radius: 18px !important;
    border: none !important;
    background: linear-gradient(135deg,#16a34a,#15803d) !important;
    color: white !important;
    font-size: 20px !important;
    font-weight: 700 !important;
}

button:hover {
    background: linear-gradient(135deg,#15803d,#166534) !important;
}

input, textarea, select {
    border-radius: 16px !important;
}

@media (max-width:768px){

    .hero {
        padding: 22px;
    }

    .hero h1 {
        font-size: 32px;
    }

    .hero p {
        font-size: 15px;
    }

    button {
        height: 54px !important;
        font-size: 17px !important;
    }
}
"""

# =========================================================
# HERO SECTION
# =========================================================
hero_html = """
<div class="hero">

<h1>🌸 AI Flower Classifier & Art Generator</h1>

<p>
Upload flower images, identify plant species instantly, and generate AI-powered artistic interpretations.
Modern responsive interface optimized for mobile and desktop devices.
</p>

</div>
"""

# =========================================================
# INTERFACE
# =========================================================
with gr.Blocks(
    css=custom_css,
    theme=gr.themes.Soft(
        primary_hue="green",
        secondary_hue="emerald"
    )
) as demo:

    gr.HTML(hero_html)

    # =====================================================
    # TOP INFO
    # =====================================================
    with gr.Row():

        with gr.Column():
            gr.Markdown("""
### ⚡ Fast AI Recognition
Instant flower classification
""")

        with gr.Column():
            gr.Markdown("""
### 🎨 AI Art Generator
Create artistic flower scenes
""")

        with gr.Column():
            gr.Markdown("""
### 📱 Mobile Responsive
Optimized modern UI
""")

    # =====================================================
    # MAIN AREA
    # =====================================================
    with gr.Row():

        # =================================================
        # LEFT SIDE
        # =================================================
        with gr.Column(scale=1):

            input_image = gr.Image(
                type="pil",
                label="🌸 Upload Flower Image"
            )

            art_style = gr.Dropdown(
                choices=[
                    "Watercolor",
                    "Oil Painting",
                    "Fantasy Art",
                    "Impressionist",
                    "Botanical Illustration"
                ],
                value="Watercolor",
                label="🎨 Art Style"
            )

            creativity_level = gr.Slider(
                minimum=10,
                maximum=100,
                value=75,
                step=5,
                label="🧠 Creativity Level"
            )

            image_quality = gr.Dropdown(
                choices=[
                    "Standard",
                    "High",
                    "Ultra"
                ],
                value="High",
                label="✨ Image Quality"
            )

            generate_btn = gr.Button(
                "🚀 Analyze & Generate Art",
                variant="primary"
            )

        # =================================================
        # RIGHT SIDE
        # =================================================
        with gr.Column(scale=1):

            output_image = gr.Image(
                label="🖌️ AI Artistic Interpretation"
            )

            wiki_output = gr.Textbox(
                label="📚 Wikipedia Article",
                lines=1
            )

            result_output = gr.Markdown()

            probability_output = gr.Label(
                label="📊 Top Predictions"
            )

    # =====================================================
    # EXAMPLES
    # =====================================================
    gr.Examples(
        examples=example_images,
        inputs=input_image
    )

    # =====================================================
    # BUTTON ACTION
    # =====================================================
    generate_btn.click(
        fn=process_image,
        inputs=[
            input_image,
            art_style,
            creativity_level,
            image_quality
        ],
        outputs=[
            output_image,
            wiki_output,
            result_output,
            probability_output
        ]
    )

# =========================================================
# LAUNCH
# =========================================================
demo.launch()