"""Sparse 1-D histogram of healpix pixel counts."""

import numpy as np

import hats.pixel_math.healpix_shim as hp


class SparseHistogram:
    """Wrapper around a naive sparse array, that is just non-zero indexes and counts.

    e.g. for a dense 1-d numpy histogram of order 0, you might see::

        [0, 4, 0, 0, 0, 0, 0, 0, 9, 0, 0]

    There are only elements at [1, 8], and they have respective values [4, 9]. You
    would create the sparse histogram like::

        SparseHistogram([1, 8], [4, 9], 0)
    """

    def __init__(self, indexes, counts, order):
        if len(indexes) != len(counts):
            raise ValueError("indexes and counts must be same length")
        self.indexes = indexes
        self.counts = counts
        self.order = order

    def to_array(self):
        """Convert the sparse array to a dense numpy array.

        Returns:
            dense 1-d numpy array.
        """
        dense = np.zeros(hp.order2npix(self.order), dtype=np.int64)
        dense[self.indexes] = self.counts
        return dense

    def to_file(self, file_name):
        """Persist the sparse array to disk.

        NB: this saves as a sparse array, and so will likely have lower space requirements
        than saving the corresponding dense 1-d numpy array.
        """
        np.savez(file_name, indexes=self.indexes, counts=self.counts, order=self.order)

    def to_dense_file(self, file_name):
        """Persist the DENSE array to disk as a numpy array."""
        with open(file_name, "wb+") as file_handle:
            file_handle.write(self.to_array().data)

    @classmethod
    def from_file(cls, file_name):
        """Read sparse histogram from a file.

        Returns:
            new sparse histogram
        """
        npzfile = np.load(file_name)
        return cls(npzfile["indexes"], npzfile["counts"], npzfile["order"])


class HistogramAggregator:
    """Utility for aggregating sparse histograms."""

    def __init__(self, order):
        self.order = order
        self.full_histogram = np.zeros(hp.order2npix(order), dtype=np.int64)

    def add(self, other):
        """Add in another sparse histogram, updating this wrapper's array.

        Args:
            other (SparseHistogram): the wrapper containing the addend
        """
        if not isinstance(other, SparseHistogram):
            raise ValueError("Both addends should be SparseHistogram.")
        if self.order != other.order:
            raise ValueError(
                "The histogram partials have incompatible sizes due to different healpix orders."
            )
        if len(other.indexes) == 0:
            return
        self.full_histogram[other.indexes] += other.counts

    def to_sparse(self):
        """Return a SparseHistogram, based on non-zero values in this aggregation."""
        indexes = self.full_histogram.nonzero()[0]
        counts = self.full_histogram[indexes]
        return SparseHistogram(indexes, counts, self.order)
