from PIL import Image, ImageDraw
import os
import io
import numpy as np
import torch
import cv2
from tqdm import tqdm
import sys
from pathlib import Path
import xarray as xr

def predict_single_tile(model, image, device, tile_size: int, th: float):
    """
    Predict a segmentation mask for a single image tile.

    This function processes a single tile of the input image through a trained model,
    applies postprocessing, and returns the predicted binary mask for that tile.

    :param model: Trained PyTorch model used for prediction.
    :type model: torch.nn.Module
    :param image: Input image (PIL Image or NumPy array).
    :type image: PIL.Image.Image | np.ndarray
    :param device: Torch device (e.g., 'cpu' or 'cuda').
    :type device: torch.device
    :param tile_size: Size of the tile (in pixels) to extract and process.
    :type tile_size: int
    :param th: Threshold value (0–1) to binarize the output probabilities.
    :type th: float
    :return: Predicted binary mask for the specified tile.
    :rtype: np.ndarray
    """
    if isinstance(image, np.ndarray):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))

    width, height = image.size
    x, y = 0, 0
    x_end = min(x + tile_size, width)
    y_end = min(y + tile_size, height)

    tile = image.crop((x, y, x_end, y_end))

    if tile_size != 256:
        tile = cv2.resize(np.array(tile), (256, 256), interpolation=cv2.INTER_LINEAR)
    else:
        tile = np.array(tile)

    tile_tensor = preprocess_tile2(tile).to(device)

    with torch.no_grad():
        output = model(tile_tensor)

    predicted_tile_mask = postprocess_mask(output, th)

    if tile_size != 256:
        predicted_tile_mask = cv2.resize(predicted_tile_mask, (tile_size, tile_size), interpolation=cv2.INTER_LINEAR)

    return predicted_tile_mask


def create_img_masks_lists(images_path: str, masks_path: str):
    """
    Create sorted lists of images and their corresponding masks.

    :param images_path: Directory containing input images.
    :type images_path: str
    :param masks_path: Directory containing corresponding mask images.
    :type masks_path: str
    :return: Sorted lists of loaded images and masks.
    :rtype: tuple[list[PIL.Image.Image], list[np.ndarray]]
    """
    images_list = sorted([Image.open(os.path.join(images_path, im)) for im in os.listdir(images_path)])
    masks_list = sorted([cv2.imread(os.path.join(masks_path, im)) for im in os.listdir(masks_path)])
    return images_list, masks_list


def preprocess_tile2(tile: np.ndarray) -> torch.Tensor:
    """
    Preprocess an image tile for model inference.

    Normalizes pixel values to [0, 1] and converts the image into a 4D PyTorch tensor.

    :param tile: Input RGB image tile.
    :type tile: np.ndarray
    :return: Normalized tensor of shape (1, 3, H, W).
    :rtype: torch.Tensor
    """
    tile = np.array(tile)
    tile = tile / 255.0
    tile = torch.from_numpy(tile).float().permute(2, 0, 1).unsqueeze(0)
    return tile


def postprocess_mask(mask: torch.Tensor, th: float) -> np.ndarray:
    """
    Postprocess model output to generate a binary mask.

    :param mask: Raw output tensor from the model.
    :type mask: torch.Tensor
    :param th: Threshold value (0–1) for binarization.
    :type th: float
    :return: Binary mask (values 0 or 1) as NumPy array.
    :rtype: np.ndarray
    """
    mask = torch.sigmoid(mask)
    mask = mask.squeeze().cpu().detach().numpy()
    mask = (mask > th).astype(np.float32)
    return mask


def predict_tiles(images_paths: list[str], model: torch.nn.Module):
    """
    Predict segmentation masks for multiple image files.

    :param images_paths: List of file paths to input images.
    :type images_paths: list[str]
    :param model: Trained PyTorch model for inference.
    :type model: torch.nn.Module
    :return: List of predicted masks for each input image.
    :rtype: list[np.ndarray]
    """
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    masks = []
    for image_path in tqdm(images_paths):
        image = Image.open(image_path)
        mask = predict_single_tile(model, image, device, tile_size=256, th=0.5)
        masks.append(mask)
    return masks


