import gradio as gr
import numpy as np
import pandas as pd
import torch
import spaces
from scipy import stats
from sklearn.decomposition import PCA
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModel, AutoTokenizer

# ============================================================
# Load BulkRNABert once at startup
# ============================================================
print("Loading BulkRNABert...")
config = AutoConfig.from_pretrained(
    "InstaDeepAI/BulkRNABert",
    trust_remote_code=True,
)
config.embeddings_layers_to_save = (4,)

tokenizer = AutoTokenizer.from_pretrained(
    "InstaDeepAI/BulkRNABert",
    trust_remote_code=True,
)

model = AutoModel.from_pretrained(
    "InstaDeepAI/BulkRNABert",
    config=config,
    trust_remote_code=True,
)
model.eval()
print(f"✓ Model loaded. Required genes: {config.n_genes}")

csv_path = hf_hub_download(
    repo_id="InstaDeepAI/BulkRNABert",
    filename="data/tcga_sample.csv",
    repo_type="model",
)
ref_df = pd.read_csv(csv_path)
ref_genes = [c for c in ref_df.columns if c != "identifier"]
print(f"✓ Reference genes: {len(ref_genes)}")


# ============================================================
# Helpers
# ============================================================
def align_genes(df, required_genes):
    aligned = pd.DataFrame(0.0, index=df.index, columns=required_genes)
    for g in required_genes:
        if g in df.columns:
            aligned[g] = df[g]
    return aligned


def get_embeddings(df, model, tokenizer, config):
    expr = df.to_numpy()
    expr = np.log10(1 + expr)

    assert expr.shape[1] == config.n_genes, \
        f"Expected {config.n_genes} genes, got {expr.shape[1]}"

    ids = tokenizer.batch_encode_plus(expr, return_tensors="pt")["input_ids"]

    with torch.no_grad():
        out = model(ids)

    return out["embeddings_4"].mean(axis=1).numpy()


# ============================================================
# GPU signal function (required for ZeroGPU startup)
# ============================================================
@spaces.GPU(duration=60)
def gpu_ping():
    """Signal function so ZeroGPU detects a registered @spaces.GPU handler."""
    return "GPU ready"


# ============================================================
# Main analysis function
# ============================================================
def analyze(tcga_file, cptac_file):
    if tcga_file is None or cptac_file is None:
        return {"error": "Please upload both TCGA and CPTAC TSV files."}

    try:
        tcga = pd.read_csv(tcga_file.name, sep="\t", index_col=0).T
        cptac = pd.read_csv(cptac_file.name, sep="\t", index_col=0).T

        tcga_a = align_genes(tcga, ref_genes)
        cptac_a = align_genes(cptac, ref_genes)

        tcga_emb = get_embeddings(tcga_a, model, tokenizer, config)
        cptac_emb = get_embeddings(cptac_a, model, tokenizer, config)

        pca = PCA(n_components=5, random_state=42)
        tcga_pca = pca.fit_transform(tcga_emb)
        cptac_pca = pca.transform(cptac_emb)

        tcga_centroid = tcga_pca.mean(axis=0)
        cptac_centroid = cptac_pca.mean(axis=0)
        centroid_dist = float(np.linalg.norm(tcga_centroid - cptac_centroid))

        tcga_d = np.linalg.norm(tcga_pca - tcga_centroid, axis=1)
        cptac_d = np.linalg.norm(cptac_pca - tcga_centroid, axis=1)

        etnsi = float((cptac_d.mean() - tcga_d.mean()) / tcga_d.std())
        _, p_val = stats.mannwhitneyu(tcga_d, cptac_d)

        X = np.vstack([tcga_pca, cptac_pca])
        y = np.array([0] * len(tcga_pca) + [1] * len(cptac_pca))

        clf = LogisticRegression(max_iter=1000, random_state=42)
        auc = cross_val_score(clf, X, y, cv=5, scoring="roc_auc")

        return {
            "status": "success",
            "tcga_samples": int(tcga_emb.shape[0]),
            "cptac_samples": int(cptac_emb.shape[0]),
            "embedding_dim": int(tcga_emb.shape[1]),
            "pca_variance_explained": round(float(pca.explained_variance_ratio_.sum()), 4),
            "centroid_distance": round(centroid_dist, 4),
            "etnsi_embedding": round(etnsi, 4),
            "etnsi_pvalue": float(p_val),
            "classifier_auc_mean": round(float(auc.mean()), 4),
            "classifier_auc_std": round(float(auc.std()), 4),
        }

    except Exception as e:
        return {"error": str(e)}


# ============================================================
# Gradio interface with hidden GPU handler registered
# ============================================================
with gr.Blocks(title="BulkRNABert UCEC Analysis") as iface:
    gr.Markdown("# BulkRNABert Novel UCEC Analysis")
    gr.Markdown(
        "Correlates a 9-gene signature with BulkRNABert embeddings across "
        "TCGA tumor and CPTAC normal samples. Outputs the embedding-based "
        "Tumor-Normal Separation Index (eTNSI) and classifier AUC."
    )

    with gr.Row():
        tcga_in = gr.File(label="TCGA expression TSV (genes × samples)")
        cptac_in = gr.File(label="CPTAC expression TSV (genes × samples)")

    run_btn = gr.Button("Run Analysis", variant="primary")
    out = gr.JSON(label="Novel Analysis Results")

    run_btn.click(fn=analyze, inputs=[tcga_in, cptac_in], outputs=out)

    # Hidden button registers the @spaces.GPU handler at startup
    gpu_btn = gr.Button("ping", visible=False)
    gpu_txt = gr.Textbox(visible=False)
    gpu_btn.click(fn=gpu_ping, inputs=[], outputs=gpu_txt)

iface.launch()