import os, sys, torch, joblib, numpy as np, gradio as gr, spaces
from PIL import Image
from torchvision import transforms
from huggingface_hub import snapshot_download
from util.pos_embed import interpolate_pos_embed
from timm.models.layers import trunc_normal_
import models_vit

print("Downloading models...")
snapshot_download(
    repo_id="janvib/retinal-weights",
    local_dir="/tmp/models",
    repo_type="model",
)
print("Downloaded.")

pth_files = [f for f in os.listdir("/tmp/models") if f.endswith(".pth")]
weights_path = f"/tmp/models/{pth_files[0]}"

model = models_vit.RETFound_mae(num_classes=2, drop_path_rate=0.2, global_pool=True)
checkpoint = torch.load(weights_path, map_location="cpu", weights_only=False)
checkpoint_model = checkpoint["model"]
state_dict = model.state_dict()
for k in ["head.weight", "head.bias"]:
    if k in checkpoint_model and checkpoint_model[k].shape != state_dict[k].shape:
        del checkpoint_model[k]
interpolate_pos_embed(model, checkpoint_model)
model.load_state_dict(checkpoint_model, strict=False)
model = model.eval()

ridge        = joblib.load("/tmp/models/retinal_age_ridge.pkl")
glaucoma_clf = joblib.load("/tmp/models/glaucoma_classifier.pkl")
dr_clf       = joblib.load("/tmp/models/dr_classifier_v2.pkl")
dr_scaler    = joblib.load("/tmp/models/dr_scaler.pkl")
print("All models loaded")

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

@spaces.GPU
def predict(image, age):
    img_tensor = transform(image).unsqueeze(0).cuda()
    model.cuda()
    with torch.no_grad():
        feats = model.forward_features(img_tensor)[:, 0, :].cpu().numpy()
    feats_scaled = dr_scaler.transform(feats)

    g_prob  = glaucoma_clf.predict_proba(feats)[0][1]
    dr_prob = dr_clf.predict_proba(feats_scaled)[0][1]
    predicted_age = float(ridge.predict(feats)[0])
    rag = predicted_age - float(age) if age else None

    if g_prob >= 0.7:
        g_flag = "HIGH - Referable Glaucoma Suspected"
        g_rec  = "Recommend full glaucoma workup: IOP, visual fields, OCT RNFL."
    elif g_prob >= 0.5:
        g_flag = "MODERATE - Glaucoma Suspect"
        g_rec  = "Consider glaucoma suspect workup."
    else:
        g_flag = "LOW - No Referable Glaucoma"
        g_rec  = "Routine follow-up."

    if dr_prob >= 0.7:
        dr_flag = "HIGH - Referable Diabetic Retinopathy"
        dr_rec  = "Recommend urgent retinal specialist referral."
    elif dr_prob >= 0.5:
        dr_flag = "MODERATE - Possible Referable DR"
        dr_rec  = "Consider retinal specialist referral."
    else:
        dr_flag = "LOW - No Referable DR"
        dr_rec  = "Routine follow-up if diabetic."

    age_line = "Predicted Retinal Age: " + str(round(predicted_age, 1)) + " yrs"
    if rag is not None:
        age_line = age_line + "\nRetinal Age Gap: " + str(round(rag, 1)) + " yrs"

    result = "\nGLAUCOMA SCREENING\n"
    result += g_flag + "\n"
    result += "Probability: " + str(round(g_prob * 100, 1)) + "%\n"
    result += "Action: " + g_rec + "\n"
    result += "\nDIABETIC RETINOPATHY SCREENING\n"
    result += dr_flag + "\n"
    result += "Probability: " + str(round(dr_prob * 100, 1)) + "%\n"
    result += "Action: " + dr_rec + "\n"
    result += "\nRETINAL AGE\n"
    result += age_line + "\n"
    result += "\nResearch use only. Not FDA-cleared. Clinical correlation required."
    return result

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(type="pil", label="Fundus Photograph"),
        gr.Number(label="Patient Age (optional)", value=60),
    ],
    outputs=gr.Textbox(label="Screening Results", lines=22),
    title="RETFound - Retinal Screening Tool",
    description="Upload a de-identified color fundus photograph to screen for glaucoma and diabetic retinopathy. Research use only. Not FDA-cleared.",
    flagging_mode="never",
)

demo.launch()