def Sliding_Window(model: torch.nn.Module, image: Image.Image, device: torch.device,
                   tile_h: int, tile_w: int, stride: int, th: float) -> np.ndarray:
    """
    Perform tiled inference with overlap (sliding window).

    Processes an entire image by dividing it into overlapping tiles,
    predicting each tile, and averaging overlapping regions.

    :param model: Trained PyTorch model.
    :type model: torch.nn.Module
    :param image: Input PIL image.
    :type image: PIL.Image.Image
    :param device: Device to run inference on ('cpu' or 'cuda').
    :type device: torch.device
    :param tile_h: Tile height in pixels.
    :type tile_h: int
    :param tile_w: Tile width in pixels.
    :type tile_w: int
    :param stride: Step size for moving the sliding window.
    :type stride: int
    :param th: Threshold to apply sigmoid binarization.
    :type th: float
    :return: Full-size probability map (averaged over overlapping tiles).
    :rtype: np.ndarray
    """
    width, height = image.size
    accum = np.zeros((height, width), dtype=np.float32)
    weight = np.zeros((height, width), dtype=np.float32)

    ys = list(range(0, max(height - tile_h + 1, 1), stride))
    xs = list(range(0, max(width - tile_w + 1, 1), stride))
    if ys[-1] != height - tile_h:
        ys.append(max(height - tile_h, 0))
    if xs[-1] != width - tile_w:
        xs.append(max(width - tile_w, 0))

    for y in ys:
        for x in xs:
            tile = image.crop((x, y, x + tile_w, y + tile_h))

            if (tile_h, tile_w) != (256, 256):
                tile_np = cv2.resize(np.array(tile), (256, 256), interpolation=cv2.INTER_LINEAR)
            else:
                tile_np = np.array(tile)

            tile_tensor = preprocess_tile2(tile_np).to(device)

            with torch.no_grad():
                output = model(tile_tensor)

            prob_tile = postprocess_mask(output, th)

            if (tile_h, tile_w) != (256, 256):
                prob_tile = cv2.resize(prob_tile, (tile_w, tile_h), interpolation=cv2.INTER_LINEAR)

            accum[y:y + tile_h, x:x + tile_w] += prob_tile
            weight[y:y + tile_h, x:x + tile_w] += 1.0

    weight = np.maximum(weight, 1e-6)
    full_prob = accum / weight
    return full_prob


def Full_Scene_Probability_Mask(
    model: torch.nn.Module,
    image_path: str,
    device: torch.device,
    tile_h: int,
    tile_w: int,
    stride: int,
    th: float = None
):
    """
    Generate a full-scene probability map or binary segmentation mask using
    tiled inference over either:
      • Standard RGB image files (JPG, PNG…)
      • NetCDF MSG Ash or Ice RGB products

    This function:
      1. Loads the input image
         - If NetCDF: detects correct RGB band names (Ash or Ice)
         - If standard image: converts to RGB
      2. Runs sliding-window inference using the trained segmentation model
      3. Optionally thresholds the probability map into a binary mask

    Parameters
    ----------
    model : torch.nn.Module
        Loaded segmentation model.
    image_path : str
        Path to input image (NetCDF or RGB).
    device : torch.device
        Device to perform inference on (CPU/GPU).
    tile_h, tile_w : int
        Tile size for sliding-window inference.
    stride : int
        Tile stride for overlap (lower = more blending, slower).
    th : float, optional
        Threshold to obtain a binary mask. If None, return raw probabilities.

    Returns
    -------
    mask_or_prob : np.ndarray
        Full-scene probability map (float32) or binary mask.
    image : PIL.Image.Image
        The reconstructed full-scene RGB image.
    """
    image_path = Path(image_path)

    # ─────────────────────────────────────────────────────────────────────
    # Load NetCDF (Satellite RGB) image
    # ─────────────────────────────────────────────────────────────────────
    if image_path.suffix == ".nc":
        with xr.open_dataset(image_path) as ds:
            
            # Variables for both supported product types
            ash_vars = ("ash_red", "ash_green", "ash_blue")
            ice_vars = ("ice_red", "ice_green", "ice_blue")

            # Identify available band set
            if all(v in ds.variables for v in ash_vars):
                variables = ash_vars
            elif all(v in ds.variables for v in ice_vars):
                variables = ice_vars
            else:
                raise ValueError(
                    f"Neither Ash nor Ice RGB variables found in {image_path}. "
                    f"Available variables: {list(ds.data_vars)}"
                )

            # Read the three RGB channels
            red   = ds[variables[0]].values
            green = ds[variables[1]].values
            blue  = ds[variables[2]].values

        # Stack into RGB format expected by the model
        rgb = np.stack([red, green, blue], axis=-1)
        image = Image.fromarray(rgb, mode="RGB")

    # ─────────────────────────────────────────────────────────────────────
    # Load standard RGB image file
    # ─────────────────────────────────────────────────────────────────────
    else:
        image = Image.open(image_path).convert("RGB")

    # ─────────────────────────────────────────────────────────────────────
    # Run inference with tiling (helps with large satellite scenes)
    # ─────────────────────────────────────────────────────────────────────
    prob = Sliding_Window(
        model=model,
        image=image,
        device=device,
        tile_h=tile_h,
        tile_w=tile_w,
        stride=stride,
        th=th
    )

    # Optionally convert probability map into binary segmentation mask
    if th is not None:
        mask = (prob >= th).astype(np.float32)
        return mask, image

    return prob, image
