app.py

import ast
import logging
import re
from typing import Dict, List, Optional, Tuple

import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from datasets import load_dataset

from sklearn.ensemble import HistGradientBoostingRegressor
from sklearn.impute import SimpleImputer
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split
from sklearn.multioutput import MultiOutputRegressor
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

# Logging configuration
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ========================= CONFIG =========================
APP_TITLE = "Quantum Noise Robustness Benchmark"
APP_SUBTITLE = (
    "Predict noisy expectation values (Z/X/Y) and errors from ideal values "
    "and circuit structure — without expensive simulation."
)

REPO_CONFIG = {
    "amplitude_damping": {
        "label": "amplitude_damping",
        "repo": "QSBench/QSBench-Amplitude-v1.0.0-demo",
    },
}

TARGET_COLS = ["error_Z_global", "error_X_global", "error_Y_global"]
IDEAL_COLS = ["ideal_expval_Z_global", "ideal_expval_X_global", "ideal_expval_Y_global"]
NOISY_COLS = ["noisy_expval_Z_global", "noisy_expval_X_global", "noisy_expval_Y_global"]

NON_FEATURE_COLS = {
    "sample_id", "sample_seed", "circuit_hash", "split", "circuit_qasm",
    "qasm_raw", "qasm_transpiled", "circuit_type_resolved", "circuit_type_requested",
    "noise_type", "noise_prob", "observable_bases", "observable_mode", "shots",
    "gpu_requested", "gpu_available", "backend_device", "precision_mode",
    "circuit_signature", "noise_label",
    *IDEAL_COLS, *NOISY_COLS, *TARGET_COLS,
    "sign_ideal_Z_global", "sign_noisy_Z_global",
    "sign_ideal_X_global", "sign_noisy_X_global",
    "sign_ideal_Y_global", "sign_noisy_Y_global",
}

SOFT_EXCLUDE_PATTERNS = ["ideal_", "noisy_", "sign_ideal_", "sign_noisy_"]

_ASSET_CACHE: Dict[str, pd.DataFrame] = {}

# ========================= HELPERS =========================

