import math
from typing import Iterable, Any

import torch

from goldener.torch_utils import get_unique_values_in_tensor


def check_x_and_y_shapes(x_shape: tuple[int, ...], y_shape: tuple[int, ...]) -> None:
    """Check compatibility of shapes of x and y tensors.

    Args:
        x_shape: Shape of the input tensor x.
        y_shape: Shape of the target tensor y.

    Raises:
        ValueError: If the shapes are incompatible. See the conditions below for details.
    Conditions:
        - If x is 1D, x and y must have the same shape.
        - If x is 2D or more, y must have a channel dimension of 1 (i.e., y_shape[1] == 1).
        - If y has a batch size greater than 1 (i.e., y_shape[0] > 1), x and y must have the same batch size (i.e., x_shape[0] == y_shape[0]).
        - If x has more than 2 dimensions, the additional (after channel) dimensions of x and y must match.
    """

    if len(x_shape) == 1:
        if x_shape != y_shape:
            raise ValueError("x and y must have the same shape when x is 1D")
    else:
        if y_shape[1] != 1:
            raise ValueError(
                "x and y must have the same channel dimension when x is 2D"
            )

        batch_y = y_shape[0]
        if batch_y > 1 and x_shape[0] != batch_y:
            raise ValueError(
                "x and y must have the same batch size when y batch size > 1"
            )

        if len(x_shape) > 2 and x_shape[2:] != y_shape[2:]:
            raise ValueError(
                "x and y must have compatible shapes when x is more than 2D"
            )


def get_ratio_list_sum(ratios: list[float]) -> float:
    """Get the sum of a list of ratios and validate it (total between 0 and 1).

    Args:
        ratios: A list of float ratios.

    Returns:
        The sum of the ratios.

    Raises:
        ValueError: If the sum of the ratios is not greater than 0 and at most 1.
    """
    ratio_sum = sum(ratios)
    if not all(ratio >= 0 for ratio in ratios):
        raise ValueError("Ratios must be non-negative.")

    if not (0 < ratio_sum <= 1.0):
        raise ValueError("Sum of ratios must be 1.0.")

    return ratio_sum


def get_ratios_for_counts(counts: list[int]) -> list[float]:
    """Get ratios for a list of counts.

    Args:
        counts: A list of integer counts.

    Returns:
        A list of float ratios corresponding to the input counts.
    """
    if len(counts) == 0:
        raise ValueError("Counts list cannot be empty.")

    if len(counts) == 1:
        return [1.0]  # single count gets all the ratio

    total = sum(counts)
    if total == 0:
        raise ValueError("Total count must be greater than 0.")

    first_ratios = [count / total for count in counts[:-1]]

    return first_ratios + [
        1.0 - sum(first_ratios)
    ]  # ensure the last ratio is adjusted to sum to 1.0


def split_sampling_among_chunks(to_split: int, chunk_sizes: list[int]) -> list[int]:
    """Split a total sampling count among chunks based on their sizes.

    Args:
        to_split: The total number of samples to split among the chunks.
        chunk_sizes: A list of integers representing the sizes of each chunk.

    Returns:
        A list of integers representing the number of samples to draw from each chunk, summing up to to_split.
    """
    if len(chunk_sizes) == 0:
        raise ValueError("At least one chunk size is required to split the sampling")

    if len(chunk_sizes) == 1:
        if to_split > chunk_sizes[0]:
            raise ValueError(
                f"Split count {to_split} cannot be greater than chunk size {chunk_sizes[0]}"
            )
        return [to_split]

    total_size = sum(chunk_sizes)
    if total_size == 0:
        raise ValueError("Total size of chunks must be greater than 0.")

    ratios = get_ratios_for_counts(chunk_sizes)
    first_split_counts = [math.floor(ratio * to_split) for ratio in ratios[:-1]]

    split_counts = first_split_counts + [
        to_split - sum(first_split_counts)
    ]  # ensure the last count is adjusted to sum to to_split

    for split_count, chunk_size in zip(split_counts, chunk_sizes):
        if split_count > chunk_size:
            raise ValueError(
                f"Split count {split_count} cannot be greater than chunk size {chunk_size}"
            )

    return split_counts


