Upload app.py

app.py

@@ -1,333 +1,684 @@
 """
-Human-in-the-Loop Feedback Tool for Rabi Oscillation Classification.
-Production-ready Gradio interface with batch processing, AI model integration,
-dynamic graphing based on curve overlap, and MongoDB feedback storage.
+Multi-User Data Annotation Platform for Rabi Oscillation Quality Classification.
+3-Tab Gradio interface: Data Ingestion → Tagging Workspace → Analytics Dashboard.
+MongoDB queue with hash-based deduplication and multi-voter consensus.
 """
 import gradio as gr
 import torch
+import torch.nn.functional as torchF
 import numpy as np
 import plotly.graph_objects as go
-import json
-import os
+import json, os, hashlib
 from datetime import datetime, timezone
-from pymongo import MongoClient
+from pymongo import MongoClient, ASCENDING
+from scipy.optimize import curve_fit
 import certifi
 
+# ── AI Model imports (RabiEstimator for seed prediction) ─────────────────
 from ai_model import (
-    RabiEstimator, predict_ai_fit, load_model, rabi_formula, DEVICE
+    RabiEstimator, predict_ai_fit, load_model, rabi_formula,
+    preprocess_sample as ai_preprocess, target_to_params, DEVICE
 )
 
-# ═══════════════════════════════════════════════════════════════════════════
+# ── Classifier imports (RabiMultiTaskNet for confidence bars) ────────────
+from model import RabiMultiTaskNet
+from preprocessor import preprocess_sample as clf_preprocess, extract_fit_params
+from config import DEVICE as CLF_DEVICE, MODEL_PATH as CLF_MODEL_PATH
+
+# ═════════════════════════════════════════════════════════════════════════
 # Configuration
-# ═══════════════════════════════════════════════════════════════════════════
+# ═════════════════════════════════════════════════════════════════════════
 MONGO_URI = "mongodb+srv://Admin:zeyl3RK8oEu6Qhz3@cluster0.iwlxabj.mongodb.net/?appName=Cluster0"
 DB_NAME = "rabi_classifier"
-COLLECTION_NAME = "physicist_feedback"
-MODEL_WEIGHTS = "rabi_model_best.pt"
+COLLECTION_NAME = "experiments"
+ESTIMATOR_WEIGHTS = "rabi_model_best.pt"
+VOTES_REQUIRED = 2
+RATING_CHOICES = ["Bad (0)", "Borderline (1)", "Perfect (2)"]
+RATING_MAP = {"Bad (0)": 0, "Borderline (1)": 1, "Perfect (2)": 2}
 
-# ═══════════════════════════════════════════════════════════════════════════
+# ═════════════════════════════════════════════════════════════════════════
 # Global state
-# ═══════════════════════════════════════════════════════════════════════════
-AI_MODEL = None
-MONGO_COLLECTION = None
+# ═════════════════════════════════════════════════════════════════════════
+AI_MODEL = None          # RabiEstimator  – seed prediction
+CLF_MODEL = None         # RabiMultiTaskNet – confidence classification
+MONGO_COL = None
+MONGO_CLIENT = None
+
 
 def init_globals():
-    """Initialize AI model and MongoDB connection at startup."""
-    global AI_MODEL, MONGO_COLLECTION
+    global AI_MODEL, CLF_MODEL, MONGO_COL, MONGO_CLIENT
+
+    # 1) Load RabiEstimator (seed → curve_fit)
+    if os.path.exists(ESTIMATOR_WEIGHTS):
+        try:
+            AI_MODEL = load_model(ESTIMATOR_WEIGHTS)
+            print(f"✓ RabiEstimator loaded from {ESTIMATOR_WEIGHTS}")
+        except Exception as e:
+            print(f"⚠ RabiEstimator load failed: {e}")
+    else:
+        print(f"⚠ Estimator weights not found: {ESTIMATOR_WEIGHTS}")
 
-    # Load AI model
-    if os.path.exists(MODEL_WEIGHTS):
+    # 2) Load RabiMultiTaskNet (classifier for confidence bars)
+    if os.path.exists(CLF_MODEL_PATH):
         try:
-            AI_MODEL = load_model(MODEL_WEIGHTS)
-            print(f"✓ AI model loaded from {MODEL_WEIGHTS} (device: {DEVICE})")
+            CLF_MODEL = RabiMultiTaskNet().to(CLF_DEVICE)
+            ckpt = torch.load(CLF_MODEL_PATH, map_location=CLF_DEVICE, weights_only=False)
+            if isinstance(ckpt, dict) and 'model_state_dict' in ckpt:
+                CLF_MODEL.load_state_dict(ckpt['model_state_dict'])
+            else:
+                CLF_MODEL.load_state_dict(ckpt)
+            CLF_MODEL.eval()
+            print(f"✓ Classifier loaded from {CLF_MODEL_PATH}")
         except Exception as e:
-            print(f"⚠ Failed to load AI model: {e}")
+            CLF_MODEL = None
+            print(f"⚠ Classifier load failed: {e}")
     else:
