import multiprocessing as mp
from collections.abc import Iterable
from concurrent.futures import ProcessPoolExecutor
from concurrent.futures.process import BrokenProcessPool
import warnings
from pathlib import Path
from typing import Union

import numpy as np
from PIL import Image

from ._check import check_image
from ._concurrent import DEFAULT_START_METHOD

MEAN = tuple[float, ...]
STD = tuple[float, ...]


def get_mean_and_std_from_images(images: Iterable[np.ndarray]) -> tuple[MEAN, STD]:
    """Get channel mean and std values for an iterable of images.

    Args:
        images: Iterable of array images.

    Returns:
        Tuple of channel mean and std values.
    """
    mean, std = [], []
    for img in images:
        m, s = _get_mean_and_std(img)
        mean.append(m)
        std.append(s)
    return tuple(np.array(mean).mean(0)), tuple(np.array(std).mean(0))


def get_mean_and_std_from_paths(
    paths: Iterable[Union[str, Path]], num_workers: int = 1
) -> tuple[MEAN, STD]:
    """Get channel mean and std values for an iterable of image paths.

    Args:
        paths: Iterable of image paths.
        num_workers: Number of image reading processes.

    Returns:
        Tuple of channel mean and std values.
    """
    if num_workers <= 1:
        return get_mean_and_std_from_images(_read_image(x) for x in paths)

    ctx = mp.get_context(DEFAULT_START_METHOD)
    try:
        with ProcessPoolExecutor(max_workers=num_workers, mp_context=ctx) as pool:
            images = pool.map(_read_image, paths)
            return get_mean_and_std_from_images(images)
    except BrokenProcessPool:
        warnings.warn(
            "Mean/std multiprocessing pool crashed; falling back to serial execution.",
            RuntimeWarning,
            stacklevel=2,
        )
        # Fall back to serial if the worker pool crashes (e.g. decoder issues).
        return get_mean_and_std_from_images(_read_image(x) for x in paths)


def _get_mean_and_std(image: np.ndarray) -> tuple[MEAN, STD]:
    """Calculate mean and standard deviation for each image channel (between [0, 1])."""
    check_image(image)
    if image.ndim == 2:  # noqa
        image = np.expand_dims(image, -1)  # noqa
    return (
        tuple([image[..., i].mean() / 255 for i in range(image.shape[-1])]),
        tuple([image[..., i].std() / 255 for i in range(image.shape[-1])]),
    )


def _read_image(path: Union[str, Path]) -> np.ndarray:
    img = Image.open(path)
    if img.mode not in ("RGB", "L"):
        img = img.convert("RGB")
    return np.array(img)
