import numpy, scipy.special, scipy.optimize, math, pylab

class GaussianFit(object):
    def __init__(self, mean, std, bins, freqs, model_freqs, chi2):
        self.mean = mean
        self.std = std
        self.bins = bins
        self.freqs = freqs
        self.model_freqs = model_freqs
        self.chi2 = chi2

    def pdf(self, x): return 1.0/(2*numpy.pi)**0.5 / self.std * numpy.exp(-0.5 * (x - self.mean)**2 / self.std**2)
    
    def plot_data(self, color1 = '#33ff33', color2='#55ff55', label=None, *args, **vargs):
        return plotpdf(self.freqs, self.bins, bars=True, label=label, color1=color1, color2=color2)

    def plot_model(self, color = 'r', label=None, *args, **vargs):
        return plotpdf(self.model_freqs, self.bins, bars=False, label=label, color1=color)

    @staticmethod
    def build(x, n=100):
        x = numpy.asarray(x).flatten()
        mean = numpy.mean(x)
        std = numpy.std(x)
        model_freqs, bins = build_model(len(x), n, mean, std)        
        freqs, bins = numpy.histogram(x, bins, new=True)
        
        return GaussianFit(mean, std, bins, freqs, model_freqs, chi2(freqs, model_freqs, 2))


    def midbins(bins):
        "return an array of middles of the bins"
        bins = numpy.asarray(bins)
        return 0.5*(bins[:-1] + bins[1:])

    def chi2(freqs, model_freqs, deg_freedom = 2):
        "calculate chi^2 for counting test. model_freqs^0.5 are used as std. devs."
        return sum((freqs - model_freqs)**2 / model_freqs) / (len(freqs) - deg_freedom)

    def pdf(probs, bins):
        "return an estimate of probability density based on probabilities and bin boundaries"
        return probs / (bins[1:] - bins[:-1])

    def plotpdf(freqs, bins, bars=True, color1 = '#33ff33', color2='#55ff55', label=None, *args, **vargs):
        """
        plot probability density histogram based on a set of frequencies/bin boundaries.
        If 'bars' is True, bars are plotted, alternating between 'color1' and 'color2',
        otherwise a line is plotted in 'color1'.
        """
        if len(bins) != len(freqs) + 1:
            raise ValueError('Mismatching bins and frequencies')

        N = numpy.sum(freqs)

        if not numpy.isfinite(bins[0]):
            bins = bins[1:]
            freqs = freqs[1:]
        
        if not numpy.isfinite(bins[-1]):
            bins = bins[:-1]
            freqs = freqs[:-1]

        left = bins[:-1]
        width = bins[1:] - bins[:-1]
        height = freqs / width / N
    #    height = pdf(freqs / N, bins)

        if bars:
            a = pylab.bar(left = left[0::2], width = width[0::2], height = height[0::2], linewidth=0, color=color1, *args, **vargs)
            if label is not None:
                vargs['label'] = label
            return a + pylab.bar(left = left[1::2], width = width[1::2], height = height[1::2], linewidth=0, color=color2, *args, **vargs)
        else:
            if label is not None:
                vargs['label'] = label
            return pylab.plot(0.5 * (bins[1:] + bins[:-1]), height, color=color1, *args, **vargs)