app.py

"""
Multi-User Data Annotation Platform for Rabi Oscillation Quality Classification.
4-Tab Gradio 6: Ingestion → Tagging → Analytics → Conflict Resolution.
MongoDB queue with hash dedup, implicit agreement, consensus, and routing.
"""
import gradio as gr
import torch
import torch.nn.functional as torchF
import numpy as np
import plotly.graph_objects as go
import json, os, hashlib, traceback
from datetime import datetime, timezone
from pymongo import MongoClient, ASCENDING
from scipy.optimize import curve_fit
from scipy.signal import find_peaks
import certifi

from ai_model import (
    RabiEstimator, predict_ai_fit, load_model, rabi_formula,
    preprocess_sample as ai_preprocess, target_to_params, DEVICE
)
from model import RabiMultiTaskNet
from config import DEVICE as CLF_DEVICE, MODEL_PATH as CLF_MODEL_PATH, SEQ_LEN

# ═════════════════════════════════════════════════════════════════════════
# Config
# ═════════════════════════════════════════════════════════════════════════
MONGO_URI = os.environ['MONGO_URI']
DB_NAME = "rabi_classifier"
COL_NAME = "experiments"
ESTIMATOR_WEIGHTS = "rabi_model_best.pt"
VOTES_REQUIRED = 1
RATING_CHOICES = ["Bad (0)", "Borderline (1)", "Perfect (2)"]
RATING_MAP = {"Bad (0)": 0, "Borderline (1)": 1, "Perfect (2)": 2}
LABEL_NAMES = ["Bad", "Borderline", "Perfect"]
BAR_COLORS = ["#ef4444", "#f59e0b", "#22c55e"]

# ═════════════════════════════════════════════════════════════════════════
# Globals
# ═════════════════════════════════════════════════════════════════════════
AI_MODEL = None
CLF_MODEL = None
MONGO_COL = None
MONGO_DB = None


def init_globals():
    global AI_MODEL, CLF_MODEL, MONGO_COL, MONGO_DB

    if os.path.exists(ESTIMATOR_WEIGHTS):
        try:
            AI_MODEL = load_model(ESTIMATOR_WEIGHTS)
            print(f"✓ RabiEstimator loaded")
        except Exception as e:
            print(f"⚠ RabiEstimator: {e}")

    if os.path.exists(CLF_MODEL_PATH):
        try:
            CLF_MODEL = RabiMultiTaskNet().to(CLF_DEVICE)
            ckpt = torch.load(CLF_MODEL_PATH, map_location=CLF_DEVICE, weights_only=False)
            state = ckpt.get('model_state_dict', ckpt) if isinstance(ckpt, dict) else ckpt
            CLF_MODEL.load_state_dict(state)
            CLF_MODEL.eval()
            print(f"✓ Classifier loaded")
        except Exception as e:
            CLF_MODEL = None
            print(f"⚠ Classifier: {e}")

    client = None
    for kw in [dict(tlsCAFile=certifi.where()),
               dict(tls=True, tlsAllowInvalidCertificates=True)]:
        try:
            client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=3000, **kw)
            client.admin.command('ping')
            break
        except Exception:
            client = None
    if client:
        MONGO_DB = client[DB_NAME]
        MONGO_COL = MONGO_DB[COL_NAME]
        MONGO_COL.create_index([("status", ASCENDING), ("num_votes", ASCENDING)])
        print(f"✓ MongoDB: {DB_NAME}.{COL_NAME}")
    else:
        print("⚠ MongoDB unavailable")


# ═════════════════════════════════════════════════════════════════════════
# Hashing
# ═════════════════════════════════════════════════════════════════════════

def compute_hash(data):
    md = data.get('measured_data', {})
    payload = json.dumps({'x': md.get('x_values', []),
                          'y': md.get('y_values', [])}, sort_keys=True)
    return hashlib.md5(payload.encode()).hexdigest()


# ═════════════════════════════════════════════════════════════════════════
# AI Seed Refinement
# ═════════════════════════════════════════════════════════════════════════

