import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import swin_t
from PIL import Image

#  Model definition
class MMIM(nn.Module):
    def __init__(self, num_classes=36):
        super(MMIM, self).__init__()
        self.backbone = swin_t(weights='IMAGENET1K_V1')
        self.backbone.head = nn.Identity()
        self.classifier = nn.Sequential(
            nn.Linear(768, 512),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(512, num_classes)
        )

    def forward(self, x):
        features = self.backbone(x)
        return self.classifier(features)

#  Load model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MMIM(num_classes=36)
checkpoint = torch.load("MMIM_best.pth", map_location=device)
filtered_checkpoint = {
    k: v for k, v in checkpoint.items() if k in model.state_dict() and model.state_dict()[k].shape == v.shape
}
model.load_state_dict(filtered_checkpoint, strict=False)
model.to(device)
model.eval()

#  Class names
class_names = [
   "Chinee apple",              # class1
    "Black grass",                   # class14
    "Charlock",                  # class15
     "Cleavers",               # class16
     "Common Chickweed",          # class17
     "Common Wheat",               # class18
     "Fat Hen",                  # class19
      "Lanthana",                  # class2
      "Loose Silky bent",           # class20
      "Maize",              # class21
    "Scentless Mayweed",                   # class22
     "Shepherds Purse",          # class23
    "Small-Flowered Cranesbill",                     # class24
    "Sugar beet",         # class25
    "Carpetweeds",           # class26
    "Crabgrass",# class27
    "Eclipta",                # class28
    "Goosegrass",               # class29
    "Negative",                # class3
    "Morningglory",                 # class30
    "Nutsedge",                   # class31
    "Palmer Amarnath",              # class32
    "Prickly Sida",                  # class33
    "Purslane",            # class34
    "Ragweed",               # class35
    "Sicklepod",                  # class36
    "SpottedSpurge",                   # class37
    "SpurredAnoda",                 # class38
    "Swinecress",             # class39
     "Parkinsonia",                # class4
    "Waterhemp",              # class40
   
    "Parthenium",                 # class5
    "Prickly acacia",               # class6
    "Rubber vine",                # class7
    "Siam weed",            # class8
   "Snake weed",  
]

#  Weed info dictionary
weed_info = {
    "Chinee apple": " Invasive shrub. Control by uprooting or herbicide treatment.",
    "Black grass": " Infests cereal crops. Remove before seed shedding.",
    "Charlock": " Common weed in oilseed crops. Responds to early herbicide.",
    "Cleavers": " Sticky climbing weed. Control before flowering.",
    "Common Chickweed": " Fast-spreading groundcover weed. Avoid soil disturbance.",
    "Common Wheat": " May appear as weed in rotation crops.",
    "Fat Hen": " Broadleaf weed. Competes heavily with crops.",
    "Lanthana": " Invasive ornamental plant, toxic to livestock.",
    "Loose Silky bent": " Grass weed affecting wheat fields.",
    "Maize": " Sometimes emerges as volunteer weed post-harvest.",
    "Scentless Mayweed": " Strong competitor in cereals. Shallow-rooted.",
    "Shepherds Purse": " Common weed in cool seasons. Heart-shaped pods.",
    "Small-Flowered Cranesbill": " Low-growing, thrives in dry areas.",
    "Sugar beet": " Appears as volunteer in crop fields.",
    "Carpetweeds": " Low mat-forming weed. Easy to remove manually.",
    "Crabgrass": " Summer annual grass. Thrives in disturbed soil.",
    "Eclipta": " Moisture-loving herbaceous weed.",
    "Goosegrass": " Mat-forming weed, tough to hand-pull.",
    "Negative": " No weed confidently detected. Please recheck input.",
    "Morningglory": " Climbing vine, chokes crops quickly.",
    "Nutsedge": " Grass-like weed with tubers. Hard to control.",
    "Palmer Amarnath": " Highly aggressive and herbicide-resistant.",
    "Prickly Sida": " Hairy, thorny stems. Requires early control.",
    "Purslane": " Succulent weed, common in warm climates.",
    "Ragweed": " Allergen-producing weed. Kill before flowering.",
    "Sicklepod": " Toxic to livestock. Control before pod set.",
    "SpottedSpurge": " Low-growing. Releases milky sap.",
    "SpurredAnoda": " Fast-growing summer annual. Common in cotton.",
    "Swinecress": " Strong odor. Grows in compacted soils.",
    "Parkinsonia": " Woody shrub. Mechanical removal advised.",
    "Waterhemp": " Fast-growing amaranth. Glyphosate-resistant strains exist.",
    "Parthenium": " Toxic and invasive. Avoid contact.",
    "Prickly acacia": " Thorny shrub. Displaces native plants.",
    "Rubber vine": " Woody climber. Toxic and invasive.",
    "Siam weed": " Highly invasive in tropical zones.",
    "Snake weed": " Woody perennial, toxic to livestock."
}

#  Transform
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

#  Prediction function
def predict(img):
    img = img.convert('RGB')
    img_tensor = transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(img_tensor)
        probs = torch.softmax(outputs, dim=1)

    conf, pred = torch.max(probs, 1)
    predicted_class = class_names[pred.item()]
    confidence = conf.item() * 100

    if predicted_class.lower() == "negative":
        label = f" Predicted as: Negative\nConfidence: {confidence:.2f}%"
    elif confidence < 60:
        label = f" Low confidence. Possibly Not a Weed\nConfidence: {confidence:.2f}%"
    else:
        label = f" Predicted class: {predicted_class}\nConfidence: {confidence:.2f}%"

    info = weed_info.get(predicted_class, " No additional info available.")
    return f"{label}\n\n Info: {info}"

#  App description
about_markdown = """
###  Weed Classifier — Swin Transformer + MMIM  
This tool predicts weed species from images using a Vision Transformer backbone trained with multi-masked image modeling.

-  Shows confidence scores  
-  Flags uncertain or non-weed predictions  
-  Displays weed info after prediction  
-  Upload an image

> Tip: Use clear, focused weed images for better results.
"""

#  Gradio Interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=gr.Textbox(label="Prediction"),
    title=" Weed Image Classifier",
    description="A Self- Spervised Learning model for weed image classification.",
    article=about_markdown
)

interface.launch()