from datetime import datetime
import gradio as gr
import requests

API_URL = "https://huggingface.co/api/models"
TIMEOUT = 20
MAX_FETCH = 200

CATEGORIES = {
    "All World Models": "",
    "World Model": "world-model",
    "World Foundation Model": "world foundation",
    "World Action Model": "world action",
    "Interactive World Model": "interactive",
    "Robotics": "robot",
    "Embodied AI": "embodied",
    "Physical AI": "physical",
    "Video World Model": "video",
    "Spatial Intelligence": "spatial",
    "Reinforcement Learning": "reinforcement",
    "Autonomous Systems": "autonomous",
    "Agents": "agent",
}

TAXONOMY = """
### 🧭 World Model Taxonomy

| Area | Core idea |
|---|---|
| **Latent World Models** | Predict future states in a compressed internal representation |
| **Predictive World Models** | Learn how environments evolve over time |
| **Generative World Models** | Generate possible future states or environments |
| **Interactive World Models** | Respond to actions and controls |
| **World Foundation Models** | General-purpose world models designed for broad reuse |
| **World Action Models** | Connect environmental representation with possible actions |
| **Embodied World Models** | Model environments for embodied agents and robots |
| **Spatial World Models** | Represent geometry, position, depth and movement |
| **Video World Models** | Use visual temporal prediction as part of world modeling |
| **Agent World Models** | Support planning and decision making for autonomous agents |

**The categories overlap.** A single project may belong to several areas at once.
"""

WHAT_ARE_WORLD_MODELS = """
### 🌍 What are World Models?

A **world model** is an AI model that learns an internal representation of an environment and how that environment changes.

Unlike a system that only recognizes the current state, a world model can try to capture **dynamics**:

```text
Current State + Action
          ↓
      World Model
          ↓
Possible Future State
```

That makes world models especially relevant to AI systems that need to:

- predict what may happen next
- simulate alternatives
- plan over multiple steps
- reason about consequences
- interact with physical or digital environments
- train from imagined experience

World models appear across reinforcement learning, robotics, autonomous driving, video generation, spatial intelligence, embodied AI, Physical AI and agentic systems.
"""

def safe_int(value):
    try:
        return int(value or 0)
    except Exception:
        return 0

def human_number(value):
    if value is None:
        return "—"
    n = safe_int(value)
    if n >= 1_000_000:
        return f"{n/1_000_000:.1f}M"
    if n >= 1_000:
        return f"{n/1_000:.1f}K"
    return str(n)

def format_date(value):
    if not value:
        return "—"
    try:
        return datetime.fromisoformat(str(value).replace("Z", "+00:00")).strftime("%Y-%m-%d")
    except Exception:
        return str(value)[:10]

def get_license(model):
    card = model.get("cardData") or {}
    lic = card.get("license")
    if isinstance(lic, list):
        lic = ", ".join(map(str, lic))
    if lic:
        return str(lic)
    for tag in model.get("tags") or []:
        if isinstance(tag, str) and tag.startswith("license:"):
            return tag.split(":", 1)[1]
    return "—"

def get_library(model):
    if model.get("library_name"):
        return str(model["library_name"])
    tags = model.get("tags") or []
    for lib in ["transformers", "diffusers", "sentence-transformers", "pytorch", "tensorflow", "mlx", "gguf"]:
        if lib in tags:
            return lib
    return "—"

def selected_tags(model):
    tags = [str(t) for t in (model.get("tags") or [])]
    cleaned = []
    for tag in tags:
        if tag.startswith(("license:", "region:", "arxiv:", "dataset:")):
            continue
        if tag in {"safetensors", "pytorch", "transformers"}:
            continue
        cleaned.append(tag)
    return ", ".join(cleaned[:7]) if cleaned else "—"

def fetch_world_models():
    queries = [
        {"filter": "world-model", "sort": "downloads", "direction": "-1", "limit": str(MAX_FETCH), "full": "true"},
        {"search": "world model", "sort": "downloads", "direction": "-1", "limit": "100", "full": "true"},
    ]
    merged = {}
    for params in queries:
        try:
            r = requests.get(
                API_URL,
                params=params,
                timeout=TIMEOUT,
                headers={"User-Agent": "World-Models-Explorer/1.0"},
            )
            r.raise_for_status()
            data = r.json()
            if isinstance(data, list):
                for model in data:
                    mid = str(model.get("id") or model.get("modelId") or "")
                    if mid:
                        merged[mid] = model
        except Exception:
            continue
    return list(merged.values())

