from flask import Flask, request, jsonify
from flask_cors import CORS
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import io
import os

# ══════════════════════════════════════════════════════════════
# Flask App Setup
# ══════════════════════════════════════════════════════════════
app = Flask(__name__)
CORS(app)  # Allow mobile app to connect

# ══════════════════════════════════════════════════════════════
# Configuration
# ══════════════════════════════════════════════════════════════
MODEL_PATH   = "resnet18_rice_v2.pth"   # ← your saved model
CLASS_NAMES  = ["ClassA-Drought", "ClassB-PestInfestation", "ClassC-Healthy"]
THRESHOLD    = 0.80                      # below this = uncertain
IMG_SIZE     = 224
device       = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ══════════════════════════════════════════════════════════════
# Load Model Once at Startup
# ══════════════════════════════════════════════════════════════
def load_model():
    model = models.resnet18(weights=None)
    model.fc = nn.Linear(model.fc.in_features, len(CLASS_NAMES))
    checkpoint = torch.load(MODEL_PATH, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model = model.to(device)
    model.eval()
    print(f"✅ Model loaded from {MODEL_PATH}")
    print(f"✅ Classes: {CLASS_NAMES}")
    return model

model = load_model()

# ══════════════════════════════════════════════════════════════
# Image Transform
# ══════════════════════════════════════════════════════════════
transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                          std=[0.229, 0.224, 0.225])
])

# ══════════════════════════════════════════════════════════════
# Routes
# ══════════════════════════════════════════════════════════════

# Health check — test if API is running
@app.route("/", methods=["GET"])
def home():
    return jsonify({
        "status"  : "running",
        "message" : "Rice Stress Detection API",
        "classes" : CLASS_NAMES,
        "model"   : "ResNet18 — 99.66% accuracy"
    })

# Main prediction endpoint
@app.route("/predict", methods=["POST"])
def predict():
    # ── Check if image was sent ─────────────────────────────
    if "image" not in request.files:
        return jsonify({
            "error": "No image provided. Send image as form-data with key 'image'"
        }), 400

    file = request.files["image"]

    # ── Check file is valid ─────────────────────────────────
    if file.filename == "":
        return jsonify({"error": "Empty filename"}), 400

    allowed = {"jpg", "jpeg", "png", "bmp", "webp"}
    ext = file.filename.rsplit(".", 1)[-1].lower()
    if ext not in allowed:
        return jsonify({"error": f"File type .{ext} not allowed. Use jpg/png"}), 400

    try:
        # ── Read and process image ──────────────────────────
        image_bytes = file.read()
        image       = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        tensor      = transform(image).unsqueeze(0).to(device)

        # ── Run prediction ──────────────────────────────────
        with torch.no_grad():
            outputs    = model(tensor)
            probs      = torch.softmax(outputs, dim=1)[0]
            confidence, predicted = torch.max(probs, 0)

        predicted_class = CLASS_NAMES[predicted.item()]
        conf_value      = confidence.item()
        is_uncertain    = conf_value < THRESHOLD

        # ── Build response ──────────────────────────────────
        all_probs = {
            CLASS_NAMES[i]: round(probs[i].item() * 100, 2)
            for i in range(len(CLASS_NAMES))
        }

        # Clean class name for display
        display_name = predicted_class.replace("ClassA-", "")\
                                      .replace("ClassB-", "")\
                                      .replace("ClassC-", "")

        response = {
            "predicted_class" : predicted_class,
            "display_name"    : display_name,
            "confidence"      : round(conf_value * 100, 2),
            "is_uncertain"    : is_uncertain,
            "all_probabilities": all_probs,
            "recommendation"  : get_recommendation(display_name, is_uncertain)
        }

        return jsonify(response), 200

    except Exception as e:
        return jsonify({"error": str(e)}), 500


# ══════════════════════════════════════════════════════════════
# Recommendation Logic
# ══════════════════════════════════════════════════════════════
def get_recommendation(class_name, is_uncertain):
    if is_uncertain:
        return "Low confidence — recommend manual inspection by agronomist"

    recommendations = {
        "Drought": (
            "Rice plant shows drought stress. "
            "Recommended actions: increase irrigation frequency, "
            "check soil moisture levels, consider mulching."
        ),
        "PestInfestation": (
            "Rice plant shows pest infestation signs. "
            "Recommended actions: inspect leaves for insects, "
            "apply appropriate pesticide, monitor surrounding plants."
        ),
        "Healthy": (
            "Rice plant appears healthy. "
            "Continue regular monitoring and maintenance."
        )
    }
    return recommendations.get(class_name, "No recommendation available")


# ══════════════════════════════════════════════════════════════
# Run Server
# ══════════════════════════════════════════════════════════════
if __name__ == "__main__":
   port = int(os.environ.get("PORT", 7860))
print("🌾 Rice Stress Detection API starting...")
app.run(debug=False, host="0.0.0.0", port=port)