-        print(f"⚠ Model weights not found at {MODEL_WEIGHTS}")
+        print(f"⚠ Classifier weights not found: {CLF_MODEL_PATH}")
+
+    # 3) MongoDB
+    try:
+        MONGO_CLIENT = MongoClient(MONGO_URI, serverSelectionTimeoutMS=3000,
+                                   tlsCAFile=certifi.where())
+        MONGO_CLIENT.admin.command('ping')
+    except Exception:
+        # Fallback: skip cert verification (common on macOS Python)
+        try:
+            MONGO_CLIENT = MongoClient(MONGO_URI, serverSelectionTimeoutMS=3000,
+                                       tls=True, tlsAllowInvalidCertificates=True)
+            MONGO_CLIENT.admin.command('ping')
+        except Exception as e:
+            MONGO_COL = None
+            print(f"⚠ MongoDB unavailable: {e}")
+            return
+
+    db = MONGO_CLIENT[DB_NAME]
+    MONGO_COL = db[COLLECTION_NAME]
+    MONGO_COL.create_index([("status", ASCENDING), ("num_votes", ASCENDING)])
+    print(f"✓ MongoDB connected: {DB_NAME}.{COLLECTION_NAME}")
+
+
+# ═════════════════════════════════════════════════════════════════════════
+# Hashing & Deduplication
+# ═════════════════════════════════════════════════════════════════════════
+
+def compute_data_hash(data: dict) -> str:
+    """MD5 hash of measured_data arrays for content-based dedup."""
+    try:
+        md = data.get('measured_data', data)
+        x = md.get('x_values', [])
+        y = md.get('y_values', [])
+    except AttributeError:
+        x, y = [], []
+    payload = json.dumps({'x': x, 'y': y}, sort_keys=True)
+    return hashlib.md5(payload.encode()).hexdigest()
+
+
+# ═════════════════════════════════════════════════════════════════════════
+# AI Seed → curve_fit Refinement
+# ═════════════════════════════════════════════════════════════════════════
+
+def refine_ai_seed(x, y, seed_params):
+    """Take AI seed [A,T,phi,C] and refine via scipy curve_fit. Fallback to seed."""
+    p0 = [seed_params['amplitude'], seed_params['T'],
+          seed_params['phase'], seed_params['offset']]
+    try:
+        popt, _ = curve_fit(
+            rabi_formula, x, y, p0=p0,
+            bounds=([0, 0.001, -np.pi*2, -1], [1, 5, np.pi*2, 2]),
+            maxfev=5000
+        )
+        return {'amplitude': popt[0], 'T': popt[1], 'phase': popt[2], 'offset': popt[3]}
+    except Exception:
+        return seed_params
+
+
+def get_ai_params(x, y):
+    """Full pipeline: AI seed → curve_fit refinement. Returns dict or None."""
+    if AI_MODEL is None:
+        return None
+    try:
+        _, seed = predict_ai_fit(x, y, AI_MODEL)
+        return refine_ai_seed(x, y, seed)
+    except Exception:
+        return None
+
+
+# ═════════════════════════════════════════════════════════════════════════
+# Classifier Confidence
+# ═════════════════════════════════════════════════════════════════════════
 
-    # Connect to MongoDB (non-blocking — app works without it)
+def get_classifier_probs(data: dict):
+    """Run RabiMultiTaskNet and return (data_probs, fit_probs) as numpy arrays of shape (3,)."""
+    if CLF_MODEL is None:
+        return None, None
     try:
-        client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=3000,
-                            tlsCAFile=certifi.where())
-        client.admin.command('ping')
-        db = client[DB_NAME]
-        MONGO_COLLECTION = db[COLLECTION_NAME]
-        print(f"✓ MongoDB connected: {DB_NAME}.{COLLECTION_NAME}")
+        signal_3ch, params_t, _, _ = clf_preprocess(data)
+        signal_in = signal_3ch.unsqueeze(0).to(CLF_DEVICE)
+        params_in = params_t.unsqueeze(0).to(CLF_DEVICE)
+        with torch.no_grad():
+            data_logits, fit_logits = CLF_MODEL(signal_in, params_in)
+        data_p = torchF.softmax(data_logits, dim=1).cpu().numpy()[0]
+        fit_p = torchF.softmax(fit_logits, dim=1).cpu().numpy()[0]
+        return data_p, fit_p
     except Exception as e:
-        MONGO_COLLECTION = None
-        print(f"⚠ MongoDB not available (will retry on submit): {e}")
-
-
-# ═══════════════════════════════════════════════════════════════════════════
-# JSON helpers
-# ═══════════════════════════════════════════════════════════════════════════
-
-def extract_fit_params(data):
-    """Extract the 4 fit parameters from JSON fitted_data."""
-    defaults = {'amplitude': 0.0, 'T': 0.05, 'phase': 0.0, 'offset': 0.0}
-    fitted = data.get('fitted_data')
-    if fitted is None:
-        return defaults
-    params_list = fitted.get('parameters')
-    if params_list is None:
-        return defaults
-    params = dict(defaults)
-    for p in params_list:
-        if p.get('name') in params:
-            params[p['name']] = float(p.get('value', 0.0))
-    return params
-
-
-# ═══════════════════════════════════════════════════════════════════════════
+        print(f"Classifier error: {e}")
+        return None, None
+
+
+def make_confidence_html(data_probs, fit_probs):
+    """Build HTML progress bars for classifier confidence."""
+    labels = ["Bad", "Borderline", "Perfect"]
+    colors = ["#ef4444", "#f59e0b", "#22c55e"]
+
+    def bar_row(name, probs):
+        if probs is None:
+            return f"<p style='color:#888;'>{name}: No classifier available</p>"
+        best = int(np.argmax(probs))
+        rows = f"<p style='font-weight:600; margin:8px 0 4px;'>{name}</p>"
+        for i, (lbl, col) in enumerate(zip(labels, colors)):
+            pct = probs[i] * 100
+            bold = "font-weight:700;" if i == best else ""
+            rows += (
+                f"<div style='display:flex;align-items:center;margin:2px 0;'>"
+                f"<span style='width:80px;font-size:13px;{bold}'>{lbl}</span>"
+                f"<div style='flex:1;background:#e5e7eb;border-radius:6px;height:18px;overflow:hidden;'>"
+                f"<div style='width:{pct:.1f}%;height:100%;background:{col};border-radius:6px;'></div></div>"
+                f"<span style='width:50px;text-align:right;font-size:13px;{bold}'>{pct:.1f}%</span></div>"
+            )
+        return rows
+
+    return (
+        "<div style='padding:12px;border:1px solid #e5e7eb;border-radius:10px;background:#fafafa;'>"
+        "<p style='font-weight:700;font-size:15px;margin:0 0 6px;'>🤖 Model Confidence</p>"
+        + bar_row("📊 Data Quality", data_probs)
+        + bar_row("📐 Fit Quality", fit_probs)
+        + "</div>"
+    )
+
+
+# ═════════════════════════════════════════════════════════════════════════
 # Plotting
-# ═══════════════════════════════════════════════════════════════════════════
+# ═════════════════════════════════════════════════════════════════════════
 
-def create_plot(data, filename, regular_params, ai_params, curves_overlap):
-    """Create Plotly plot with raw data, Regular Fit (blue), AI Fit (green)."""
+def create_plot(data, title_label, regular_params, ai_params, overlap):
     x = np.array(data['measured_data']['x_values'])
     y = np.array(data['measured_data']['y_values'])
-    x_dense = np.linspace(x.min(), x.max(), 500)
+    xd = np.linspace(x.min(), x.max(), 500)
 
-    y_regular = rabi_formula(
-        x_dense, regular_params['amplitude'], regular_params['T'],
-        regular_params['phase'], regular_params['offset']
-    )
+    yr = rabi_formula(xd, regular_params['amplitude'], regular_params['T'],
+                      regular_params['phase'], regular_params['offset'])
 
     fig = go.Figure()
+    fig.add_trace(go.Scatter(x=x, y=y, mode='markers', name='Raw Data',
+                             marker=dict(color='rgba(52,152,219,0.6)', size=5)))
+    fig.add_trace(go.Scatter(x=xd, y=yr, mode='lines', name='Regular Fit',
+                             line=dict(color='#2980b9', width=2.5)))
 
-    # Raw data points
-    fig.add_trace(go.Scatter(
-        x=x, y=y, mode='markers', name='Raw Data',
-        marker=dict(color='rgba(52, 152, 219, 0.7)', size=6),
-    ))
-
-    # Regular Fit — Blue
-    fig.add_trace(go.Scatter(
-        x=x_dense, y=y_regular, mode='lines', name='Regular Fit (JSON)',
-        line=dict(color='#2980b9', width=2.5),
-    ))
-
-    # AI Fit — Green
     if ai_params is not None:
-        y_ai = rabi_formula(
-            x_dense, ai_params['amplitude'], ai_params['T'],
-            ai_params['phase'], ai_params['offset']
-        )
-        if curves_overlap:
-            fig.add_trace(go.Scatter(
-                x=x_dense, y=y_ai, mode='lines',
-                name='AI Fit (overlapping)',
-                line=dict(color='#27ae60', width=3.5, dash='dot'),
-            ))
+        ya = rabi_formula(xd, ai_params['amplitude'], ai_params['T'],
+                          ai_params['phase'], ai_params['offset'])
+        if overlap:
+            fig.add_trace(go.Scatter(x=xd, y=ya, mode='lines', name='AI Fit (overlap)',
+                                     line=dict(color='#27ae60', width=3, dash='dot')))
         else:
-            fig.add_trace(go.Scatter(
-                x=x_dense, y=y_ai, mode='lines', name='AI Fit',
-                line=dict(color='#27ae60', width=2.5),
-            ))
-
-    title = f"Rabi Oscillation — {filename}"
-    if curves_overlap:
-        title += " (OVERLAP)"
+            fig.add_trace(go.Scatter(x=xd, y=ya, mode='lines', name='AI Fit',
+                                     line=dict(color='#27ae60', width=2.5)))
 
+    tag = " ✅ OVERLAP" if overlap else ""
     fig.update_layout(
-        title=title,
-        xaxis_title='Drive Amplitude (a.u.)',
-        yaxis_title='Signal (a.u.)',
-        height=550,
-        margin=dict(l=50, r=30, t=60, b=50),
-        legend=dict(x=0.01, y=0.99),
-        font=dict(size=14),
+        title=f"Rabi — {title_label}{tag}",
+        xaxis_title='Drive Amplitude (a.u.)', yaxis_title='Signal (a.u.)',
+        height=500, margin=dict(l=50, r=30, t=50, b=50),
+        legend=dict(x=0.01, y=0.99), font=dict(size=13),
     )
     return fig
 
 