def filter_batch_from_indices(
    batch: dict[str, Any],
    to_remove: set[int],
    index_key: str = "idx",
) -> dict[str, Any]:
    """Filter a batch dictionary to only include items at specified indices.

    Args:
        batch: A dictionary representing a batch of data (each key corresponds to stacked information).
        to_remove: A set of indices to filter out from the batch.
        index_key: The key in the batch dictionary that contains the indices.

    Returns:
        A filtered batch dictionary without the specified indices.
    """
    keep_in_batch = [
        idx_position
        for idx_position, idx_value in enumerate(batch[index_key])
        if (idx_value.item() if isinstance(idx_value, torch.Tensor) else idx_value)
        not in to_remove
    ]
    if not keep_in_batch:
        return {}  # all samples already described

    def filter_batched_values(
        batched_value: list | torch.Tensor,
    ) -> list | torch.Tensor:
        """Inner function to remove already described samples from the batch."""
        filtered = [
            value
            for idx_value, value in enumerate(batched_value)
            if idx_value in keep_in_batch
        ]
        if isinstance(batched_value, torch.Tensor):
            return torch.stack(filtered, dim=0)
        else:
            return filtered

    return {
        key: filter_batched_values(batched_value)
        for key, batched_value in batch.items()
    }


def get_size_and_sampling_count_per_chunk(
    total_size: int, sampling_size: int, max_chunk_size: int
) -> tuple[list[int], list[int]]:
    """Get sizes and sampling sizes per chunk.

    Args:
        total_size: Total size of the data to sample from.
        sampling_size: Total number of samples to draw.
        max_chunk_size: Maximum size of each chunk.

    Returns:
        A tuple of two lists:
            - A list of integers representing the size of each chunk.
            - A list of integers representing the number of samples to draw from each chunk.

    Raises:
        ValueError: If sampling_size is greater than or equal to total_size.
    """
    if sampling_size >= total_size:
        raise ValueError("sampling_size must be less than or equal to total_size")

    if max_chunk_size >= total_size:
        return [total_size], [sampling_size]  # single chunk

    chunk_count = math.ceil(total_size / max_chunk_size)
    chunk_size = total_size // chunk_count
    chunk_sampling_size = sampling_size // chunk_count

    chunk_sizes = [chunk_size] * (chunk_count - 1)
    chunk_sizes.append(total_size - sum(chunk_sizes))
    chunk_sampling_sizes = [chunk_sampling_size] * (chunk_count - 1)
    chunk_sampling_sizes.append(sampling_size - sum(chunk_sampling_sizes))

    return chunk_sizes, chunk_sampling_sizes


def check_sampling_size(
    sampling_size: int | float, total_size: int | None = None, force_max: bool = False
) -> None:
    """Check the validity of the sampling size.

    Args:
        sampling_size: The sampling size to check (can be int or float).
        total_size: The total size of the data (Optional).
        force_max: Whether the sampling size is expected to be the maximum possible (1.0 for float, total for int)

    Raises:
        ValueError: If the sampling size is invalid based on its type and total size.
        If sampling_size is a float, it must be in the range (0.0, 1.0].
        If sampling_size is an int, it must be in the range (0, total_size].
    """
    if isinstance(sampling_size, float):
        if force_max:
            if not math.isclose(sampling_size, 1.0, rel_tol=1e-9, abs_tol=0.0):
                raise ValueError("Sampling size as float must be equal to 1.0")
        else:
            if not 0 < sampling_size <= 1.0:
                raise ValueError(
                    "Sampling size as float must be greater than 0.0 and at most 1.0"
                )

    if isinstance(sampling_size, int) and total_size is not None:
        if force_max:
            if sampling_size != total_size:
                raise ValueError(
                    "Sampling size as int must be equal to the total number of samples"
                )
        else:
            if not 0 < sampling_size <= total_size:
                raise ValueError(
                    "Sampling size as int must be greater than 0 and less or equal than the total number of samples"
                )


def check_all_same_type(iterable: Iterable[Any]) -> None:
    """Check if all elements in an iterable are of the same type.

    Args:
        iterable: An iterable containing elements to check.

    Raises:
        TypeError: If any element in the iterable is not of the expected type.
    """
    first_type = type(next(iter(iterable)))
    different = []
    for item in iterable:
        if not isinstance(item, first_type):
            different.append(type(item))

    if different:
        raise TypeError(
            f"All elements must be of the same type {first_type}. Found different types: {different}"
        )


