from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
from PIL import Image
import io

app = FastAPI(title="AgriVision API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Treatment recommendations ──────────────────────────────────
TREATMENTS = {
    "Armyworm": {
        "type": "Pest",
        "treatment": "Apply insecticides such as chlorpyrifos or lambda-cyhalothrin. Use pheromone traps for monitoring. Encourage natural predators like parasitic wasps."
    },
    "Grasshopper": {
        "type": "Pest",
        "treatment": "Apply malathion or carbaryl-based insecticides in early morning. Use barriers and biological controls such as Nosema locustae."
    },
    "aphids": {
        "type": "Pest",
        "treatment": "Spray neem oil or insecticidal soap. Introduce ladybugs as natural predators. Use systemic insecticides like imidacloprid for severe infestations."
    },
    "bean_rust": {
        "type": "Disease",
        "treatment": "Apply fungicides containing mancozeb or chlorothalonil. Remove and destroy infected leaves. Avoid overhead irrigation to reduce leaf wetness."
    },
    "beans_angular_leaf_spot": {
        "type": "Disease",
        "treatment": "Use copper-based fungicides. Plant resistant varieties. Practice crop rotation and remove infected plant debris after harvest."
    },
    "beans_anthracnose": {
        "type": "Disease",
        "treatment": "Apply fungicides such as azoxystrobin or mancozeb. Use certified disease-free seeds. Avoid working in fields when plants are wet."
    },
    "beans_healthy": {
        "type": "Healthy",
        "treatment": "No treatment needed. Maintain good agronomic practices including proper spacing, irrigation, and balanced fertilisation."
    },
    "cutworm": {
        "type": "Pest",
        "treatment": "Apply soil insecticides like chlorpyrifos around the base of plants. Use physical collars around seedlings. Bait with insecticide-treated bran."
    },
    "maize_common_rust": {
        "type": "Disease",
        "treatment": "Apply fungicides such as propiconazole or azoxystrobin at early infection stages. Plant rust-resistant maize varieties. Monitor fields regularly."
    },
    "maize_healthy": {
        "type": "Healthy",
        "treatment": "No treatment needed. Continue regular scouting and maintain optimal soil nutrition and irrigation practices."
    },
    "maize_leaf_blight": {
        "type": "Disease",
        "treatment": "Apply fungicides containing mancozeb or strobilurins. Plant tolerant varieties. Ensure proper plant spacing for good air circulation."
    },
    "maize_streak_virus": {
        "type": "Disease",
        "treatment": "No direct cure. Control leafhopper vectors with insecticides. Remove and destroy infected plants. Plant virus-resistant maize varieties."
    },
    "potato healthy": {
        "type": "Healthy",
        "treatment": "No treatment needed. Maintain proper hilling, irrigation scheduling, and balanced NPK fertilisation."
    },
    "potato_early_blight": {
        "type": "Disease",
        "treatment": "Apply fungicides such as chlorothalonil or mancozeb every 7–10 days. Remove lower infected leaves. Avoid excessive nitrogen fertilisation."
    },
    "potato_late_blight": {
        "type": "Disease",
        "treatment": "Apply systemic fungicides like metalaxyl or cymoxanil immediately. Destroy infected plant material. Avoid overhead irrigation and improve field drainage."
    },
    "rice_blast": {
        "type": "Disease",
        "treatment": "Apply tricyclazole or isoprothiolane fungicides. Use blast-resistant rice varieties. Avoid excessive nitrogen application which promotes susceptibility."
    },
    "rice_brown_spots": {
        "type": "Disease",
        "treatment": "Apply fungicides such as iprodione or propiconazole. Improve soil fertility especially potassium levels. Use certified healthy seeds."
    },
    "rice_healthy": {
        "type": "Healthy",
        "treatment": "No treatment needed. Maintain proper water management and balanced fertilisation for continued healthy growth."
    },
    "rice_leaf_blight": {
        "type": "Disease",
        "treatment": "Apply copper-based bactericides or streptomycin. Drain fields periodically. Avoid high nitrogen rates and plant resistant varieties."
    },
    "rice_tungro": {
        "type": "Disease",
        "treatment": "No cure available. Control green leafhopper vectors with insecticides. Remove infected plants immediately. Plant tungro-resistant varieties."
    },
    "stem_borer": {
        "type": "Pest",
        "treatment": "Apply carbofuran granules into leaf whorls. Use biological controls like Trichogramma parasitoids. Remove and destroy egg masses and infected stems."
    },
    "thrips": {
        "type": "Pest",
        "treatment": "Apply spinosad or abamectin-based insecticides. Use reflective mulches to deter thrips. Introduce predatory mites as biological control."
    },
    "tomato_early_blight": {
        "type": "Disease",
        "treatment": "Apply fungicides containing chlorothalonil or copper oxychloride. Remove infected lower leaves. Mulch around plants to prevent soil splash."
    },
    "tomato_healthy": {
        "type": "Healthy",
        "treatment": "No treatment needed. Maintain staking, pruning, and regular fertilisation for continued healthy growth."
    },
    "tomato_late_blight": {
        "type": "Disease",
        "treatment": "Apply metalaxyl or cymoxanil fungicides immediately. Remove and destroy all infected material. Improve air circulation and avoid wetting foliage."
    },
    "tomato_mosaic_virus": {
        "type": "Disease",
        "treatment": "No cure available. Remove and destroy infected plants immediately. Disinfect tools with bleach solution. Control aphid vectors and use virus-free seeds."
    },
    "tomato_septoria_spots": {
        "type": "Disease",
        "treatment": "Apply fungicides such as mancozeb or azoxystrobin. Remove infected leaves. Avoid overhead watering and practice crop rotation."
    },
    "tomato_yellow_leaf_curl_virus": {
        "type": "Disease",
        "treatment": "No cure available. Control whitefly vectors with imidacloprid or reflective mulches. Remove infected plants. Use resistant tomato varieties."
    },
    "weevil": {
        "type": "Pest",
        "treatment": "Apply pyrethroid insecticides. Store harvested grain in airtight containers. Use hermetic storage bags and apply diatomaceous earth as a physical control."
    },
}

# ── Severity function ──────────────────────────────────────────
def get_severity(confidence: float, crop_type: str) -> str:
    if crop_type == "Healthy":
        return "None"
    if confidence >= 0.90:
        return "High"
    elif confidence >= 0.70:
        return "Medium"
    else:
        return "Low"

# ── Name formatter ─────────────────────────────────────────────
def format_class_name(name: str) -> str:
    return name.replace("_", " ").title()

# ── Load model once at startup ─────────────────────────────────
model = YOLO("best.pt")

# ── Routes ─────────────────────────────────────────────────────
@app.get("/")
def home():
    return {"status": "AgriVision API is running"}

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    contents = await file.read()
    image = Image.open(io.BytesIO(contents)).convert("RGB")

    results = model.predict(image, imgsz=224, verbose=False)
    probs   = results[0].probs

    top1_index      = probs.top1
    top1_class      = model.names[top1_index]
    top1_confidence = round(float(probs.top1conf), 4)

    # ── Invalid image detection ────────────────────────────────
    # Get all top 5 probabilities
    top5_probs = [float(probs.data[i]) for i in probs.top5]

    # How much the top prediction dominates over the rest
    top1_prob  = top5_probs[0]
    top2_prob  = top5_probs[1]
    top3_prob  = top5_probs[2]

    # The gap between 1st and 2nd — a confident real crop image
    # will have a clear winner. A random image spreads probability evenly.
    top1_top2_gap = top1_prob - top2_prob

    # Entropy-like spread across top 3
    top3_spread = top1_prob - top3_prob

    # Reject if EITHER of these is true:
    # 1. Very low confidence (model has no idea)
    # 2. Low confidence AND the top predictions are bunched together
    #    (model is guessing randomly across classes)
    is_invalid = (
        top1_confidence < 0.50                          # extremely low confidence
        or (top1_confidence < 0.55 and top1_top2_gap < 0.20)  # low + no clear winner
        or (top1_confidence < 0.55 and top3_spread < 0.25)    # low + very spread out
    )

    if is_invalid:
        raise HTTPException(
            status_code=422,
            detail={
                "error": "Invalid image",
                "message": "The uploaded image does not appear to be a crop or pest. Please upload a clear photo of a crop leaf or pest.",
                "confidence": f"{round(top1_confidence * 100, 2)}%"
            }
        )

    # ── Valid scan ─────────────────────────────────────────────
    info     = TREATMENTS.get(top1_class, {
        "type": "Unknown",
        "treatment": "No treatment data available for this class."
    })
    severity = get_severity(top1_confidence, info["type"])

    top3 = [
        {
            "class":      format_class_name(model.names[i]),
            "confidence": f"{round(float(probs.data[i]) * 100, 2)}%"
        }
        for i in probs.top5[:3]
    ]

    return {
        "prediction":   format_class_name(top1_class),
        "type":         info["type"],
        "confidence":   f"{round(top1_confidence * 100, 2)}%",
        "severity":     severity,
        "treatment":    info["treatment"],
        "top3":         top3
    }