app/main.py

import streamlit as st
import os, sys
import numpy as np

# ── Page Config ──
st.set_page_config(
    page_title="PETIMOT Explorer",
    page_icon="🧬",
    layout="wide",
    initial_sidebar_state="expanded",
)

# ── Ensure PETIMOT is importable ──
PETIMOT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PETIMOT_ROOT not in sys.path:
    sys.path.insert(0, PETIMOT_ROOT)

# ── Custom CSS ──
st.markdown("""
<style>
    /* ─── Dark Theme ─── */
    .stApp { background-color: #0f0d1a; }
    .block-container { padding-top: 1rem; max-width: 1200px; }

    /* Sidebar */
    section[data-testid="stSidebar"] {
        background: linear-gradient(180deg, #1a1730 0%, #0f0d1a 100%);
        border-right: 1px solid #2d2b55;
    }

    /* Headers */
    h1, h2, h3 { color: #c4b5fd !important; }

    /* Metric cards */
    [data-testid="stMetric"] {
        background: linear-gradient(135deg, #1e1b4b 0%, #312e81 100%);
        border: 1px solid #4338ca;
        border-radius: 16px;
        padding: 16px 20px;
        box-shadow: 0 4px 20px rgba(99, 102, 241, 0.15);
        transition: transform 0.2s ease, box-shadow 0.2s ease;
    }
    [data-testid="stMetric"]:hover {
        transform: translateY(-2px);
        box-shadow: 0 8px 30px rgba(99, 102, 241, 0.25);
    }
    [data-testid="stMetricLabel"] { color: #a5b4fc !important; font-size: 0.85rem !important; }
    [data-testid="stMetricValue"] { color: #e0e7ff !important; font-weight: 700 !important; }

    /* Dataframe */
    .stDataFrame { border-radius: 12px; overflow: hidden; }

    /* Tabs */
    .stTabs [data-baseweb="tab"] {
        background-color: #1e1b4b;
        border-radius: 8px 8px 0 0;
        color: #a5b4fc;
    }
    .stTabs [data-baseweb="tab"][aria-selected="true"] {
        background-color: #312e81;
        color: white;
    }

    /* ─── Hero Section ─── */
    .hero-container {
        background: linear-gradient(135deg, #1e1b4b 0%, #312e81 30%, #4338ca 60%, #6366f1 100%);
        border-radius: 24px;
        padding: 48px 40px;
        margin-bottom: 2rem;
        position: relative;
        overflow: hidden;
        box-shadow: 0 8px 40px rgba(99, 102, 241, 0.3);
    }
    .hero-container::before {
        content: '';
        position: absolute;
        top: -50%;
        right: -20%;
        width: 60%;
        height: 200%;
        background: radial-gradient(circle, rgba(139, 92, 246, 0.15) 0%, transparent 70%);
        animation: heroGlow 6s ease-in-out infinite alternate;
    }
    @keyframes heroGlow {
        0% { transform: translate(0, 0) scale(1); opacity: 0.5; }
        100% { transform: translate(-10%, 10%) scale(1.2); opacity: 1; }
    }
    .hero-title {
        font-size: 3rem;
        font-weight: 800;
        color: #ffffff !important;
        margin: 0 0 8px 0;
        letter-spacing: -0.02em;
        position: relative;
        z-index: 1;
    }
    .hero-subtitle {
        font-size: 1.3rem;
        color: #c4b5fd;
        margin: 0 0 24px 0;
        font-weight: 400;
        position: relative;
        z-index: 1;
    }
    .hero-badge {
        display: inline-block;
        background: rgba(255,255,255,0.1);
        backdrop-filter: blur(10px);
        border: 1px solid rgba(255,255,255,0.15);
        border-radius: 999px;
        padding: 6px 16px;
        font-size: 0.85rem;
        color: #e0e7ff;
        margin-right: 8px;
        margin-bottom: 8px;
    }

    /* ─── Feature Cards ─── */
    .feature-card {
        background: linear-gradient(135deg, #1e1b4b 0%, #1a1730 100%);
        border: 1px solid #312e81;
        border-radius: 16px;
        padding: 24px;
        height: 100%;
        transition: all 0.3s ease;
        box-shadow: 0 2px 12px rgba(0,0,0,0.2);
    }
    .feature-card:hover {
        border-color: #6366f1;
        box-shadow: 0 4px 24px rgba(99, 102, 241, 0.2);
        transform: translateY(-2px);
    }
    .feature-icon {
        font-size: 2.2rem;
        margin-bottom: 12px;
    }
    .feature-title {
        color: #e0e7ff !important;
        font-size: 1.15rem;
        font-weight: 700;
        margin-bottom: 8px;
    }
    .feature-desc {
        color: #94a3b8;
        font-size: 0.9rem;
        line-height: 1.5;
    }

    /* ─── Leaderboard cards ─── */
    .leader-row {
        display: flex;
        align-items: center;
        gap: 12px;
        padding: 10px 16px;
        background: rgba(30, 27, 75, 0.6);
        border-radius: 10px;
        margin-bottom: 6px;
        border-left: 3px solid #6366f1;
        transition: background 0.2s;
    }
    .leader-row:hover { background: rgba(49, 46, 129, 0.6); }
    .leader-rank {
        color: #6366f1;
        font-weight: 800;
        font-size: 1rem;
        min-width: 28px;
    }
    .leader-name {
        color: #e0e7ff;
        font-weight: 600;
        flex: 1;
        font-family: 'SF Mono', 'Fira Code', monospace;
        font-size: 0.85rem;
    }
    .leader-val {
        color: #a5b4fc;
        font-weight: 500;
        font-size: 0.85rem;
    }

    /* ─── Animations ─── */
    @keyframes fadeInUp {
        from { opacity: 0; transform: translateY(20px); }
        to { opacity: 1; transform: translateY(0); }
    }
    .animate-in { animation: fadeInUp 0.6s ease-out forwards; }
    .animate-delay-1 { animation-delay: 0.1s; }
    .animate-delay-2 { animation-delay: 0.2s; }
    .animate-delay-3 { animation-delay: 0.3s; }
</style>
""", unsafe_allow_html=True)

