import numpy as np
import scipy.stats as stats
from scipy.interpolate import PchipInterpolator
import matplotlib.pyplot as plt


class Histogram:
    def __init__(self, sample, bin_size, confidence_level=0.95):
        self.sample = sample
        self.bin_size = bin_size
        self.confidence_level = confidence_level
        self.bins = int((np.amax(sample) - np.amin(sample)) // bin_size + 1)
        self.__min = np.amin(sample) // bin_size - 1
        self.__shift = self.__min * bin_size + bin_size / 2
        self.__shifted_sample = sample - self.__shift
        self._x, self._y = self.__bin(self.__shifted_sample, bin_size, self.bins)

    def __bin(self, sample, bin_size, bins):
        assert bin_size > 0
        assert bins > 0
        assert type(bins) == int
        assert type(bin_size) in [int, float]

        x = np.linspace(0, bins * bin_size, bins + 1, endpoint=True)
        y = np.zeros(bins + 1)

        for s in sample:
            index = int((s // bin_size))
            if index >= bins:  # Ensure index is within bounds
                index = bins - 1
            y[index] += 1

        return x, y

    @property
    def x(self):
        return self._x + self.__shift + self.bin_size / 2

    @property
    def y(self):
        return self._y / (np.sum(self._y) * self.bin_size)
    
    @property
    def e(self):
        return self.error_estimate(self._y) / (np.sum(self._y) * self.bin_size)
                                       
    @property
    def counts(self):
        return self._y

    @property
    def counts_error(self):
        return self.error_estimate(self._y)
    
    def poisson_confidence_interval(self, k):
        alpha = 1 - self.confidence_level
        lower_bound = 0.5 * stats.chi2.ppf(alpha / 2, 2 * k)
        upper_bound = 0.5 * stats.chi2.ppf(1 - alpha / 2, 2 * (k + 1))
        
        if type(k) == np.ndarray:
            lower_bound[k == 0] = 0
        return lower_bound, upper_bound
    
    def error_estimate(self, k):
        lower_bound, upper_bound = self.poisson_confidence_interval(k)
        return (upper_bound - lower_bound) / 2
    
    def FWHM(self, resolution):
        y = PchipInterpolator(self.x, self.y)

        max_value = np.amax(self.y)
        half_max = max_value / 2

        peak_x = self.x[np.argmax(self.y)]

        min_x = np.amin(self.x)
        max_x = np.amax(self.x)

        left = bisection(self.x, y, half_max, resolution, min_x, peak_x)
        right = bisection(self.x, y, half_max, resolution, peak_x, max_x)

        return np.abs(right - left)
    
    def plot(self, **kwargs):
        
        x = self.x if 'x' not in kwargs else kwargs.pop('x')
        
        plt.rcParams["font.family"] = "monospace"
        plt.errorbar(x=self.x, 
                     y=self.y, 
                     yerr=self.e, 
                     fmt="s", 
                     color='black' if "color" not in kwargs else kwargs.pop("color"), 
                     zorder=1 if "zorder" not in kwargs else kwargs.pop("zorder"),
                     markerfacecolor='none' if "markerfacecolor" not in kwargs else kwargs.pop("markerfacecolor"),
                     label="Data" if "label" not in kwargs else kwargs.pop("label"))

    def show(self):
        plt.legend()
        plt.show()
        
def bisection(x, y, value, tolerance, left, right):
    assert len(x) > 1
    assert tolerance > 0

    if right - left <= tolerance:
        return (left + right) / 2

    mid = (left + right) / 2

    if y(mid) < value < y(right) or y(mid) > value > y(right):
        return bisection(x, y, value, tolerance, mid, right)
    else:
        return bisection(x, y, value, tolerance, left, mid)


if __name__ == '__main__':
    sample = np.random.normal(0, 1, 10000000)
    hist = Histogram(sample, 0.1)
    hist.plot(label='Data 1')
    hist.show()
    print(hist.FWHM(1e-5))