def get_sampling_count_from_size(
    sampling_size: int | float, total_size: int | None = None
) -> int:
    """Get the sampling count from the sampling size.

    Args:
        sampling_size: The sampling size (can be int or float). If int, it represents the exact number of samples
        and it is required to be less than total_size. If float, it represents the fraction of total size and it
        is required to be in the range (0.0, 1.0].
        total_size: The total size of the data (Optional, required if sampling_size is float).

    Returns:
        The calculated sampling count as an integer.

    Raises:
        ValueError: If the sampling size is invalid based on its type and total size.
    """
    if isinstance(sampling_size, int):
        if sampling_size <= 0:
            raise ValueError("Sampling size as int must be greater than 0.")

        if total_size is not None and not (sampling_size < total_size):
            raise ValueError(
                "Sampling size as int must be less than the total number of samples."
            )

        return sampling_size

    if not (0 < sampling_size <= 1.0):
        raise ValueError(
            "Sampling size as float must be greater than 0.0 and at most 1.0."
        )

    if total_size is None:
        raise ValueError("Total size must be provided when sampling size is a float.")

    return math.ceil(sampling_size * total_size)


def transform_batch_from_multiple_to_binarized_targets(
    batch: dict[str, Any],
    target_key: str,
    label_key: str | None,
    target_to_label: dict[tuple[int, ...], str],
    exclude_full_zero: bool = False,
) -> dict[str, Any]:
    """Transform a batch with multiple targets into a batch with binarized targets for each label.

    Args:
        batch: A dictionary representing a batch of data, where the target key contains a target with multiple
            values corresponding to the labels in target_to_label.
        target_key: The key in the batch dictionary that contains the target.
        label_key: The key in the batch dictionary that containes the labels.
        target_to_label: A mapping from unique target tuples to label values.
        exclude_full_zero: Whether to exclude the all-zero target from the transformation (default: False).

    Returns:
        A new batch dictionary where the target key contains binarized targets for each label
        and the label key contains the corresponding labels,
        with other batch elements duplicated accordingly.
    """
    multi_target = batch[target_key]

    multi_target_shape = multi_target.shape
    bin_target_shape = list(multi_target_shape)
    bin_target_shape[1] = 1

    # create a new binarized target from the target
    # containing a different value for each label
    target_flattened = multi_target.movedim(1, -1).reshape(-1, multi_target_shape[1])
    target_per_label = {}
    for unique_target in get_unique_values_in_tensor(target_flattened, -1):
        if exclude_full_zero and (unique_target == 0).all():
            continue  # skip the all-zero target if specified

        unique_target_tuple = tuple(unique_target.tolist())
        if unique_target_tuple not in target_to_label:
            raise ValueError(
                f"Unique target {unique_target} not found in target_to_label mapping."
            )
        label = target_to_label[unique_target_tuple]

        new_target = (
            (target_flattened == unique_target.unsqueeze(0))
            .all(dim=1)
            .to(torch.uint8)
            .reshape(bin_target_shape)
        )

        target_per_label[label] = new_target

    if not target_per_label:
        raise ValueError(
            "No valid targets found after applying exclude_full_zero filter."
        )

    # duplicate the batch element for each label/target and
    # insert the corresponding binarized target/label in the batch alongside them.
    new_batch: dict[str, torch.Tensor | list[Any]] = {}
    n_values = len(target_per_label)
    for batch_key, batch_value in batch.items():
        if batch_key not in (target_key, label_key):
            if isinstance(batch_value, torch.Tensor):
                new_batch[batch_key] = torch.cat([batch_value] * n_values, dim=0)
            else:
                if not isinstance(batch_value, list):
                    raise ValueError(
                        f"Batch value for key {batch_key} must be a list or a torch.Tensor. Got {type(batch_value)}."
                    )
                new_batch[batch_key] = batch_value * n_values
        else:
            if batch_key == target_key:
                new_batch[target_key] = torch.cat(
                    [target for target in target_per_label.values()],
                    dim=0,
                )
            elif label_key is not None and batch_key == label_key:
                new_batch[label_key] = list(target_per_label.keys())

    # add label if not already there
    if label_key is not None and label_key not in new_batch:
        new_batch[label_key] = list(target_per_label.keys())

    return new_batch
