app.py

# MARS-LOOP Interactive Demo (Gradio + Plotly)
# --------------------------------------------------
# Features:
# - Hero section with big rotating Mars + "Iniciemos simulación" button
# - Slide 1: Jezero-like crater with blue trash points and an orange bot
#            The bot collects points following an AI-optimized route (k-means sorting + 2-opt path)
# - Slide 2: "Inside the robot" process animation (Sort -> Shred+Wash -> Pelletize -> Mix -> Form/Print)
#            with live KPIs (mass recovery, water recovery, crew time)
# - Slide 3: Clean crater "after" view
#
# Dependencies: gradio, plotly, numpy, pandas
# Run: pip install gradio plotly numpy pandas
#      python app.py
#
# This file is self-contained (no external images needed).

import math
import numpy as np
import pandas as pd
import plotly.graph_objects as go
import gradio as gr
from dataclasses import dataclass

# --------------------------
# Helpers: simple k-means
# --------------------------
def kmeans(X, k=3, iters=15, seed=42):
    rng = np.random.RandomState(seed)
    # choose random points as initial centroids
    idx = rng.choice(len(X), size=k, replace=False)
    C = X[idx].copy()
    for _ in range(iters):
        # assign
        d = ((X[:, None, :] - C[None, :, :]) ** 2).sum(axis=2)  # (n,k)
        labels = d.argmin(axis=1)
        # update
        for j in range(k):
            pts = X[labels == j]
            if len(pts) > 0:
                C[j] = pts.mean(axis=0)
    return labels, C

# --------------------------
# Helpers: route optimization
# --------------------------
def nn_route(points):
    """Nearest-neighbor heuristic route through all points (returns indices order)."""
    n = len(points)
    if n == 0:
        return []
    remaining = set(range(n))
    order = [0]
    remaining.remove(0)
    while remaining:
        last = order[-1]
        # choose nearest
        best = min(remaining, key=lambda j: np.linalg.norm(points[j] - points[last]))
        order.append(best)
        remaining.remove(best)
    return order

def two_opt(points, order, iters=100):
    """2-opt improvement on an existing route order."""
    n = len(order)
    if n < 4:
        return order
    def route_len(ordr):
        return sum(np.linalg.norm(points[ordr[i]] - points[ordr[(i+1) % n]]) for i in range(n-1))
    best = order[:]
    best_len = route_len(best)
    improved = True
    loops = 0
    while improved and loops < iters:
        improved = False
        loops += 1
        for i in range(1, n-2):
            for k in range(i+1, n-1):
                new_order = best[:i] + best[i:k+1][::-1] + best[k+1:]
                new_len = route_len(new_order)
                if new_len < best_len:
                    best, best_len = new_order, new_len
                    improved = True
        if not improved:
            break
    return best

# --------------------------
# Data classes
# --------------------------
@dataclass
class Scenario:
    seed: int
    points: np.ndarray  # (n,2) trash points
    labels: np.ndarray  # material classes per point
    order: list         # route visiting order
    unit_mass: float    # kg per trash item
    # KPIs
    batch_mass: float
    water_used: float
    water_recov: float
    water_loss: float
    useful_mass: float
    metals_reuse: float
    mass_recovery_pct: float
    crew_time_min: float

# --------------------------
# Scenario generation
# --------------------------
def generate_scenario(n_points=60, seed=42, crater_radius=20.0):
    rng = np.random.RandomState(seed)
    # Sample points in an ellipse/valley to evoke crater
    theta = rng.uniform(0, 2*np.pi, size=n_points)
    r = crater_radius * np.sqrt(rng.uniform(0, 1, size=n_points))  # denser center
    # elliptical distortion
    a, b = 1.0, 0.65
    x = a * r * np.cos(theta)
    y = b * r * np.sin(theta) - 2.5  # slight offset
    P = np.stack([x, y], axis=1)

    # Fake spectral features by mixing position + random
    f1 = (x - x.min()) / (x.max() - x.min() + 1e-6)
    f2 = (y - y.min()) / (y.max() - y.min() + 1e-6)
    f3 = rng.beta(2, 5, size=n_points)  # polymer-likeness
    X = np.stack([f1, f2, f3], axis=1)

    labels, _ = kmeans(X, k=3, iters=20, seed=seed)

    # Route planning on points (AI optimizer)
    order = two_opt(P, nn_route(P))

    # KPI model (toy but consistent)
    unit_mass = 1.8  # kg per trash item (avg packaging/textile piece)
    batch_mass = n_points * unit_mass
    # water use per kg + recovery
    water_perkg = 0.7
    water_used = water_perkg * batch_mass
    water_recovery = 0.93
    water_recov = water_used * water_recovery
    water_loss = water_used - water_recov

    # mass conversion to useful composite including regolith (r=0.3), efficiencies product
    eff = 0.95 * 0.97 * 0.96 * 0.95 * 0.95
    polymer_mass_after = batch_mass * eff
    # regolith fraction in final (r)
    rfrac = 0.30
    final_mass = polymer_mass_after / (1 - rfrac) * 0.98  # 2% trim loss fudge
    metals_reuse = 12.0  # kg available as frames reuse, constant per batch
    useful_mass = final_mass
    mass_recovery_pct = (useful_mass + metals_reuse) / batch_mass * 100.0

    # crew time estimation
    crew_time = max(6.0, 8.0 + 0.04 * batch_mass)

    return Scenario(
        seed=seed, points=P, labels=labels, order=order, unit_mass=unit_mass,
        batch_mass=batch_mass, water_used=water_used, water_recov=water_recov, water_loss=water_loss,
        useful_mass=useful_mass, metals_reuse=metals_reuse, mass_recovery_pct=mass_recovery_pct,
        crew_time_min=crew_time
    )