def refine_ai_seed(x, y, seed):
    p0 = [seed['amplitude'], seed['T'], seed['phase'], seed['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 dict(amplitude=popt[0], T=popt[1], phase=popt[2], offset=popt[3])
    except Exception:
        return seed


def get_ai_params(x, y):
    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
# ═════════════════════════════════════════════════════════════════════════

def extract_fit_params(data):
    defaults = {'amplitude': 0.0, 'T': 0.05, 'phase': 0.0, 'offset': 0.0}
    fitted = data.get('fitted_data')
    if not fitted:
        return defaults
    plist = fitted.get('parameters')
    if not plist:
        return defaults
    p = dict(defaults)
    for item in plist:
        if item.get('name') in p:
            p[item['name']] = float(item.get('value', 0.0))
    return p


def classify_with_params(data, fit_params):
    if CLF_MODEL is None:
        return None, None
    try:
        x = np.array(data['measured_data']['x_values'], dtype=np.float64)
        y = np.array(data['measured_data']['y_values'], dtype=np.float64)
        xu = np.linspace(x.min(), x.max(), SEQ_LEN)
        yr = np.interp(xu, x, y)
        yf = rabi_formula(xu, fit_params['amplitude'], fit_params['T'],
                          fit_params['phase'], fit_params['offset'])
        ymin, ymax = yr.min(), yr.max()
        span = max(ymax - ymin, 1e-10)
        yr_n = ((yr - ymin) / span).astype(np.float32)
        yf_n = ((yf - ymin) / span).astype(np.float32)
        res = yr_n - yf_n
        rm = max(np.abs(res).max(), 1e-10)
        res_n = (res / rm).astype(np.float32)

        sig = torch.tensor(np.stack([yr_n, yf_n, res_n]), dtype=torch.float32)
        par = torch.tensor([fit_params['amplitude'], fit_params['T'],
                            fit_params['phase'], fit_params['offset']], dtype=torch.float32)
        with torch.no_grad():
            dl, fl = CLF_MODEL(sig.unsqueeze(0).to(CLF_DEVICE),
                               par.unsqueeze(0).to(CLF_DEVICE))
        return (torchF.softmax(dl, dim=1).cpu().numpy()[0].tolist(),
                torchF.softmax(fl, dim=1).cpu().numpy()[0].tolist())
    except Exception as e:
        print(f"Classifier err: {e}")
        return None, None


def compute_model_predictions(data, reg_params, ai_params, overlap):
    dp, rfp = classify_with_params(data, reg_params)
    afp = None
    if ai_params and not overlap:
        _, afp = classify_with_params(data, ai_params)
    preds = {'data_probs': dp, 'regular_fit_probs': rfp, 'ai_fit_probs': afp}
    classes = {}
    if dp:   classes['data_quality'] = int(np.argmax(dp))
    if rfp:  classes['regular_fit_quality'] = int(np.argmax(rfp))
    if afp:  classes['ai_fit_quality'] = int(np.argmax(afp))
    elif rfp: classes['ai_fit_quality'] = int(np.argmax(rfp))
    preds['model_classes'] = classes
    return preds


# ═════════════════════════════════════════════════════════════════════════
# Confidence HTML
# ═════════════════════════════════════════════════════════════════════════

def _bar_section(title, probs):
    if probs is None:
        return f"<p style='color:var(--body-text-color-subdued);margin:4px 0;'>{title}: N/A</p>"
    best = int(np.argmax(probs))
    s = f"<p style='font-weight:600;margin:10px 0 4px;color:var(--body-text-color);'>{title}</p>"
    for i in range(3):
        pct = probs[i] * 100
        bold = "font-weight:700;" if i == best else ""
        s += (f"<div style='display:flex;align-items:center;margin:2px 0;'>"
              f"<span style='width:85px;font-size:13px;color:var(--body-text-color);{bold}'>{LABEL_NAMES[i]}</span>"
              f"<div style='flex:1;background:var(--border-color-primary);border-radius:6px;height:18px;overflow:hidden;'>"
              f"<div style='width:{pct:.1f}%;height:100%;background:{BAR_COLORS[i]};border-radius:6px;'>"
              f"</div></div>"
              f"<span style='width:52px;text-align:right;font-size:13px;color:var(--body-text-color);{bold}'>{pct:.1f}%</span></div>")
    return s


def make_confidence_html(mp, overlap):
    if not mp:
        return "<p style='color:var(--body-text-color-subdued)'>No model predictions.</p>"
    h = "<p style='font-weight:700;font-size:15px;margin:0 0 4px;color:var(--body-text-color);'>🤖 Model Confidence</p>"
    h += _bar_section("📊 Data Quality", mp.get('data_probs'))
    if overlap:
        h += _bar_section("📐 Fit Quality", mp.get('regular_fit_probs'))
    else:
        h += _bar_section("📐 Regular Fit (Blue)", mp.get('regular_fit_probs'))
        h += _bar_section("📐 AI Fit (Green)", mp.get('ai_fit_probs'))
    return (f"<div style='padding:12px;border:1px solid var(--border-color-primary);"
            f"border-radius:10px;background:var(--background-fill-secondary);'>{h}</div>")


# ═════════════════════════════════════════════════════════════════════════
# Plotting
# ═════════════════════════════════════════════════════════════════════════

def create_plot(data, label, reg_p, ai_p, overlap, flip_reg=None, flip_ai=None):
    try:
        x = np.array(data['measured_data']['x_values'])
        y = np.array(data['measured_data']['y_values'])
        xd = np.linspace(x.min(), x.max(), 500)
        yr = rabi_formula(xd, reg_p['amplitude'], reg_p['T'], reg_p['phase'], reg_p['offset'])

        fig = go.Figure()
        fig.add_trace(go.Scatter(x=x, y=y, mode='markers', name='Raw Data',
                                 marker=dict(color='#e67e22', size=6, opacity=0.85)))
        fig.add_trace(go.Scatter(x=xd, y=yr, mode='lines', name='Regular Fit',
                                 line=dict(color='#2980b9', width=2.5)))
        if ai_p:
            ya = rabi_formula(xd, ai_p['amplitude'], ai_p['T'], ai_p['phase'], ai_p['offset'])
            st = dict(color='#27ae60', width=3.5, dash='dot') if overlap else dict(color='#27ae60', width=2.5)
            fig.add_trace(go.Scatter(x=xd, y=ya, mode='lines',
                                     name='AI Fit (overlap)' if overlap else 'AI Fit', line=st))

        # Flip-point markers
        if flip_reg:
            fig.add_trace(go.Scatter(
                x=[flip_reg['x']], y=[flip_reg['y']],
                mode='markers',
                name=f"Reg Flip ({flip_reg['type']}): x={flip_reg['x']:.4f}",
                marker=dict(color='#e74c3c', size=12, symbol='x'),
                showlegend=True))
            fig.add_vline(x=flip_reg['x'], line_dash='dash',
                          line_color='#e74c3c', line_width=2)
        if flip_ai and not overlap:
            fig.add_trace(go.Scatter(
                x=[flip_ai['x']], y=[flip_ai['y']],
                mode='markers',
                name=f"AI Flip ({flip_ai['type']}): x={flip_ai['x']:.4f}",
                marker=dict(color='#2ecc71', size=12, symbol='diamond'),
                showlegend=True))
            fig.add_vline(x=flip_ai['x'], line_dash='dot',
                          line_color='#2ecc71', line_width=2)

        tag = " ✅ OVERLAP" if overlap else ""
        fig.update_layout(title=dict(text=f"{label}{tag}", x=0.5, xanchor='center'),
                          xaxis_title='Drive Amplitude (a.u.)', yaxis_title='Signal (a.u.)',
                          autosize=True, height=600,
                          margin=dict(l=40, r=10, t=50, b=100),
                          legend=dict(orientation='h', x=0.5, y=-0.18,
                                      xanchor='center', yanchor='top',
                                      bgcolor='rgba(0,0,0,0)', borderwidth=0,
                                      font=dict(size=11, color='#e0e0e0')),
                          font=dict(size=13),
                          xaxis=dict(range=[x.min(), x.max()]),
                          yaxis=dict(autorange=True),
                          paper_bgcolor='rgba(0,0,0,0)', plot_bgcolor='rgba(0,0,0,0)')
        return fig
    except Exception as e:
        print(f"Plot error: {e}")
        return go.Figure().update_layout(title="Error creating plot", height=600)


def find_flip_point(data, fit_params):
    """Find the first physical extremum (peak or valley) on the SMOOTH fit curve.
    Filters out initial hardware micro-bumps using prominence and x-exclusion zone.
    Returns {'x': float, 'y': float, 'type': 'max'|'min'} or None.
    """
    try:
        x = np.array(data['measured_data']['x_values'])
        xd = np.linspace(x.min(), x.max(), 500)
        yf = rabi_formula(xd, fit_params['amplitude'], fit_params['T'],
                          fit_params['phase'], fit_params['offset'])

        # Prominence threshold: 10% of signal amplitude to reject micro-bumps
        y_range = yf.max() - yf.min()
        prom_thresh = y_range * 0.10 if y_range > 1e-10 else 0

        # Exclusion zone: skip extrema in the first 5% of the x-range
        x_range = xd[-1] - xd[0]
        x_min_cutoff = xd[0] + x_range * 0.05

        peaks, props_p = find_peaks(yf, prominence=prom_thresh)
        valleys, props_v = find_peaks(-yf, prominence=prom_thresh)

        candidates = []
        for idx in peaks:
            if xd[idx] >= x_min_cutoff:
                candidates.append((xd[idx], yf[idx], 'max'))
        for idx in valleys:
            if xd[idx] >= x_min_cutoff:
                candidates.append((xd[idx], yf[idx], 'min'))

        if not candidates:
            return None
        candidates.sort(key=lambda c: c[0])
        first = candidates[0]
        return {'x': float(first[0]), 'y': float(first[1]), 'type': first[2]}
    except Exception as e:
        print(f"find_flip_point error: {e}")
        return None


def _plot_fingerprint(flip_reg, flip_ai, overlap):
    """Create a hashable fingerprint of the plot's visual state.
    The plot only changes when flip-point markers or overlap status change;
    raw data and fit curves are constant for a given experiment."""
    fr = (flip_reg['x'], flip_reg['y'], flip_reg['type']) if flip_reg else None
    fa = (flip_ai['x'], flip_ai['y'], flip_ai['type']) if flip_ai else None
    return (fr, fa, overlap)



# ═════════════════════════════════════════════════════════════════════════
# Consensus, Routing & State Machine
# ═════════════════════════════════════════════════════════════════════════

def _extract_scores(ev):
    s = {'data': ev['data_quality']['score']}
    fq = ev['fit_quality']
    if 'combined' in fq:
        s['reg_fit'] = fq['combined']['score']
        s['ai_fit'] = fq['combined']['score']
    else:
        s['reg_fit'] = fq['regular_fit']['score']
        s['ai_fit'] = fq['ai_fit']['score']
    return s


def check_consensus_and_route(doc_id):
    """Single-labeler mode: use the last evaluation as the consensus."""
    doc = MONGO_COL.find_one({'_id': doc_id})
    if not doc or len(doc.get('evaluations', [])) < 1:
        return
    last_eval = doc['evaluations'][-1]
    scores = _extract_scores(last_eval)
    consensus = {'data_quality': scores['data'],
                 'regular_fit_quality': scores['reg_fit'],
                 'ai_fit_quality': scores['ai_fit']}
    MONGO_COL.update_one({'_id': doc_id},
        {'$set': {'status': 'completed', 'final_consensus': consensus}})
    route_completed(doc_id)


def _max_deviation(consensus, model_classes):
    """Calculate the maximum abs difference between human consensus and model prediction."""
    max_diff = 0
    for key in ('data_quality', 'regular_fit_quality', 'ai_fit_quality'):
        h = consensus.get(key)
        m = model_classes.get(key)
        if h is not None and m is not None:
            max_diff = max(max_diff, abs(h - m))
    return max_diff


def route_completed(doc_id):
    """Route completed doc to granular collections based on human-model deviation."""
    if MONGO_DB is None:
        return
    doc = MONGO_COL.find_one({'_id': doc_id})
    if not doc:
        return
    consensus = doc.get('final_consensus', {})
    mc = (doc.get('model_predictions') or {}).get('model_classes', {})
    overlap = doc.get('curves_overlap', True)
    routed = {k: v for k, v in doc.items() if k != '_id'}
    routed['original_id'] = doc['_id']

    def _route(collection):
        try:
            MONGO_DB[collection].update_one(
                {'original_id': doc['_id']}, {'$set': routed}, upsert=True)
        except Exception as e:
            print(f"Route {collection} err: {e}")

    # --- Classification routing (granular) ---
    dev = _max_deviation(consensus, mc)
    if dev == 0:
        _route('experiments_clf_success')
    elif dev >= 2:
        _route('experiments_clf_catastrophic_failure')
    else:
        _route('experiments_clf_failed')

    # --- Fit-level routing (diverging fits) ---
    if not overlap:
        hd = consensus.get('data_quality')
        hr = consensus.get('regular_fit_quality')
        ha = consensus.get('ai_fit_quality')
        data_ok = hd is not None and hd >= 1

        if data_ok:
            if hr is not None and hr <= 1:
                _route('experiments_reg_fit_failed')
            if ha is not None and ha <= 1:
                _route('experiments_ai_fit_failed')
            if hr is not None and hr <= 1 and ha is not None and ha <= 1:
                _route('experiments_both_fits_failed')
            if hr is not None and hr >= 1 and ha == 0:
                _route('experiments_seed_borderline_fail')

        if ha == 2 and hr is not None and hr <= 1:
            _route('experiments_seed_success_reg_failed')
        if hr == 2 and ha is not None and ha <= 1:
            _route('experiments_reg_success_seed_failed')

    # --- Flip function failure routing ---
    evals = doc.get('evaluations', [])
    reg_flip_fail = False
    ai_flip_fail = False
    for ev in evals:
        fp = ev.get('flip_point') or {}
        if fp.get('approved') is False:
            reg_flip_fail = True
        fpd = ev.get('flip_point_dual') or {}
        if fpd.get('reg', {}).get('approved') is False:
            reg_flip_fail = True
        if fpd.get('ai', {}).get('approved') is False:
            ai_flip_fail = True
    if reg_flip_fail:
        _route('experiments_flip_reg_failed')
    if ai_flip_fail:
        _route('experiments_flip_ai_failed')
    if reg_flip_fail and ai_flip_fail:
        _route('experiments_flip_both_failed')


# ═════════════════════════════════════════════════════════════════════════
# Tab 1 — Data Ingestion
# ═════════════════════════════════════════════════════════════════════════

def ingest_files(files):
    if MONGO_COL is None:
        return "❌ MongoDB not connected."
    if not files:
        return "⚠ No files uploaded."
    added, duped, errors = 0, 0, 0
    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)
            x = np.array(data['measured_data']['x_values'])
            y = np.array(data['measured_data']['y_values'])
            h = compute_hash(data)
            reg_p = extract_fit_params(data)
            ai_p = get_ai_params(x, y)
            overlap = True
            if ai_p:
                yr = rabi_formula(x, reg_p['amplitude'], reg_p['T'], reg_p['phase'], reg_p['offset'])
                ya = rabi_formula(x, ai_p['amplitude'], ai_p['T'], ai_p['phase'], ai_p['offset'])
                sig_range = max(np.ptp(yr), np.ptp(ya), 1e-15)
                max_diff = np.max(np.abs(yr - ya))
                overlap = bool(max_diff < sig_range * 0.05)
            mp = compute_model_predictions(data, reg_p, ai_p, overlap)
            doc = {
                'raw_data': data, 'filename': os.path.basename(fpath),
                'regular_fit_params': reg_p, 'ai_fit_params': ai_p,
                'curves_overlap': overlap, 'model_predictions': mp,
                'status': 'pending', 'num_votes': 0,
                'tagged_by': [], 'evaluations': [],
                'inserted_at': datetime.now(timezone.utc).isoformat(),
            }
            r = MONGO_COL.update_one({'_id': h}, {'$setOnInsert': doc}, upsert=True)
            added += 1 if r.upserted_id else 0
            duped += 0 if r.upserted_id else 1
        except Exception as e:
            errors += 1
            print(f"Ingest err: {e}\n{traceback.format_exc()}")
    parts = []
    if added: parts.append(f"**{added}** new")
    if duped: parts.append(f"**{duped}** duplicates skipped")
    if errors: parts.append(f"**{errors}** errors")
    return "### Ingestion Complete\n" + " · ".join(parts)


# ═════════════════════════════════════════════════════════════════════════
# Tab 2 — Tagging Workspace (29 outputs)
# ═════════════════════════════════════════════════════════════════════════
# Output order:
#  0-4:   status, plot, info, ov_msg, conf_html
#  5-6:   fit_single_col, fit_dual_col
#  7-10:  data_r, fit_s_r, fit_reg_r, fit_ai_r
#  11-14: c_data, c_single, c_reg, c_ai
#  15-17: doc_id, overlap, preds
#  18-21: flip_info, flip_tag, flip_col, flip_state  (single/overlap mode)
#  22:    exp_data
#  23-28: flip_reg_info, flip_reg_tag, flip_ai_info, flip_ai_tag,
#         flip_dual_col, flip_dual_state  (dual/diverging mode)
#  29-30: flip_reg_sub, flip_ai_sub
#  31:    plot_fp  (plot fingerprint state for skip optimization)

def _empty_ws(msg=""):
    return (msg,
            None, "", "", "<p style='color:var(--body-text-color-subdued)'>No data loaded.</p>",
            gr.Column(visible=True), gr.Column(visible=False),
            None, None, None, None,
            "", "", "", "",
            None, False, None,
            "", None, gr.Column(visible=False), None,
            None,
            "", None, "", None, gr.Column(visible=False), None,
            gr.Column(visible=True), gr.Column(visible=True),
            None)


def fetch_next(username):
    if not username or not username.strip():
        return _empty_ws("⚠ Enter your username first.")
    username = username.strip().title()
    if MONGO_COL is None:
        return _empty_ws("❌ MongoDB not connected.")

    try:
        # Smart queue: prioritize experiments flagged needs_second_opinion
        doc = MONGO_COL.find_one(
            {'status': 'pending', 'tagged_by': {'$nin': [username]}},
            sort=[('needs_second_opinion', -1), ('num_votes', -1)])
    except Exception as e:
        return _empty_ws(f"❌ DB query failed: {e}")

    if not doc:
        return _empty_ws("🎉 **Queue empty!** No pending experiments for you.")

    try:
        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)
        mp = doc.get('model_predictions')
        fname = doc.get('filename', str(doc['_id'])[:12])

        # Store experiment context for dynamic recalculation
        exp_data = {'raw_data': data, 'reg_params': reg_p,
                    'ai_params': ai_p, 'filename': fname}

        # Gating: flip-point when model says data+fit ≥ Borderline
        mc = (mp or {}).get('model_classes', {})
        model_data_ok = mc.get('data_quality', 0) >= 1
        model_fit_ok = mc.get('regular_fit_quality', 0) >= 1
        model_ai_fit_ok = mc.get('ai_fit_quality', 0) >= 1

        flip_reg = None
        flip_ai = None
        if model_data_ok:
            if model_fit_ok:
                flip_reg = find_flip_point(data, reg_p)
            if not overlap and ai_p and model_ai_fit_ok:
                flip_ai = find_flip_point(data, ai_p)

        fig = create_plot(data, fname, reg_p, ai_p, overlap,
                          flip_reg=flip_reg, flip_ai=flip_ai)
        conf = make_confidence_html(mp, overlap)
        ov_msg = ("🔗 **Curves overlap** — Rate fit once."
                  if overlap else
                  "↔️ **Curves diverge** — Rate Regular (Blue) and AI (Green) separately.")
        info = f"**{fname}** · votes: {doc['num_votes']}/{VOTES_REQUIRED} · `{str(doc['_id'])[:12]}…`"

        # Visibility
        show_fit = model_data_ok
        show_flip_single = flip_reg is not None and overlap and show_fit
        show_flip_dual = (not overlap) and show_fit and (
            flip_reg is not None or flip_ai is not None)
        # Fit-specific flip sub-section visibility (initial load: both visible)
        show_flip_reg_sub = True
        show_flip_ai_sub = True

        # Single flip info
        flip_info = ""
        if overlap and flip_reg:
            flip_info = f"Flip Point: x = {flip_reg['x']:.5f}, y = {flip_reg['y']:.4f} ({flip_reg['type']})"
        elif overlap and show_fit:
            flip_info = "No flip point detected in the fit curve."

        # Dual flip info
        flip_reg_info = ""
        flip_ai_info = ""
        if not overlap:
            if flip_reg:
                flip_reg_info = f"Reg Flip: x = {flip_reg['x']:.5f}, y = {flip_reg['y']:.4f} ({flip_reg['type']})"
            elif show_fit:
                flip_reg_info = "No flip point detected."
            if flip_ai:
                flip_ai_info = f"AI Flip: x = {flip_ai['x']:.5f}, y = {flip_ai['y']:.4f} ({flip_ai['type']})"
            elif show_fit:
                flip_ai_info = "No flip point detected."

        dual_state = {'reg': flip_reg, 'ai': flip_ai} if not overlap else None

        return (f"Loaded: {fname}",
                fig, info, ov_msg, conf,
                gr.Column(visible=overlap and show_fit),
                gr.Column(visible=(not overlap) and show_fit),
                None, None, None, None,
                "", "", "", "",
                doc['_id'], overlap, mp,
                flip_info, None, gr.Column(visible=show_flip_single),
                flip_reg if overlap else None,
                exp_data,
                flip_reg_info, None, flip_ai_info, None,
                gr.Column(visible=show_flip_dual), dual_state,
                gr.Column(visible=show_flip_reg_sub),
                gr.Column(visible=show_flip_ai_sub),
                None)
    except Exception as e:
        print(f"fetch_next render error: {e}\n{traceback.format_exc()}")
        return _empty_ws(f"❌ Render failed: {e}")


