import math


def cluster_size(n_obs, n_cluster, size_flexibility=0):
    """
    Compute the number of items allowed in each cluster. Allow a flexibility to create almost same size clusters.

    Args:
        n_obs: int, n_samples value.
        n_cluster: int, The number of clusters to form as well as the number of centroids to generate.
        size_flexibility: int | float (default 0), flexibility allowed for the size of clusters.

            float numbers under 1 represent the percentage of variation allowed for cluster's size.
            Example: 0.2 for ±20% variation per cluster.

            integer numbers over 1 represent the number of observation variation allowed for cluster's size.
            Example: 5 for ±5 items per cluster.

            0 represents a fixed sized without any variation allowed.

    Returns:
        n_max: int, maximum size of a cluster. Needed in initialization part. (Unused atm)
        n_min: int, minimum size of a cluster. Needed in initialization part.
        n_min_flex: int, minimum size of a cluster (including the flexibility). Needed in iteration part.
        n_max_flex: int, maximum size of a cluster (including the flexibility). Needed in iteration part.

    """

    # Get the estimated size of clusters (total number of items / number of clusters)
    if (n_obs % n_cluster == 0) and (size_flexibility == 0):
        # Case of equal sized clusters
        n_min = n_min_flex = n_obs / n_cluster
        n_max = n_max_flex = n_obs / n_cluster

    else:
        # Case of unequal sized clusters
        n_min = math.floor(n_obs / n_cluster)
        n_max = math.ceil(n_obs / n_cluster)

        if size_flexibility >= 1:
            n_min_flex = max(n_min - int(abs(size_flexibility)), 1)
            n_max_flex = min(n_max + int(abs(size_flexibility)), n_obs - 1)
        else:
            n_min_flex = max(n_min - math.floor(abs(size_flexibility) * math.floor(n_obs / n_cluster)), 1)
            n_max_flex = min(n_max + math.floor(abs(size_flexibility) * math.floor(n_obs / n_cluster)), n_obs - 1)

    return int(n_max), int(n_min), int(n_max_flex), int(n_min_flex)