import torch
import gradio as gr
import numpy as np
from PIL import Image
from open_clip import create_model_from_pretrained, get_tokenizer

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
MODEL_NAME = "microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224"

# ---------------- Load BiomedCLIP ----------------
print("Loading BiomedCLIP...")
clip_model, preprocess = create_model_from_pretrained(f"hf-hub:{MODEL_NAME}")
tokenizer = get_tokenizer(f"hf-hub:{MODEL_NAME}")
clip_model = clip_model.to(DEVICE).eval()

# ---------------- Load Checkpoint ----------------
print("Loading Fusion Classifier checkpoint...")
checkpoint = torch.load("fusion_classifier.pth", map_location=DEVICE)

# Some checkpoints store {"state_dict": ...}
if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
    state = checkpoint["state_dict"]
else:
    state = checkpoint

# ---------------- Determine Architecture From Checkpoint ----------------
# Required keys must exist
layer1 = state["net.0.weight"].shape   # [out, in] e.g. [1024, 1024]
layer2 = state["net.3.weight"].shape   # [out, in] e.g. [512, 1024]
layer3 = state["net.6.weight"].shape   # [out, in] e.g. [C, 512]

INPUT_DIM = layer1[1]
H1 = layer1[0]
H2 = layer2[0]
NUM_CLASSES = layer3[0]

print(f"Recovered model structure:")
print(f"Input: {INPUT_DIM}  -> {H1} -> {H2} -> {NUM_CLASSES}")

# ---------------- Load Class Names ----------------
try:
    classes = checkpoint["classes"]
except:
    try:
        classes = checkpoint["label_classes"]
    except:
        print("No class names in checkpoint, loading label_classes.npy")
        try:
            classes = np.load("label_classes.npy", allow_pickle=True)
        except:
            print("No label_classes.npy found. Using index labels.")
            classes = [f"class_{i}" for i in range(NUM_CLASSES)]

classes = list(classes)

# Safety: if class names count mismatches, fix length
if len(classes) != NUM_CLASSES:
    print("Class count mismatch. Adjusting names.")
    classes = classes[:NUM_CLASSES] if len(classes) > NUM_CLASSES else classes + [
        f"class_{i}" for i in range(len(classes), NUM_CLASSES)
    ]

# ---------------- Build EXACT SAME MODEL ----------------
class FusionClassifier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.net = torch.nn.Sequential(
            torch.nn.Linear(INPUT_DIM, H1),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(H1, H2),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(H2, NUM_CLASSES)
        )

    def forward(self, x):
        return self.net(x)

classifier = FusionClassifier().to(DEVICE)
classifier.load_state_dict(state)
classifier.eval()

# ---------------- Prediction ----------------
def predict(image, symptoms):
    if image is None:
        return "Please upload an image"

    if not symptoms or symptoms.strip() == "":
        symptoms = "no symptom information provided"

    img = preprocess(image).unsqueeze(0).to(DEVICE)
    text_tok = tokenizer([symptoms]).to(DEVICE)

    with torch.no_grad():
        img_feat = clip_model.encode_image(img)
        txt_feat = clip_model.encode_text(text_tok)

        img_feat /= img_feat.norm(dim=-1, keepdim=True)
        txt_feat /= txt_feat.norm(dim=-1, keepdim=True)

        fusion = torch.cat([img_feat, txt_feat], dim=1)

        logits = classifier(fusion)
        probs = torch.softmax(logits, dim=1).cpu().numpy()[0]

        idx = probs.argmax()
        conf = float(probs[idx])

        label = str(classes[idx]) if idx < len(classes) else "Unknown"

        return f"Prediction: {label}  | Confidence: {conf:.2f}"

# ---------------- UI ----------------
iface = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Upload Skin Image"),
        gr.Textbox(label="Describe symptoms (itching, duration, pain, etc.)")
    ],
    outputs="text",
    title="BioLens — Multimodal Skin Condition AI",
    description="Upload image + symptoms. Not medical advice."
)

iface.launch()