def _resolve_rating(user_val, model_preds, pred_key):
    if user_val is not None:
        score = RATING_MAP.get(user_val)
        if score is not None:
            return {'score': score, 'source': 'manual_correction'}
    if model_preds and 'model_classes' in model_preds:
        cls = model_preds['model_classes'].get(pred_key)
        if cls is not None:
            return {'score': int(cls), 'source': 'implicit_agreement'}
    return None


def submit_feedback(username, doc_id, overlap, model_preds,
                    data_r, c_data, fit_s, fit_reg, fit_ai, c_single, c_reg, c_ai,
                    flip_tag, flip_state,
                    flip_reg_tag, flip_ai_tag, flip_dual_state):
    if not username or not username.strip():
        return _empty_ws("⚠ Enter your username.")
    username = username.strip().title()
    if doc_id is None:
        return _empty_ws("⚠ No experiment loaded. Fetch first.")
    if MONGO_COL is None:
        return _empty_ws("❌ MongoDB not connected.")

    dq = _resolve_rating(data_r, model_preds, 'data_quality')
    if dq is None:
        return ("⚠ Rate **Data Quality** — no model prediction available.",) + _empty_ws()[1:]

    eval_doc = {
        'username': username, 'data_quality': dq,
        'data_comment': c_data or '',
        'timestamp': datetime.now(timezone.utc).isoformat(),
    }

    if overlap:
        fq = _resolve_rating(fit_s, model_preds, 'regular_fit_quality')
        if fq is None:
            return ("⚠ Rate **Fit Quality** — no model prediction.",) + _empty_ws()[1:]
        eval_doc['fit_quality'] = {'combined': fq}
        eval_doc['fit_comment'] = {'combined': c_single or ''}
    else:
        fqr = _resolve_rating(fit_reg, model_preds, 'regular_fit_quality')
        fqa = _resolve_rating(fit_ai, model_preds, 'ai_fit_quality')
        if fqr is None or fqa is None:
            return ("⚠ Rate both fits — no model prediction.",) + _empty_ws()[1:]
        eval_doc['fit_quality'] = {'regular_fit': fqr, 'ai_fit': fqa}
        eval_doc['fit_comment'] = {'regular_fit': c_reg or '', 'ai_fit': c_ai or ''}

    # Flip-point tagging
    if overlap:
        if flip_state and flip_tag:
            approved = flip_tag == '✅ Correct'
            eval_doc['flip_point'] = {'detected': flip_state, 'approved': approved}
        else:
            eval_doc['flip_point'] = None
        eval_doc['flip_point_dual'] = None
    else:
        eval_doc['flip_point'] = None
        dual_fp = {}
        if flip_dual_state:
            freg = flip_dual_state.get('reg')
            fai = flip_dual_state.get('ai')
            if freg and flip_reg_tag:
                dual_fp['reg'] = {'detected': freg,
                                  'approved': flip_reg_tag == '✅ Correct'}
            if fai and flip_ai_tag:
                dual_fp['ai'] = {'detected': fai,
                                 'approved': flip_ai_tag == '✅ Correct'}
        eval_doc['flip_point_dual'] = dual_fp or None

    # Determine if this evaluation contains any manual corrections
    all_sources = [dq['source']]
    if overlap:
        all_sources.append(eval_doc['fit_quality']['combined']['source'])
    else:
        all_sources.append(eval_doc['fit_quality']['regular_fit']['source'])
        all_sources.append(eval_doc['fit_quality']['ai_fit']['source'])
    has_correction = any(s == 'manual_correction' for s in all_sources)
    all_implicit = all(s == 'implicit_agreement' for s in all_sources)

    try:
        update_ops = {
            '$push': {'evaluations': eval_doc},
            '$inc': {'num_votes': 1},
            '$addToSet': {'tagged_by': username},
        }
        if has_correction:
            update_ops.setdefault('$set', {})['needs_second_opinion'] = True
        # Single flip failure flag
        if overlap and flip_state and flip_tag == '❌ Wrong':
            update_ops.setdefault('$set', {})['flip_function_failed'] = True
        # Dual flip failure flags (single-veto)
        if not overlap and flip_dual_state:
            rf = flip_dual_state.get('reg') and flip_reg_tag == '❌ Wrong'
            af = flip_dual_state.get('ai') and flip_ai_tag == '❌ Wrong'
            if rf:
                update_ops.setdefault('$set', {})['flip_reg_failed'] = True
            if af:
                update_ops.setdefault('$set', {})['flip_ai_failed'] = True
            if rf and af:
                update_ops.setdefault('$set', {})['flip_both_failed'] = True

        MONGO_COL.update_one({'_id': doc_id}, update_ops)

        # --- Single-labeler mode: 1 vote = completed ---
        scores = _extract_scores(eval_doc)
        consensus = {'data_quality': scores['data'],
                     'regular_fit_quality': scores['reg_fit'],
                     'ai_fit_quality': scores['ai_fit']}
        MONGO_COL.update_one({'_id': doc_id},
            {'$set': {'status': 'completed', 'final_consensus': consensus}})
        route_completed(doc_id)
        status = "✅ Feedback saved!"
    except Exception as e:
        status = f"❌ Save error: {e}"

    nxt = fetch_next(username)
    return (f"{status} {nxt[0]}",) + nxt[1:]