def load_guide_content() -> str:
    """Read the GUIDE.md file from the root directory."""
    try:
        with open("GUIDE.md", "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        return "### ⚠️ GUIDE.md not found in the root directory."

def safe_parse(value):
    """Safely parse stringified Python literals."""
    if isinstance(value, str):
        try:
            return ast.literal_eval(value)
        except Exception:
            return value
    return value

def adjacency_features(adj_value) -> Dict[str, float]:
    """Derive graph statistics from an adjacency matrix."""
    parsed = safe_parse(adj_value)
    if not isinstance(parsed, list) or len(parsed) == 0:
        return {
            "adj_edge_count": np.nan,
            "adj_density": np.nan,
            "adj_degree_mean": np.nan,
            "adj_degree_std": np.nan,
        }
    try:
        arr = np.array(parsed, dtype=float)
        n = arr.shape[0]
        edge_count = float(np.triu(arr, k=1).sum())
        possible_edges = float(n * (n - 1) / 2)
        density = edge_count / possible_edges if possible_edges > 0 else np.nan
        degrees = arr.sum(axis=1)
        return {
            "adj_edge_count": edge_count,
            "adj_density": density,
            "adj_degree_mean": float(np.mean(degrees)),
            "adj_degree_std": float(np.std(degrees)),
        }
    except Exception:
        return {
            "adj_edge_count": np.nan,
            "adj_density": np.nan,
            "adj_degree_mean": np.nan,
            "adj_degree_std": np.nan,
        }

def qasm_features(qasm_value) -> Dict[str, float]:
    """Extract lightweight text statistics from QASM."""
    if not isinstance(qasm_value, str) or not qasm_value.strip():
        return {
            "qasm_length": np.nan,
            "qasm_line_count": np.nan,
            "qasm_gate_keyword_count": np.nan,
            "qasm_measure_count": np.nan,
        }
    text = qasm_value
    lines = [line for line in text.splitlines() if line.strip()]
    gate_keywords = re.findall(
        r"\b(cx|h|x|y|z|rx|ry|rz|u1|u2|u3|u|swap|cz|ccx|rxx|ryy|rzz)\b",
        text,
        flags=re.IGNORECASE,
    )
    measure_count = len(re.findall(r"\bmeasure\b", text, flags=re.IGNORECASE))
    return {
        "qasm_length": float(len(text)),
        "qasm_line_count": float(len(lines)),
        "qasm_gate_keyword_count": float(len(gate_keywords)),
        "qasm_measure_count": float(measure_count),
    }

def enrich_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    """Add derived numeric features and compute error targets."""
    df = df.copy()
    if "adjacency" in df.columns:
        adj_df = df["adjacency"].apply(adjacency_features).apply(pd.Series)
        df = pd.concat([df, adj_df], axis=1)
    qasm_source = "qasm_transpiled" if "qasm_transpiled" in df.columns else "qasm_raw"
    if qasm_source in df.columns:
        qasm_df = df[qasm_source].apply(qasm_features).apply(pd.Series)
        df = pd.concat([df, qasm_df], axis=1)

    for basis in ["Z", "X", "Y"]:
        ideal_col = f"ideal_expval_{basis}_global"
        noisy_col = f"noisy_expval_{basis}_global"
        error_col = f"error_{basis}_global"
        if ideal_col in df.columns and noisy_col in df.columns:
            df[error_col] = df[noisy_col] - df[ideal_col]
    return df

def load_single_dataset() -> pd.DataFrame:
    """Fetch and cache the dataset."""
    key = "amplitude_damping"
    if key not in _ASSET_CACHE:
        logger.info("Loading dataset: %s", key)
        ds = load_dataset(REPO_CONFIG[key]["repo"])
        df = pd.DataFrame(ds["train"])
        df = enrich_dataframe(df)
        _ASSET_CACHE[key] = df
    return _ASSET_CACHE[key]

def get_available_feature_columns(df: pd.DataFrame) -> List[str]:
    """Retrieve filtered list of numerical feature columns."""
    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
    features = []
    for col in numeric_cols:
        if col in NON_FEATURE_COLS:
            continue
        if any(pattern in col for pattern in SOFT_EXCLUDE_PATTERNS):
            continue
        features.append(col)
    return sorted(features)

def default_feature_selection(features: List[str]) -> List[str]:
    """Provide a curated list of default structural features."""
    preferred = [
        "gate_entropy", "adj_density", "adj_degree_mean", "adj_degree_std",
        "depth", "total_gates", "cx_count", "two_qubit_gates",
        "qasm_length", "qasm_line_count", "qasm_gate_keyword_count",
    ]
    selected = [f for f in preferred if f in features]
    return selected[:10] if selected else features[:10]

def make_regression_figure(
    y_true: np.ndarray, 
    y_pred: np.ndarray, 
    ideal_vals: np.ndarray, 
    noisy_vals: np.ndarray, 
    basis: str
) -> plt.Figure:
    """Generate diagnostic regression plots including physics emulation."""
    fig, axs = plt.subplots(1, 3, figsize=(20, 6))

    # 1. Error Prediction (Predicted vs True)
    axs[0].scatter(y_true, y_pred, alpha=0.6, s=15, color='#3498db')
    min_v, max_v = min(y_true.min(), y_pred.min()), max(y_true.max(), y_pred.max())
    axs[0].plot([min_v, max_v], [min_v, max_v], 'r--', lw=2)
    axs[0].set_xlabel("True Error")
    axs[0].set_ylabel("Predicted Error")
    axs[0].set_title(f"{basis} Error: Predicted vs True")
    axs[0].grid(True, alpha=0.3)

    # 2. Residual Distribution
    residuals = y_true - y_pred
    axs[1].hist(residuals, bins=50, alpha=0.7, color="#2ecc71", edgecolor="black")
    axs[1].axvline(0, color="red", linestyle="--")
    axs[1].set_xlabel("Residual")
    axs[1].set_ylabel("Count")
    axs[1].set_title(f"{basis} Error Residuals")
    axs[1].grid(True, alpha=0.3)

    # 3. Physics Emulation (Ideal vs Noisy Expectation Values)
    pred_noisy_vals = ideal_vals + y_pred
    
    axs[2].scatter(ideal_vals, noisy_vals, alpha=0.4, s=15, label="Actual Noisy (Simulated)", color="#95a5a6")
    axs[2].scatter(ideal_vals, pred_noisy_vals, alpha=0.6, s=15, label="Predicted Noisy (ML)", color="#e74c3c")
    axs[2].plot([-1, 1], [-1, 1], 'k--', lw=1, alpha=0.7, label="No Noise Limit")
    axs[2].set_xlabel("Ideal Expectation Value")
    axs[2].set_ylabel("Noisy Expectation Value")
    axs[2].set_title(f"Physics Emulation: {basis} Basis Shift")
    axs[2].legend()
    axs[2].grid(True, alpha=0.3)

    fig.tight_layout()
    return fig

def train_regressor(
    feature_columns: List[str],
    test_size: float,
    max_iter: int,
    max_depth: float,
    random_state: float,
) -> Tuple[Optional[plt.Figure], str, Optional[plt.Figure], Optional[plt.Figure]]:
    """Train multi-output regressor and return metrics with plots."""
    if not feature_columns:
        return None, "### ❌ Please select at least one feature.", None, None

    df = load_single_dataset()
    required_cols = feature_columns + TARGET_COLS + IDEAL_COLS + NOISY_COLS
    train_df = df.dropna(subset=required_cols).copy()

    if len(train_df) < 50:
        return None, "### ❌ Not enough rows after filtering missing values.", None, None

    X = train_df[feature_columns]
    y = train_df[TARGET_COLS]

    seed = int(random_state)
    depth = int(max_depth) if max_depth and int(max_depth) > 0 else None

    # Track indices to extract ideal and noisy arrays for the test set later
    indices = np.arange(len(train_df))
    idx_train, idx_test = train_test_split(indices, test_size=test_size, random_state=seed)

    X_train, X_test = X.iloc[idx_train], X.iloc[idx_test]
    y_train, y_test = y.iloc[idx_train], y.iloc[idx_test]
    
    ideal_test = train_df[IDEAL_COLS].iloc[idx_test].values
    noisy_test = train_df[NOISY_COLS].iloc[idx_test].values

    model = Pipeline([
        ("imputer", SimpleImputer(strategy="median")),
        ("scaler", StandardScaler()),
        ("regressor", MultiOutputRegressor(
            HistGradientBoostingRegressor(
                max_iter=int(max_iter),
                max_depth=depth,
                random_state=seed,
                learning_rate=0.1,
                min_samples_leaf=1,
            )
        ))
    ])

    model.fit(X_train, y_train)
    y_pred = model.predict(X_test)

    mae = mean_absolute_error(y_test, y_pred, multioutput="raw_values")
    rmse = np.sqrt(mean_squared_error(y_test, y_pred, multioutput="raw_values"))
    r2 = r2_score(y_test, y_pred, multioutput="raw_values")

    metrics_text = (
        "### Regression Results\n\n"
        f"**Rows used:** {len(train_df):,}\n"
        f"**Test size:** {test_size:.0%}\n\n"
        f"**Z-error** — MAE: {mae[0]:.5f} | RMSE: {rmse[0]:.5f} | R²: {r2[0]:.4f}\n"
        f"**X-error** — MAE: {mae[1]:.5f} | RMSE: {rmse[1]:.5f} | R²: {r2[1]:.4f}\n"
        f"**Y-error** — MAE: {mae[2]:.5f} | RMSE: {rmse[2]:.5f} | R²: {r2[2]:.4f}\n"
    )

    # Generate figures passing ideal and true noisy data
    fig_z = make_regression_figure(y_test.iloc[:, 0].values, y_pred[:, 0], ideal_test[:, 0], noisy_test[:, 0], "Z")
    fig_x = make_regression_figure(y_test.iloc[:, 1].values, y_pred[:, 1], ideal_test[:, 1], noisy_test[:, 1], "X")
    fig_y = make_regression_figure(y_test.iloc[:, 2].values, y_pred[:, 2], ideal_test[:, 2], noisy_test[:, 2], "Y")

    return fig_z, metrics_text, fig_x, fig_y

# ======================= EXPLORER FUNCTIONS =======================

def build_dataset_profile(df: pd.DataFrame) -> str:
    """Generate Markdown summary of the loaded dataset."""
    return (
        f"### Dataset profile\n\n"
        f"**Rows:** {len(df):,} \n"
        f"**Columns:** {len(df.columns):,} \n"
        f"**Classes / Noise:** amplitude_damping"
    )

def refresh_explorer(dataset_key: str, split_name: str):
    """Update Explorer tab components."""
    df = load_single_dataset()
    splits = df["split"].dropna().unique().tolist() if "split" in df.columns else ["train"]
    if not splits:
        splits = ["train"]
    if split_name not in splits:
        split_name = splits[0]

    filtered = df[df["split"] == split_name] if "split" in df.columns else df
    display_df = filtered.head(12).copy()

    raw_qasm = display_df["qasm_raw"].iloc[0] if not display_df.empty and "qasm_raw" in display_df.columns else "// N/A"
    transpiled_qasm = display_df["qasm_transpiled"].iloc[0] if not display_df.empty and "qasm_transpiled" in display_df.columns else "// N/A"

    profile_box = build_dataset_profile(df)
    summary_box = (
        f"### Split summary\n\n"
        f"**Dataset:** `{dataset_key}` \n"
        f"**Label:** `amplitude_damping` \n"
        f"**Available splits:** {', '.join(splits)} \n"
        f"**Preview rows:** {len(display_df)}"
    )

    return (
        gr.update(choices=splits, value=split_name),
        display_df,
        raw_qasm,
        transpiled_qasm,
        profile_box,
        summary_box,
    )

# ========================= INTERFACE =========================

CUSTOM_CSS = """
.gradio-container {
    max-width: 1400px !important;
}
footer {
    margin-top: 1rem;
}
"""

with gr.Blocks(title=APP_TITLE) as demo:
    gr.Markdown(f"# 🌌 {APP_TITLE}")
    gr.Markdown(APP_SUBTITLE)

    with gr.Tabs():
        with gr.TabItem("🔎 Explorer"):
            dataset_dropdown = gr.Dropdown(
                list(REPO_CONFIG.keys()),
                value="amplitude_damping",
                label="Dataset",
            )
            split_dropdown = gr.Dropdown(
                ["train"],
                value="train",
                label="Split",
            )
            profile_box = gr.Markdown(value="### Loading dataset...")
            summary_box = gr.Markdown(value="### Loading split summary...")
            explorer_df = gr.Dataframe(label="Preview", interactive=False)

            with gr.Row():
                raw_qasm = gr.Code(label="Raw QASM", language=None)
                transpiled_qasm = gr.Code(label="Transpiled QASM", language="python")

        with gr.TabItem("🧠 Regression Training"):
            feature_picker = gr.CheckboxGroup(
                label="Input features (circuit structure + topology)",
                choices=[],
                value=[],
            )
            test_size = gr.Slider(0.1, 0.4, value=0.25, step=0.05, label="Test Split")
            max_iter = gr.Slider(100, 800, value=400, step=50, label="Max Iterations")
            max_depth = gr.Slider(3, 25, value=12, step=1, label="Max Depth")
            seed = gr.Number(value=42, precision=0, label="Random Seed")

            run_btn = gr.Button("🚀 Train Multi-Output Regressor", variant="primary")

            with gr.Row():
                plot_z = gr.Plot(label="Z Error Metrics")
                plot_x = gr.Plot(label="X Error Metrics")
            plot_y = gr.Plot(label="Y Error Metrics")
            metrics = gr.Markdown()

        with gr.TabItem("📖 Guide"):
            gr.Markdown(load_guide_content())

    gr.Markdown("---")
    gr.Markdown(
        "### 🔗 Links\n"
        "[Website](https://qsbench.github.io) | "
        "[Hugging Face](https://huggingface.co/QSBench) | "
        "[GitHub](https://github.com/QSBench)"
    )

    # ======================= CALLBACKS =======================
    def sync_features(dataset_key):
        """Update available feature choices when dataset changes."""
        df = load_single_dataset()
        features = get_available_feature_columns(df)
        defaults = default_feature_selection(features)
        return gr.update(choices=features, value=defaults)

    dataset_dropdown.change(
        refresh_explorer,
        [dataset_dropdown, split_dropdown],
        [split_dropdown, explorer_df, raw_qasm, transpiled_qasm, profile_box, summary_box],
    )
    split_dropdown.change(
        refresh_explorer,
        [dataset_dropdown, split_dropdown],
        [split_dropdown, explorer_df, raw_qasm, transpiled_qasm, profile_box, summary_box],
    )
    dataset_dropdown.change(sync_features, [dataset_dropdown], [feature_picker])

    run_btn.click(
        train_regressor,
        [feature_picker, test_size, max_iter, max_depth, seed],
        [plot_z, metrics, plot_x, plot_y],
    )

    demo.load(
        refresh_explorer,
        [dataset_dropdown, split_dropdown],
        [split_dropdown, explorer_df, raw_qasm, transpiled_qasm, profile_box, summary_box],
    )
    demo.load(sync_features, [dataset_dropdown], [feature_picker])

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(), css=CUSTOM_CSS)