"""
PlantField - 3-Stage ResNet50 Classification Inference API
Endpoint: POST /classification

Stage 1: Binary       - Plant vs Non-Plant
Stage 2: Binary       - Healthy vs Diseased  (hanya jika Plant)
Stage 3: Multi-class  - Downy Mildew, Powdery Mildew, Septoria Blight, Viral
                        (hanya jika Diseased)
"""

import io
import os
import logging
from pathlib import Path

import torch
import torch.nn as nn
from torchvision import models, transforms
from torchvision.models import ResNet50_Weights
from safetensors.torch import load_file
from PIL import Image
import numpy as np

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn

# ── Logging ──────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ── Konfigurasi ───────────────────────────────────────────────────────────────
IMG_SIZE = 224
DEVICE   = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Threshold cascade (sama dengan notebook training)
STAGE1_THRESHOLD = 0.5   # p_plant  > threshold  → lanjut Stage 2
STAGE2_THRESHOLD = 0.5   # p_diseased > threshold → lanjut Stage 3

# Label kelas
STAGE1_CLASSES  = ["Non-Plant", "Plant"]          # index 0 = Non-Plant, 1 = Plant
STAGE2_CLASSES  = ["Healthy", "Diseased"]         # index 0 = Healthy,   1 = Diseased
DISEASE_CLASSES = [                               # 4 kelas penyakit (Stage 3)
    "Downy_mildew_on_lettuce",
    "Powdery_mildew_on_lettuce",
    "Septoria_blight_on_lettuce",
    "Viral",
]
DISEASE_LABEL_MAP = {
    "Downy_mildew_on_lettuce":    "Downy Mildew",
    "Powdery_mildew_on_lettuce":  "Powdery Mildew",
    "Septoria_blight_on_lettuce": "Septoria Blight",
    "Viral":                      "Viral",
}

# Path model — file safetensors dari notebook training
MODEL_DIR          = Path(__file__).parent
STAGE1_MODEL_PATH  = MODEL_DIR / "ResNet50_S1_best.safetensors"
STAGE2_MODEL_PATH  = MODEL_DIR / "ResNet50_S2_best.safetensors"
STAGE3_MODEL_PATH  = MODEL_DIR / "ResNet50_S3_best.safetensors"


# ── Arsitektur Model ──────────────────────────────────────────────────────────
def build_resnet50(num_classes: int) -> nn.Module:
    """
    ResNet50 dengan identical custom head sesuai notebook training.
    Head: BN -> Linear(256) -> ReLU -> Dropout(0.5)
              -> Linear(128) -> ReLU -> Dropout(0.3)
              -> Linear(num_classes)
    """
    model = models.resnet50(weights=None)
    in_features = model.fc.in_features
    model.fc = nn.Sequential(
        nn.BatchNorm1d(in_features),
        nn.Linear(in_features, 256),
        nn.ReLU(inplace=True),
        nn.Dropout(0.5),
        nn.Linear(256, 128),
        nn.ReLU(inplace=True),
        nn.Dropout(0.3),
        nn.Linear(128, num_classes),
    )
    return model


# ── Preprocessing ─────────────────────────────────────────────────────────────
val_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE), Image.LANCZOS),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    ),
])


# ── Load Models ───────────────────────────────────────────────────────────────
logger.info(f"Loading 3-stage models on device: {DEVICE}")

stage1_model = build_resnet50(num_classes=2)   # Plant / Non-Plant
stage1_model.load_state_dict(load_file(str(STAGE1_MODEL_PATH)))
stage1_model.to(DEVICE).eval()
logger.info("Stage 1 (Plant/Non-Plant) model loaded")

stage2_model = build_resnet50(num_classes=2)   # Healthy / Diseased
stage2_model.load_state_dict(load_file(str(STAGE2_MODEL_PATH)))
stage2_model.to(DEVICE).eval()
logger.info("Stage 2 (Healthy/Diseased) model loaded")

stage3_model = build_resnet50(num_classes=4)   # Jenis penyakit (4 kelas)
stage3_model.load_state_dict(load_file(str(STAGE3_MODEL_PATH)))
stage3_model.to(DEVICE).eval()
logger.info("Stage 3 (Disease type) model loaded")


# ── FastAPI App ────────────────────────────────────────────────────────────────
app = FastAPI(
    title="PlantField Classification API",
    description=(
        "3-Stage ResNet50 cascade classifier untuk deteksi penyakit selada. "
        "Stage 1: Plant/Non-Plant → Stage 2: Healthy/Diseased → Stage 3: Jenis penyakit."
    ),
    version="2.0.0",
)

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


# ── Helper Functions ──────────────────────────────────────────────────────────
def preprocess_image(image_bytes: bytes) -> torch.Tensor:
    """Preprocess gambar dari bytes ke tensor siap inference."""
    try:
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        tensor = val_transform(image)
        return tensor.unsqueeze(0)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Gagal memproses gambar: {str(e)}")


@torch.no_grad()
def run_inference(model: nn.Module, tensor: torch.Tensor) -> tuple:
    """Jalankan inference dan kembalikan prediksi, confidence, dan probabilitas."""
    tensor = tensor.to(DEVICE)
    output = model(tensor)
    probs  = torch.softmax(output, dim=1)[0]
    pred   = probs.argmax().item()
    conf   = probs[pred].item()
    return pred, conf, probs.cpu().numpy().tolist()


