import tempfile
import base64
import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import io
from PIL import Image
from matplotlib.backends.backend_pdf import PdfPages

import os
_current_dir = os.path.dirname(os.path.abspath(__file__))
_image_path = os.path.join(_current_dir, "figures", "active_learning_performance.png")
_image_path_logo = os.path.join(_current_dir, "figures", "scarse_banner.png")

import scarse

# ── Global state ──────────────────────────────────────────────────────────────
VALID_AA = set("ACDEFGHIKLMNPQRSTVWY") # Valid amino acids
_train_file_path: str | None = None    # path to a temp CSV — needed by scarse.train()
_predict_file_path: str | None = None  # path to the uploaded prediction file
_trained: bool = False                 # flag so predict tab knows training has run
_predictions = None                    # store CV predictions from training
_target_list = []                      # List of unique endpoints
_problem_type = "Regression"           # Problem type
_class_labels = {}                     # target 
_metrics_cache = {}                    # stores raw metrics dict for PDF generation
_metrics_html_cache = ""
_prediction_file_path_result  = None

# ── Helpers ───────────────────────────────────────────────────────────────────

def _read_df(path: str) -> pd.DataFrame:
    """Read CSV or Excel from a file path."""
    if path.endswith((".xlsx", ".xls")):
        return pd.read_excel(path)
    return pd.read_csv(path, sep=None, engine="python")


def _df_to_temp_csv(df: pd.DataFrame) -> str:
    """Write a DataFrame to a temporary CSV and return the path."""
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv")
    df.to_csv(tmp.name, index=False)
    return tmp.name

    
def clear_predict_ui():
    return gr.update(visible=False), "⏳ Running predictions, please wait..."


def make_download_visible(file_path):
    if file_path is None:
        return gr.update(visible=False)
    return gr.update(value=file_path, visible=True)


def get_model_status_ready():
    return (
        "<div style='padding:10px;background:#dcfce7;border:1px solid #86efac;"
        "border-radius:8px;color:#166534'>"
        "✅ Model trained and ready for prediction."
        "</div>"
    )


def show_download_button():
    if _prediction_file_path_result:
        return gr.update(value=_prediction_file_path_result, visible=True)
    return gr.update(visible=False)
    

def get_model_status_not_ready():
    return (
        "<div style='padding:10px;background:#fef9c3;border:1px solid #fbbf24;"
        "border-radius:8px;color:#92400e'>"
        "⚠️ No model trained yet — please train a model in the Train Model tab first."
        "</div>"
    )