def explore(query, category, library_filter, license_filter, limit):
    models = fetch_world_models()
    q = (query or "").strip().lower()
    cat = CATEGORIES.get(category, "").lower()
    lib_q = (library_filter or "").strip().lower()
    lic_q = (license_filter or "").strip().lower()

    filtered = []
    for model in models:
        mid = str(model.get("id") or model.get("modelId") or "")
        tags = [str(t) for t in (model.get("tags") or [])]
        library = get_library(model)
        license_name = get_license(model)
        pipeline = str(model.get("pipeline_tag") or "—")
        haystack = " ".join([mid, library, license_name, pipeline] + tags).lower()

        if q and q not in haystack:
            continue
        if cat and cat not in haystack:
            continue
        if lib_q and lib_q not in library.lower():
            continue
        if lic_q and lic_q not in license_name.lower():
            continue
        filtered.append(model)

    filtered.sort(
        key=lambda m: (safe_int(m.get("downloads")), safe_int(m.get("likes"))),
        reverse=True
    )

    rows = []
    for model in filtered[:int(limit)]:
        mid = str(model.get("id") or model.get("modelId") or "")
        rows.append([
            mid,
            human_number(model.get("downloads")),
            human_number(model.get("likes")),
            str(model.get("pipeline_tag") or "—"),
            get_library(model),
            get_license(model),
            format_date(model.get("lastModified")),
            selected_tags(model),
            f"https://huggingface.co/{mid}",
        ])

    if not models:
        summary = (
            "**The Hugging Face Hub could not be reached right now.** "
            "Use Refresh later. The explanatory sections of the Explorer remain available."
        )
    else:
        summary = (
            f"**{len(rows)} models shown** · {len(filtered)} matches · "
            f"{len(models)} world-model-related repositories discovered. "
            "Downloads and likes are discovery signals, not quality scores."
        )
    return rows, summary

CSS = """
.gradio-container {max-width: 1500px !important;}
#hero {
    padding: 1.5rem 1.7rem;
    border: 1px solid var(--border-color-primary);
    border-radius: 18px;
    margin-bottom: 1rem;
}
#hero h1 {margin-bottom: .25rem;}
#hero p {font-size: 1.02rem;}
"""

with gr.Blocks(css=CSS, title="World Models Explorer") as demo:
    gr.Markdown("""
<div id="hero">

# 🌍 World Models Explorer

**Explore world models for simulation, prediction, planning, agents, robotics and Physical AI.**

Discover models and projects from the Hugging Face ecosystem while learning how the major world-model categories fit together.

</div>
""")

    with gr.Accordion("What are World Models?", open=True):
        gr.Markdown(WHAT_ARE_WORLD_MODELS)

    with gr.Row():
        query = gr.Textbox(label="Search", placeholder="e.g. Genie, Cosmos, Dreamer, robotics, video, agent…", scale=2)
        category = gr.Dropdown(choices=list(CATEGORIES.keys()), value="All World Models", label="Category", scale=1)
        library_filter = gr.Textbox(label="Library", placeholder="e.g. diffusers", scale=1)
        license_filter = gr.Textbox(label="License", placeholder="e.g. apache", scale=1)

    with gr.Row():
        limit = gr.Slider(minimum=10, maximum=100, value=50, step=10, label="Maximum results")
        refresh = gr.Button("Refresh Explorer", variant="primary")

    summary = gr.Markdown()

    table = gr.Dataframe(
        headers=["Model", "Downloads", "Likes", "Pipeline", "Library", "License", "Updated", "Tags", "URL"],
        datatype=["str"] * 9,
        interactive=False,
        wrap=True,
        label="World-model ecosystem",
    )

    with gr.Accordion("World Model Taxonomy", open=False):
        gr.Markdown(TAXONOMY)

    gr.Markdown("""
### How to use the Explorer

1. **Discover** relevant models and projects.
2. **Filter** by category, library, license or keyword.
3. **Open the original model card** before making technical decisions.
4. **Verify architecture, training data, license, modalities and limitations.**
5. **Benchmark on your own task** before production deployment.

---

### Planned World Models ecosystem

**World Models Explorer** → discover the field  
**World Model Registry** → structured metadata  
**World Model Benchmark** → evaluation and benchmarks  
**World Model Landscape** → visual ecosystem map

---

**Cooperation, research and ecosystem partnerships:**  
[agenten@magenta.de](mailto:agenten@magenta.de)
""")

    inputs = [query, category, library_filter, license_filter, limit]

    refresh.click(fn=explore, inputs=inputs, outputs=[table, summary])
    for component in [query, library_filter, license_filter]:
        component.submit(fn=explore, inputs=inputs, outputs=[table, summary])
    category.change(fn=explore, inputs=inputs, outputs=[table, summary])
    limit.change(fn=explore, inputs=inputs, outputs=[table, summary])
    demo.load(fn=explore, inputs=inputs, outputs=[table, summary])

if __name__ == "__main__":
    demo.launch()