# --------------------------
# Figures
# --------------------------
def hero_html():
    # Rotating planet built with pure CSS (no external image)
    # Large, soft gradients evoke the Red Planet.
    return gr.HTML(
        '''
        <style>
        .hero-wrap{height:78vh;display:flex;flex-direction:column;align-items:center;justify-content:center;background:#12090a;}
        .mars{
            width: 420px; height: 420px; border-radius:50%;
            background: radial-gradient( circle at 35% 30%,
                 #ffb199 0%, #e06045 35%, #b63a27 55%, #5a1e16 75%, #2a1010 100% );
            box-shadow: 0 0 80px rgba(255,90,50,0.35), inset -30px -40px 80px rgba(0,0,0,0.6);
            position: relative; animation: spin 18s linear infinite;
        }
        .mars:before{
            content:""; position:absolute; inset:0; border-radius:50%;
            background: radial-gradient(circle at 70% 65%, rgba(255,255,255,0.12), rgba(0,0,0,0) 40%);
            filter: blur(1px);
        }
        @keyframes spin{ from{transform: rotate(0deg)} to{transform: rotate(360deg)} }
        h1{font-size:48px; letter-spacing:10px; color:#ffe6df; margin:28px 0 6px; font-weight:300;}
        .sub{color:#ffb8a9; letter-spacing:4px; margin-bottom:22px;}
        .btn-start{
            background:#e06045; color:#170d0d; border:none; padding:14px 26px; border-radius:999px;
            font-weight:700; letter-spacing:1px; cursor:pointer; transition: all .2s ease;
        }
        .btn-start:hover{ transform: translateY(-1px); filter: brightness(1.05); }
        </style>
        <div class="hero-wrap">
            <div class="mars"></div>
            <h1>MARTE</h1>
            <div class="sub">PLANETA ROJO</div>
            <button id="start-sim" class="btn-start">Iniciemos simulación</button>
            <script>
            // Bridge click to Gradio event by triggering a hidden button if exists
            setTimeout(()=>{
                const btn = document.querySelector('button.sr-only-start');
                const real = document.getElementById('start-sim');
                if(btn && real){ real.onclick = ()=> btn.click(); }
            }, 500);
            </script>
        </div>
        '''
    )