def generate_pdf_report():
    if not _predictions or not _metrics_cache:
        return None

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", prefix="scarse_report_")

    with PdfPages(tmp.name) as pdf:

        # ── Page 1: Metrics table ──────────────────────────────────────
        n_targets = len(_metrics_cache)
        fig, axes = plt.subplots(n_targets, 1, figsize=(8.27, max(4, n_targets * 3)))
        if n_targets == 1:
            axes = [axes]

        fig.suptitle("SCARSE Training Report", fontsize=16, fontweight="bold",
                     color="#4f46e5", y=1.02)

        for ax, (target, scores) in zip(axes, _metrics_cache.items()):
            ax.axis("off")
            ax.set_title(f"Target: {target}", fontsize=11, fontweight="bold",
                         color="#4f46e5", loc="left", pad=10)

            rows = [[k, f"{v:.4f}"] for k, v in scores.items()]
            tbl = ax.table(cellText=rows, colLabels=["Metric", "Value"],
                           loc="center", cellLoc="left")
            tbl.auto_set_font_size(False)
            tbl.set_fontsize(10)
            tbl.scale(1, 1.6)

            for j in range(2):
                tbl[0, j].set_facecolor("#4f46e5")
                tbl[0, j].set_text_props(color="white", fontweight="bold")
            for i in range(1, len(rows) + 1):
                for j in range(2):
                    tbl[i, j].set_facecolor("#f8fafc" if i % 2 == 0 else "white")

        plt.tight_layout()
        pdf.savefig(fig, bbox_inches="tight")
        plt.close(fig)

        # ── One plot page per target ───────────────────────────────────
        for target in _target_list:
            data = _predictions[target]
            y_true = list(data["y_vall_all"])
            y_pred = list(data["y_pred_all"])

            fig, ax = plt.subplots(figsize=(8, 6))

            if _problem_type == "Classification":
                if target in _class_labels:
                    labels = _class_labels[target]
                    y_true = [labels[int(v)] if int(v) < len(labels) else str(v) for v in y_true]
                    y_pred = [labels[int(v)] if int(v) < len(labels) else str(v) for v in y_pred]

                true_counts = pd.Series(y_true).value_counts()
                pred_counts = pd.Series(y_pred).value_counts()
                df_counts = pd.DataFrame({"True": true_counts, "Predicted": pred_counts}).fillna(0)
                df_counts.plot(kind="bar", ax=ax, color=["#4f46e5", "#a5b4fc"])
                ax.set_title(f"Class Distribution — {target}", fontsize=13, fontweight="bold")
                ax.set_xlabel("Class")
                ax.set_ylabel("Count")
                plt.xticks(rotation=45, ha="right")

            else:
                ax.scatter(y_true, y_pred, alpha=0.7, color="#4f46e5")
                min_v = min(min(y_true), min(y_pred))
                max_v = max(max(y_true), max(y_pred))
                ax.plot([min_v, max_v], [min_v, max_v], "r--",
                        linewidth=1.5, label="Perfect prediction")
                ax.set_title(f"True vs Predicted — {target}", fontsize=13, fontweight="bold")
                ax.set_xlabel("True values")
                ax.set_ylabel("Predicted values")
                ax.legend()

            plt.tight_layout()
            pdf.savefig(fig, bbox_inches="tight")
            plt.close(fig)

    return tmp.name


def show_report_button():
    path = generate_pdf_report()
    if path:
        return gr.update(value=path, visible=True)
    return gr.update(visible=False)
    

def _metrics_to_html(metrics: dict, prob_strategy: str, folds: int) -> str:
    """Render the metrics dict as a clean HTML table block."""
    title = f"{folds}-Fold Cross-Validation Training Performance"

    blocks = [f"<h3 style='margin-bottom:16px;color:#4f46e5'>{title}</h3>"]

    for target, scores in metrics.items():
        rows = "".join(
            f"<tr>"
            f"<td style='padding:8px 16px;border-bottom:1px solid #e5e7eb;font-weight:600;"
            f"color:#6b7280;text-transform:uppercase;font-size:0.78rem'>{k}</td>"
            f"<td style='padding:8px 16px;border-bottom:1px solid #e5e7eb;font-weight:700;"
            f"color:#111827'>{v:.4f}</td>"
            f"</tr>"
            for k, v in scores.items()
        )
        blocks.append(
            f"""
            <div style='margin-bottom:24px'>
              <div style='border-left:4px solid #4f46e5;padding-left:10px;
                          font-size:1rem;font-weight:600;color:#4f46e5;margin-bottom:10px'>
                Target: {target}
              </div>
              <table style='border-collapse:collapse;width:100%;
                            background:#fff;border:1px solid #e5e7eb;border-radius:8px;overflow:hidden'>
                <thead>
                  <tr style='background:#f8fafc'>
                    <th style='padding:8px 16px;text-align:left;color:#6b7280;
                               font-size:0.8rem;text-transform:uppercase'>Metric</th>
                    <th style='padding:8px 16px;text-align:left;color:#6b7280;
                               font-size:0.8rem;text-transform:uppercase'>Value</th>
                  </tr>
                </thead>
                <tbody>{rows}</tbody>
              </table>
            </div>
            """
        )

    return "".join(blocks)


