"""
Inference utilities - Fish Freshness Classifier
"""

import os
import json
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
from torchvision import transforms, models

# ── Constants ────────────────────────────────────────────────────────────────
IMG_SIZE    = 224
NUM_CLASSES = 4
MODEL_PATH  = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "models", "best_model_fx.pth"))
CLASS_JSON  = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "models", "class_names.json"))

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

GRADE_RULES = {
    # (eye_result, gill_result) -> grade
    ("fresh",     "fresh"):     ("A", "Sangat Segar",  "#22c55e"),
    ("fresh",     "non-fresh"): ("B", "Cukup Segar",   "#f59e0b"),
    ("non-fresh", "fresh"):     ("B", "Cukup Segar",   "#f59e0b"),
    ("non-fresh", "non-fresh"): ("C", "Tidak Segar",   "#ef4444"),
}

GRADE_DESC = {
    "A": "Ikan sangat segar, layak ekspor premium. Mata jernih dan insang merah cerah.",
    "B": "Ikan cukup segar, layak ekspor standar. Salah satu indikator menunjukkan penurunan kesegaran.",
    "C": "Ikan tidak segar, tidak memenuhi syarat ekspor. Mata dan insang menunjukkan tanda tidak segar.",
}

# ── Transform ─────────────────────────────────────────────────────────────────
infer_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])

# ── Model loader (singleton cache) ───────────────────────────────────────────
_model_cache = None
_classes_cache = None

def load_model():
    global _model_cache, _classes_cache
    if _model_cache is not None:
        return _model_cache, _classes_cache

    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(
            f"Model tidak ditemukan di: {MODEL_PATH}\n"
            "Jalankan train.py terlebih dahulu."
        )

    # Load class names
    with open(CLASS_JSON) as f:
        classes = json.load(f)

    # Build model
    model = models.mobilenet_v3_large(weights=None)  # Ganti dengan DEFAULT jika ada akses internet
    in_features = model.classifier[3].in_features
    model.classifier[3] = nn.Sequential(
        nn.Dropout(p=0.4),
        nn.Linear(in_features, NUM_CLASSES),
    )
    model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
    model.to(DEVICE)
    model.eval()

    _model_cache = model
    _classes_cache = classes
    return model, classes


def predict_single(image_pil):
    """
    Predict freshness from a single cropped image (eye OR gill).
    Returns:
        class_name  : str  e.g. 'eye-fresh'
        confidence  : float
        all_probs   : dict {class_name: prob}
    """
    model, classes = load_model()
    tensor = infer_transform(image_pil.convert("RGB")).unsqueeze(0).to(DEVICE)
    with torch.no_grad():
        logits = model(tensor)
        probs  = torch.softmax(logits, dim=1)[0].cpu().numpy()

    idx        = int(probs.argmax())
    class_name = classes[idx]
    all_probs  = {classes[i]: float(probs[i]) for i in range(len(classes))}
    return class_name, float(probs[idx]), all_probs


def determine_grade(eye_class, gill_class):
    """
    Map eye + gill prediction to export grade.
    eye_class / gill_class: 'eye-fresh', 'eye-non-fresh', 'gill-fresh', 'gill-non-fresh'
    Returns: grade, label, color, description
    """
    # Extract just 'fresh' or 'non-fresh'
    eye_status  = "fresh" if "non" not in eye_class  else "non-fresh"
    gill_status = "fresh" if "non" not in gill_class else "non-fresh"

    grade, label, color = GRADE_RULES[(eye_status, gill_status)]
    desc = GRADE_DESC[grade]
    return grade, label, color, desc, eye_status, gill_status


def classify_fish(eye_image_pil, gill_image_pil):
    """
    Full pipeline: predict eye + gill → determine grade.
    Returns a result dict.
    """
    eye_class,  eye_conf,  eye_probs  = predict_single(eye_image_pil)
    gill_class, gill_conf, gill_probs = predict_single(gill_image_pil)

    grade, label, color, desc, eye_status, gill_status = determine_grade(eye_class, gill_class)

    return {
        "grade":       grade,
        "label":       label,
        "color":       color,
        "description": desc,
        "eye": {
            "class":      eye_class,
            "status":     eye_status,
            "confidence": eye_conf,
            "probs":      eye_probs,
        },
        "gill": {
            "class":      gill_class,
            "status":     gill_status,
            "confidence": gill_conf,
            "probs":      gill_probs,
        },
    }


def is_model_ready():
    return os.path.exists(MODEL_PATH) and os.path.exists(CLASS_JSON)