def reset_ratings():
    return None, None, None, None, None, None, None


def update_workspace_state(data_val, fit_s_val, fit_reg_val, fit_ai_val,
                           model_preds, overlap, exp_data, prev_fp):
    """Centralized state recalculation. Called when user changes any rating.
    Recomputes flip points, redraws plot, recalculates all visibility.

    Visibility rules:
      - Explicit "Bad (0)" on DATA → hide ALL fit + flip sections
      - Explicit "Bad (0)" on a FIT → hide only THAT fit's flip section
      - None (cleared) or Borderline/Perfect → show sections
    Flip-point gating:
      - Uses effective quality (user > model) for computation
    Plot optimization:
      - Uses _plot_fingerprint to detect if the plot's visual state changed.
      - Returns gr.skip() for the plot output when unchanged, saving bandwidth.

    Outputs 16 items: (plot, fs_col, fd_col,
      flip_info, flip_tag, flip_col, flip_state,
      flip_reg_info, flip_reg_tag, flip_ai_info, flip_ai_tag,
      flip_dual_col, flip_dual_state,
      flip_reg_sub, flip_ai_sub, new_fp)
    """
    _empty = (None,
              gr.Column(visible=True), gr.Column(visible=False),
              "", None, gr.Column(visible=False), None,
              "", None, "", None, gr.Column(visible=False), None,
              gr.Column(visible=True), gr.Column(visible=True),
              None)
    if exp_data is None:
        return _empty

    raw_data = exp_data['raw_data']
    reg_params = exp_data['reg_params']
    ai_params = exp_data.get('ai_params')
    fname = exp_data.get('filename', '?')
    mc = (model_preds or {}).get('model_classes', {})

    # ── UI VISIBILITY: only explicit "Bad" hides sections ──
    ui_show = data_val != 'Bad (0)'

    # ── Fit-specific UI visibility for flip sections ──
    # Explicit "Bad" on a fit hides only THAT fit's flip section
    fit_s_ui_ok = fit_s_val != 'Bad (0)'
    fit_reg_ui_ok = fit_reg_val != 'Bad (0)'
    fit_ai_ui_ok = fit_ai_val != 'Bad (0)'

    # ── FLIP-POINT GATING: uses effective quality for computation ──
    if data_val is not None:
        eff_data = RATING_MAP.get(data_val, -1)
    else:
        eff_data = mc.get('data_quality', 0)
    data_gate = eff_data >= 1

    if overlap:
        eff_fit = RATING_MAP.get(fit_s_val, -1) if fit_s_val is not None else mc.get('regular_fit_quality', 0)
    else:
        eff_fit = RATING_MAP.get(fit_reg_val, -1) if fit_reg_val is not None else mc.get('regular_fit_quality', 0)
    fit_gate = eff_fit >= 1

    eff_ai_fit = RATING_MAP.get(fit_ai_val, -1) if fit_ai_val is not None else mc.get('ai_fit_quality', 0)
    ai_fit_gate = eff_ai_fit >= 1

    # ── Compute flip points (gated on effective quality) ──
    flip_reg = None
    flip_ai = None
    if data_gate:
        if fit_gate:
            flip_reg = find_flip_point(raw_data, reg_params)
        if not overlap and ai_params and ai_fit_gate:
            flip_ai = find_flip_point(raw_data, ai_params)

    # ── Plot optimization: skip rebuild if visual state unchanged ──
    new_fp = _plot_fingerprint(flip_reg, flip_ai, overlap)
    if new_fp == prev_fp:
        fig = gr.skip()
    else:
        fig = create_plot(raw_data, fname, reg_params, ai_params, overlap,
                          flip_reg=flip_reg, flip_ai=flip_ai)

    # ── Section visibility ──
    show_fit_single = overlap and ui_show
    show_fit_dual = (not overlap) and ui_show

    # Single flip: visible only if data OK AND fit not Bad
    show_flip_single = (flip_reg is not None and overlap
                        and ui_show and fit_s_ui_ok)

    # Dual flip: outer container visible if at least one sub-section is active
    show_flip_reg_sub = (flip_reg is not None and ui_show and fit_reg_ui_ok)
    show_flip_ai_sub = (flip_ai is not None and ui_show and fit_ai_ui_ok)
    show_flip_dual = (not overlap) and (show_flip_reg_sub or show_flip_ai_sub)

    # ── Info strings ──
    flip_info = ""
    if overlap and flip_reg:
        flip_info = f"Flip Point: x = {flip_reg['x']:.5f}, y = {flip_reg['y']:.4f} ({flip_reg['type']})"
    elif overlap and ui_show and data_gate and fit_gate:
        flip_info = "No flip point detected in the fit curve."

    flip_reg_info = ""
    flip_ai_info = ""
    if not overlap and ui_show:
        if flip_reg and fit_reg_ui_ok:
            flip_reg_info = f"Reg Flip: x = {flip_reg['x']:.5f}, y = {flip_reg['y']:.4f} ({flip_reg['type']})"
        elif data_gate and fit_gate and fit_reg_ui_ok:
            flip_reg_info = "No flip point detected."
        if flip_ai and fit_ai_ui_ok:
            flip_ai_info = f"AI Flip: x = {flip_ai['x']:.5f}, y = {flip_ai['y']:.4f} ({flip_ai['type']})"
        elif data_gate and ai_fit_gate and fit_ai_ui_ok:
            flip_ai_info = "No flip point detected."

    dual_state = {'reg': flip_reg, 'ai': flip_ai} if not overlap else None

    return (fig,
            gr.Column(visible=show_fit_single),
            gr.Column(visible=show_fit_dual),
            flip_info, None,
            gr.Column(visible=show_flip_single),
            flip_reg if overlap else None,
            flip_reg_info, None, flip_ai_info, None,
            gr.Column(visible=show_flip_dual),
            dual_state,
            gr.Column(visible=show_flip_reg_sub),
            gr.Column(visible=show_flip_ai_sub),
            new_fp)


