import gradio as gr
import numpy as np
import pandas as pd
import torch
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 model ---
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()

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"✓ Model ready. Required genes: {len(ref_genes)}")


def align_genes(df, required):
    out = pd.DataFrame(0.0, index=df.index, columns=required)
    for g in required:
        if g in df.columns:
            out[g] = df[g]
    return out


def embed(df):
    expr = np.log10(1 + df.to_numpy())
    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()


def analyze(tcga_file, cptac_file):
    """Novel analysis: signature vs embeddings."""
    if tcga_file is None or cptac_file is None:
        return {"error": "Upload both files"}

    # Load
    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

    # Align and embed
    tcga_a = align_genes(tcga, ref_genes)
    cptac_a = align_genes(cptac, ref_genes)
    tcga_emb = embed(tcga_a)
    cptac_emb = embed(cptac_a)

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

    # Novel metric: eTNSI
    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)

    # Classifier
    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": float(pca.explained_variance_ratio_.sum()),
        "centroid_distance": centroid_dist,
        "etnsi_embedding": etnsi,
        "etnsi_pvalue": float(p_val),
        "classifier_auc_mean": float(auc.mean()),
        "classifier_auc_std": float(auc.std()),
    }


iface = gr.Interface(
    fn=analyze,
    inputs=[
        gr.File(label="TCGA expression TSV (genes × samples)"),
        gr.File(label="CPTAC expression TSV (genes × samples)"),
    ],
    outputs=gr.JSON(label="Novel Results"),
    title="BulkRNABert Novel Analysis",
    description="Correlates 9-gene signature with BulkRNABert embeddings. "
                "Outputs eTNSI and classifier AUC — novel metrics for UCEC.",
)

iface.launch(server_name="0.0.0.0", server_port=7860)