import numpy as np
import pickle

class histEq:
    def __init__(self,X_sample,levels=10000,nbins=10000, path=[]): #X_sample is a list of one or more images sampled from the distribution, levels is the ourput norm range
        if len(path) > 0:
            self.load(path)
        else:
            self.nbins = nbins
            self.init(X_sample,levels,nbins)

    def init(self,X_sample,levels,nbins):
        #format X
        X_f = np.concatenate(X_sample)
        self.N =  np.prod(X_f.shape) #number of pixels
        self.L = levels #grey levels to eq to

        hist, bins = np.histogram(X_f.ravel(), nbins)
        self.hist = hist/len(X_sample) #normalize to thenumber of images
        self.bins = bins + ((max(bins)-min(bins))/len(bins))/2 #center of each bin
        self.bins = self.bins[:nbins]
        self.cdf = np.cumsum(self.hist)
        self.cdf_min = np.min(self.cdf)

    def equalize(self,image):
        out = np.interp(image.flat, self.bins, self.cdf)
        return out.reshape(image.shape)

    def dequalize(self,image):
        out = np.interp(image.flat, self.cdf, self.bins)
        return out.reshape(image.shape)

    def save(self,path=None):
        if path is None:
            return [self.bins,self.cdf]
        with open(path, 'wb') as f:
            pickle.dump([self.bins,self.cdf], f)

    def load(self,path):
        with open(path, 'rb') as f:
            data = pickle.load(f)
        self.bins = data[0]
        self.cdf = data[1]