# ═════════════════════════════════════════════════════════════════════════
# Tab 3 — Analytics
# ═════════════════════════════════════════════════════════════════════════

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'})
    conflicts = MONGO_COL.count_documents({'status': 'conflict'})

    pipe = [{'$unwind': '$evaluations'},
            {'$group': {'_id': None, 'n': {'$sum': 1},
                        'p': {'$sum': {'$cond': [{'$eq': ['$evaluations.data_quality.score', 2]}, 1, 0]}}}}]
    agg = list(MONGO_COL.aggregate(pipe))
    n_ev = agg[0]['n'] if agg else 0
    n_p = agg[0]['p'] if agg else 0
    pct = (n_p / n_ev * 100) if n_ev else 0

    tp = [{'$unwind': '$evaluations'},
          {'$group': {'_id': '$evaluations.username', 'c': {'$sum': 1}}},
          {'$sort': {'c': -1}}, {'$limit': 10}]
    taggers = list(MONGO_COL.aggregate(tp))
    tr = "\n".join(f"| {t['_id']} | {t['c']} |" for t in taggers) or "| — | — |"

    clf_s = MONGO_DB['experiments_clf_success'].count_documents({}) if MONGO_DB is not None else 0
    clf_f = MONGO_DB['experiments_clf_failed'].count_documents({}) if MONGO_DB is not None else 0
    clf_c = MONGO_DB['experiments_clf_catastrophic_failure'].count_documents({}) if MONGO_DB is not None else 0
    seed_w = MONGO_DB['experiments_seed_success_reg_failed'].count_documents({}) if MONGO_DB is not None else 0
    reg_w = MONGO_DB['experiments_reg_success_seed_failed'].count_documents({}) if MONGO_DB is not None else 0
    # New granular collections
    reg_ff = MONGO_DB['experiments_reg_fit_failed'].count_documents({}) if MONGO_DB is not None else 0
    ai_ff = MONGO_DB['experiments_ai_fit_failed'].count_documents({}) if MONGO_DB is not None else 0
    both_ff = MONGO_DB['experiments_both_fits_failed'].count_documents({}) if MONGO_DB is not None else 0
    seed_bf = MONGO_DB['experiments_seed_borderline_fail'].count_documents({}) if MONGO_DB is not None else 0
    flip_rf = MONGO_DB['experiments_flip_reg_failed'].count_documents({}) if MONGO_DB is not None else 0
    flip_af = MONGO_DB['experiments_flip_ai_failed'].count_documents({}) if MONGO_DB is not None else 0
    flip_bf = MONGO_DB['experiments_flip_both_failed'].count_documents({}) if MONGO_DB is not None else 0

    return f"""### 📈 Platform Statistics
| Metric | Count |
|--------|-------|
| Total Experiments | **{total}** |
| Pending | **{pending}** |
| Completed | **{completed}** |
| Conflicts (legacy) | **{conflicts}** |
| Evaluations | **{n_ev}** |

### 🎯 Data Quality — Model vs Human Agreement
**{pct:.1f}%** rated Perfect ({n_p}/{n_ev})

### 📦 Classification Routing
| Collection | Count |
|------------|-------|
| ✅ clf_success | {clf_s} |
| ⚠️ clf_failed | {clf_f} |
| 💥 clf_catastrophic_failure | {clf_c} |

### 🔧 Fit-Level Failures (Diverging)
| Collection | Count |
|------------|-------|
| 🟦 reg_fit_failed | {reg_ff} |
| 🟩 ai_fit_failed | {ai_ff} |
| 🔴 both_fits_failed | {both_ff} |
| 🌱 seed_borderline_fail | {seed_bf} |
| 🟩 seed_success_reg_failed | {seed_w} |
| 🟦 reg_success_seed_failed | {reg_w} |

### 📍 Flip Function Failures
| Collection | Count |
|------------|-------|
| 📍 flip_reg_failed | {flip_rf} |
| 📍 flip_ai_failed | {flip_af} |
| 📍 flip_both_failed | {flip_bf} |

### 🏆 Top Annotators
| Username | Evals |
|----------|-------|
{tr}"""


# ═════════════════════════════════════════════════════════════════════════
# Tab 4 — Conflict Resolution (14 outputs)
# ═════════════════════════════════════════════════════════════════════════
# Output order:
#  0: status, 1: plot, 2: details_md,
#  3: concede_col, 4: tiebreak_col, 5: tb_fit_s_col, 6: tb_fit_d_col,
#  7: tb_data, 8: tb_fit_s, 9: tb_fit_r, 10: tb_fit_a,
#  11: doc_id_state, 12: overlap_state, 13: evals_state

def _empty_cr(msg=""):
    return (msg, None, "",
            gr.Column(visible=False), gr.Column(visible=False),
            gr.Column(visible=True), gr.Column(visible=False),
            None, None, None, None,
            None, False, None)


def _format_votes(evals, overlap):
    def _lbl(ev, key):
        fq = ev.get('fit_quality', {})
        if key == 'data':
            d = ev.get('data_quality', {})
            return f"{LABEL_NAMES[d['score']]} *({d['source'][:3]})*"
        if key == 'combined' and 'combined' in fq:
            d = fq['combined']
            return f"{LABEL_NAMES[d['score']]} *({d['source'][:3]})*"
        if key in fq:
            d = fq[key]
            return f"{LABEL_NAMES[d['score']]} *({d['source'][:3]})*"
        return "—"

    if len(evals) >= 2:
        e1, e2 = evals[0], evals[1]
        rows = f"| Data Quality | {_lbl(e1,'data')} | {_lbl(e2,'data')} |\n"
        if overlap:
            rows += f"| Fit Quality | {_lbl(e1,'combined')} | {_lbl(e2,'combined')} |\n"
        else:
            rows += f"| Regular Fit | {_lbl(e1,'regular_fit')} | {_lbl(e2,'regular_fit')} |\n"
            rows += f"| AI Fit | {_lbl(e1,'ai_fit')} | {_lbl(e2,'ai_fit')} |\n"

        return f"""### ⚔️ Conflicting Evaluations
| Aspect | {e1['username']} | {e2['username']} |
|--------|---------|---------|
{rows}
*(man = manual, imp = implicit)*"""
    elif len(evals) == 1:
        e1 = evals[0]
        rows = f"| Data Quality | {_lbl(e1,'data')} |\n"
        if overlap:
            rows += f"| Fit Quality | {_lbl(e1,'combined')} |\n"
        else:
            rows += f"| Regular Fit | {_lbl(e1,'regular_fit')} |\n"
            rows += f"| AI Fit | {_lbl(e1,'ai_fit')} |\n"

        return f"""### Single Evaluation
| Aspect | {e1['username']} |
|--------|---------|
{rows}
*(man = manual, imp = implicit)*"""
    else:
        return "### No evaluations recorded."


def fetch_conflict(username):
    if not username or not username.strip():
        return _empty_cr("⚠ Enter your username.")
    username = username.strip().title()
    if MONGO_COL is None:
        return _empty_cr("❌ MongoDB not connected.")
    doc = MONGO_COL.find_one({'status': 'conflict'})
    if not doc:
        return _empty_cr("🎉 No conflicts to resolve!")

    try:
        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)
        fname = doc.get('filename', str(doc['_id'])[:12])
        evals = doc.get('evaluations', [])[-2:]
        fig = create_plot(data, f"CONFLICT — {fname}", reg_p, ai_p, overlap)
        details = _format_votes(evals, overlap)
        is_original = username in doc.get('tagged_by', [])

        return (f"Conflict: {fname}", fig, details,
                gr.Column(visible=is_original),
                gr.Column(visible=not is_original),
                gr.Column(visible=overlap),
                gr.Column(visible=not overlap),
                None, None, None, None,
                doc['_id'], overlap, evals)
    except Exception as e:
        return _empty_cr(f"❌ Error: {e}")


def concede_vote(username, doc_id, evals_state):
    if not username or not doc_id or not evals_state:
        return _empty_cr("⚠ Invalid state.")
    username = username.strip().title()
    other = next((ev for ev in evals_state if ev.get('username') != username), None)
    if not other:
        return _empty_cr("⚠ Could not find peer's vote.")
    scores = _extract_scores(other)
    consensus = {'data_quality': scores['data'],
                 'regular_fit_quality': scores['reg_fit'],
                 'ai_fit_quality': scores['ai_fit']}
    try:
        MONGO_COL.update_one({'_id': doc_id}, {'$set': {
            'status': 'completed', 'final_consensus': consensus,
            'resolved_by': 'concession', 'conceded_by': username}})
        route_completed(doc_id)
    except Exception as e:
        return _empty_cr(f"❌ Error: {e}")
    nxt = fetch_conflict(username)
    return (f"✅ Conceded. {nxt[0]}",) + nxt[1:]


def submit_tiebreaker(username, doc_id, overlap, evals_state,
                      tb_data, tb_fit_s, tb_fit_r, tb_fit_a):
    if not username or not doc_id:
        return _empty_cr("⚠ Invalid state.")
    username = username.strip().title()
    d = RATING_MAP.get(tb_data)
    if d is None:
        return _empty_cr("⚠ Rate Data Quality.")
    consensus = {'data_quality': d}
    if overlap:
        f = RATING_MAP.get(tb_fit_s)
        if f is None:
            return _empty_cr("⚠ Rate Fit Quality.")
        consensus['regular_fit_quality'] = f
        consensus['ai_fit_quality'] = f
    else:
        fr, fa = RATING_MAP.get(tb_fit_r), RATING_MAP.get(tb_fit_a)
        if fr is None or fa is None:
            return _empty_cr("⚠ Rate both fits.")
        consensus['regular_fit_quality'] = fr
        consensus['ai_fit_quality'] = fa
    try:
        MONGO_COL.update_one({'_id': doc_id}, {'$set': {
            'status': 'completed', 'final_consensus': consensus,
            'resolved_by': 'tiebreaker', 'tiebreaker_by': username}})
        route_completed(doc_id)
    except Exception as e:
        return _empty_cr(f"❌ Error: {e}")
    nxt = fetch_conflict(username)
    return (f"✅ Tie broken. {nxt[0]}",) + nxt[1:]


