"""
Forge MIP Embedding Explorer — HuggingFace Spaces app.
Single MIP: upload a .mps/.lp file → learned instance embedding + stats.
Multi-MIP:  upload several files → interactive 2-D scatter plot (PCA / t-SNE).
Embedding type: Forge pre-trained model (mip_to_embeddings).
"""

import os
import tempfile
import traceback
from pathlib import Path

import gradio as gr
import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import spaces
from sklearn.decomposition import PCA

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_ACCEPTED_EXTS = (".mps", ".lp", ".mps.gz", ".lp.gz")

SAMPLE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "samples")
MODEL_PKL = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                          "models", "forge_pretrain_trained.pkl")
TRAIN_CONFIG_YAML = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                  "configs", "train_config.yaml")


def _extract_label(filename: str) -> str:
    """Best-effort: pull the problem-type token from DMIPLIB filenames."""
    stem = Path(filename).stem
    if stem.endswith(".mps") or stem.endswith(".lp"):
        stem = Path(stem).stem
    parts = stem.split("_")
    return parts[1] if len(parts) > 1 else stem


def _clean_embed_matrix(mat: np.ndarray) -> np.ndarray:
    mat = np.where(np.isfinite(mat), mat, 0.0)
    col_sum = mat.sum(axis=0, keepdims=True)
    return mat / (col_sum + 1e-10)


def _build_embedder():
    from forge.embeddings import Forge
    return Forge(train_config_yaml=TRAIN_CONFIG_YAML)


def _mip_stats(path: str) -> dict:
    """Read basic MIP stats directly from the Gurobi model attributes."""
    import gurobipy as gp
    model = gp.read(path)
    type_map = {
        (False, False, False): "LP",
        (True, False, False): "MILP",
        (False, True, False): "QP",
        (True, True, False): "MIQP",
        (False, False, True): "QCP",
        (True, False, True): "MIQCP",
    }
    prob_type = type_map.get((bool(model.IsMIP), bool(model.IsQP), bool(model.IsQCP)), "MILP")
    return {
        "probtype": prob_type,
        "n_vars": int(model.NumVars),
        "n_constr": int(model.NumConstrs),
        "n_nzcnt": int(model.NumNZs),
        "num_b_variables": int(model.NumBinVars),
        "num_i_variables": int(model.NumIntVars) - int(model.NumBinVars),
        "num_c_variables": int(model.NumVars) - int(model.NumIntVars),
    }


def _embed_file(path: str, forge):
    from forge.pipeline import mip_to_embeddings
    from forge.utils import Constants

    with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as tmp:
        output_pkl = tmp.name
    try:
        result = mip_to_embeddings(
            forge=forge,
            input_forge_pkl=MODEL_PKL,
            model_type=Constants.FORGE_PRE_TRAIN,
            input_mips=path,
            input_mip_instances_file=None,
            output_mip_to_embeddings_pkl=output_pkl,
            instance_embedding_only=True,
        )
    finally:
        if os.path.exists(output_pkl):
            os.remove(output_pkl)

    if not result:
        return None
    key = next(iter(result))
    return result[key]


# ---------------------------------------------------------------------------
# Single-MIP
# ---------------------------------------------------------------------------

@spaces.GPU(duration=60)
def embed_single(mip_file):
    if mip_file is None:
        return None, None, None, "⚠ Upload a MIP instance first."

    path = mip_file if isinstance(mip_file, str) else mip_file.name

    try:
        forge = _build_embedder()
        emb = _embed_file(path, forge)
        if emb is None:
            return None, None, None, "⚠ Could not process the file."

        stats = _mip_stats(path)

        stats_df = pd.DataFrame({
            "Metric": [
                "Problem type",
                "Variables",
                "Constraints",
                "Non-zeros",
                "Binary vars",
                "Integer vars",
                "Continuous vars",
            ],
            "Value": [
                stats["probtype"],
                stats["n_vars"],
                stats["n_constr"],
                stats["n_nzcnt"],
                stats["num_b_variables"],
                stats["num_i_variables"],
                stats["num_c_variables"],
            ],
        })

        vals = np.where(np.isfinite(emb.instance_embedding),
                        emb.instance_embedding, 0.0)
        names = [f"dim_{i}" for i in range(len(vals))]

        top_n = min(40, len(names))
        top_idx = np.argsort(np.abs(vals))[-top_n:][::-1]
        top_names = [names[i] for i in top_idx]
        top_vals = [float(vals[i]) for i in top_idx]

        fig = go.Figure(
            go.Bar(
                x=top_names,
                y=top_vals,
                marker=dict(
                    color=top_vals,
                    colorscale="RdBu",
                    cmid=0,
                    showscale=True,
                ),
            )
        )
        fig.update_layout(
            title=f"Top-{top_n} Instance Embedding Dimensions — {Path(path).name}",
            xaxis_title="Embedding dimension",
            yaxis_title="Value",
            xaxis_tickangle=-55,
            height=480,
            margin=dict(b=180),
        )

        csv_df = pd.DataFrame({"dimension": names, "value": emb.instance_embedding})

        return (
            stats_df,
            fig,
            csv_df,
            f"✓ Processed {Path(path).name}  "
            f"({stats['n_vars']} vars, {stats['n_constr']} constraints)",
        )

    except Exception as exc:
        tb = traceback.format_exc()
        msg = (
            f"❌ {exc}\n\n"
            "Make sure gurobipy is installed and a valid Gurobi licence is available.\n"
            "For small instances the size-limited free licence works out of the box.\n"
            "For larger instances set GRB_WLSACCESSID / GRB_WLSSECRET / GRB_LICENSEID "
            "as Space secrets.\n\n"
            f"Traceback:\n{tb}"
        )
        return None, None, None, msg