# ── Sidebar ──
with st.sidebar:
    st.markdown("""
    <div style="text-align:center; padding: 16px 0;">
        <div style="font-size: 3.5rem;">🧬</div>
        <div style="font-size: 1.6rem; font-weight: 800; color: #c4b5fd; letter-spacing: -0.02em;">PETIMOT</div>
        <div style="font-size: 0.8rem; color: #94a3b8; margin-top: 4px;">Protein Motion from Sparse Data</div>
        <div style="font-size: 0.7rem; color: #6366f1; margin-top: 2px;">SE(3)-Equivariant GNNs</div>
    </div>
    """, unsafe_allow_html=True)
    st.divider()

    # Global settings
    st.markdown("### ⚙️ Settings")

    weights_dir = os.path.join(PETIMOT_ROOT, "weights")
    pt_files = []
    if os.path.isdir(weights_dir):
        for root, dirs, files in os.walk(weights_dir):
            for f in files:
                if f.endswith(".pt"):
                    pt_files.append(os.path.join(root, f))

    if pt_files:
        selected_weights = st.selectbox(
            "Model weights",
            pt_files,
            format_func=lambda x: os.path.basename(x),
            key="weights"
        )
    else:
        selected_weights = None
        st.warning("No weights found in `weights/`")

    st.divider()
    st.markdown("""
    **Links**
    - [Paper](https://arxiv.org/abs/2504.02839)
    - [GitHub](https://github.com/PhyloSofS-Team/PETIMOT)
    - [Data](https://figshare.com/s/ab400d852b4669a83b64)
    """)
    st.caption("GPL-3.0 · Lombard, Grudinin & Laine")

