import gradio as gr
import torch
import torch.nn.functional as F
from torchvision import transforms
from huggingface_hub import hf_hub_download
import numpy as np
from PIL import Image

class TinySSLBase(torch.nn.Module):
    def __init__(self, out_dim=256):
        super().__init__()
        self.patch_embed = torch.nn.Conv2d(3, 256, kernel_size=3, stride=4, padding=1)
        self.proj = torch.nn.Linear(256, out_dim)
        encoder_layer = torch.nn.TransformerEncoderLayer(d_model=out_dim, nhead=4, dim_feedforward=512, batch_first=True, dropout=0.1)
        self.encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=4)
        self.cls_token = torch.nn.Parameter(torch.randn(1, 1, out_dim) * 0.02)
        self.out_dim = out_dim

    def forward(self, x):
        B = x.shape[0]
        x = self.patch_embed(x).flatten(2).transpose(1, 2)
        x = self.proj(x)
        cls = self.cls_token.expand(B, -1, -1)
        x = torch.cat([cls, x], dim=1)
        x = self.encoder(x)
        return {"cls": x[:, 0], "patches": x[:, 1:]}

class TinySSLTiny(torch.nn.Module):
    def __init__(self, out_dim=128):
        super().__init__()
        self.patch_embed = torch.nn.Conv2d(3, 128, kernel_size=3, stride=4, padding=1)
        self.proj = torch.nn.Linear(128, out_dim)
        encoder_layer = torch.nn.TransformerEncoderLayer(d_model=out_dim, nhead=4, dim_feedforward=256, batch_first=True, dropout=0.1)
        self.encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=2)
        self.out_dim = out_dim

    def forward(self, x):
        x = self.patch_embed(x).flatten(2).transpose(1, 2)
        x = self.proj(x)
        x = self.encoder(x)
        cls = x.mean(dim=1)
        return {"cls": cls, "patches": x}

class TinySSLCNN(torch.nn.Module):
    def __init__(self, out_dim=256):
        super().__init__()
        self.blocks = torch.nn.Sequential(
            torch.nn.Conv2d(3, 32, 3, stride=2, padding=1), torch.nn.BatchNorm2d(32), torch.nn.ReLU(inplace=True),
            torch.nn.Conv2d(32, 64, 3, stride=2, padding=1), torch.nn.BatchNorm2d(64), torch.nn.ReLU(inplace=True),
            torch.nn.Conv2d(64, 128, 3, stride=2, padding=1), torch.nn.BatchNorm2d(128), torch.nn.ReLU(inplace=True),
            torch.nn.Conv2d(128, 256, 3, stride=2, padding=1), torch.nn.BatchNorm2d(256), torch.nn.ReLU(inplace=True),
        )
        self.head = torch.nn.Linear(256, out_dim)
        self.out_dim = out_dim

    def forward(self, x):
        feat = self.blocks(x)
        cls = self.head(feat.mean(dim=(2, 3)))
        patches = feat.flatten(2).transpose(1, 2)
        return {"cls": cls, "patches": patches}

MODEL_MAP = {"base": TinySSLBase, "tiny": TinySSLTiny, "cnn": TinySSLCNN}
DATASET_INFO = {"flowers102": 102, "oxford_pets": 37, "eurosat": 10, "breastmnist": 2}
DATASET_LABELS = {
    "flowers102": [f"Flower {i}" for i in range(102)],
    "eurosat": ["AnnualCrop", "Forest", "HerbaceousVegetation", "Highway", "Industrial", "Pasture", "PermanentCrop", "Residential", "River", "SeaLake"],
    "oxford_pets": [f"Pet Class {i}" for i in range(37)],
    "breastmnist": ["Malignant", "Benign"],
}
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

eval_t = transforms.Compose([
    transforms.Resize(256), transforms.CenterCrop(224),
    transforms.ToTensor(), transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
])

loaded_models = {}

def load_model(model_name, dataset):
    key = f"{model_name}_{dataset}"
    if key in loaded_models:
        return loaded_models[key]
    ckpt_path = hf_hub_download(repo_id="emran696996966/TinySSL", filename=f"checkpoints/{key}.pt")
    model_cls = MODEL_MAP[model_name]
    model = model_cls()
    model.load_state_dict(torch.load(ckpt_path, weights_only=False)["model"])
    model.eval()
    num_classes = DATASET_INFO[dataset]
    head = torch.nn.Linear(model.out_dim, num_classes)
    loaded_models[key] = (model, head)
    return model, head

def predict(image, model_name, dataset):
    if image is None:
        return "Please upload an image."
    img = Image.fromarray(np.array(image)).convert("RGB")
    x = eval_t(img).unsqueeze(0)
    model, head = load_model(model_name, dataset)
    with torch.no_grad():
        feat = model(x)["cls"]
        logits = head(feat)
        probs = F.softmax(logits, dim=-1)[0]
    top5 = torch.topk(probs, min(5, probs.shape[-1]))
    labels = DATASET_LABELS.get(dataset, [f"Class {i}" for i in range(DATASET_INFO[dataset])])
    lines = []
    for prob, idx in zip(top5.values, top5.indices):
        label = labels[idx.item()] if idx.item() < len(labels) else f"Class {idx.item()}"
        lines.append(f"{label}: {prob.item()*100:.1f}%")
    return "\n".join(lines)

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Image(label="Upload Image"),
        gr.Dropdown(choices=["base", "tiny", "cnn"], value="base", label="Model"),
        gr.Dropdown(choices=["flowers102", "oxford_pets", "eurosat", "breastmnist"], value="flowers102", label="Dataset"),
    ],
    outputs=gr.Textbox(label="Top-5 Predictions"),
    title="TinySSL: Distilling DINOv2 into Tiny Vision Models",
    description="2M-param students trained on DINOv2 features. Upload an image to classify.",
)

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