import gradio as gr
import torch
import torch.nn as nn
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModel

# --- Load BLIP model from Hugging Face Hub ---
processor = BlipProcessor.from_pretrained("Mewish/blip_medical")
blip_model = BlipForConditionalGeneration.from_pretrained("Mewish/blip_medical").to("cpu")

# --- Load BioClinicalBERT backbone ---
tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
bert_model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")

# --- Define BiLSTM classifier ---
class BiLSTMClassifier(nn.Module):
    def __init__(self, bert_model, hidden_dim=256, num_classes=7, dropout=0.5):
        super().__init__()
        self.bert = bert_model
        self.lstm = nn.LSTM(768, hidden_dim, batch_first=True, bidirectional=True)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(hidden_dim*2, num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        embeddings = outputs.last_hidden_state
        lstm_out, _ = self.lstm(embeddings)
        pooled = torch.mean(lstm_out, dim=1)
        pooled = self.dropout(pooled)
        return self.fc(pooled)

# --- Load classifier weights from Hugging Face Hub ---
classifier = BiLSTMClassifier(bert_model).to("cpu")
state_dict_url = "https://huggingface.co/Mewish/skin_cancer_classifier/resolve/main/classifier.pt"
classifier.load_state_dict(torch.hub.load_state_dict_from_url(state_dict_url, map_location="cpu"))
classifier.eval()

# --- Classes ---
classes = [
    "actinic keratosis",
    "basal cell carcinoma",
    "dermatofibroma",
    "nevus",
    "pigmented benign keratosis",
    "squamous cell carcinoma",
    "vascular lesion"
]

# --- Prediction function ---
def predict(image):
    # Step 1: BLIP caption
    inputs = processor(images=image, return_tensors="pt").to("cpu")
    generated_ids = blip_model.generate(**inputs, max_length=50)
    caption = processor.decode(generated_ids[0], skip_special_tokens=True)

    # Step 2: Classifier prediction
    text_inputs = tokenizer(caption, max_length=50, padding="max_length",
                            truncation=True, return_tensors="pt")
    input_ids = text_inputs["input_ids"].to("cpu")
    attention_mask = text_inputs["attention_mask"].to("cpu")

    with torch.no_grad():
        outputs = classifier(input_ids, attention_mask)
        probs = torch.softmax(outputs, dim=1).cpu().numpy()[0]
        pred_class = classes[probs.argmax()]
        confidence = probs.max()

    return caption, pred_class, float(confidence)

# --- Gradio Interface ---
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[
        gr.Textbox(label="Generated Caption"),
        gr.Textbox(label="Predicted Class"),
        gr.Number(label="Confidence Score")
    ],
    title="Skin Cancer AI Agent",
    description="Upload a lesion image to generate a caption and predict the cancer type."
)

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