def update_plot(selected_target):
    if not _predictions or selected_target is None:
        return None

    data = _predictions[selected_target]
    y_true = data["y_vall_all"]
    y_pred = data["y_pred_all"]

    fig, ax = plt.subplots()

    if _problem_type == "Classification":
        # Map encoded integers back to original class names
        if selected_target in _class_labels:
            labels = _class_labels[selected_target]
            y_true = [labels[int(v)] if int(v) < len(labels) else str(v) for v in y_true]
            y_pred = [labels[int(v)] if int(v) < len(labels) else str(v) for v in y_pred]

        true_counts = pd.Series(y_true).value_counts()
        pred_counts = pd.Series(y_pred).value_counts()

        df_counts = pd.DataFrame({
            "True": true_counts,
            "Predicted": pred_counts
        }).fillna(0)

        df_counts.plot(kind="bar", ax=ax)
        ax.set_title(f"Class Distribution ({selected_target})")
        ax.set_xlabel("Class")
        ax.set_ylabel("Count")
        plt.xticks(rotation=45, ha="right")

    else:
        ax.scatter(y_true, y_pred, alpha=0.7)
        min_v = min(min(y_true), min(y_pred))
        max_v = max(max(y_true), max(y_pred))
        ax.plot([min_v, max_v], [min_v, max_v], 'r--', label="Perfect prediction")
        ax.legend()
        ax.set_title(f"True vs Predicted ({selected_target})")
        ax.set_xlabel("True values")
        ax.set_ylabel("Predicted values")

    plt.tight_layout()
    return fig


def _make_pred_plot(predictions: dict):
    if not predictions:
        return None

    images = []

    for target, data in predictions.items():
        y_true = data["y_vall_all"]
        y_pred = data["y_pred_all"]

        fig, ax = plt.subplots(figsize=(5, 4))

        ax.scatter(y_true, y_pred, alpha=0.7)

        min_v = min(min(y_true), min(y_pred))
        max_v = max(max(y_true), max(y_pred))
        ax.plot([min_v, max_v], [min_v, max_v], 'r--', label="Perfect prediction")
        ax.legend()

        ax.set_title(f"True vs Predicted ({target})")
        ax.set_xlabel("True values")
        ax.set_ylabel("Predicted values")

        # --- convert fig → image ---
        buf = io.BytesIO()
        fig.savefig(buf, format="png", bbox_inches="tight", dpi=300)
        buf.seek(0)

        img = Image.open(buf).convert("RGB")
        images.append(img)

        plt.close(fig)

    return images

def validate_sequences(df, sequence_col):

    # -------------------------------------------------
    # sequence column selected
    # -------------------------------------------------

    if not sequence_col:
        raise gr.Error("Please select a sequence column.")

    # -------------------------------------------------
    # column exists
    # -------------------------------------------------

    if sequence_col not in df.columns:
        raise gr.Error(
            f"Sequence column '{sequence_col}' not found."
        )

    # -------------------------------------------------
    # NaN check
    # -------------------------------------------------

    if df[sequence_col].isna().any():
        raise gr.Error(
            f"Sequence column '{sequence_col}' contains missing values."
        )

    # -------------------------------------------------
    # amino acid validation
    # -------------------------------------------------

    invalid_rows = []

    for idx, seq in enumerate(df[sequence_col].astype(str)):

        seq = seq.strip().upper()

        # empty sequence
        if seq == "":
            invalid_rows.append(
                f"Row {idx + 1}: empty sequence"
            )
            continue

        invalid_chars = set(seq) - VALID_AA

        if invalid_chars:

            invalid_rows.append(
                f"Row {idx + 1}: invalid amino acids "
                f"{sorted(invalid_chars)}"
            )

    if invalid_rows:

        preview = "<br>".join(invalid_rows[:10])

        if len(invalid_rows) > 10:
            preview += (
                f"<br>... and {len(invalid_rows)-10} more rows"
            )

        raise gr.Error(
            f"Invalid amino acid sequences detected:"
            f"<br><br>{preview}"
        )


