import os, io
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import matplotlib.pyplot as plt
import pandas as pd
import gradio as gr

# ----------------------------
# Device
# ----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"

# ----------------------------
# Labels
# ----------------------------
breeds = [
    "Alambadi", "Amritmahal", "Ayrshire", "Banni", "Bargur", "Bhadawari", "Brown_Swiss",
    "Dangi", "Deoni", "Gir", "Guernsey", "Hallikar", "Hariana", "Holstein_Friesian",
    "Jaffrabadi", "Jersey", "Kangayam", "Kankrej", "Kasargod", "Kenkatha", "Kherigarh",
    "Khillari", "Krishna_Valley", "Malnad_gidda", "Mehsana", "Murrah", "Nagori", "Nagpuri",
    "Nili_Ravi", "Nimari", "Ongole", "Pulikulam", "Rathi", "Red_Dane", "Red_Sindhi",
    "Sahiwal", "Surti", "Tharparkar", "Toda", "Umblachery", "Vechur"
]
num_classes = len(breeds)

# ----------------------------
# Model
# ----------------------------
model = models.efficientnet_b0(weights=None)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
state = torch.load("bovine_model.pth", map_location=device)
model.load_state_dict(state)
model.to(device).eval()

# ----------------------------
# Preprocessing
# ----------------------------
val_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225])
])

# ----------------------------
# Predict
# ----------------------------
def predict_image(img_path: str):
    # Keep original filename for outputs
    base_name = os.path.basename(img_path)
    stem, _ = os.path.splitext(base_name)

    img = Image.open(img_path).convert("RGB")
    input_tensor = val_transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        outputs = model(input_tensor)
        probs = torch.nn.functional.softmax(outputs, dim=1)[0]

    top_prob, top_idx = torch.max(probs, dim=0)
    conf = float(top_prob.item()) * 100.0
    predicted_breed = breeds[int(top_idx.item())]

    # Annotate image (title overlay)
    fig, ax = plt.subplots()
    ax.imshow(img)
    ax.set_title(f"{predicted_breed} ({conf:.2f}%)")
    ax.axis("off")
    annotated_name = f"{predicted_breed}_{conf:.2f}pct_{stem}.png"
    plt.savefig(annotated_name, format="png", bbox_inches="tight", pad_inches=0.1, dpi=150)
    plt.close(fig)

    # CSV output
    df = pd.DataFrame([{
        "breed": predicted_breed,
        "confidence_percent": f"{conf:.2f}%",
        "filename": base_name
    }])
    csv_name = f"{stem}_prediction.csv"
    df.to_csv(csv_name, index=False)

    # Return:
    # 1) predicted breed (text)
    # 2) confidence (%) (text)
    # 3) file (CSV)
    # 4) file (annotated image with breed+confidence in filename)
    return predicted_breed, f"{conf:.2f}%", csv_name, annotated_name


# ----------------------------
# UI
# ----------------------------
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="filepath", label="Upload Cattle/Buffalo Image"),
    outputs=[
        gr.Textbox(label="Predicted Breed"),
        gr.Textbox(label="Prediction Confidence (%)"),
        gr.File(label="Download CSV"),
        gr.File(label="Download Annotated Image")
    ],
    title="Indian Cattle/Buffalo Breed Detection",
    description="Upload an image → get predicted breed, confidence score, CSV, and an annotated image file named with the predicted breed and confidence."
)

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