-# ═══════════════════════════════════════════════════════════════════════════
-# Core processing
-# ═══════════════════════════════════════════════════════════════════════════
+# ═════════════════════════════════════════════════════════════════════════
+# Tab 1 — Data Ingestion
+# ═════════════════════════════════════════════════════════════════════════
 
-def process_experiment(data, filename):
-    x = np.array(data['measured_data']['x_values'])
-    y = np.array(data['measured_data']['y_values'])
+def ingest_files(files):
+    if MONGO_COL is None:
+        return "❌ MongoDB is not connected. Cannot ingest data."
+    if not files:
+        return "⚠ No files uploaded."
 
-    regular_params = extract_fit_params(data)
-    ai_params = None
-    overlap = True
-    model_status_msg = "Awaiting model..."
-
-    if AI_MODEL is not None:
+    added, duped, errors = 0, 0, 0
+    for f in files:
         try:
-            _, ai_params = predict_ai_fit(x, y, AI_MODEL)
-            y_reg = rabi_formula(x, regular_params['amplitude'], regular_params['T'],
-                                 regular_params['phase'], regular_params['offset'])
-            y_ai = rabi_formula(x, ai_params['amplitude'], ai_params['T'],
-                                ai_params['phase'], ai_params['offset'])
-            overlap = np.allclose(y_reg, y_ai, rtol=1e-3, atol=1e-3)
-            model_status_msg = "✅ AI Prediction successful."
-        except Exception as e:
-            model_status_msg = f"❌ AI prediction failed: {e}"
-            ai_params = None
+            fpath = f.name if hasattr(f, 'name') else str(f)
+            with open(fpath, 'r') as fp:
+                data = json.load(fp)
+
+            x = np.array(data['measured_data']['x_values'])
+            y = np.array(data['measured_data']['y_values'])
+            doc_hash = compute_data_hash(data)
+
+            # Compute fits at ingestion time
+            reg_params = extract_fit_params(data)
+            ai_params = get_ai_params(x, y)
+
+            # Overlap detection
             overlap = True
-    else:
-        model_status_msg = "⚠ No AI model loaded. Only regular fit shown."
-
-    fig = create_plot(data, filename, regular_params, ai_params, overlap)
-    return fig, overlap, regular_params, ai_params, model_status_msg
-
-def display_experiment(batch_data, idx):
-    if not batch_data or idx < 0 or idx >= len(batch_data):
-        return (
-            None, "📂 Upload JSON files to begin", "", "Waiting for data...",
-            gr.update(visible=True), gr.update(visible=False),
-            None, "", None, None, None, "", "", "",
-            batch_data, idx, False
-        )
+            if ai_params is not None:
+                yr = rabi_formula(x, reg_params['amplitude'], reg_params['T'],
+                                  reg_params['phase'], reg_params['offset'])
+                ya = rabi_formula(x, ai_params['amplitude'], ai_params['T'],
+                                  ai_params['phase'], ai_params['offset'])
+                overlap = bool(np.allclose(yr, ya, rtol=1e-3, atol=1e-3))
+
+            doc = {
+                'raw_data': data,
+                'filename': os.path.basename(fpath),
+                'regular_fit_params': reg_params,
+                'ai_fit_params': ai_params,
+                'curves_overlap': overlap,
+                'status': 'pending',
+                'num_votes': 0,
+                'tagged_by': [],
+                'evaluations': [],
+                'inserted_at': datetime.now(timezone.utc).isoformat(),
+            }
+
+            result = MONGO_COL.update_one(
+                {'_id': doc_hash},
+                {'$setOnInsert': doc},
+                upsert=True
+            )
+            if result.upserted_id is not None:
+                added += 1
+            else:
+                duped += 1
+        except Exception as e:
+            errors += 1
+            print(f"Ingest error: {e}")
+
+    parts = []
+    if added: parts.append(f"**{added}** new experiments added")
+    if duped: parts.append(f"**{duped}** duplicates skipped")
+    if errors: parts.append(f"**{errors}** errors")
+    return "### Ingestion Complete\n" + " · ".join(parts)
 
-    exp = batch_data[idx]
-    fig, overlap, reg_p, ai_p, mod_stat = process_experiment(exp['data'], exp['filename'])
-    progress = f"Experiment {idx + 1} / {len(batch_data)} — {exp['filename']}"
+
+# ═════════════════════════════════════════════════════════════════════════
+# Tab 2 — Tagging Workspace
+# ═════════════════════════════════════════════════════════════════════════
+
+def fetch_next(username):
+    """Fetch the next pending experiment that this user hasn't tagged yet."""
+    if not username or not username.strip():
+        return _empty_workspace("⚠ Please enter your username first.")
+    username = username.strip()
+
+    if MONGO_COL is None:
+        return _empty_workspace("❌ MongoDB not connected.")
+
+    doc = MONGO_COL.find_one(
+        {'status': 'pending', 'tagged_by': {'$nin': [username]}},
+        sort=[('num_votes', ASCENDING)]
+    )
+
+    if doc is None:
+        return _empty_workspace("🎉 **Queue empty!** No pending experiments for you.")
+
+    data = doc['raw_data']
+    reg_p = doc.get('regular_fit_params') or extract_fit_params(data)
+    ai_p = doc.get('ai_fit_params')
+    overlap = doc.get('curves_overlap', True)
+    doc_id = doc['_id']
+    fname = doc.get('filename', doc_id[:12])
+
+    fig = create_plot(data, fname, reg_p, ai_p, overlap)
+    data_probs, fit_probs = get_classifier_probs(data)
+    conf_html = make_confidence_html(data_probs, fit_probs)
 
     if overlap:
-        overlap_info = "🔗 Curves overlap (Regular and AI solutions match). Rate once."
+        overlap_msg = "🔗 **Curves overlap** — fits converged to same solution. Rate once."
     else:
-        overlap_info = "↔️ Curves diverge (Regular = Blue, AI = Green). Rate each separately."
+        overlap_msg = "↔️ **Curves diverge** — Regular (blue) vs AI (green). Rate each."
+
+    info = f"**{fname}** · votes: {doc['num_votes']} · id: `{doc_id[:12]}…`"
 
     return (
-        fig, progress, overlap_info, mod_stat,
-        gr.update(visible=overlap),
-        gr.update(visible=not overlap),
-        None, "", None, None, None, "", "", "",
-        batch_data, idx, overlap
+        fig, info, overlap_msg, conf_html, "",
+        gr.update(visible=overlap), gr.update(visible=not overlap),
+        "Perfect (2)", "Perfect (2)", "Perfect (2)", "Perfect (2)",
+        "", "", "",
+        doc_id, overlap
     )
 
-def on_upload(files):
-    if not files: return display_experiment([], 0)
-    batch_data = []
-    for f in files:
-        try:
-            fpath = f.name if hasattr(f, 'name') else str(f)
-            with open(fpath, 'r') as fp: data = json.load(fp)
-            batch_data.append({'data': data, 'filename': os.path.basename(fpath)})
-        except: pass
-    return display_experiment(batch_data, 0)
-
-def on_prev(batch, idx): return display_experiment(batch, max(0, idx - 1))
-def on_next(batch, idx): return display_experiment(batch, min(len(batch) - 1, idx + 1) if batch else 0)
-
-def on_submit(batch, idx, overlap, data_rate, data_comment, fit_sing, fit_reg, fit_ai, c_sing, c_reg, c_ai):
-    if not batch: return ("⚠ No experiment.",) + display_experiment([], 0)
-    if data_rate is None: return ("⚠ Rate Data Quality.",) + display_experiment(batch, idx)
-    if overlap and fit_sing is None: return ("⚠ Rate Fit Quality.",) + display_experiment(batch, idx)
-    if not overlap and (fit_reg is None or fit_ai is None): return ("⚠ Rate both fits.",) + display_experiment(batch, idx)
-
-    exp = batch[idx]
-    x = np.array(exp['data']['measured_data']['x_values'])
-    y = np.array(exp['data']['measured_data']['y_values'])
-    ai_params = None
-    if AI_MODEL:
-        try: _, ai_params = predict_ai_fit(x, y, AI_MODEL)
-        except: pass
-
-    doc = {
-        'filename': exp['filename'],
-        'raw_data': exp['data'],
-        'ai_prediction': ai_params,
-        'curves_overlap': bool(overlap),
-        'data_quality_rating': data_rate,
+
+def _empty_workspace(msg):
+    return (
+        None, msg, "", "<p style='color:#888'>No data loaded.</p>", "",
+        gr.update(visible=True), gr.update(visible=False),
+        "Perfect (2)", "Perfect (2)", "Perfect (2)", "Perfect (2)",
+        "", "", "",
+        None, False
+    )
+
+
+def submit_feedback(username, doc_id, overlap,
+                    data_rate, data_comment,
+                    fit_single, fit_reg, fit_ai,
+                    comment_single, comment_reg, comment_ai):
+    """Validate, write evaluation to MongoDB, fetch next."""
+    if not username or not username.strip():
+        return ("⚠ Enter your username.",) + _empty_workspace("")
+    username = username.strip()
+
+    if doc_id is None:
+        return ("⚠ No experiment loaded. Fetch one first.",) + _empty_workspace("")
+
+    if MONGO_COL is None:
+        return ("❌ MongoDB not connected.",) + _empty_workspace("")
+
+    # Map ratings
+    dr = RATING_MAP.get(data_rate)
+    if dr is None:
+        return ("⚠ Please rate **Data Quality**.",) + _empty_workspace("")
+
+    evaluation = {
+        'username': username,
+        'data_quality': dr,
         'data_comment': data_comment or '',
         'timestamp': datetime.now(timezone.utc).isoformat(),
     }
+
     if overlap:
-        doc['fit_quality_rating'] = {'combined': fit_sing}
-        doc['fit_comments'] = {'combined': c_sing or ''}
+        fr = RATING_MAP.get(fit_single)
+        if fr is None:
+            return ("⚠ Please rate **Fit Quality**.",) + _empty_workspace("")
+        evaluation['fit_quality'] = {'combined': fr}
+        evaluation['fit_comment'] = {'combined': comment_single or ''}
+    else:
+        fr_reg = RATING_MAP.get(fit_reg)
+        fr_ai = RATING_MAP.get(fit_ai)
+        if fr_reg is None or fr_ai is None:
+            return ("⚠ Please rate **both** fits.",) + _empty_workspace("")
+        evaluation['fit_quality'] = {'regular_fit': fr_reg, 'ai_fit': fr_ai}
+        evaluation['fit_comment'] = {'regular_fit': comment_reg or '', 'ai_fit': comment_ai or ''}
+
+    try:
+        MONGO_COL.update_one(
+            {'_id': doc_id},
+            {
+                '$push': {'evaluations': evaluation},
+                '$inc': {'num_votes': 1},
+                '$addToSet': {'tagged_by': username},
+            }
+        )
+        # Check if we reached the vote threshold
+        updated = MONGO_COL.find_one({'_id': doc_id})
+        if updated and updated.get('num_votes', 0) >= VOTES_REQUIRED:
+            MONGO_COL.update_one({'_id': doc_id}, {'$set': {'status': 'completed'}})
+
+        status = "✅ Feedback saved!"
+    except Exception as e:
+        status = f"❌ Error: {e}"
+
+    # Fetch next
+    nxt = fetch_next(username)
+    return (status,) + nxt
+
+
+def mark_perfect_and_next(username, doc_id, overlap):
+    """Shortcut: submit all ratings as Perfect(2) and fetch next."""
+    return submit_feedback(
+        username, doc_id, overlap,
+        "Perfect (2)", "",
+        "Perfect (2)", "Perfect (2)", "Perfect (2)",
+        "", "", ""
+    )
+
+
+# ═════════════════════════════════════════════════════════════════════════
+# Tab 3 — Analytics Dashboard
+# ═════════════════════════════════════════════════════════════════════════
+
+def refresh_analytics():
+    if MONGO_COL is None:
+        return "❌ MongoDB not connected."
+
+    total = MONGO_COL.count_documents({})
+    pending = MONGO_COL.count_documents({'status': 'pending'})
+    completed = MONGO_COL.count_documents({'status': 'completed'})
+
+    # Model success rate: % of evaluations with data_quality == 2
+    pipeline = [
+        {'$unwind': '$evaluations'},
+        {'$group': {
+            '_id': None,
+            'total_evals': {'$sum': 1},
+            'perfect_data': {'$sum': {'$cond': [{'$eq': ['$evaluations.data_quality', 2]}, 1, 0]}},
+        }}
+    ]
+    agg = list(MONGO_COL.aggregate(pipeline))
+
+    if agg:
+        total_evals = agg[0]['total_evals']
+        perfect_data = agg[0]['perfect_data']
+        success_pct = (perfect_data / total_evals * 100) if total_evals > 0 else 0
     else:
-        doc['fit_quality_rating'] = {'regular_fit': fit_reg, 'ai_fit': fit_ai}
-        doc['fit_comments'] = {'regular_fit': c_reg or '', 'ai_fit': c_ai or ''}
+        total_evals, perfect_data, success_pct = 0, 0, 0
 
-    status = "Feedback saved!"
-    if MONGO_COLLECTION is not None:
-        try: MONGO_COLLECTION.insert_one(doc)
-        except Exception as e: status = f"MongoDB error: {e}"
-    else: status = "⚠ MongoDB disconnected. Not saved."
+    # Top taggers
+    tagger_pipeline = [
+        {'$unwind': '$evaluations'},
+        {'$group': {'_id': '$evaluations.username', 'count': {'$sum': 1}}},
+        {'$sort': {'count': -1}},
+        {'$limit': 10}
+    ]
+    taggers = list(MONGO_COL.aggregate(tagger_pipeline))
+    tagger_rows = "\n".join([f"| {t['_id']} | {t['count']} |" for t in taggers]) if taggers else "| — | — |"
 
-    nxt = idx + 1 if idx < len(batch) - 1 else idx
-    status += " Loaded next." if idx < len(batch) - 1 else " Last experiment!"
-    return (status,) + display_experiment(batch, nxt)
+    md = f"""### 📈 Platform Statistics
 
+| Metric | Count |
+|--------|-------|
+| Total Experiments | **{total}** |
+| Pending | **{pending}** |
+| Completed (≥{VOTES_REQUIRED} votes) | **{completed}** |
+| Total Evaluations | **{total_evals}** |
 
-RATING_CHOICES = ["Bad (0)", "Borderline (1)", "Perfect (2)"]
+---
+
+### 🎯 Model Success Rate
+**{success_pct:.1f}%** of evaluations rated Data Quality as Perfect (Class 2)
+({perfect_data} / {total_evals})
+
+---
+
+### 🏆 Top Annotators
+
+| Username | Evaluations |
+|----------|-------------|
+{tagger_rows}
+"""
+    return md
+
+
+# ═════════════════════════════════════════════════════════════════════════
+# Custom CSS
+# ═════════════════════════════════════════════════════════════════════════
+
+CUSTOM_CSS = """
+.gradio-container { max-width: 1100px !important; margin: 0 auto !important; }
+
+/* Rating radio colors */
+.rating-radio label { padding: 10px 20px !important; border-radius: 10px !important;
+    margin: 3px 6px !important; font-weight: 600 !important; font-size: 14px !important;
+    transition: all 0.2s !important; cursor: pointer !important; }
+
+/* Bad = Red */
+.rating-radio label:nth-child(1) {
+    border: 2px solid #ef4444 !important; color: #dc2626 !important;
+    background: rgba(239,68,68,0.06) !important; }
+.rating-radio label:nth-child(1).selected,
+.rating-radio label:nth-child(1):has(input:checked) {
+    background: #ef4444 !important; color: #fff !important; }
+
+/* Borderline = Orange */
+.rating-radio label:nth-child(2) {
+    border: 2px solid #f59e0b !important; color: #d97706 !important;
+    background: rgba(245,158,11,0.06) !important; }
+.rating-radio label:nth-child(2).selected,
+.rating-radio label:nth-child(2):has(input:checked) {
+    background: #f59e0b !important; color: #fff !important; }
+
+/* Perfect = Green */
+.rating-radio label:nth-child(3) {
+    border: 2px solid #22c55e !important; color: #16a34a !important;
+    background: rgba(34,197,94,0.06) !important; }
+.rating-radio label:nth-child(3).selected,
+.rating-radio label:nth-child(3):has(input:checked) {
+    background: #22c55e !important; color: #fff !important; }
+
+/* Submit button */
+.submit-btn { font-size: 16px !important; padding: 14px 28px !important;
+    border-radius: 12px !important; }
+.perfect-btn { font-size: 14px !important; padding: 14px 28px !important;
+    border-radius: 12px !important; background: #22c55e !important;
+    color: white !important; border: none !important; }
+"""
+
+
+# ═════════════════════════════════════════════════════════════════════════
+# Build UI
+# ═════════════════════════════════════════════════════════════════════════
 
 def build_ui():
     init_globals()
-    with gr.Blocks(title="Rabi Feedback") as demo:
-        batch_data = gr.State([])
-        batch_index = gr.State(0)
-        overlap_state = gr.State(False)
-
-        gr.Markdown("# ⚛️ Rabi Oscillation — Physicist Feedback Tool")
-        gr.Markdown("Upload JSON experiments, review the quality of the data and fits, and submit your ratings.")
-        
-        with gr.Row():
-            file_upload = gr.File(label="Upload JSON Files", file_count="multiple", file_types=[".json"])
-        
-        with gr.Row():
-            prev_btn = gr.Button("← Previous")
-            progress_md = gr.Markdown("📂 Upload JSON files to begin")
-            next_btn = gr.Button("Next →")
-
-        plot = gr.Plot()
-        
-        with gr.Row():
-            overlap_info = gr.Markdown("")
-            model_status = gr.Markdown("Awaiting model...")
-
-        gr.Markdown("---")
-        
-        with gr.Row():
-            with gr.Column():
-                gr.Markdown("### 1. Data Quality")
-                data_rating = gr.Radio(choices=RATING_CHOICES, label="Rate raw data visually:")
-                comment_data = gr.Textbox(label="Data notes", lines=2)
-            
-            with gr.Column(visible=True) as fit_single_col:
-                gr.Markdown("### 2. Fit Quality (Combined)")
-                fit_single_rating = gr.Radio(choices=RATING_CHOICES, label="Rate the unified fit:")
-                comment_single = gr.Textbox(label="Fit notes", lines=2)
-                
-        with gr.Column(visible=False) as fit_dual_col:
-            gr.Markdown("### 2. Fit Quality (Diverging)")
-            with gr.Row():
-                with gr.Column():
-                    fit_regular_rating = gr.Radio(choices=RATING_CHOICES, label="Regular Fit (Blue):")
-                    comment_regular = gr.Textbox(label="Notes on blue curve", lines=2)
-                with gr.Column():
-                    fit_ai_rating = gr.Radio(choices=RATING_CHOICES, label="AI Fit (Green):")
-                    comment_ai = gr.Textbox(label="Notes on green curve", lines=2)
-
-        gr.Markdown("---")
-        submit_btn = gr.Button("Submit Feedback & Next", variant="primary")
-        status_md = gr.Markdown("")
-
-        outputs = [
-            plot, progress_md, overlap_info, model_status,
-            fit_single_col, fit_dual_col,
-            data_rating, comment_data,
-            fit_single_rating, fit_regular_rating, fit_ai_rating,
-            comment_single, comment_regular, comment_ai,
-            batch_data, batch_index, overlap_state
-        ]
-
-        file_upload.change(on_upload, [file_upload], outputs)
-        prev_btn.click(on_prev, [batch_data, batch_index], outputs)
-        next_btn.click(on_next, [batch_data, batch_index], outputs)
-        submit_btn.click(on_submit, [
-            batch_data, batch_index, overlap_state,
-            data_rating, comment_data,
-            fit_single_rating, fit_regular_rating, fit_ai_rating,
-            comment_single, comment_regular, comment_ai
-        ], [status_md] + outputs)
+
+    with gr.Blocks(title="Rabi Annotation Platform") as demo:
+
+        gr.Markdown("# ⚛️ Rabi Oscillation — Data Annotation Platform")
+
+        with gr.Tabs():
+
+            # ─────────────────────────────────────────────────────
+            # TAB 1: Data Ingestion
+            # ─────────────────────────────────────────────────────
+            with gr.Tab("📥 Data Ingestion"):
+                gr.Markdown("Upload JSON experiment files. Duplicates (identical measured data) are automatically skipped.")
+                file_upload = gr.File(
+                    label="Upload JSON Experiments",
+                    file_count="multiple", file_types=[".json"]
+                )
+                ingest_btn = gr.Button("Process & Upload to Database", variant="primary")
+                ingest_result = gr.Markdown("")
+
+                ingest_btn.click(ingest_files, [file_upload], [ingest_result])
+
+            # ─────────────────────────────────────────────────────
+            # TAB 2: Tagging Workspace
+            # ─────────────────────────────────────────────────────
+            with gr.Tab("🏷️ Tagging Workspace"):
+
+                # State
+                current_doc_id = gr.State(None)
+                overlap_state = gr.State(False)
+
+                # User ID + Fetch
+                with gr.Row():
+                    username_box = gr.Textbox(label="Your Username", placeholder="e.g. alice",
+                                             scale=2, max_lines=1)
+                    fetch_btn = gr.Button("🔄 Fetch Next Experiment", variant="primary", scale=1)
+
+                status_md = gr.Markdown("")
+
+                # Plot + Confidence side by side
+                with gr.Row():
+                    plot = gr.Plot(label="Signal & Fit Curves", scale=3)
+                    with gr.Column(scale=1):
+                        exp_info = gr.Markdown("")
+                        overlap_info = gr.Markdown("")
+                        confidence_html = gr.HTML("<p style='color:#888'>Fetch an experiment to see predictions.</p>")
+
+                gr.Markdown("---")
+
+                # Ratings
+                with gr.Row():
+                    # Data Quality — always visible
+                    with gr.Column():
+                        gr.Markdown("### 📊 Data Quality")
+                        data_rating = gr.Radio(choices=RATING_CHOICES, value="Perfect (2)",
+                                               label="Rate raw data:", elem_classes=["rating-radio"])
+                        comment_data = gr.Textbox(label="Data notes", lines=2)
+
+                    # Fit Quality — SINGLE (overlap)
+                    with gr.Column(visible=True) as fit_single_col:
+                        gr.Markdown("### 📐 Fit Quality (Combined)")
+                        fit_single_rating = gr.Radio(choices=RATING_CHOICES, value="Perfect (2)",
+                                                     label="Rate the fit:", elem_classes=["rating-radio"])
+                        comment_single = gr.Textbox(label="Fit notes", lines=2)
+
+                # Fit Quality — DUAL (diverge)
+                with gr.Column(visible=False) as fit_dual_col:
+                    gr.Markdown("### 📐 Fit Quality (Diverging)")
+                    with gr.Row():
+                        with gr.Column():
+                            gr.Markdown("#### 🟦 Regular Fit (Blue)")
+                            fit_regular_rating = gr.Radio(choices=RATING_CHOICES, value="Perfect (2)",
+                                                          label="Regular Fit:", elem_classes=["rating-radio"])
+                            comment_regular = gr.Textbox(label="Regular fit notes", lines=2)
+                        with gr.Column():
+                            gr.Markdown("#### 🟩 AI Fit (Green)")
+                            fit_ai_rating = gr.Radio(choices=RATING_CHOICES, value="Perfect (2)",
+                                                      label="AI Fit:", elem_classes=["rating-radio"])
+                            comment_ai = gr.Textbox(label="AI fit notes", lines=2)
+
+                gr.Markdown("---")
+
+                with gr.Row():
+                    submit_btn = gr.Button("🚀 Submit Feedback & Next", variant="primary",
+                                           elem_classes=["submit-btn"], scale=2)
+                    perfect_btn = gr.Button("✅ Mark as Perfect & Next",
+                                            elem_classes=["perfect-btn"], scale=1)
+
+                # ── Wire Tab 2 ───────────────────────────────────
+                fetch_outputs = [
+                    plot, exp_info, overlap_info, confidence_html, status_md,
+                    fit_single_col, fit_dual_col,
+                    data_rating, fit_single_rating, fit_regular_rating, fit_ai_rating,
+                    comment_data, comment_single, comment_ai,
+                    current_doc_id, overlap_state
+                ]
+
+                fetch_btn.click(fetch_next, [username_box], fetch_outputs)
+
+                submit_btn.click(
+                    submit_feedback,
+                    [username_box, current_doc_id, overlap_state,
+                     data_rating, comment_data,
+                     fit_single_rating, fit_regular_rating, fit_ai_rating,
+                     comment_single, comment_regular, comment_ai],
+                    [status_md] + fetch_outputs
+                )
+
+                perfect_btn.click(
+                    mark_perfect_and_next,
+                    [username_box, current_doc_id, overlap_state],
+                    [status_md] + fetch_outputs
+                )
+
+            # ─────────────────────────────────────────────────────
+            # TAB 3: Analytics Dashboard
+            # ─────────────────────────────────────────────────────
+            with gr.Tab("📊 Analytics"):
+                gr.Markdown("View aggregate statistics from all annotator submissions.")
+                refresh_btn = gr.Button("🔄 Refresh Statistics", variant="primary")
+                analytics_md = gr.Markdown("Click **Refresh** to load stats.")
+                refresh_btn.click(refresh_analytics, [], [analytics_md])
 
     return demo
 
+
+# ═════════════════════════════════════════════════════════════════════════
+# Launch
+# ═════════════════════════════════════════════════════════════════════════
+
 if __name__ == '__main__':
     demo = build_ui()
-    demo.launch()
+    demo.launch(css=CUSTOM_CSS)