#!/usr/bin/env python

"""
Lorentzian peak model

References::
    None
"""

from numpy import array, pi, asarray, arange, sqrt
from numpy import random


def lorentzian(E, Io, Eo, Gamma, A=0, B=0, C=0):
    """
    lorentzian with quadratic background::

        I = Io  (Gamma/2*pi) / ( (E-Eo)^2 + (Gamma/2)^2 ) + (A + B E + C E^2)
    """
    E = asarray(E)  # force E to be a numpy array
    return A + (B + C * E) * E + Io * (Gamma / 2 / pi) / ((E - Eo) ** 2 + (Gamma / 2) ** 2)


def simulate_events(params, xmin, xmax, npts=4000):
    """Generate a lorentzian dataset of npts between [min,max] from given params"""

    def gensample(F, xmin, xmax):
        a = arange(xmin, xmax, (xmax - xmin) / 200.0)
        ymin = 0
        ymax = F(a).max()
        while 1:
            t1 = random.random() * (xmax - xmin) + xmin
            t2 = random.random() * (ymax - ymin) + ymin
            t3 = F(t1)
            if t2 < t3:
                return t1

    fwd = lambda x: lorentzian(x, **params)
    return array([gensample(fwd, xmin, xmax) for _ in range(npts)])


def simulate_histogram(pars, Emin, Emax, dE, npts=4000):
    events = simulate_events(pars, Emin, Emax)
    E, I = histogram(events, dE, Emin, Emax)
    # print min(events),max(events)
    dI = sqrt(I)
    data = {"x": (E[1:] + E[:-1]) / 2.0, "y": I, "dy": dI}
    return data


def demo_data():
    # integrated intensity = 4000
    # center = 6400
    # width gamma = 180
    # background = -(E-6340)^2/1000 + 10
    bgC = 6340.0
    bgW = 10000.0
    A = 35 - bgC**2 / bgW
    B = 2.0 * bgC / bgW
    C = -1.0 / bgW
    # A,B,C = 0,0,0
    N = 4000
    pars = {"Eo": 6500.0, "Gamma": 180.0, "Io": 20 * N, "A": A, "B": B, "C": C}
    Emin, Emax = 6000, 6700
    dE = 20

    return pars, simulate_histogram(pars, Emin, Emax, dE, N)


# probably shouldn't be in here...
from numpy import histogram as numpyhisto


def histogram(data, binwidth, xmin, xmax):
    """
    generate bin-centered histogram of provided data

    return bins of given binwidth (and histogram) generated between [xmin,xmax]
    """
    edges = arange(xmin, xmax + binwidth * 0.9999999, binwidth)
    centers = edges + (0.5 * binwidth)
    histo, _ = numpyhisto(data, bins=edges)
    # print data.size, sum(histo), edges[0], edges[-1], min(data),max(data)
    return centers, histo


coeff, data = demo_data()


def demo():
    import pylab

    x, y, dy = data["x"], data["y"], data["dy"]
    A, B, C = coeff["A"], coeff["B"], coeff["C"]
    Io, Eo, Gamma = coeff["Io"], coeff["Eo"], coeff["Gamma"]
    pylab.errorbar(x, y, yerr=dy, label="data")
    pylab.plot(x, pylab.polyval([C, B, A], x), label="background")
    pylab.plot(x, lorentzian(x, Eo=Eo, Io=Io, Gamma=Gamma), label="peak")
    pylab.plot(x, lorentzian(x, **coeff), label="peak+bkg")
    pylab.legend()
    pylab.show()


if __name__ == "__main__":
    demo()