def validate_training_targets(
    df,
    target_cols,
    problem_type
):

    # -------------------------------------------------
    # target selected
    # -------------------------------------------------

    if not target_cols:
        raise gr.Error(
            "Please select at least one target column."
        )

    # -------------------------------------------------
    # existence + NaN
    # -------------------------------------------------

    for col in target_cols:

        if col not in df.columns:

            raise gr.Error(
                f"Target column '{col}' not found."
            )

        if df[col].isna().any():

            raise gr.Error(
                f"Target column '{col}' contains missing values."
            )

    # -------------------------------------------------
    # regression validation
    # -------------------------------------------------

    if problem_type == "Regression":

        for col in target_cols:

            if not pd.api.types.is_numeric_dtype(df[col]):

                raise gr.Error(
                    f"Regression target '{col}' "
                    f"must contain numeric values."
                )

            if df[col].nunique() < 2:

                raise gr.Error(
                    f"Regression target '{col}' "
                    f"contains fewer than 2 unique values."
                )

    # -------------------------------------------------
    # classification validation
    # -------------------------------------------------

    elif problem_type == "Classification":

        for col in target_cols:

            n_unique = df[col].nunique()

            if n_unique < 2:

                raise gr.Error(
                    f"Classification target '{col}' "
                    f"contains only one class."
                    f"(requires at least 2 classes)"
                )

            if n_unique == len(df):

                raise gr.Error(
                    f"Classification target '{col}' "
                    f"contains all unique values "
                    f"(looks like an ID column)."
                )

            class_counts = df[col].value_counts()

            tiny_classes = class_counts[class_counts < 2]

            if len(tiny_classes) > 0:

                raise gr.Error(
                    f"Classification target '{col}' "
                    f"contains classes with fewer than "
                    f"2 samples: {list(tiny_classes.index)}"
                )
        
# ── File-upload callbacks ─────────────────────────────────────────────────────

def load_train_file(file):
    """Parse uploaded training file and populate column selectors."""
    global _train_file_path
    if file is None:
        return gr.update(choices=[], value=None), gr.update(choices=[], value=[])
    try:
        df = _read_df(file.name)
        # scarse.train() expects a file path, so persist Excel as temp CSV
        if file.name.endswith((".xlsx", ".xls")):
            _train_file_path = _df_to_temp_csv(df)
        else:
            _train_file_path = file.name
        cols = list(df.columns)
        return (
            gr.update(choices=cols, value=cols[0] if cols else None),
            gr.update(choices=cols, value=[]),
        )
    except Exception as e:
        raise gr.Error(f"Could not read training file: {str(e)}")

def load_predict_file(file):
    """Parse uploaded prediction file and populate sequence column selector."""
    global _predict_file_path
    if file is None:
        return gr.update(choices=[], value=None)
    try:
        df = _read_df(file.name)
        # Convert Excel to CSV so ModelOptimization can read with pd.read_csv
        if file.name.endswith((".xlsx", ".xls")):
            _predict_file_path = _df_to_temp_csv(df)
        else:
            _predict_file_path = file.name
        cols = list(df.columns)
        return gr.update(choices=cols, value=cols[0] if cols else None)
    except Exception as e:
        raise gr.Error(f"Could not read prediction file: {str(e)}")

def clear_training_ui():
    return None, gr.update(choices=[], value=None), "", "", False

# ── Training callback ─────────────────────────────────────────────────────────

def train_model(
    sequence_col,
    target_cols,
    prob_strategy,
    random_seed,
    cv_folds,
    progress=gr.Progress(track_tqdm=True),
):
    global _trained, _predictions, _target_list, _problem_type, _metrics_html_cache, _metrics_cache, _prediction_file_path_result, _class_labels
    _problem_type = prob_strategy
    _trained = False
    _predictions = None
    _metrics_html_cache = ""
    _metrics_cache = {}
    _prediction_file_path_result  = None
    _target_list = []                   
    _class_labels = {}  

    progress(0, desc="Initialising...")

    if not sequence_col:
        raise gr.Error("Please select a sequence column.")
    if not target_cols:
        raise gr.Error("Please select at least one target column.")
    if _train_file_path is None:
        raise gr.Error("No training file loaded.")

    try:
        df_temp = _read_df(_train_file_path)
        if prob_strategy != "Classification":
            for col in target_cols:
                df_temp[col] = df_temp[col].astype(float)
        if prob_strategy == "Classification":
            for col in target_cols:
                _class_labels[col] = sorted(df_temp[col].astype(str).unique())

        if len(df_temp) < int(cv_folds):
            raise gr.Error(
                f"Dataset only has {len(df_temp)} rows but {int(cv_folds)} CV folds were requested. "
                f"Please reduce CV folds to at most {len(df_temp)}."
            )

        validate_sequences(df_temp, sequence_col)
        validate_training_targets(df_temp, target_cols, prob_strategy)

        prob_bool = prob_strategy == "Classification"
        progress(0, desc="Initialising...")

        metrics, predictions = scarse.train(
            data_path=_train_file_path,
            classification=prob_bool,
            seq_col=sequence_col,
            score_col=list(target_cols),
            random_seed=int(random_seed),
            folds=int(cv_folds),
            optuna_print=False,
            progress=progress,
        )

        _trained = True
        _predictions = predictions
        _target_list = list(predictions.keys())
        _metrics_cache = metrics
        _metrics_html_cache = _metrics_to_html(metrics, prob_strategy, int(cv_folds))

        #return gr.update(choices=_target_list, value=_target_list[0] if _target_list else None)
        return "✅ Model trained successfully!", True

    except Exception as e:
        raise gr.Error(f"Training error: {str(e)}")