# ---------------------------------------------------------------------------
# Multi-MIP
# ---------------------------------------------------------------------------

@spaces.GPU(duration=120)
def embed_multi(mip_files, method, label_mode, max_files=10):
    if not mip_files:
        return None, None, "⚠ Upload at least 2 MIP files."

    paths = [f if isinstance(f, str) else f.name for f in mip_files]

    if len(paths) > max_files:
        paths = paths[:max_files]

    try:
        embedder = _build_embedder()
        embeddings, names, labels, errors = [], [], [], []

        for p in paths:
            try:
                emb = _embed_file(p, embedder)
                if emb is None:
                    errors.append(f"{Path(p).name}: no embedding returned")
                    continue
                embeddings.append(emb.instance_embedding)
                name = Path(p).name
                names.append(name)
                labels.append(_extract_label(name) if label_mode == "problem type" else name)
            except Exception as exc:
                errors.append(f"{Path(p).name}: {exc}")

        if len(embeddings) < 2:
            return None, None, "⚠ Need ≥ 2 valid embeddings.\n" + "\n".join(errors)

        mat = _clean_embed_matrix(np.array(embeddings))

        dim_label = ""

        if method == "PCA":
            reducer = PCA(n_components=2, random_state=42)
            coords = reducer.fit_transform(mat)
            ev = reducer.explained_variance_ratio_
            dim_label = f"PCA  (PC1 {ev[0]:.1%}, PC2 {ev[1]:.1%})"

        elif method == "t-SNE":
            from sklearn.manifold import TSNE
            perp = max(5, min(30, len(embeddings) - 1))
            coords = TSNE(n_components=2, perplexity=perp, random_state=42,
                          max_iter=1000).fit_transform(mat)
            dim_label = f"t-SNE  (perplexity={perp})"

        else:  # UMAP
            try:
                import umap
                coords = umap.UMAP(n_components=2, random_state=42).fit_transform(mat)
                dim_label = "UMAP"
            except ImportError:
                reducer = PCA(n_components=2, random_state=42)
                coords = reducer.fit_transform(mat)
                ev = reducer.explained_variance_ratio_
                dim_label = f"PCA (UMAP unavailable)  PC1 {ev[0]:.1%} PC2 {ev[1]:.1%}"

        df_plot = pd.DataFrame({
            "Dim 1": coords[:, 0],
            "Dim 2": coords[:, 1],
            "label": labels,
            "file": names,
        })

        unique_labels = df_plot["label"].unique()
        palette = px.colors.qualitative.Plotly + px.colors.qualitative.Set2
        color_map = {lbl: palette[i % len(palette)]
                     for i, lbl in enumerate(sorted(unique_labels))}

        fig = go.Figure()
        for lbl in sorted(unique_labels):
            sub = df_plot[df_plot["label"] == lbl]
            fig.add_trace(go.Scatter(
                x=sub["Dim 1"],
                y=sub["Dim 2"],
                mode="markers+text",
                text=sub["file"] if len(paths) <= 15 else None,
                textposition="top center",
                marker=dict(size=10, color=color_map[lbl]),
                name=lbl,
                hovertemplate="<b>%{customdata}</b><br>Dim1=%{x:.3f}  Dim2=%{y:.3f}",
                customdata=sub["file"],
            ))

        fig.update_layout(
            title=f"MIP Embedding Space — {dim_label}",
            xaxis_title="Dim 1",
            yaxis_title="Dim 2",
            legend_title="Label",
            height=560,
        )

        summary = pd.DataFrame({
            "File": names,
            "Label": labels,
            "Dim 1": coords[:, 0].round(4),
            "Dim 2": coords[:, 1].round(4),
        })

        status = f"✓ Embedded {len(embeddings)} MIPs  |  {dim_label}"
        if errors:
            status += f"\n⚠ {len(errors)} failed:\n" + "\n".join(errors)

        return fig, summary, status

    except Exception as exc:
        tb = traceback.format_exc()
        return None, None, f"❌ {exc}\n\n{tb}"


# ---------------------------------------------------------------------------
# Sample-data helper
# ---------------------------------------------------------------------------