def crater_animation(scn: Scenario, show_classes=True):
    P = scn.points
    labels = scn.labels
    order = scn.order

    # Colors by class
    palette = np.array(["deepskyblue", "dodgerblue", "lightskyblue"]) if show_classes else np.array(["deepskyblue"]*3)
    colors = palette[labels]

    # Robot initial pos at first point
    fig = go.Figure()

    # crater / landscape (stylized)
    t = np.linspace(0, 2*np.pi, 200)
    a, b = 22, 14
    fig.add_trace(go.Scatter(x=a*np.cos(t), y=b*np.sin(t)-2.5, mode="lines", line=dict(width=2), name="Crater Edge", opacity=0.25))

    # trash points initial
    fig.add_trace(go.Scatter(x=P[:,0], y=P[:,1], mode="markers",
                             marker=dict(size=9, color=colors, opacity=0.95, line=dict(width=0)),
                             name="Trash"))

    # robot
    start = P[order[0]]
    fig.add_trace(go.Scatter(x=[start[0]], y=[start[1]], mode="markers",
                             marker=dict(size=18, color="orange", symbol="triangle-up"), name="Bot"))

    frames = []
    # Animate visiting each point (3 substeps per edge)
    sub = 3
    trash_opacity = np.ones(len(P)) * 0.95
    bx, by = start[0], start[1]

    for idx in range(1, len(order)):
        p0 = P[order[idx-1]]
        p1 = P[order[idx]]
        for s in range(sub):
            tfrac = (s+1)/sub
            x = p0[0]*(1-tfrac) + p1[0]*tfrac
            y = p0[1]*(1-tfrac) + p1[1]*tfrac
            frames.append(go.Frame(data=[
                # crater edge again
                go.Scatter(x=a*np.cos(t), y=b*np.sin(t)-2.5, mode="lines", line=dict(width=2), opacity=0.25, showlegend=False),
                # trash with current opacity
                go.Scatter(x=P[:,0], y=P[:,1], mode="markers",
                           marker=dict(size=9, color=colors, opacity=trash_opacity, line=dict(width=0)), showlegend=False),
                # bot position
                go.Scatter(x=[x], y=[y], mode="markers", marker=dict(size=18, color="orange", symbol="triangle-up"), showlegend=False)
            ]))

        # when arriving, mark this point as collected (fade it)
        trash_opacity[order[idx]] = 0.1

    fig.update(frames=frames)
    fig.update_layout(
        title="Jezero Crater • Trash Collection",
        xaxis=dict(visible=False), yaxis=dict(visible=False, scaleanchor="x", scaleratio=1),
        height=520, plot_bgcolor="rgba(18,9,10,1)", paper_bgcolor="rgba(18,9,10,1)",
        font=dict(color="#ffe6df"),
        updatemenus=[dict(type="buttons", x=0.02, y=0.96, buttons=[
            dict(label="Play", method="animate", args=[None]),
            dict(label="Pause", method="animate",
                 args=[[None], {"mode":"immediate","frame":{"duration":0,"redraw":False},
                                "transition":{"duration":0}}])
        ])],
        showlegend=False
    )
    return fig

def process_animation(scn: Scenario):
    stages = ["SORT", "SHRED+WASH", "PELLETIZE", "MIX (REGOLITH)", "FORM/PRINT"]
    x = [0, 1, 2, 3, 4]
    y = [0]*5

    # Base layout
    fig = go.Figure()
    # boxes
    for i, s in enumerate(stages):
        fig.add_trace(go.Scatter(
            x=[i], y=[0], mode="markers+text",
            marker=dict(size=140, symbol="square", color="#2b1a1a", line=dict(color="#623a35", width=2)),
            text=[s], textfont=dict(color="#ffb8a9"), textposition="middle center", showlegend=False
        ))
    # progress dot (animated)
    frames = []
    nframes = 60
    for f in range(nframes):
        idx = int((f / nframes) * len(stages))
        idx = min(idx, len(stages)-1)
        frames.append(go.Frame(data=[
            go.Scatter(x=[i for i in x], y=[0]*5, mode="markers+text",
                       marker=dict(size=140, symbol="square", color=["#3e2321" if j<=idx else "#2b1a1a" for j in range(5)],
                                   line=dict(color="#623a35", width=2)),
                       text=stages, textfont=dict(color="#ffb8a9"), textposition="middle center", showlegend=False),
            go.Scatter(x=[idx], y=[0.55], mode="markers", marker=dict(size=18, color="orange", symbol="triangle-up"), showlegend=False)
        ]))

    # KPIs as annotations
    kp = scn
    ann = [
        dict(x=0, y=-0.9, text=f"Batch waste: {kp.batch_mass:.1f} kg", showarrow=False, font=dict(color="#ffe6df")),
        dict(x=1, y=-0.9, text=f"Water used: {kp.water_used:.1f} L • Recovered: {kp.water_recov:.1f} L", showarrow=False, font=dict(color="#ffe6df")),
        dict(x=2, y=-0.9, text=f"Useful composite: {kp.useful_mass:.1f} kg", showarrow=False, font=dict(color="#ffe6df")),
        dict(x=3, y=-0.9, text=f"Metals → reuse: {kp.metals_reuse:.1f} kg", showarrow=False, font=dict(color="#ffe6df")),
        dict(x=4, y=-0.9, text=f"Mass recovery: {kp.mass_recovery_pct:.1f}% • Crew time ≤ {kp.crew_time_min:.0f} min", showarrow=False, font=dict(color="#ffe6df")),
    ]

    fig.update(frames=frames)
    fig.update_layout(
        title="Inside the Bot • MARS-LOOP Process",
        xaxis=dict(visible=False, range=[-0.5, 4.5]),
        yaxis=dict(visible=False, range=[-1.2, 1.2]),
        height=420, plot_bgcolor="rgba(18,9,10,1)", paper_bgcolor="rgba(18,9,10,1)",
        font=dict(color="#ffe6df"),
        updatemenus=[dict(type="buttons", x=0.02, y=0.96, buttons=[
            dict(label="Play", method="animate", args=[None]),
            dict(label="Pause", method="animate",
                 args=[[None], {"mode":"immediate","frame":{"duration":0,"redraw":False},
                                "transition":{"duration":0}}])
        ])],
        annotations=ann
    )
    return fig

