import torch
import cv2
import numpy as np
from pathlib import Path
from PIL import Image
from torchvision import transforms
import segmentation_models_pytorch as smp

try:
    # Python >= 3.9
    from importlib.resources import files
except ImportError:
    # Python < 3.9 fallback (generally not recommended)
    from importlib_resources import files


class BladeSegmentor:
    """
    Wind turbine blade segmentation inference wrapper.

    External API (stable):
        segment(image_path, output_path, threshold=0.5)

    Notes:
    - Model architecture and preprocessing strictly match train.py
    - Output is a binary mask (uint8, values in {0, 255})
    """

    # torchvision.transforms.Resize expects (height, width)
    IMAGE_HEIGHT = 640
    IMAGE_WIDTH = 960

    def __init__(
        self,
        weight_path: str | None = None,
        device: str | None = None
    ):
        """
        Args:
            weight_path:
                - None: load built-in package weight (for pip-installed usage)
                - str : load weight from specified path
            device:
                - None: auto select cuda if available else cpu
                - "cpu" / "cuda"
        """
        # -------- device --------
        self.device = torch.device(
            device if device else ("cuda" if torch.cuda.is_available() else "cpu")
        )

        # -------- resolve weight path --------
        if weight_path is None:
            # load packaged weight: blade_seg/weights/best_model.pth
            try:
                weight_path = files("blade_seg.weights") / "best_model.pth"
            except Exception as e:
                raise RuntimeError(
                    "Default weight not found inside package. "
                    "If you are running from source, please pass weight_path explicitly."
                ) from e
        else:
            weight_path = Path(weight_path)

        if not Path(weight_path).exists():
            raise FileNotFoundError(f"Weight file not found: {weight_path}")

        # -------- model definition (IDENTICAL to train.py) --------
        self.model = smp.Segformer(
            encoder_name="mit_b0",
            encoder_weights=None,   # inference: do NOT load imagenet weights
            in_channels=3,
            classes=1,
            activation=None
        ).to(self.device)

        # -------- load weights --------
        state_dict = torch.load(weight_path, map_location=self.device)
        self.model.load_state_dict(state_dict)
        self.model.eval()

        # -------- preprocessing --------
        self.transform = transforms.Compose([
            transforms.Resize((self.IMAGE_HEIGHT, self.IMAGE_WIDTH)),
            transforms.ToTensor()
        ])

    @torch.no_grad()
    def segment(
        self,
        image_path: str,
        output_path: str,
        threshold: float = 0.5
    ):
        """
        Perform blade segmentation on a single image.

        Args:
            image_path:
                Path to input RGB image.
            output_path:
                Path to output binary mask image (PNG recommended).
            threshold:
                Probability threshold for binarization (default = 0.5).

        Output:
            - Saved mask image (uint8, 0 or 255)
        """
        image_path = Path(image_path)
        if not image_path.exists():
            raise FileNotFoundError(f"Input image not found: {image_path}")

        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # -------- load image --------
        img = Image.open(image_path).convert("RGB")
        orig_w, orig_h = img.size  # PIL uses (width, height)

        # -------- preprocess --------
        x = self.transform(img).unsqueeze(0).to(self.device)

        # -------- inference --------
        logits = self.model(x)              # (1, 1, H, W)
        prob = torch.sigmoid(logits)[0, 0]  # (H, W)

        # -------- binarize --------
        mask = (prob > threshold).cpu().numpy().astype(np.uint8) * 255

        # -------- resize back to original size --------
        mask = cv2.resize(
            mask,
            (orig_w, orig_h),               # OpenCV expects (width, height)
            interpolation=cv2.INTER_NEAREST
        )

        # -------- save result (check success) --------
        ok = cv2.imwrite(str(output_path), mask)
        if not ok:
            raise RuntimeError(f"Failed to write output mask to {output_path}")