# ── Hero Section ──
st.markdown("""
<div class="hero-container">
    <p class="hero-title">🧬 PETIMOT Explorer</p>
    <p class="hero-subtitle">
        Explore protein motion predictions at scale — 36K+ proteins analyzed with SE(3)-Equivariant Graph Neural Networks
    </p>
    <div>
        <span class="hero-badge">🔬 36K+ Proteins</span>
        <span class="hero-badge">🧠 SE(3)-Equivariant</span>
        <span class="hero-badge">📊 4 Motion Modes</span>
        <span class="hero-badge">⚡ CPU Inference</span>
    </div>
</div>
""", unsafe_allow_html=True)

# ── Data Status ──
from app.utils.download import check_data_status, ensure_weights
from app.utils.data_loader import find_predictions_dir, load_prediction_index

status = check_data_status(PETIMOT_ROOT)

# ── Quick Search ──
st.markdown("### 🔍 Quick Search")
quick_search = st.text_input(
    "Search proteins by name",
    placeholder="e.g. 1ake, 4ake, lysozyme...",
    label_visibility="collapsed",
    key="home_search",
)

if quick_search:
    st.session_state["explorer_search"] = quick_search
    st.info(f"🔍 Navigate to the **Explorer** page to see results for **\"{quick_search}\"**")

# ── Metrics Row ──
st.markdown("### 📊 Dataset Overview")
col1, col2, col3, col4 = st.columns(4)
with col1:
    n_pred = status['predictions']
    pred_display = "36,675" if n_pred < 0 else f"{n_pred:,}"
    st.metric("🧬 Proteins", pred_display,
              delta="✅" if status['has_predictions'] else "⚠️ No data")
with col2:
    st.metric("🎯 Ground Truth", f"{status['ground_truth']:,}",
              delta="✅" if status['has_gt'] else "Not loaded")
with col3:
    st.metric("⚖️ Model Weights", "4.7M params",
              delta="✅" if status['has_weights'] else "Missing")
with col4:
    st.metric("🔮 Motion Modes", "4 per protein",
              delta="Normal Modes" if status['has_predictions'] else "—")

# ── Feature Cards ──
st.markdown("---")
col1, col2, col3 = st.columns(3)

with col1:
    st.markdown("""
    <div class="feature-card animate-in animate-delay-1">
        <div class="feature-icon">🔍</div>
        <div class="feature-title">Explorer</div>
        <div class="feature-desc">
            Browse pre-computed predictions for 36K+ proteins. Filter by sequence length,
            displacement, and view 3D motion visualizations with interactive controls.
        </div>
    </div>
    """, unsafe_allow_html=True)

with col2:
    st.markdown("""
    <div class="feature-card animate-in animate-delay-2">
        <div class="feature-icon">🔮</div>
        <div class="feature-title">Inference</div>
        <div class="feature-desc">
            Predict motion modes for any protein structure. Upload a PDB file or fetch
            from RCSB. Runs on CPU in 5–30 seconds.
        </div>
    </div>
    """, unsafe_allow_html=True)

with col3:
    st.markdown("""
    <div class="feature-card animate-in animate-delay-3">
        <div class="feature-icon">📊</div>
        <div class="feature-title">Statistics</div>
        <div class="feature-desc">
            Dataset-wide analysis with interactive charts: displacement distributions,
            correlation heatmaps, leaderboards, and length-stratified analysis.
        </div>
    </div>
    """, unsafe_allow_html=True)

