import math
from typing import Iterable, Any, TypeVar

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.

    Raises:
        ValueError: If the input parameters are invalid. See the conditions below for details.
    Conditions:
        - At least one chunk size must be provided.
        - If only one chunk size is provided, to_split must not be greater than that chunk size.
        - The total size of the chunks must be greater than 0.
    """
    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)
    counts = get_sampling_count_from_ratios(
        ratios={
            i: ratio for i, ratio in enumerate(ratios)
        },  # add dummy keys for the ratios
        sampling_size=to_split,
        force_non_zero=False,
    )
    counts = dict(sorted(counts.items()))  # reorder in the chunk order

    return list(counts.values())


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_indices_with_labels(
    batch: dict[str, Any],
    label_key: str,
    exclude_labels: set[str],
    index_key: str = "idx",
) -> set[int]:
    """Get the set of index values for batch items whose sharing label values with exclude_labels.

    Args:
        batch: A dictionary representing a batch of data. It must contain the keys specified by label_key and index_key,
        where the values are lists or tensors of the same length. Label values can be strings or iterables of strings.
        label_key: The key in the batch dictionary that contains the labels.
        exclude_labels: A set of label strings to exclude.
        index_key: The key in the batch dictionary that contains the indices.

    Returns:
        A set of index values corresponding to items with excluded labels.
    """
    indices = set()
    for idx_value, label_value in zip(batch[index_key], batch[label_key]):
        labels = (
            {
                label_value,
            }
            if isinstance(label_value, str)
            else label_value
        )
        if exclude_labels.intersection(labels):
            indices.add(
                int(idx_value.item())
                if isinstance(idx_value, torch.Tensor)
                else idx_value
            )

    return indices


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_labels: set[str] | None = None,
    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_labels: An optional set of labels to exclude from the transformation (default: None).
        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]

        if exclude_labels is not None and label in exclude_labels:
            continue

        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_labels = []
                for label, target in target_per_label.items():
                    new_labels += [label] * len(target)
                new_batch[label_key] = new_labels

    # 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


T = TypeVar("T")


def get_sampling_count_from_ratios(
    ratios: dict[T, float], sampling_size: int, force_non_zero: bool = False
) -> dict[T, int]:
    """Get the sampling count for each key from the ratios.

    Args:
        ratios: A dictionary mapping keys to their corresponding ratios (values between 0 and 1).
            the initial counts are computed by multiplying the sampling size by the ratios
            and taking the floor of the result. If the sum of the initial counts is less than
            the sampling size, the remaining samples are distributed one by one starting from the
            most represented key until the total count matches the sampling size.
        sampling_size: The total number of samples to draw.
        force_non_zero: Whether to force a non-zero sampling count for each key (default: False).
            If activated, the least represented keys will have at least 1 sample,
            and the most represented keys will have their count reduced accordingly to maintain the total sampling size.
    Returns:
        A dictionary mapping keys to their corresponding sampling counts (integers).
            The dictionary is sorted by count in ascending order.
    """
    if not math.isclose(sum(ratios.values()), 1.0, rel_tol=1e-9, abs_tol=0.0):
        raise ValueError("Ratios must sum to 1.0")

    # compute the count per key and make sure they sum to the sampling size by
    # adjusting each count starting from the most represented key until the total count matches the sampling size
    counts = {key: math.floor(sampling_size * ratio) for key, ratio in ratios.items()}

    count_total = sum(counts.values())
    if count_total < sampling_size:
        counts = dict(sorted(counts.items(), key=lambda x: x[1], reverse=True))
        to_distribute = sampling_size - count_total
        loop_count = math.ceil(to_distribute / len(counts))
        for _ in range(loop_count):
            for key in counts.keys():
                counts[key] += 1
                to_distribute -= 1
                if to_distribute == 0:
                    break

    assert sum(counts.values()) == sampling_size

    if not force_non_zero:
        return dict(sorted(counts.items(), key=lambda x: x[1]))

    # If force_non_zero is activated, adjust counts to ensure no key has a count of zero
    # first, the zero counts are identified and incremented by 1, while keeping track of how many were incremented
    added = 0
    for key, count in counts.items():
        if count == 0:
            added += 1
            counts[key] += 1

    # then, for each added count, the counts are sorted in descending order
    # and the highest count is decremented by 1 to maintain the total sampling size
    for _ in range(added):
        counts = dict(sorted(counts.items(), key=lambda x: x[1], reverse=True))
        highest_count_key = next(iter(counts.keys()))
        counts[highest_count_key] -= 1
        if counts[highest_count_key] == 0:
            raise ValueError(
                f"While trying to adjust counts, the count for key '{highest_count_key}' became zero."
            )

    return dict(sorted(counts.items(), key=lambda x: x[1]))
