import spaces
import os
import torch
import torch.nn as nn
import gradio as gr
from transformers import BigBirdTokenizer, BigBirdModel
from huggingface_hub import hf_hub_download

# --- 1. PyTorch Architecture Definition ---
SCALAR_COLS = [
    "bail_type_mapped", "days_in_custody_available", "days_in_custody",
    "age_available_mapped", "min_age", "max_age", "median_age",
    "shap_sum_pos", "shap_sum_neg", "shap_max_pos", "shap_min_neg",
    "shap_pos_count", "shap_neg_count", "shap_l1_total",
    "shap_top3_abs_sum", "xgb_pred"
]

class BigBirdScalarClassifier(nn.Module):
    def __init__(self, bigbird_name="google/bigbird-roberta-large", scalar_dim=16):
        super().__init__()
        self.bigbird = BigBirdModel.from_pretrained(bigbird_name)
        hidden = self.bigbird.config.hidden_size
        self.att_vector = nn.Parameter(torch.randn(hidden))
        self.drop_text1 = nn.Dropout(0.1)
        self.proj = nn.Linear(hidden, 512)
        self.ln = nn.LayerNorm(512)
        self.drop_text2 = nn.Dropout(0.2)
        self.classifier = nn.Linear(512 + scalar_dim, 1)

    def attention_pool(self, H, attention_mask):
        scores = torch.matmul(H, self.att_vector)
        scores = scores.masked_fill(attention_mask == 0, -1e4)
        alpha = torch.softmax(scores, dim=1)
        return torch.sum(H * alpha.unsqueeze(-1), dim=1)

    def forward(self, input_ids, attention_mask, scalars):
        out = self.bigbird(input_ids=input_ids, attention_mask=attention_mask, return_dict=True)
        H = out.last_hidden_state
        pooled = self.attention_pool(H, attention_mask)
        x_text = self.drop_text1(pooled)
        x_text = self.proj(x_text)
        x_text = self.ln(x_text)
        x_text = self.drop_text2(x_text)
        x = torch.cat([x_text, scalars], dim=1)
        return self.classifier(x).squeeze(1)

# --- 2. Global Model Initialization ---
# ZeroGPU handles the CUDA allocation automatically behind the scenes
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = BigBirdTokenizer.from_pretrained("google/bigbird-roberta-large")

# Download directly from your HF model repository
model_path = hf_hub_download(repo_id="puspeshsrivastava/ibpsv0.2", filename="best_model_full.pt")

classifier = BigBirdScalarClassifier(scalar_dim=len(SCALAR_COLS)).to(device)
checkpoint = torch.load(model_path, map_location=device)
state_dict = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint
classifier.load_state_dict(state_dict)
classifier.eval()

# --- 3. Inference Logic (with ZeroGPU Decorator) ---
@spaces.GPU
def predict_bail_outcome(
    case_details, bail_type_mapped, days_in_custody_available, days_in_custody,
    age_available_mapped, min_age, max_age, median_age,
    shap_sum_pos, shap_sum_neg, shap_max_pos, shap_min_neg,
    shap_pos_count, shap_neg_count, shap_l1_total, shap_top3_abs_sum, xgb_pred
):
    try:
        encoded = tokenizer(
            case_details, padding=True, truncation=True, max_length=1200, return_tensors="pt"
        )
        input_ids = encoded["input_ids"].to(device)
        attention_mask = encoded["attention_mask"].to(device)

        scalars_list = [
            float(bail_type_mapped), float(days_in_custody_available), float(days_in_custody),
            float(age_available_mapped), float(min_age), float(max_age), float(median_age),
            float(shap_sum_pos), float(shap_sum_neg), float(shap_max_pos), float(shap_min_neg),
            float(shap_pos_count), float(shap_neg_count), float(shap_l1_total), float(shap_top3_abs_sum), float(xgb_pred)
        ]
        scalars_tensor = torch.tensor([scalars_list], dtype=torch.float32).to(device)

        with torch.no_grad():
            logit = classifier(input_ids, attention_mask, scalars_tensor)
            prob = float(torch.sigmoid(logit).item())
            pred_label = 1 if prob >= 0.5 else 0

        return {
            "prediction": pred_label,
            "probability": round(prob, 4),
            "outcome_label": "Bail Granted" if pred_label == 1 else "Bail Denied"
        }
    except Exception as e:
        return {"error": str(e)}

# --- 4. Gradio UI & Auto-API Mapping ---
inputs = [
    gr.Textbox(label="Case Details (Narrative)", lines=4),
    gr.Number(label="Bail Type (1=Regular, 2=Anticipatory)", value=1),
    gr.Number(label="Days in Custody Available (1=Yes, 0=No)", value=1),
    gr.Number(label="Days in Custody", value=0.0),
    gr.Number(label="Age Available (1=Yes, 0=No)", value=1),
    gr.Number(label="Min Age", value=35.0),
    gr.Number(label="Max Age", value=35.0),
    gr.Number(label="Median Age", value=35.0),
    gr.Number(label="SHAP Sum Pos", value=0.0),
    gr.Number(label="SHAP Sum Neg", value=0.0),
    gr.Number(label="SHAP Max Pos", value=0.0),
    gr.Number(label="SHAP Min Neg", value=0.0),
    gr.Number(label="SHAP Pos Count", value=0.0),
    gr.Number(label="SHAP Neg Count", value=0.0),
    gr.Number(label="SHAP L1 Total", value=0.0),
    gr.Number(label="SHAP Top 3 Abs Sum", value=0.0),
    gr.Number(label="XGB Pred Baseline", value=0.5),
]

demo = gr.Interface(
    fn=predict_bail_outcome,
    inputs=inputs,
    outputs=gr.JSON(label="Prediction Response"),
    title="IBPS BigBird Predictor (ZeroGPU)",
    description="Interactive UI & API for the Indian Bail Prediction System."
)

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