# ── Featured Proteins ──
if status['has_predictions']:
    pred_dir = find_predictions_dir(PETIMOT_ROOT)
    if pred_dir:
        try:
            df = load_prediction_index(pred_dir)
            if not df.empty and len(df) > 0:
                st.markdown("---")
                st.markdown("### 🏆 Featured Proteins")

                col_flex, col_rigid = st.columns(2)

                with col_flex:
                    st.markdown("**🔴 Most Flexible** (highest mean displacement)")
                    top5 = df.nlargest(5, "mean_disp_m0")
                    html_rows = ""
                    for i, (_, row) in enumerate(top5.iterrows()):
                        html_rows += f"""
                        <div class="leader-row">
                            <span class="leader-rank">#{i+1}</span>
                            <span class="leader-name">{row['name']}</span>
                            <span class="leader-val">{row['mean_disp_m0']:.3f} Å · {int(row['seq_len'])} res</span>
                        </div>"""
                    st.markdown(html_rows, unsafe_allow_html=True)

                with col_rigid:
                    st.markdown("**🔵 Most Rigid** (lowest mean displacement)")
                    bot5 = df.nsmallest(5, "mean_disp_m0")
                    html_rows = ""
                    for i, (_, row) in enumerate(bot5.iterrows()):
                        html_rows += f"""
                        <div class="leader-row">
                            <span class="leader-rank">#{i+1}</span>
                            <span class="leader-name">{row['name']}</span>
                            <span class="leader-val">{row['mean_disp_m0']:.3f} Å · {int(row['seq_len'])} res</span>
                        </div>"""
                    st.markdown(html_rows, unsafe_allow_html=True)

                # ── Sparkline overview ──
                st.markdown("---")
                st.markdown("### 📈 At a Glance")

                import plotly.graph_objects as go

                col_s1, col_s2, col_s3 = st.columns(3)

                def make_sparkline(values, title, color, unit=""):
                    fig = go.Figure()
                    fig.add_trace(go.Histogram(
                        x=values, nbinsx=40,
                        marker_color=color,
                        marker_line_width=0,
                        opacity=0.85,
                    ))
                    fig.update_layout(
                        template="plotly_dark",
                        height=140,
                        margin=dict(l=0, r=0, t=30, b=0),
                        paper_bgcolor="rgba(0,0,0,0)",
                        plot_bgcolor="rgba(0,0,0,0)",
                        showlegend=False,
                        title=dict(text=f"<b>{title}</b>", font=dict(size=13, color="#a5b4fc"), x=0.02),
                        xaxis=dict(showgrid=False, showticklabels=True, color="#6366f1",
                                   tickfont=dict(size=9)),
                        yaxis=dict(showgrid=False, showticklabels=False),
                    )
                    return fig

                with col_s1:
                    fig = make_sparkline(df['seq_len'], "Sequence Length", "#6366f1")
                    st.plotly_chart(fig, use_container_width=True, key="spark_len")
                with col_s2:
                    fig = make_sparkline(df['mean_disp_m0'], "Mean Displacement (Å)", "#10b981")
                    st.plotly_chart(fig, use_container_width=True, key="spark_mean")
                with col_s3:
                    fig = make_sparkline(df['max_disp_m0'], "Max Displacement (Å)", "#f59e0b")
                    st.plotly_chart(fig, use_container_width=True, key="spark_max")

        except Exception as e:
            st.warning(f"Could not load featured proteins: {e}")

# ── Auto-download if missing ──
if not status['has_weights']:
    st.divider()
    st.warning("⚠️ Model weights not found.")
    if st.button("⬇️ Download weights from Figshare (18 MB)", type="primary"):
        with st.spinner("Downloading..."):
            wt = ensure_weights(PETIMOT_ROOT)
        if wt:
            st.success(f"✅ Weights downloaded: {os.path.basename(wt)}")
            st.rerun()
        else:
            st.error("Download failed. Please manually download from "
                     "[Figshare](https://figshare.com/s/ab400d852b4669a83b64) "
                     "and place in `weights/`")

if not status['has_predictions'] and status['has_weights']:
    st.info("💡 No pre-computed predictions yet. Use the **Inference** page to predict "
            "individual proteins, or run batch inference from the Colab notebook.")