import os, sys, torch, gradio as gr
from PIL import Image
from torchvision import transforms
from huggingface_hub import hf_hub_download

# Clone PanDerm source saat startup (sekali saja di container)
PANDERM_DIR = "/tmp/PanDerm"
if not os.path.exists(PANDERM_DIR):
    os.system(f"git clone https://github.com/SiyuanYan1/PanDerm.git {PANDERM_DIR}")
sys.path.insert(0, os.path.join(PANDERM_DIR, "classification"))
sys.path.insert(0, os.path.join(PANDERM_DIR, "classification", "models"))

import importlib.util
_spec = importlib.util.spec_from_file_location(
    "modeling_finetune",
    os.path.join(PANDERM_DIR, "classification", "models", "modeling_finetune.py"),
)
_mod = importlib.util.module_from_spec(_spec)
sys.modules["modeling_finetune"] = _mod   # ← TAMBAHKAN baris ini
_spec.loader.exec_module(_mod)

from timm.models import create_model

CLASSES = ["akiec", "bcc", "bkl", "df", "mel", "nv", "vasc"]
LABELS = {
    "akiec": "Actinic Keratosis (suspicious)",
    "bcc":   "Basal Cell Carcinoma (malignant)",
    "bkl":   "Benign Keratosis (low risk)",
    "df":    "Dermatofibroma (low risk)",
    "mel":   "Melanoma (malignant)",
    "nv":    "Melanocytic Nevus (low risk)",
    "vasc":  "Vascular Lesion (low risk)",
}

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# Download checkpoint dari model repo
ckpt_path = hf_hub_download(
    repo_id="farelfebryan/panderm-ham10000",
    filename="checkpoint-best.pth",
)

# Build model persis seperti training Anda
model = create_model(
    "PanDerm_Large_FT",
    pretrained=False,
    num_classes=7,
    drop_path_rate=0.2,
    use_mean_pooling=True,
    init_scale=0.001,
    use_rel_pos_bias=False,
    use_abs_pos_emb=True,
    init_values=0.1,
    sin_pos_emb=True,
)

ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
state = ckpt.get("model", ckpt.get("state_dict", ckpt))
state = {k.replace("module.", ""): v for k, v in state.items()}
missing, unexpected = model.load_state_dict(state, strict=False)
print(f"Loaded checkpoint — missing: {len(missing)}, unexpected: {len(unexpected)}")
if missing: print("  missing examples:", missing[:5])
if unexpected: print("  unexpected examples:", unexpected[:5])

model.to(device).eval()

# Transform sesuai training: img_size=224, imagenet mean/std
tfm = transforms.Compose([
    transforms.Resize(256, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

@torch.no_grad()
def predict(img: Image.Image):
    if img is None:
        return {}
    x = tfm(img.convert("RGB")).unsqueeze(0).to(device)
    logits = model(x)
    probs = torch.softmax(logits, dim=1)[0].cpu().tolist()
    return {LABELS[CLASSES[i]]: float(probs[i]) for i in range(7)}

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload dermoscopy image"),
    outputs=gr.Label(num_top_classes=3, label="Prediction"),
    title="PanDerm Skin Lesion Classifier (HAM10000)",
    description=(
        "7-class skin lesion classifier based on PanDerm-Large, fine-tuned on "
        "HAM10000 with SDXL DoRA synthetic augmentation. "
        "Classes: akiec, bcc, bkl, df, mel, nv, vasc."
    ),
    examples=None,
    flagging_mode="never",
)

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