# ═════════════════════════════════════════════════════════════════════════
# CSS & JS
# ═════════════════════════════════════════════════════════════════════════

CSS = """
.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }

.rating-radio label { padding: 10px 16px !important; border-radius: 10px !important;
    margin: 3px 4px !important; font-weight: 600 !important; font-size: 14px !important;
    transition: all 0.2s !important; cursor: pointer !important; }

.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; }

.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; }

.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; }


"""

# JavaScript to enable click-to-toggle on radio buttons
TOGGLE_JS = """
() => {
    function setup() {
        document.querySelectorAll('.rating-radio').forEach(function(group) {
            if (group.dataset.toggleBound) return;
            group.dataset.toggleBound = '1';
            group.querySelectorAll('input[type="radio"]').forEach(function(radio) {
                var wasChecked = false;
                radio.addEventListener('mousedown', function() {
                    wasChecked = this.checked;
                });
                radio.addEventListener('click', function(e) {
                    if (wasChecked) {
                        this.checked = false;
                        this.dispatchEvent(new Event('input', {bubbles: true}));
                        this.dispatchEvent(new Event('change', {bubbles: true}));
                    }
                });
            });
        });
    }
    setup();
    new MutationObserver(setup).observe(document.body, {childList: true, subtree: true});
}
"""


# ═════════════════════════════════════════════════════════════════════════
# Build UI
# ═════════════════════════════════════════════════════════════════════════

def build_ui():
    init_globals()

    with gr.Blocks(title="Rabi Annotation Platform") as demo:
        gr.Markdown("# ⚛️ Rabi Oscillation — Data Annotation Platform")

        with gr.Tabs():

            # ──────────── TAB 1: INGESTION ────────────
            with gr.Tab("📥 Data Ingestion"):
                gr.Markdown("Upload JSON files. Identical measured data is automatically deduplicated.")
                ig_files = gr.File(label="Upload JSON Experiments",
                                   file_count="multiple", file_types=[".json"])
                ig_btn = gr.Button("Process & Upload to Database", variant="primary")
                ig_result = gr.Markdown("")
                ig_btn.click(ingest_files, [ig_files], [ig_result])

            # ──────────── TAB 2: TAGGING WORKSPACE ────────────
            with gr.Tab("🏷️ Tagging Workspace"):
                ws_doc = gr.State(None)
                ws_overlap = gr.State(False)
                ws_preds = gr.State(None)

                with gr.Row():
                    ws_user = gr.Textbox(label="Username", placeholder="e.g. alice",
                                         scale=2, max_lines=1)
                    ws_fetch = gr.Button("🔄 Fetch Next Experiment", variant="primary", scale=1)

                ws_status = gr.Markdown("")

                # Graph — full width
                ws_plot = gr.Plot()

                # Experiment info + model confidence below graph
                with gr.Row():
                    with gr.Column(scale=1):
                        ws_info = gr.Markdown("")
                        ws_ov_msg = gr.Markdown("")
                    with gr.Column(scale=2):
                        ws_conf = gr.HTML("<p style='color:#888'>Fetch an experiment to begin.</p>")

                gr.Markdown("---")
                gr.Markdown("💡 **Leave a rating empty = accept model prediction.** "
                            "Select only to correct the model.")

                ws_exp_data = gr.State(None)
                ws_flip_state = gr.State(None)
                ws_plot_fp = gr.State(None)

                ws_reset = gr.Button("🔄 Clear All Ratings", size="sm")

                # Data Quality
                gr.Markdown("### 📊 Data Quality")
                ws_dr = gr.Radio(choices=RATING_CHOICES, value=None,
                                 label="Rate raw data:", elem_classes=["rating-radio"])
                ws_cd = gr.Textbox(label="Data notes (optional)", lines=2)

                # Fit Quality — SINGLE (overlap)
                with gr.Column(visible=True) as ws_fs_col:
                    gr.Markdown("### 📐 Fit Quality (Combined)")
                    ws_fs = gr.Radio(choices=RATING_CHOICES, value=None,
                                     label="Rate the fit:", elem_classes=["rating-radio"])
                    ws_cs = gr.Textbox(label="Fit notes (optional)", lines=2)

                # Fit Quality — DUAL (diverge)
                with gr.Column(visible=False) as ws_fd_col:
                    gr.Markdown("### 📐 Fit Quality (Diverging)")
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("#### 🟦 Regular Fit")
                            ws_fr = gr.Radio(choices=RATING_CHOICES, value=None,
                                             label="Regular Fit:", elem_classes=["rating-radio"])
                            ws_cr = gr.Textbox(label="Regular fit notes", lines=2)
                        with gr.Column():
                            gr.Markdown("#### 🟩 AI Fit")
                            ws_fa = gr.Radio(choices=RATING_CHOICES, value=None,
                                             label="AI Fit:", elem_classes=["rating-radio"])
                            ws_ca = gr.Textbox(label="AI fit notes", lines=2)

                # Flip-Point Detection — SINGLE (overlap)
                with gr.Column(visible=False) as ws_flip_col:
                    gr.Markdown("### 🎯 Flip-Point Detection")
                    ws_flip_info = gr.Textbox(label="Detected Flip Point",
                                              interactive=False, lines=1)
                    ws_flip_tag = gr.Radio(
                        choices=["✅ Correct", "❌ Wrong"], value=None,
                        label="Did the function find the correct physical flip?",
                        elem_classes=["rating-radio"])

                # Flip-Point Detection — DUAL (diverging fits)
                ws_flip_dual_state = gr.State(None)
                with gr.Column(visible=False) as ws_flip_dual_col:
                    gr.Markdown("### 🎯 Flip-Point Detection (Diverging Fits)")
                    with gr.Row():
                        with gr.Column(visible=True) as ws_flip_reg_sub:
                            gr.Markdown("#### 🟦 Regular Fit Flip")
                            ws_flip_reg_info = gr.Textbox(
                                label="Regular Flip", interactive=False, lines=1)
                            ws_flip_reg_tag = gr.Radio(
                                choices=["✅ Correct", "❌ Wrong"], value=None,
                                label="Regular flip correct?",
                                elem_classes=["rating-radio"])
                        with gr.Column(visible=True) as ws_flip_ai_sub:
                            gr.Markdown("#### 🟩 AI Fit Flip")
                            ws_flip_ai_info = gr.Textbox(
                                label="AI Flip", interactive=False, lines=1)
                            ws_flip_ai_tag = gr.Radio(
                                choices=["✅ Correct", "❌ Wrong"], value=None,
                                label="AI flip correct?",
                                elem_classes=["rating-radio"])

                gr.Markdown("---")
                ws_submit = gr.Button("🚀 Submit Feedback & Next", variant="primary", size="lg")

                ws_reset.click(reset_ratings, [],
                    [ws_dr, ws_fs, ws_fr, ws_fa,
                     ws_flip_tag, ws_flip_reg_tag, ws_flip_ai_tag]).then(
                    update_workspace_state,
                    [ws_dr, ws_fs, ws_fr, ws_fa,
                     ws_preds, ws_overlap, ws_exp_data, ws_plot_fp],
                    [ws_plot, ws_fs_col, ws_fd_col,
                     ws_flip_info, ws_flip_tag, ws_flip_col, ws_flip_state,
                     ws_flip_reg_info, ws_flip_reg_tag,
                     ws_flip_ai_info, ws_flip_ai_tag,
                     ws_flip_dual_col, ws_flip_dual_state,
                     ws_flip_reg_sub, ws_flip_ai_sub, ws_plot_fp])

                # Centralized state recalculation on any rating change
                _state_inputs = [ws_dr, ws_fs, ws_fr, ws_fa,
                                 ws_preds, ws_overlap, ws_exp_data, ws_plot_fp]
                _state_outputs = [ws_plot, ws_fs_col, ws_fd_col,
                                  ws_flip_info, ws_flip_tag, ws_flip_col, ws_flip_state,
                                  ws_flip_reg_info, ws_flip_reg_tag,
                                  ws_flip_ai_info, ws_flip_ai_tag,
                                  ws_flip_dual_col, ws_flip_dual_state,
                                  ws_flip_reg_sub, ws_flip_ai_sub, ws_plot_fp]
                ws_dr.input(update_workspace_state, _state_inputs, _state_outputs)
                ws_fs.input(update_workspace_state, _state_inputs, _state_outputs)
                ws_fr.input(update_workspace_state, _state_inputs, _state_outputs)
                ws_fa.input(update_workspace_state, _state_inputs, _state_outputs)

                # Output list — 32 components
                ws_outs = [ws_status, ws_plot, ws_info, ws_ov_msg, ws_conf,
                           ws_fs_col, ws_fd_col,
                           ws_dr, ws_fs, ws_fr, ws_fa,
                           ws_cd, ws_cs, ws_cr, ws_ca,
                           ws_doc, ws_overlap, ws_preds,
                           ws_flip_info, ws_flip_tag, ws_flip_col, ws_flip_state,
                           ws_exp_data,
                           ws_flip_reg_info, ws_flip_reg_tag,
                           ws_flip_ai_info, ws_flip_ai_tag,
                           ws_flip_dual_col, ws_flip_dual_state,
                           ws_flip_reg_sub, ws_flip_ai_sub,
                           ws_plot_fp]

                ws_fetch.click(fetch_next, [ws_user], ws_outs)
                ws_submit.click(submit_feedback,
                    [ws_user, ws_doc, ws_overlap, ws_preds,
                     ws_dr, ws_cd, ws_fs, ws_fr, ws_fa, ws_cs, ws_cr, ws_ca,
                     ws_flip_tag, ws_flip_state,
                     ws_flip_reg_tag, ws_flip_ai_tag, ws_flip_dual_state],
                    ws_outs)

            # ──────────── TAB 3: ANALYTICS ────────────
            with gr.Tab("📊 Analytics"):
                gr.Markdown("Aggregate statistics from all annotations.")
                an_btn = gr.Button("🔄 Refresh Statistics", variant="primary")
                an_md = gr.Markdown("Click **Refresh** to load.")
                an_btn.click(refresh_analytics, [], [an_md])

            # ──────────── TAB 4: CONFLICT RESOLUTION ────────────
            with gr.Tab("⚔️ Conflicts"):
                cr_doc = gr.State(None)
                cr_overlap = gr.State(False)
                cr_evals = gr.State(None)

                with gr.Row():
                    cr_user = gr.Textbox(label="Username", placeholder="e.g. alice",
                                         scale=2, max_lines=1)
                    cr_fetch = gr.Button("🔄 Fetch Conflict", variant="primary", scale=1)

                cr_status = gr.Markdown("")
                cr_plot = gr.Plot()
                cr_details = gr.Markdown("")

                # Concede section (original voter)
                with gr.Column(visible=False) as cr_concede_col:
                    gr.Markdown("### You are an original voter.")
                    gr.Markdown("Accept your peer's vote as the final consensus.")
                    cr_concede_btn = gr.Button("🤝 Concede / Accept Peer's Vote", variant="secondary")

                # Tiebreaker section (3rd party)
                with gr.Column(visible=False) as cr_tb_col:
                    gr.Markdown("### You are the tie-breaker. Your vote decides.")
                    cr_tb_dr = gr.Radio(choices=RATING_CHOICES, value=None,
                                       label="Data Quality:", elem_classes=["rating-radio"])

                    with gr.Column(visible=True) as cr_tb_fs_col:
                        cr_tb_fsr = gr.Radio(choices=RATING_CHOICES, value=None,
                                             label="Fit Quality:", elem_classes=["rating-radio"])

                    with gr.Column(visible=False) as cr_tb_fd_col:
                        with gr.Row():
                            cr_tb_frr = gr.Radio(choices=RATING_CHOICES, value=None,
                                                 label="Regular Fit:", elem_classes=["rating-radio"])
                            cr_tb_far = gr.Radio(choices=RATING_CHOICES, value=None,
                                                 label="AI Fit:", elem_classes=["rating-radio"])

                    cr_tb_btn = gr.Button("⚖️ Submit Tie-Breaker Vote", variant="primary")

                # Output list — 14 components
                cr_outs = [cr_status, cr_plot, cr_details,
                           cr_concede_col, cr_tb_col, cr_tb_fs_col, cr_tb_fd_col,
                           cr_tb_dr, cr_tb_fsr, cr_tb_frr, cr_tb_far,
                           cr_doc, cr_overlap, cr_evals]

                cr_fetch.click(fetch_conflict, [cr_user], cr_outs)
                cr_concede_btn.click(concede_vote, [cr_user, cr_doc, cr_evals], cr_outs)

                def _handle_tb(u, did, ov, evs, d, fs, fr, fa):
                    return submit_tiebreaker(u, did, ov, evs, d, fs, fr, fa)

                cr_tb_btn.click(_handle_tb,
                    [cr_user, cr_doc, cr_overlap, cr_evals,
                     cr_tb_dr, cr_tb_fsr, cr_tb_frr, cr_tb_far],
                    cr_outs)

    return demo


