class HistogramClustering(skl.base.BaseEstimator, skl.base.TransformerMixin):
    """Template class for HistogramClustering (HC)

    Attributes:
        centroids (np.ndarray): Array of centroid distributions p(y|c) with shape (n_clusters, n_bins).

    Parameters:
        n_clusters (int): Number of clusters (textures).
        n_bins (int): Number of bins used to discretize the range of pixel values found in input image X.
        window_size (int): Size of the window used to compute the local histograms for each pixel.
                           Should be an odd number larger or equal to 3.
        random_state (int): Random seed.
        estimation (str): Whether to use Maximum a Posteriori ("MAP") or
                          Deterministic Annealing ("DA") estimation.
    """

    def __init__(self, n_clusters=10, n_bins=64, window_size=7, random_state=42, estimation="MAP"):
        self.n_clusters = n_clusters
        self.n_bins =n_bins
        self.window_size = window_size
        self.random_state = random_state
        self.estimation = estimation
        # Add more parameters, if necessary.

    def fit(self, X):
        """Compute HC for input image X

        Compute centroids.

        Args:
            X (np.ndarray): Input array with shape (height, width)

        Returns:
            self
        """

        np.histogram(X, bins=np.arange(self.n_bins))

        # First calculate the bins

        normalized_X = X.copy()

        normalized_X *= self.n_bins
        # normalized_X = np.ceil(normalized_X)

        N = np.prod(X.shape)

        n = np.zeros((N, self.n_bins))

        indices = [(x, y) for x, y in np.ndindex(X.shape)]

        pad = self.window_size // 2
        indices = np.add(indices, pad)
        Xmax = normalized_X.max()
        padded_X = np.pad(normalized_X, pad, 'constant', constant_values=normalized_X.max() * 10)


        i = 0
        for idx in indices:
            histogram = np.histogram(padded_X[idx[0]-pad:idx[0]+pad, idx[1]-pad:idx[1]+pad], bins=self.n_bins, range=(padded_X.min(), Xmax))
            n[i, ] = histogram[0]
            i += 1

        if self.estimation == "MAP":

            # These are constant w.r.t the iterative algorithm
            n_x = np.sum(n, axis=1).reshape(-1, 1)

            # p_x = n_x / np.sum(n_x)

            p_y_x = n / n_x

            p_y_c = np.random.multivariate_normal(np.zeros(self.n_bins * self.n_clusters), np.eye(self.n_bins * self.n_clusters)).reshape(self.n_bins, self.n_clusters)
            p_y_c /= np.sum(p_y_c, axis=1).reshape(-1, 1)

            c_x = np.random.randint(0, self.n_clusters, size=(N, 1))
            c_x_old = c_x

            p_y_c_old = p_y_c

            iter = 0
            while True:
                # estimate p(y|c):w
                mask = np.zeros((N, self.n_clusters))
                mask[np.arange(N).reshape(-1, 1), c_x_old] = 1
                for c in range(self.n_clusters):
                    inner_sum = n_x / np.sum(n_x * mask[:, c].reshape(-1, 1), axis=0, keepdims=True)
                    row = np.sum( inner_sum * p_y_x * mask[:, c].reshape(-1, 1), axis=0)

                    p_y_c[:, c] = row

                cluster_assignments = np.zeros((N, self.n_clusters))
                for c in range(self.n_clusters):
                    # m = 1e16
                    cluster_assignments[:, c] = -np.sum(p_y_x * np.log(p_y_c_old[:, c] + np.finfo(float).eps), axis=1)

                c_x = np.argmin(cluster_assignments, axis=1).reshape(-1, 1)

                diff = np.sum(p_y_c - p_y_c_old)
                if iter > 10:
                    break
                c_x_old = c_x
                p_y_c_old = p_y_c
                iter += 1

                print('\rFinished iter {}, diff={}'.format(iter, diff), end='', flush=True)


            self.p_y_x = p_y_x
            self.p_y_c = p_y_c
            self.c_x = c_x

            self.centroids = p_y_c

            # Code for Maximum a Posteriori estimation

        elif self.estimation == "DA":
            raise NotImplementedError()

            # Code for Deterministic Annealing estimation

        return self

    def predict(self, X):
        """Predict cluster assignments for each pixel in image X.

        Args:
            X (np.ndarray): Input array with shape (height, width)

        Returns:
            C (np.ndarray): Assignment map (height, width)
        """
        check_is_fitted(self, ["centroids"])

        # Your code goes here

        return C

    def generate(self, C):
        """Generate a sample image X from a texture label map C.

        The entries of C are integers from the set {1,...,n_clusters}. They represent the texture labels
        of each pixel. Given the texture labels, a sample image X is generated by sampling
        the value of each pixel from the fitted p(y|c).

        Args:
            C (np.ndarray): Input array with shape (height, width)

        Returns:
            X (np.ndarray): Sample image (height, width)
        """
        check_is_fitted(self, ["centroids"])

        # Your code goes here

        return X