def load_sample_files():
    """Return all files from the bundled samples directory."""
    if not os.path.isdir(SAMPLE_DIR):
        return [], "⚠ Sample data directory not found."

    from collections import defaultdict
    by_type = defaultdict(list)
    for f in os.listdir(SAMPLE_DIR):
        if any(f.lower().endswith(e) for e in _ACCEPTED_EXTS):
            lbl = _extract_label(f)
            by_type[lbl].append(os.path.join(SAMPLE_DIR, f))

    selected = []
    for lbl in sorted(by_type):
        selected.extend(sorted(by_type[lbl]))

    if not selected:
        return [], "⚠ No MIP files found in samples directory."

    selected = selected[:10]
    return selected, f"Loaded {len(selected)} sample files (capped at 10 for free-tier GPU)."


def use_sample_data(method, label_mode):
    # Not decorated with @spaces.GPU itself — it delegates to embed_multi,
    # which is already decorated and will request the GPU when called.
    #
    # NOTE: ZeroGPU runs the decorated call in a separate worker process and
    # pickles the arguments to send them over. embed_multi already accepts
    # plain path strings (`f if isinstance(f, str) else f.name`), so pass
    # strings directly rather than wrapping them in a locally-defined class
    # (local/inner classes can't be pickled, which is what caused the
    # PicklingError).
    paths, msg = load_sample_files()
    if not paths:
        return None, None, msg

    fig, table, status = embed_multi(paths, method, label_mode)
    return fig, table, msg + "\n" + (status or "")


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

_DESCRIPTION = """
# 
For more details, visit [Forge: Foundational Optimization Representation from Graph Embeddings](https://skadio.github.io/forge/)
"""

with gr.Blocks(title="Forge MIP Embeddings", theme=gr.themes.Soft()) as demo:

    gr.Markdown(_DESCRIPTION)

    with gr.Tab("Single MIP"):
        gr.Markdown("Upload a MIP instace to visualize its Forge embedding.")
        with gr.Row():
            single_file = gr.File(
                label="MIP instance (.mps / .lp / .mps.gz / .lp.gz)",
                file_types=[".mps", ".lp", ".gz"],
            )
            single_btn = gr.Button("Generate Forge Embedding", variant="primary", scale=0)

        single_status = gr.Textbox(label="Status", lines=3, interactive=False)

        with gr.Row():
            stats_df = gr.DataFrame(label="MIP Statistics", interactive=False)

        single_plot = gr.Plot(label="Top Instance Features")
        full_emb_df = gr.DataFrame(label="Full Embedding Vector", interactive=False,
                                   wrap=True)

        single_btn.click(
            fn=embed_single,
            inputs=single_file,
            outputs=[stats_df, single_plot, full_emb_df, single_status],
        )

    with gr.Tab("Multiple MIPs"):
        gr.Markdown(
            "Upload MIP instances to visualize their Forge embeddings. "
            "Max 10 instances on this free-tier CPU. Alternatively, download this app and run on your GPU."
        )

        with gr.Row():
            multi_files = gr.File(
                label="MIP files (multiple)",
                file_count="multiple",
                file_types=[".mps", ".lp", ".gz"],
            )

        with gr.Row():
            method_radio = gr.Radio(
                ["PCA", "t-SNE", "UMAP"],
                value="PCA",
                label="Dimensionality reduction",
            )
            label_mode = gr.Radio(
                ["problem type", "filename"],
                value="problem type",
                label="Colour by",
            )
            multi_btn = gr.Button("Embed & Visualise", variant="primary", scale=0)

        multi_status = gr.Textbox(label="Status", lines=3, interactive=False)
        scatter_plot = gr.Plot(label="2D Forge Embedding Space")
        coords_df = gr.DataFrame(label="Embedding Coordinates", interactive=False)

        multi_btn.click(
            fn=embed_multi,
            inputs=[multi_files, method_radio, label_mode],
            outputs=[scatter_plot, coords_df, multi_status],
        )

        gr.Markdown("---\n### Demo with bundled MIP instances")
        gr.Markdown(
            "No files to upload? Run on MIP instances bundled within this space."
        )
        with gr.Row():
            sample_method = gr.Radio(
                ["PCA", "t-SNE", "UMAP"],
                value="PCA",
                label="Reduction method",
            )
            sample_label = gr.Radio(
                ["problem type", "filename"],
                value="problem type",
                label="Colour by",
            )
            sample_btn = gr.Button("Demo on Sample Data", variant="primary", scale=0)

        sample_status = gr.Textbox(label="Status", lines=2, interactive=False)
        sample_plot = gr.Plot(label="Sample Embedding Space")
        sample_df = gr.DataFrame(label="Sample Coordinates", interactive=False)

        sample_btn.click(
            fn=use_sample_data,
            inputs=[sample_method, sample_label],
            outputs=[sample_plot, sample_df, sample_status],
        )

    gr.Markdown(
        "---\n"
        "[Forge Homepage](https://skadio.github.io/forge/)"
    )

if __name__ == "__main__":
    demo.launch()