# ═════════════════════════════════════════════════════════════════════════
# Validation Tests (run with: python app.py --test)
# ═════════════════════════════════════════════════════════════════════════

def run_validation_tests():
    """Verify MLOps logic: fast-track, catastrophic routing, smart queue."""
    print("=" * 60)
    print("VALIDATION TESTS")
    print("=" * 60)
    passed = 0
    failed = 0

    def check(name, condition):
        nonlocal passed, failed
        if condition:
            print(f"  ✅ {name}")
            passed += 1
        else:
            print(f"  ❌ {name}")
            failed += 1

    # --- 1. _max_deviation ---
    print("\n1. Catastrophic Failure Detection (_max_deviation)")
    check("Perfect match = 0",
          _max_deviation({'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2},
                         {'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2}) == 0)
    check("Off-by-1 = 1",
          _max_deviation({'data_quality': 2, 'regular_fit_quality': 1, 'ai_fit_quality': 2},
                         {'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2}) == 1)
    check("Catastrophic (2→0) = 2",
          _max_deviation({'data_quality': 0, 'regular_fit_quality': 2, 'ai_fit_quality': 2},
                         {'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2}) == 2)
    check("Catastrophic (0→2) = 2",
          _max_deviation({'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2},
                         {'data_quality': 0, 'regular_fit_quality': 2, 'ai_fit_quality': 2}) == 2)
    check("Mixed: worst wins (data diff=2, fit diff=1) = 2",
          _max_deviation({'data_quality': 0, 'regular_fit_quality': 1, 'ai_fit_quality': 2},
                         {'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2}) == 2)
    check("Missing keys handled gracefully",
          _max_deviation({'data_quality': 2}, {'regular_fit_quality': 0}) == 0)

    # --- 2. Single-Labeler Mode ---
    print("\n2. Single-Labeler Mode")
    check("VOTES_REQUIRED == 1", VOTES_REQUIRED == 1)

    # --- 3. _resolve_rating implicit vs manual ---
    print("\n3. Implicit Agreement vs Manual Correction")
    mp = {'model_classes': {'data_quality': 2, 'regular_fit_quality': 1, 'ai_fit_quality': 0}}
    r1 = _resolve_rating(None, mp, 'data_quality')
    check("None input → implicit_agreement, score=2",
          r1 == {'score': 2, 'source': 'implicit_agreement'})
    r2 = _resolve_rating("Bad (0)", mp, 'data_quality')
    check("'Bad (0)' input → manual_correction, score=0",
          r2 == {'score': 0, 'source': 'manual_correction'})

    # --- 4. Smart Queue needs_second_opinion ---
    print("\n4. Smart Queue Prioritization")
    sources_all_implicit = ['implicit_agreement', 'implicit_agreement']
    sources_with_correction = ['implicit_agreement', 'manual_correction']
    check("All implicit → no needs_second_opinion flag",
          not any(s == 'manual_correction' for s in sources_all_implicit))
    check("Has correction → needs_second_opinion = True",
          any(s == 'manual_correction' for s in sources_with_correction))

    # --- 5. Username normalization ---
    print("\n5. Username Normalization")
    check("'aDI mASKIL' → 'Adi Maskil'", 'aDI mASKIL'.strip().title() == 'Adi Maskil')
    check("'bob' → 'Bob'", 'bob'.strip().title() == 'Bob')
    check("'ALICE SMITH' → 'Alice Smith'", 'ALICE SMITH'.strip().title() == 'Alice Smith')

    # --- 6. Flip-Point Detection ---
    print("\n6. Flip-Point Detection (find_flip_point)")
    # Single-peak sine
    mock_data = {'measured_data': {'x_values': list(np.linspace(0, 1, 100)),
                                   'y_values': list(np.sin(2 * np.pi * np.linspace(0, 1, 100)))}}
    mock_params = {'amplitude': 1.0, 'T': 1.0, 'phase': 0.0, 'offset': 0.0}
    fp = find_flip_point(mock_data, mock_params)
    check("Sine wave → finds a flip point", fp is not None)
    check("Flip type is 'max' or 'min'", fp is not None and fp['type'] in ('max', 'min'))
    check("Flip x is within data range", fp is not None and 0 <= fp['x'] <= 1)

    # Flat line → no flip
    flat_data = {'measured_data': {'x_values': [0, 0.5, 1], 'y_values': [0.5, 0.5, 0.5]}}
    flat_params = {'amplitude': 0.0, 'T': 1.0, 'phase': 0.0, 'offset': 0.5}
    fp_flat = find_flip_point(flat_data, flat_params)
    check("Flat line → no flip point", fp_flat is None)

    # Multi-peak → returns leftmost
    mp_data = {'measured_data': {'x_values': list(np.linspace(0, 2, 200)),
                                  'y_values': list(np.sin(4 * np.pi * np.linspace(0, 2, 200)))}}
    mp_params = {'amplitude': 1.0, 'T': 0.5, 'phase': 0.0, 'offset': 0.0}
    fp_mp = find_flip_point(mp_data, mp_params)
    check("Multi-peak → returns leftmost extremum", fp_mp is not None and fp_mp['x'] < 0.5)

    # --- 7. Smart Queue: 1-vote before 0-vote ---
    print(f"\n7. Smart Queue Priority (1-vote before 0-vote)")
    from pymongo import ASCENDING
    # Simulate the sort used in fetch_next
    docs = [
        {'_id': 'A', 'needs_second_opinion': False, 'num_votes': 0},
        {'_id': 'B', 'needs_second_opinion': True,  'num_votes': 1},
        {'_id': 'C', 'needs_second_opinion': False, 'num_votes': 1},
        {'_id': 'D', 'needs_second_opinion': False, 'num_votes': 0},
    ]
    sort_keys = [('needs_second_opinion', -1), ('num_votes', -1)]
    sorted_docs = sorted(docs,
        key=lambda d: tuple(-d.get(k, 0) if v == -1 else d.get(k, 0)
                            for k, v in sort_keys))
    check("needs_second_opinion doc served first",
          sorted_docs[0]['_id'] == 'B')
    check("1-vote doc served before 0-vote doc",
          sorted_docs[1]['_id'] == 'C' and sorted_docs[1]['num_votes'] == 1)
    check("0-vote docs served last",
          sorted_docs[2]['num_votes'] == 0 and sorted_docs[3]['num_votes'] == 0)

    # --- 8. Plot Fingerprint / gr.skip() Optimization ---
    print("\n8. Plot Fingerprint Optimization (_plot_fingerprint)")
    fp_a = _plot_fingerprint(None, None, True)
    fp_b = _plot_fingerprint(None, None, True)
    check("Same inputs → identical fingerprint", fp_a == fp_b)
    fp_with_flip = _plot_fingerprint({'x': 0.25, 'y': 0.5, 'type': 'max'}, None, True)
    check("Flip vs no-flip → different fingerprint", fp_with_flip != fp_a)
    fp_overlap_t = _plot_fingerprint(None, None, True)
    fp_overlap_f = _plot_fingerprint(None, None, False)
    check("Overlap True vs False → different fingerprint", fp_overlap_t != fp_overlap_f)
    # Verify update_workspace_state returns gr.skip() when fingerprint unchanged
    mock_exp = {
        'raw_data': {'measured_data': {'x_values': list(np.linspace(0, 1, 50)),
                                       'y_values': list(np.sin(np.linspace(0, 1, 50)))}},
        'reg_params': {'amplitude': 0.5, 'T': 1.0, 'phase': 0.0, 'offset': 0.0},
        'ai_params': None, 'filename': 'test'
    }
    mock_mp = {'model_classes': {'data_quality': 2, 'regular_fit_quality': 2}}
    # First call: prev_fp=None → must draw (not skip)
    result1 = update_workspace_state('Perfect (2)', 'Perfect (2)', None, None,
                                     mock_mp, True, mock_exp, None)
    check("First call (prev_fp=None) → plot is NOT gr.skip()",
          not isinstance(result1[0], type(gr.skip())))
    new_fp = result1[-1]  # last element is the new fingerprint
    check("First call returns a fingerprint tuple", isinstance(new_fp, tuple))
    # Second call: same ratings → fingerprint unchanged → gr.skip()
    result2 = update_workspace_state('Perfect (2)', 'Perfect (2)', None, None,
                                     mock_mp, True, mock_exp, new_fp)
    is_skip = (result2[0].__class__.__name__ == 'skip'
               or result2[0] is gr.skip()  # Gradio skip sentinel
               or repr(result2[0]) == repr(gr.skip()))
    check("Second call (same ratings) → plot IS gr.skip()", is_skip)
    # Third call: change rating to cross gate → fingerprint changes → redraw
    result3 = update_workspace_state('Bad (0)', 'Perfect (2)', None, None,
                                     mock_mp, True, mock_exp, new_fp)
    check("Gate-crossing rating change → plot is NOT gr.skip()",
          not isinstance(result3[0], type(gr.skip()))
          or repr(result3[0]) != repr(gr.skip()))

    print(f"\n{'=' * 60}")
    print(f"RESULTS: {passed} passed, {failed} failed")
    print(f"{'=' * 60}")
    return failed == 0


def run_mongodb_routing_test():
    """E2E test: inject mock data → route → verify collections exist in Atlas."""
    print("=" * 60)
    print("MONGODB E2E ROUTING TEST")
    print("=" * 60)

    # 1. Connect using same logic as init_globals
    client = None
    for kw in [dict(tlsCAFile=certifi.where()),
               dict(tls=True, tlsAllowInvalidCertificates=True)]:
        try:
            client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000, **kw)
            client.admin.command('ping')
            break
        except Exception:
            client = None
    if not client:
        print("❌ Cannot connect to MongoDB. Aborting.")
        return False

    db = client[DB_NAME]
    test_col = db['experiments_test_mock']
    print(f"✓ Connected to {DB_NAME}")

    # 2. Clean up previous test data
    test_col.delete_many({})
    for col_name in ['experiments_clf_success', 'experiments_clf_failed',
                     'experiments_clf_catastrophic_failure',
                     'experiments_seed_success_reg_failed',
                     'experiments_reg_success_seed_failed']:
        db[col_name].delete_many({'original_id': {'$regex': '^e2e_test_'}})
    print("✓ Cleaned previous test data")

    # 3. Insert mock documents for each routing scenario
    test_cases = [
        {
            'name': 'Catastrophic Failure (model=2, human=0)',
            '_id': 'e2e_test_catastrophic',
            'model_predictions': {
                'model_classes': {'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2},
                'data_probs': [0.01, 0.02, 0.97],
                'regular_fit_probs': [0.02, 0.03, 0.95],
                'ai_fit_probs': [0.01, 0.04, 0.95],
            },
            'final_consensus': {'data_quality': 0, 'regular_fit_quality': 0, 'ai_fit_quality': 0},
            'curves_overlap': True,
            'status': 'completed',
            'expected_collection': 'experiments_clf_catastrophic_failure',
        },
        {
            'name': 'Minor Failure (model=2, human=1)',
            '_id': 'e2e_test_minor_fail',
            'model_predictions': {
                'model_classes': {'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2},
                'data_probs': [0.01, 0.02, 0.97],
                'regular_fit_probs': [0.02, 0.03, 0.95],
                'ai_fit_probs': [0.01, 0.04, 0.95],
            },
            'final_consensus': {'data_quality': 1, 'regular_fit_quality': 2, 'ai_fit_quality': 2},
            'curves_overlap': True,
            'status': 'completed',
            'expected_collection': 'experiments_clf_failed',
        },
        {
            'name': 'Perfect Match (model=human)',
            '_id': 'e2e_test_success',
            'model_predictions': {
                'model_classes': {'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2},
                'data_probs': [0.01, 0.02, 0.97],
                'regular_fit_probs': [0.02, 0.03, 0.95],
                'ai_fit_probs': [0.01, 0.04, 0.95],
            },
            'final_consensus': {'data_quality': 2, 'regular_fit_quality': 2, 'ai_fit_quality': 2},
            'curves_overlap': True,
            'status': 'completed',
            'expected_collection': 'experiments_clf_success',
        },
    ]

    # Insert all test docs into the test collection
    for tc in test_cases:
        doc = {k: v for k, v in tc.items() if k not in ('name', 'expected_collection')}
        doc['raw_data'] = {'measured_data': {'x_values': [0, 1], 'y_values': [0, 1]}}
        doc['evaluations'] = []
        doc['tagged_by'] = []
        doc['num_votes'] = 1
        test_col.insert_one(doc)
    print(f"✓ Inserted {len(test_cases)} mock documents into experiments_test_mock")

    # 4. Run routing on each mock doc using the REAL route_completed logic
    # We temporarily point MONGO_COL and MONGO_DB to our test setup
    global MONGO_COL, MONGO_DB
    original_col = MONGO_COL
    original_db = MONGO_DB
    MONGO_COL = test_col
    MONGO_DB = db

    print("\nRouting documents...")
    for tc in test_cases:
        route_completed(tc['_id'])
        print(f"  → Routed: {tc['name']}")

    # 5. Verify each document landed in the correct collection
    print("\nChecking for collection creation...")
    all_passed = True
    for tc in test_cases:
        target = tc['expected_collection']
        found = db[target].find_one({'original_id': tc['_id']})
        if found:
            print(f"  ✅ {target} contains '{tc['name']}'")
        else:
            print(f"  ❌ {target} MISSING '{tc['name']}'")
            all_passed = False

    # 6. List all collections to confirm they exist in Atlas
    print("\nCollections in database:")
    for name in sorted(db.list_collection_names()):
        count = db[name].count_documents({})
        print(f"  📁 {name} ({count} docs)")

    # Restore globals
    MONGO_COL = original_col
    MONGO_DB = original_db

    print(f"\n{'=' * 60}")
    print(f"E2E RESULT: {'ALL PASSED ✅' if all_passed else 'FAILURES DETECTED ❌'}")
    print(f"{'=' * 60}")
    return all_passed


if __name__ == '__main__':
    import sys
    if '--test' in sys.argv:
        success = run_validation_tests()
        sys.exit(0 if success else 1)
    elif '--e2e' in sys.argv:
        success = run_mongodb_routing_test()
        sys.exit(0 if success else 1)
    else:
        demo = build_ui()
        _user = os.environ.get('APP_USER', 'qarakal')
        _pass = os.environ.get('APP_PASS', 'lab2026')
        demo.launch(css=CSS, js=TOGGLE_JS, auth=(_user, _pass))