def clean_crater(P):
    # stylized empty crater
    fig = go.Figure()
    t = np.linspace(0, 2*np.pi, 200)
    a, b = 22, 14
    fig.add_trace(go.Scatter(x=a*np.cos(t), y=b*np.sin(t)-2.5, mode="lines",
                             line=dict(width=2), opacity=0.28, name="Crater"))
    fig.add_annotation(x=0, y=0, text="Área limpia ✅", showarrow=False, font=dict(size=28, color="#b4ffb4"))
    fig.update_layout(
        title="Jezero Crater • After Cleaning",
        xaxis=dict(visible=False), yaxis=dict(visible=False, scaleanchor="x", scaleratio=1),
        height=520, plot_bgcolor="rgba(18,9,10,1)", paper_bgcolor="rgba(18,9,10,1)",
        font=dict(color="#ffe6df"), showlegend=False
    )
    return fig

# --------------------------
# Gradio App
# --------------------------
with gr.Blocks(title="MARS-LOOP Interactive", theme=gr.themes.Soft(primary_hue="red")) as demo:
    gr.HTML("<style> .gradio-container {max-width: 980px !important;} </style>")
    state_scn = gr.State()   # Scenario object
    state_slide = gr.State(0)

    # HERO
    hero = hero_html()
    hidden_start_bridge = gr.Button("Iniciar", visible=False, elem_classes=["sr-only-start"])

    # SLIDE 1
    with gr.Group(visible=False) as slide1:
        gr.Markdown("### 🏜️ Simulación: Jezero Crater (detección y recolección de desechos)")
        with gr.Row():
            show_classes = gr.Checkbox(value=True, label="Mostrar clasificación automática de materiales (IA)")
            seed_in = gr.Slider(1, 9999, value=42, step=1, label="Semilla")
            regen = gr.Button("🔄 Regenerar escenario")
        crater_plot = gr.Plot()
        next1 = gr.Button("Siguiente ➜ Proceso interno")

    # SLIDE 2
    with gr.Group(visible=False) as slide2:
        gr.Markdown("### ⚙️ Proceso dentro del robot (MARS-LOOP)")
        process_plot = gr.Plot()
        next2 = gr.Button("Siguiente ➜ Cráter limpio")

    # SLIDE 3
    with gr.Group(visible=False) as slide3:
        gr.Markdown("### ✅ Resultado: Cráter limpio")
        clean_plot = gr.Plot()
        reset = gr.Button("Reiniciar")

    # ------------------
    # Callbacks
    # ------------------
    def on_start():
        scn = generate_scenario(seed=42)
        fig = crater_animation(scn, show_classes=True)
        return (
            gr.update(visible=False),    # hide hero
            gr.update(visible=True),     # show slide1
            scn, 1,                      # state scenario + slide
            gr.update(value=fig)         # crater plot
        )

    def on_regen(seed, show):
        scn = generate_scenario(seed=int(seed))
        fig = crater_animation(scn, show_classes=bool(show))
        return scn, gr.update(value=fig)

    def go_next1(scn: Scenario):
        fig = process_animation(scn)
        return gr.update(visible=False), gr.update(visible=True), gr.update(value=fig), 2

    def go_next2(scn: Scenario):
        fig = clean_crater(scn.points)
        return gr.update(visible=False), gr.update(visible=True), gr.update(value=fig), 3

    def on_reset():
        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), 0

    # wire
    hidden_start_bridge.click(on_start, outputs=[hero, slide1, state_scn, state_slide, crater_plot])
    regen.click(on_regen, inputs=[seed_in, show_classes], outputs=[state_scn, crater_plot])
    next1.click(go_next1, inputs=[state_scn], outputs=[slide1, slide2, process_plot, state_slide])
    next2.click(go_next2, inputs=[state_scn], outputs=[slide2, slide3, clean_plot, state_slide])
    reset.click(on_reset, outputs=[slide3, slide2, hero, state_slide])

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