from java.awt import Color
from jhplot  import *
from jhplot.io  import *
from jhplot.fit  import *
import jhplot

c1 = HPlot("Canvas",600,400)
c1.setGTitle("NegativeBinomial fit")
c1.visible()
c1.setNameX("P_{n}")
c1.setNameY("Probability")
c1.setRange(0,60,0.00001,0.5)
c1.setLegendPos(0.6,0.9,"NDC")
c1.setMarginLeft(90)
c1.setLogScale(1,1)

# generate NBD data
from jhplot.math.num.random import NegativeBinomialRandomVariable
h1 = H1D('Data',50, 0, 100)
r=NegativeBinomialRandomVariable(20,0.5)
for i in range(5000):
      h1.fill(r.nextRandomVariable(),1.0/5000)
p1=P1D("data",h1, 0)
p1.setColor(Color.blue)
c1.draw(p1)

f=HFitter("leastsquares") # prepare fits
print f.getFitMethod()
func=jhplot.fit.NegativeBinomial()
print func.parameterNames()
func.setParameter("norm",1)
func.setParameter("numbertrials",21)
func.setParameter("probsuccess",0.5)
f.setFunc(func)
f.setParRange("probsuccess",0.00001, 0.9999)
f.setParRange("numertrials",10, 1000000)
f.setParStep("numbertrials",1)
f.setParStep("probsuccess",0.01)
f.setRange(0,80)
f.fit(p1)

ff=f.getFittedFunc() # get fitted results
r=f.getResult()
Pars   = r.fittedParameters()
Errors = r.errors()
Names  = r.fittedParameterNames()
print "Fit results:"
for i in range(ff.numberOfParameters()):
  print Names[i]+" : "+str(Pars[i])+" +- "+str(Errors[i])

mess='χ^{2}/ndf='+str(round(r.quality()*r.ndf()))
mess=mess+" / "+str(r.ndf()) 
lab=HLabel(mess, 0.8, 0.69, 'NDC')
c1.add(lab) 

# draw function
f2 = F1D("NegativeBinomial fit",ff,0,80)
f2.setPoints(80) # since disribution discrite, make smaller steps
f2.setPenWidth(1)
c1.draw(f2)
print 'Quality=',r.quality(), ' NDF=',r.ndf()