"""
predictor.py
============
Class DRPredictor chính đóng gói toàn bộ quy trình load model, nạp trọng số,
tiền xử lý ảnh và suy luận (Inference) phân loại mức độ bệnh Võng mạc Tiểu đường.
"""

from __future__ import annotations
import os
import json
from typing import Dict, List, Union, Any
import torch
import numpy as np
from PIL import Image

try:
    from .model import ResNet50_DR
    from .preprocessing import prepare_image_tensor, load_image, full_preprocess_pipeline
except ImportError:
    from model import ResNet50_DR
    from preprocessing import prepare_image_tensor, load_image, full_preprocess_pipeline


class DRPredictor:
    """
    Predictor đính kèm mô hình AI ResNet-50.
    
    Ví dụ sử dụng:
        >>> from modelAI_ResNet50 import DRPredictor
        >>> predictor = DRPredictor()
        >>> result = predictor.predict("path/to/fundus_image.png")
        >>> print(result["class_name"], result["confidence"])
    """
    def __init__(
        self,
        weights_path: str | None = None,
        config_path: str | None = None,
        device: str | None = None,
    ):
        base_dir = os.path.dirname(os.path.abspath(__file__))

        # Load file cấu hình config.json
        if config_path is None:
            config_path = os.path.join(base_dir, "config.json")
        
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"Không tìm thấy file cấu hình: {config_path}")
        
        with open(config_path, "r", encoding="utf-8") as f:
            self.config = json.load(f)

        # Xác định thiết bị tính toán (CUDA / CPU)
        if device is None:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(device)

        # Cấu hình đường dẫn trọng số .pth
        if weights_path is None:
            weights_name = self.config.get("default_weights", "resnet50_baseline_fold1.pth")
            weights_path = os.path.join(base_dir, weights_name)

        if not os.path.exists(weights_path):
            raise FileNotFoundError(f"Không tìm thấy file trọng số: {weights_path}")

        # Khởi tạo khung mô hình ResNet50
        self.model = ResNet50_DR(
            num_classes=self.config["num_classes"],
            drop_rate=self.config.get("drop_rate", 0.3),
            pretrained=False,
        )

        # Load trọng số PyTorch checkpoint
        try:
            checkpoint = torch.load(weights_path, map_location=self.device, weights_only=False)
        except TypeError:
            checkpoint = torch.load(weights_path, map_location=self.device)
        if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
            state_dict = checkpoint["model_state_dict"]
        elif isinstance(checkpoint, dict) and "state_dict" in checkpoint:
            state_dict = checkpoint["state_dict"]
        else:
            state_dict = checkpoint

        self.model.load_state_dict(state_dict)
        self.model.to(self.device)
        self.model.eval()

        self.class_names: Dict[str, str] = self.config["class_names"]
        self.img_size = tuple(self.config["img_size"])
        self.mean = tuple(self.config["mean"])
        self.std = tuple(self.config["std"])

    def predict(
        self,
        image_input: Union[str, bytes, Image.Image, np.ndarray],
        use_ben_graham: bool = True,
    ) -> Dict[str, Any]:
        """
        Dự đoán mức độ DR cho một ảnh duy nhất.

        Parameters
        ----------
        image_input : Đường dẫn file (str), Data Bytes, PIL Image, hoặc NumPy BGR Array.
        use_ben_graham : Áp dụng lọc Ben Graham trước khi suy luận.

        Returns
        -------
        Dict chứa class_id, class_name, confidence (%) và xác suất cho từng class.
        """
        tensor_img = prepare_image_tensor(
            image_input=image_input,
            target_size=self.img_size,
            mean=self.mean,
            std=self.std,
            use_ben_graham=use_ben_graham,
        )
        batch_tensor = tensor_img.unsqueeze(0).to(self.device)

        with torch.no_grad():
            outputs = self.model(batch_tensor)
            probs = torch.softmax(outputs, dim=1)[0].cpu().numpy()

        pred_class = int(np.argmax(probs))
        confidence = float(probs[pred_class])

        probabilities_dict = {
            self.class_names.get(str(i), f"Class {i}"): float(probs[i])
            for i in range(len(probs))
        }

        return {
            "class_id": pred_class,
            "class_name": self.class_names.get(str(pred_class), f"Class {pred_class}"),
            "confidence": confidence,
            "probabilities": probabilities_dict,
        }

    def predict_batch(
        self,
        image_inputs: List[Union[str, bytes, Image.Image, np.ndarray]],
        use_ben_graham: bool = True,
    ) -> List[Dict[str, Any]]:
        """
        Dự đoán đồng thời theo danh sách ảnh (Batch Inference).
        """
        if not image_inputs:
            return []

        tensors = [
            prepare_image_tensor(
                img,
                target_size=self.img_size,
                mean=self.mean,
                std=self.std,
                use_ben_graham=use_ben_graham,
            )
            for img in image_inputs
        ]
        batch_tensor = torch.stack(tensors, dim=0).to(self.device)

        with torch.no_grad():
            outputs = self.model(batch_tensor)
            probs_batch = torch.softmax(outputs, dim=1).cpu().numpy()

        results = []
        for probs in probs_batch:
            pred_class = int(np.argmax(probs))
            confidence = float(probs[pred_class])
            probabilities_dict = {
                self.class_names.get(str(i), f"Class {i}"): float(probs[i])
                for i in range(len(probs))
            }
            results.append({
                "class_id": pred_class,
                "class_name": self.class_names.get(str(pred_class), f"Class {pred_class}"),
                "confidence": confidence,
                "probabilities": probabilities_dict,
            })

        return results