# ── Endpoints ─────────────────────────────────────────────────────────────────
@app.get("/")
async def root():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "model":  "ResNet50 3-Stage Cascade Classification",
        "version": "2.0.0",
        "description": "PlantField - Lettuce Disease Detection API",
        "pipeline": {
            "stage1": "Plant vs Non-Plant",
            "stage2": "Healthy vs Diseased (hanya jika Plant)",
            "stage3": "Jenis penyakit spesifik (hanya jika Diseased)",
        },
        "endpoints": {
            "classification": "POST /classification - Upload gambar untuk klasifikasi"
        },
    }


@app.post("/classification")
async def classify_plant(file: UploadFile = File(...)):
    """
    Endpoint klasifikasi penyakit tanaman selada dengan cascade 3-stage.

    Pipeline:
      1. Stage 1 — Plant vs Non-Plant
         • Jika Non-Plant → berhenti, kembalikan hasil
      2. Stage 2 — Healthy vs Diseased  (hanya jika terdeteksi Plant)
         • Jika Healthy → berhenti, kembalikan hasil
      3. Stage 3 — Klasifikasi jenis penyakit (hanya jika Diseased)

    Response fields:
      - is_plant            : true/false
      - stage              : stage terakhir yang dieksekusi (1, 2, atau 3)
      - final_label         : label akhir hasil klasifikasi
      - final_confidence    : confidence label akhir (0.0 – 1.0)
      - stage1              : hasil detail Stage 1
      - stage2              : hasil detail Stage 2 (null jika tidak dieksekusi)
      - stage3              : hasil detail Stage 3 (null jika tidak dieksekusi)
    """
    # ── Validasi file ──────────────────────────────────────────────────────────
    if not file.content_type or not file.content_type.startswith("image/"):
        raise HTTPException(
            status_code=400,
            detail=f"File harus berupa gambar. Content-type: {file.content_type}"
        )

    image_bytes = await file.read()
    if len(image_bytes) == 0:
        raise HTTPException(status_code=400, detail="File gambar kosong")

    tensor = preprocess_image(image_bytes)

    # ── Stage 1: Plant vs Non-Plant ────────────────────────────────────────────
    s1_pred, s1_conf, s1_probs = run_inference(stage1_model, tensor)
    # s1_probs[0] = Non-Plant, s1_probs[1] = Plant
    p_plant = s1_probs[1]  # probabilitas "Plant"
    is_plant = p_plant > STAGE1_THRESHOLD

    stage1_detail = {
        "prediction": "Plant" if is_plant else "Non-Plant",
        "confidence": round(s1_conf, 4),
        "probabilities": {
            "Non-Plant": round(s1_probs[0], 4),
            "Plant":     round(s1_probs[1], 4),
        },
    }

    if not is_plant:
        # Bukan tanaman — tidak perlu lanjut
        result = {
            "is_plant":         False,
            "stage":            1,
            "final_label":      "Non-Plant",
            "final_confidence": round(s1_conf, 4),
            "stage1":           stage1_detail,
            "stage2":           None,
            "stage3":           None,
        }
        logger.info(
            f"[{file.filename}] Stage1=Non-Plant ({s1_conf:.3f}) → SELESAI"
        )
        return JSONResponse(content=result)

    # ── Stage 2: Healthy vs Diseased ───────────────────────────────────────────
    s2_pred, s2_conf, s2_probs = run_inference(stage2_model, tensor)
    # s2_probs[0] = Healthy, s2_probs[1] = Diseased
    p_diseased = s2_probs[1]
    is_diseased = p_diseased > STAGE2_THRESHOLD

    stage2_detail = {
        "prediction": "Diseased" if is_diseased else "Healthy",
        "confidence": round(s2_conf, 4),
        "probabilities": {
            "Healthy":  round(s2_probs[0], 4),
            "Diseased": round(s2_probs[1], 4),
        },
    }

    if not is_diseased:
        # Tanaman sehat
        result = {
            "is_plant":         True,
            "stage":            2,
            "final_label":      "Healthy",
            "final_confidence": round(s2_conf, 4),
            "stage1":           stage1_detail,
            "stage2":           stage2_detail,
            "stage3":           None,
        }
        logger.info(
            f"[{file.filename}] Stage1=Plant ({s1_conf:.3f}) | "
            f"Stage2=Healthy ({s2_conf:.3f}) → SELESAI"
        )
        return JSONResponse(content=result)

    # ── Stage 3: Jenis Penyakit ────────────────────────────────────────────────
    s3_pred, s3_conf, s3_probs = run_inference(stage3_model, tensor)
    disease_raw    = DISEASE_CLASSES[s3_pred]
    disease_pretty = DISEASE_LABEL_MAP[disease_raw]

    stage3_detail = {
        "prediction": disease_pretty,
        "confidence": round(s3_conf, 4),
        "probabilities": {
            DISEASE_LABEL_MAP[cls]: round(prob, 4)
            for cls, prob in zip(DISEASE_CLASSES, s3_probs)
        },
    }

    result = {
        "is_plant":         True,
        "stage":            3,
        "final_label":      disease_pretty,
        "final_confidence": round(s3_conf, 4),
        "stage1":           stage1_detail,
        "stage2":           stage2_detail,
        "stage3":           stage3_detail,
    }

    logger.info(
        f"[{file.filename}] Stage1=Plant ({s1_conf:.3f}) | "
        f"Stage2=Diseased ({s2_conf:.3f}) | "
        f"Stage3={disease_pretty} ({s3_conf:.3f}) → SELESAI"
    )
    return JSONResponse(content=result)


# ── Main ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=7860,
        reload=False,
    )