def get_training_metrics():
    return _metrics_html_cache

def get_training_status():
    return "✅ Model trained successfully!"

def get_target_selector():
    """Called after training to populate the dropdown."""
    return gr.update(choices=_target_list, value=_target_list[0] if _target_list else None)


# ── Prediction callback ───────────────────────────────────────────────────────

def predict_and_download(sequence_col, trained, progress=gr.Progress(track_tqdm=True)):
    global _predict_file_path, _prediction_file_path_result

    if not trained:
        raise gr.Error("No model has been trained yet. Please go to the Train Model tab first.")
    if _predict_file_path is None:
        raise gr.Error("No prediction data loaded.")
    if not sequence_col:
        raise gr.Error("Please select a sequence column.")

    try:
        df = _read_df(_predict_file_path)
        validate_sequences(df, sequence_col)

        df_pred = scarse.pred(
            data_path=_predict_file_path,
            seq_col=sequence_col,
            progress=progress,
        )

        first_pred_col = [c for c in df_pred.columns if c.startswith("pred_")]
        if first_pred_col and pd.api.types.is_numeric_dtype(df_pred[first_pred_col[0]]):
            df_pred = df_pred.sort_values(by=first_pred_col[0], ascending=False)

        out_tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="scarse_predictions_")
        df_pred.to_csv(out_tmp.name, index=False)
        #return out_tmp.name  # single output
        _prediction_file_path_result = out_tmp.name
        return "✅ Predictions complete! Click the button below to download."

    except Exception as e:
        raise gr.Error(f"Prediction error: {str(e)}")


def get_predict_status():
    return f"✅ Predictions ready!"

# ── UI ────────────────────────────────────────────────────────────────────────

CSS = """
body { font-family: 'Inter', sans-serif; }
#title  { text-align:center; margin-bottom: 4px; }
#subtitle { text-align:center; color: #6b7280; margin-bottom: 24px; }
.status-box textarea { font-size: 0.9rem !important; }
"""

with gr.Blocks(title="SCARSE") as demo:

    trained_state = gr.State(False)
    file_path_state = gr.State(None) 
    #gr.Markdown("# 💊 SCARSE", elem_id="title")
    #gr.Markdown(
    #    "Sequence-based low-resource peptide property prediction.  \n"
    #    "Upload your peptide dataset, train a model, and download predictions.",
    #    elem_id="subtitle",
    #)
    with open(_image_path_logo, "rb") as f:
        encoded = base64.b64encode(f.read()).decode()

    gr.HTML(f"""
    <div style="display:flex; justify-content:center;">
        <img
            src="data:image/png;base64,{encoded}"
            style="width:800px; display:block; pointer-events:none;"
        />
    </div>
    """)

    with gr.Tabs():

        # ── Train tab ──────────────────────────────────────────────────────────
        with gr.Tab("🎛️ Train Model"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Training Data")
                    train_file = gr.File(
                        label="Upload CSV / Excel",
                        file_types=[".csv", ".xlsx", ".xls"],
                    )
                    train_seq_col = gr.Dropdown(
                        label="Sequence Column",
                        choices=[],
                        interactive=True,
                    )
                    train_target_cols = gr.CheckboxGroup(
                        label="Target Column(s)",
                        choices=[],
                        interactive=True,
                    )

                with gr.Column(scale=1):
                    gr.Markdown("### Training Settings")
                    prob_strategy = gr.Radio(
                        choices=["Regression", "Classification"],
                        value="Regression",
                        label="Problem Type",
                    )
                    random_seed = gr.Number(
                        value=42,
                        label="Random Seed",
                        precision=0,
                        info="Random seed controls reproducibility. Same seed = same results.",
                    )
                    cv_folds = gr.Number(
                        value=10,
                        label="CV Folds",
                        minimum=2,
                        maximum=20,
                        precision=0,
                        info="Number of cross-validation splits. Higher = more robust but slower (we recommend to use the default value of 10, but can be between 2 and 20).",
                    )
            gr.Markdown(
                """
                <div style="margin-top:10px; margin-bottom:10px; color:gray; font-size:0.9rem">
                ⏱️ <b>Estimated runtime</b><br>
                • Regression (~100 sequences): ~3-4 min<br>
                • Classification (~100 sequences): ~15-25 min<br>
                <span style="font-size:0.8rem">(Runtime will vary depending on dataset and dataset size. If running locally, this can run significantly faster depending on hardware)</span>
                </div>
                """
                )
            
            train_btn    = gr.Button("⚡ Train Model", variant="primary")
            status_train = gr.Textbox(label="Status", interactive=False, elem_classes="status-box")
            metrics_html = gr.HTML(label="Results")
            target_selector = gr.Dropdown(
                label="Select Target",
                #choices=[],
                #interactive=True
            )

            pred_plot = gr.Plot(label="True vs Predicted")

            report_download = gr.DownloadButton(
                label="📄 Download Report (PDF)",
                visible=False,
                variant="primary",
            )

            train_file.change(
                load_train_file,
                inputs=train_file,
                outputs=[train_seq_col, train_target_cols],
            )

            train_btn.click(
                clear_training_ui,
                inputs=[],
                outputs=[pred_plot, target_selector, metrics_html, status_train, trained_state],
            ).then(
                train_model,
                inputs=[train_seq_col, train_target_cols, prob_strategy, random_seed, cv_folds],
                outputs=[status_train, trained_state],
            ).then(
                get_target_selector,
                inputs=[],
                outputs=[target_selector],
            ).then(
                get_training_metrics,
                inputs=[],
                outputs=[metrics_html],
            ).then(
                show_report_button,     
                inputs=[],
                outputs=[report_download],
            )
            target_selector.change(
                update_plot,
                inputs=target_selector,
                outputs=pred_plot
            )

        # ── Predict tab ────────────────────────────────────────────────────────
        with gr.Tab("🔮 Predict") as predict_tab:
            file_path_state = gr.State(None)
            model_status_indicator = gr.HTML(
                value="<div style='padding:10px;background:#fef9c3;border:1px solid #fbbf24;"
                      "border-radius:8px;color:#92400e'>"
                      "⚠️ No model trained yet — please train a model in the Train Model tab first."
                      "</div>"
            )
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Prediction Input")
                    predict_file = gr.File(
                        label="Upload CSV / Excel",
                        file_types=[".csv", ".xlsx", ".xls"],
                    )
                    predict_seq_col = gr.Dropdown(
                        label="Sequence Column",
                        choices=[],
                        interactive=True,
                    )

                with gr.Column(scale=1):
                    gr.Markdown("### Output")
                    #status_predict   = gr.Textbox(label="Status", interactive=False, elem_classes="status-box")
                    #predict_download = gr.File(label="Download Predictions CSV", interactive=False)

                    status_predict   = gr.Textbox(label="Status", interactive=False, elem_classes="status-box")
                    predict_download = gr.DownloadButton( 
                        label="📥 Download Predictions CSV",
                        visible=False,
                        variant="primary"
                    )


            gr.Markdown(
                """
                <small style="color:gray">
                ⏱️ Estimated runtime:  
                • ~5-7 minutes for predicting on 500 sequences, using 100 sequences for training  
                (Runtime will vary depending on input size and amino acid sequence length.)
                </small>
                """
                )
            
            predict_btn = gr.Button("📥 Predict", variant="primary")

            predict_file.change(
                load_predict_file,
                inputs=predict_file,
                outputs=[predict_seq_col],
            )

            predict_btn.click(
                clear_predict_ui,
                inputs=[],
                outputs=[predict_download, status_predict],
            ).then(
                predict_and_download,
                inputs=[predict_seq_col, trained_state],
                outputs=[status_predict],      
            ).then(
                show_download_button,
                inputs=[],
                outputs=[predict_download],
            )

        predict_tab.select(       
            fn=lambda trained: get_model_status_ready() if trained else get_model_status_not_ready(),
            inputs=[trained_state],
            outputs=[model_status_indicator],
        )

        # ── About tab ──────────────────────────────────────────────────────────
        # ── About tab ──────────────────────────────────────────────────────────
        with gr.Tab("ℹ️ Help"):
            gr.Markdown(
                """
                # SCARSE - User Documentation
                
                ### 1. Introduction
                Welcome to the user documentation for SCARSE: Small-sample Classification And Regression Solution for low-resource peptide Engineering. This application allows researchers to train SCARSE on peptide sequences without writing code. It utilizes the ESM-2 650M foundation model to extract features from sequences and optimizes the top model (GPR (regression) / ExtraTrees (classification)) for your specific target variables.
                
                ---
                
                ### 2. Data Preparation
                Proper formatting is crucial for the AI framework to understand your data.
                * **Format**: Your data must be in CSV (Comma Separated Values).
                * **Sequence Column**: Must contain one column with amino acid sequences. It should be given in the one-letter amino acid sequence notation (ACDEFGHIKLMNPQRSTVWY).
                * **Target Columns**: Must contain one or more columns with either numerical values or categorical values you want to predict.
                * **Clean Data**: Ensure there are no empty rows or missing values in your data.
                * **Note**: The file is allowed to contain more columns than just the sequence and targets.
                * **Selection**: You will later select which column corresponds to each within the application.
                * See the tutorials of SCARSE on the GitHub page for examples how the CSV file should look like.

                * **Privacy & Data Handling**:
                    * Uploaded datasets, generated embeddings, trained models, and prediction files are **not stored permanently** by this application. 
                    * All uploaded data is processed temporarily in memory or ephemeral session storage only for the duration of the active session. 
                    * No datasets or trained models are used for further training, shared with third parties, or retained after the session ends.
                
                ---
                ### 3. Training the Model
                Follow these steps to train your custom model:
                1. Go to the **Train Model** card.
                2. **Upload Dataset**: Click "Click to Upload".
                3. **Select Sequence Column**: Choose the column containing your amino acid sequences from the dropdown menu.
                4. **Select Target Columns**: Check the boxes for the endpoints/properties you wish to predict.
                5. **Problem Type**: Select either regression (should be numeric targets, such as 4.6, 3.34, 34.5...) or classification (should be class targets, such as High, Medium, Low...).
                6. **Requirements**: You can select multiple columns, but they must be of the same problem type (all regression-based or all classification-based). The predictions made on each target column are independent of each other, meaning one model is trained per target column.
                7. **Mixed Types**: If you have one column for regression and one for classification, you must start by predicting for one and then later select and predict for the next.
                8. **Seed (Optional)**: Define which random seed to use for the training process.
                9. **CV Folds (Optional)**: Define the number of folds used for cross-validation.
                10. **Train**: Click the "Train Model" button.
                11. **Report**: Once complete, a Performance Report will appear at the bottom of the screen showing metrics for each target from the CV.

                ---
                ### 4. Understanding the Metrics

                #### Regression Metrics

                | Metric | What it measures | Score range |
                |--------|-----------------|-------------|
                | **MSE** (Mean Squared Error) | Average squared difference between predicted and true values. Penalises large errors heavily. Lower is better. | 0 = perfect |
                | **RMSE** (Root Mean Squared Error) | Same as MSE but in the same units as your target, making it easier to interpret (e.g. if predicting IC50 in nM, RMSE is also in nM). Lower is better. | 0 = perfect |
                | **MAE** (Mean Absolute Error) | Average absolute difference between predicted and true values. More robust to outliers than RMSE. Lower is better. | 0 = perfect |
                | **R²** (R-squared) | How much of the variation in the target the model explains. A model that always predicts the mean scores **0** — anything above 0 means the model has learned something. Negative values mean the model is worse than just predicting the mean. | 0 = no better than predicting the mean, 1 = perfect |
                | **Spearman correlation** | Whether the model correctly **ranks** sequences relative to each other, regardless of exact predicted values. A random model scores around **0**. | −1 to 1, higher is better |
        
                ---
            
                #### Classification Metrics
            
                | Metric | What it measures | Score range |
                |--------|-----------------|-------------|
                | **Accuracy** | Fraction of sequences correctly classified. Can be misleading if classes are imbalanced — a model that always predicts the majority class can still score high. | 0 to 1, higher is better |
                | **Balanced Accuracy** | Like accuracy, but gives equal weight to each class regardless of size. More reliable than accuracy for imbalanced data. | 0 to 1, higher is better |
                | **F1 weighted** | Balance between precision (are positive predictions correct?) and recall (are positives found?), averaged across classes weighted by their frequency. A random model scores close to **0**. | 0 to 1, higher is better |
                | **MCC** (Matthews Correlation Coefficient) | Overall agreement between predicted and actual classes. Unlike accuracy, it only scores high when the model performs well on **all** classes. A random model scores around **0**, a perfect model scores **1**, and a completely wrong model scores **−1**. | −1 to 1, higher is better |
                | **ROC AUC** | How well the model separates classes across all possible decision thresholds. A random classifier scores **0.5** (equivalent to random guessing), a perfect model scores **1**. | 0.5 = random, 1 = perfect |
            
                ---
                ### 5. Correlate to active learning end-point performance
                Below we illustrate the relation between CV R² score and end-point active learning performance.
                The y-axis display how many times better performance SCARSE guided active learning delivers compared to random sampling when looking at the accumulation of top 10% of peptides.
                By comparing the CV R² score of your data to the corresponding figure below for your dataset size one can get and indication of how suitable your data is combined with SCARSE to perform active learning peptide engineering.
                Note that this can only be used as a guide to evaluate regression problem performance.
                """)
            gr.Image(value=_image_path, width=700)   
            gr.Markdown(
                """
                
                ### 6. Generating Predictions
                Once a model is trained, you can predict values for new sequences:
                1. Go to the **Predict** card.
                2. **Upload Input**: Click "Click to Upload" and upload the sequences you wish to predict on, using the trained model.
                3. **Select Sequence Column**: Choose the column in the new file that contains the sequences you want to predict on.
                4. **Predict**: Click the "Predict & Download" button, then press the download link to download predictions.
                
                ---

                ### 7. Notes and best practices

                * **Call order matters**  
                  You must run training before making predictions, as the trained model is stored internally.
                
                * **Data quality is important**  
                  * Ensure no missing or empty sequences and endpoint/target values  
                  * Use consistent formatting (standard amino acid codes(ACDEFGHIKLMNPQRSTVWY))
                
                * **Small datasets are supported**  
                  SCARSE have been evaluated for datasets as small as ~20 samples, but performance generally improves with more data.

                ---
                
                ### 8. Troubleshooting
                * **"Error reading file"**: Ensure your CSV file is not currently open in Excel while trying to upload it.
                * **“Error when defining multiple target columns”**: Make sure that you only select targets of the same problem type (all regression-based or all classification-based).
                
                ---
                
                ### 9. Citation
                If using this application, please cite:
                **Andrekson L, Rydbergh R, Mercado R, Wenzel M. AI-guided discovery for low-resource peptide engineering using evolutionary scale modeling. bioRxiv. 2026. https://doi.org/10.64898/2026.06.25.734678.**
                                
                GitHub of SCARSE: https://github.com/LeoAnd00/SCARSE                 
                """
                )

            


if __name__ == "__main__":
    demo.launch(
        css=CSS,
        server_name="0.0.0.0",
    )